Do not pass opt type in hstu umia st publish (#5125)

Summary:
Pull Request resolved: #5125

X-link: meta-pytorch/torchrec#3544

X-link: https://github.com/facebookresearch/FBGEMM/pull/2127

Previously if kvzch table enable PARTIAL_ROWWISE_ADAM opt type. It will pass PARTIAL_ROWWISE_ADAM to all sharder as fused param, which will let sharder init opt with PARTIAL_ROWWISE_ADAM and will cause OOM issue. This diff is changing only pass opt type PARTIAL_ROWWISE_ADAM to KVZCH tbe and avoid OOM issue.

Reviewed By: steven1327

Differential Revision: D86787539

fbshipit-source-id: eeeed6a449e8ea130fc684b58e6f15e5f7418e3a

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -11,7 +11,7 @@
 
 import enum
 from dataclasses import dataclass
-from typing import NamedTuple, Optional
+from typing import FrozenSet, NamedTuple, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -249,6 +249,8 @@ class KVZCHParams(NamedTuple):
     eviction_policy: EvictionPolicy = EvictionPolicy()
     embedding_cache_mode: bool = False
     load_ckpt_without_opt: bool = False
+    optimizer_type_for_st: Optional[str] = None
+    optimizer_state_dtypes_for_st: Optional[FrozenSet[Tuple[str, int]]] = None
 
     def validate(self) -> None:
         assert len(self.bucket_offsets) == len(self.bucket_sizes), (
@@ -274,6 +276,10 @@ class KVZCHTBEConfig(NamedTuple):
     threshold_calculation_bucket_num: Optional[int] = 1000000  # 1M
     # When true, we only save weight to kvzch backend and not optimizer state.
     load_ckpt_without_opt: bool = False
+    # [DO NOT USE] This is for st publish only, do not set it in your config
+    optimizer_type_for_st: Optional[str] = None
+    # [DO NOT USE] This is for st publish only, do not set it in your config
+    optimizer_state_dtypes_for_st: Optional[FrozenSet[Tuple[str, int]]] = None
 
 
 class BackendType(enum.IntEnum):

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -249,6 +249,30 @@ def __init__(
                 assert self.optimizer in [
                     OptimType.EXACT_ROWWISE_ADAGRAD
                 ], f"only EXACT_ROWWISE_ADAGRAD supports embedding cache mode, but got {self.optimizer}"
+            if self.is_st_publish:
+                if (
+                    # pyre-ignore [16]
+                    self.kv_zch_params.optimizer_type_for_st
+                    == OptimType.PARTIAL_ROWWISE_ADAM.value
+                ):
+                    self.optimizer = OptimType.PARTIAL_ROWWISE_ADAM
+                    logging.info(
+                        f"Override optimizer type with {self.optimizer=} for st publish"
+                    )
+                if (
+                    # pyre-ignore [16]
+                    self.kv_zch_params.optimizer_state_dtypes_for_st
+                    is not None
+                ):
+                    optimizer_state_dtypes = {}
+                    for k, v in dict(
+                        self.kv_zch_params.optimizer_state_dtypes_for_st
+                    ).items():
+                        optimizer_state_dtypes[k] = SparseType.from_int(v)
+                    self.optimizer_state_dtypes = optimizer_state_dtypes
+                    logging.info(
+                        f"Override optimizer_state_dtypes with {self.optimizer_state_dtypes=} for st publish"
+                    )
 
         self.pooling_mode = pooling_mode
         self.bounds_check_mode_int: int = bounds_check_mode.value