Backout

Summary:
X-link: facebookresearch/FBGEMM#887

Backout D68055168 as it seems to break pyper and causes S498612.

Differential Revision: D70996903

Author: spcyppt

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -447,7 +447,7 @@ def generate() -> None:
                     ssd_optimizers.append(optim)
 
             BackwardSplitGenerator.generate_backward_split(
-                ssd_tensors=ssd_tensors, aux_args=aux_args, **optimizer
+                ssd_tensors=ssd_tensors, **optimizer
             )
         BackwardSplitGenerator.generate_rocm_backward_split()
 

fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu

@@ -601,7 +601,6 @@ Tensor {{ embedding_cuda_op }}(
 
     {%- if "learning_rate" in args.split_kernel_arg_names %}
     // convert `learning rate` to float since `learning rate` is float in kernels
-    TORCH_CHECK(learning_rate_tensor.is_cpu(), "learning_rate_tensor tensor needs to be on CPU. Ensure learning_rate_tensor is on CPU or contact FBGEMM team if you get this error.")
     const float learning_rate = learning_rate_tensor.item<float>();
     {%- endif %}
 

fbgemm_gpu/codegen/training/python/split_embedding_codegen_lookup_invoker.template

@@ -597,7 +597,7 @@ class SplitTableBatchedEmbeddingBagsCodegen(nn.Module):
     """
 
     embedding_specs: List[Tuple[int, int, EmbeddingLocation, ComputeDevice]]
-    optimizer_args: invokers.lookup_args.OptimizerArgsPT2
+    optimizer_args: invokers.lookup_args.OptimizerArgs
     lxu_cache_locations_list: List[Tensor]
     lxu_cache_locations_empty: Tensor
     timesteps_prefetched: List[int]
@@ -926,13 +926,6 @@ def __init__(  # noqa C901
             "feature_dims",
             torch.tensor(feature_dims, device="cpu", dtype=torch.int64),
         )
-        (_info_B_num_bits, _info_B_mask) = torch.ops.fbgemm.get_infos_metadata(
-            self.D_offsets,  # unused tensor
-            1,  # max_B
-            T,  # T
-        )
-        self.info_B_num_bits: int = _info_B_num_bits
-        self.info_B_mask: int = _info_B_mask
 
         # A flag for indicating whether all embedding tables are placed in the
         # same locations
@@ -1077,9 +1070,6 @@ def __init__(  # noqa C901
             # which should not be effective when CounterBasedRegularizationDefinition
             # and CowClipDefinition are not used
             counter_halflife = -1
-        learning_rate_tensor = torch.tensor(
-            learning_rate, device=torch.device("cpu"), dtype=torch.float
-        )
 
         # TO DO: Enable this on the new interface
         # learning_rate_tensor = torch.tensor(
@@ -1095,12 +1085,12 @@ def __init__(  # noqa C901
                 "`use_rowwise_bias_correction` is only supported for OptimType.ADAM",
             )
 
-        self.optimizer_args = invokers.lookup_args.OptimizerArgsPT2(
+        self.optimizer_args = invokers.lookup_args.OptimizerArgs(
             stochastic_rounding=stochastic_rounding,
             gradient_clipping=gradient_clipping,
             max_gradient=max_gradient,
             max_norm=max_norm,
-            learning_rate_tensor=learning_rate_tensor,
+            learning_rate=learning_rate,
             eps=eps,
             beta1=beta1,
             beta2=beta2,
@@ -1885,7 +1875,7 @@ def forward(  # noqa: C901
             if len(self.lxu_cache_locations_list) == 0
             else self.lxu_cache_locations_list.pop(0)
         )
-        common_args = invokers.lookup_args.CommonArgsPT2(
+        common_args = invokers.lookup_args.CommonArgs(
             placeholder_autograd_tensor=self.placeholder_autograd_tensor,
             # pyre-fixme[6]: For 2nd argument expected `Tensor` but got
             #  `Union[Module, Tensor]`.
@@ -1932,8 +1922,6 @@ def forward(  # noqa: C901
             is_experimental=self.is_experimental,
             use_uniq_cache_locations_bwd=self.use_uniq_cache_locations_bwd,
             use_homogeneous_placements=self.use_homogeneous_placements,
-            info_B_num_bits=self.info_B_num_bits,
-            info_B_mask=self.info_B_mask,
         )
 
         if self.optimizer == OptimType.NONE:
@@ -2046,6 +2034,7 @@ def forward(  # noqa: C901
                     momentum1,
                     momentum2,
                     iter_int,
+                    self.use_rowwise_bias_correction,
                     row_counter=(
                         row_counter if self.use_rowwise_bias_correction else None
                     ),
@@ -2931,7 +2920,7 @@ def _set_learning_rate(self, lr: float) -> float:
         Helper function to script `set_learning_rate`.
         Note that returning None does not work.
         """
-        self.optimizer_args.learning_rate_tensor.fill_(lr)
+        self.optimizer_args = self.optimizer_args._replace(learning_rate=lr)
         return 0.0
 
     @torch.jit.ignore
@@ -3446,22 +3435,6 @@ def prepare_inputs(
             offsets, batch_size_per_feature_per_rank
         )
 
-        vbe = vbe_metadata.B_offsets is not None
-        # TODO: assert vbe_metadata.B_offsets.numel() - 1 == T
-        # T = self.D_offsets.numel() - 1
-        # vbe_metadata.B_offsets causes jit to fail for cogwheel forward compatibility test
-        # max_B = int(vbe_metadata.max_B) if vbe else int(offsets.numel() - 1 / T)
-
-        # TODO:  max_B <= self.info_B_mask
-        # cannot use assert as it breaks pt2 compile for dynamic shape
-        # Need to use torch._check for dynamic shape and cannot construct fstring, use constant string.
-        # cannot use lambda as it fails jit script.
-        # torch._check is not supported in jitscript
-        # torch._check(
-        #     max_B <= self.info_B_mask,
-        #     "Not enough infos bits to accommodate T and B.",
-        # )
-
         # TODO: remove this and add an assert after updating
         # bounds_check_indices to support different indices type and offset
         # type
@@ -3489,6 +3462,7 @@ def prepare_inputs(
                 per_sample_weights = per_sample_weights.float()
 
         if self.bounds_check_mode_int != BoundsCheckMode.NONE.value:
+            vbe = vbe_metadata.B_offsets is not None
             # Compute B info and VBE metadata for bounds_check_indices only if
             # VBE and bounds check indices v2 are used
             if vbe and self.use_bounds_check_v2:
@@ -3502,7 +3476,11 @@ def prepare_inputs(
                 assert isinstance(
                     output_offsets_feature_rank, Tensor
                 ), "output_offsets_feature_rank must be tensor"
-
+                info_B_num_bits, info_B_mask = torch.ops.fbgemm.get_infos_metadata(
+                    B_offsets,  # unused tensor
+                    vbe_metadata.max_B,
+                    B_offsets.numel() - 1,  # T
+                )
                 row_output_offsets, b_t_map = torch.ops.fbgemm.generate_vbe_metadata(
                     B_offsets,
                     B_offsets_rank_per_feature,
@@ -3511,11 +3489,13 @@ def prepare_inputs(
                     self.max_D,
                     self.is_nobag,
                     vbe_metadata.max_B_feature_rank,
-                    self.info_B_num_bits,
+                    info_B_num_bits,
                     offsets.numel() - 1,  # total_B
                 )
             else:
                 b_t_map = None
+                info_B_num_bits = -1
+                info_B_mask = -1
 
             torch.ops.fbgemm.bounds_check_indices(
                 self.rows_per_table,
@@ -3527,8 +3507,8 @@ def prepare_inputs(
                 B_offsets=vbe_metadata.B_offsets,
                 max_B=vbe_metadata.max_B,
                 b_t_map=b_t_map,
-                info_B_num_bits=self.info_B_num_bits,
-                info_B_mask=self.info_B_mask,
+                info_B_num_bits=info_B_num_bits,
+                info_B_mask=info_B_mask,
                 bounds_check_version=self.bounds_check_version,
             )
 

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -228,12 +228,6 @@ def __init__(
             torch.tensor(feature_dims, device="cpu", dtype=torch.int64),
         )
 
-        (self.info_B_num_bits, self.info_B_mask) = torch.ops.fbgemm.get_infos_metadata(
-            self.D_offsets,  # unused tensor
-            1,  # max_B
-            T,  # T
-        )
-
         assert cache_sets > 0
         element_size = weights_precision.bit_rate() // 8
         assert (
@@ -550,15 +544,12 @@ def __init__(
             )
         cowclip_regularization = CowClipDefinition()
 
-        learning_rate_tensor = torch.tensor(
-            learning_rate, device=torch.device("cpu"), dtype=torch.float
-        )
-        self.optimizer_args = invokers.lookup_args_ssd.OptimizerArgsPT2(
+        self.optimizer_args = invokers.lookup_args_ssd.OptimizerArgs(
             stochastic_rounding=stochastic_rounding,
             gradient_clipping=gradient_clipping,
             max_gradient=max_gradient,
             max_norm=max_norm,
-            learning_rate_tensor=learning_rate_tensor,
+            learning_rate=learning_rate,
             eps=eps,
             beta1=beta1,
             beta2=beta2,
@@ -1639,7 +1630,7 @@ def forward(
             vbe_metadata.max_B,
         )
 
-        common_args = invokers.lookup_args_ssd.CommonArgsPT2(
+        common_args = invokers.lookup_args_ssd.CommonArgs(
             placeholder_autograd_tensor=self.placeholder_autograd_tensor,
             output_dtype=self.output_dtype,
             dev_weights=self.weights_dev,
@@ -1674,8 +1665,6 @@ def forward(
             },
             # pyre-fixme[6]: Expected `lookup_args_ssd.VBEMetadata` but got `lookup_args.VBEMetadata`
             vbe_metadata=vbe_metadata,
-            info_B_num_bits=self.info_B_num_bits,
-            info_B_mask=self.info_B_mask,
         )
 
         self.timesteps_prefetched.pop(0)
@@ -1826,7 +1815,7 @@ def _set_learning_rate(self, lr: float) -> float:
         Helper function to script `set_learning_rate`.
         Note that returning None does not work.
         """
-        self.optimizer_args.learning_rate_tensor.fill_(lr)
+        self.optimizer_args = self.optimizer_args._replace(learning_rate=lr)
         return 0.0
 
     def flush(self) -> None:

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -82,10 +82,8 @@
     "use_cache": st.booleans(),
     "cache_algorithm": st.sampled_from(CacheAlgorithm),
     "use_cpu": use_cpu_strategy(),
-    "output_dtype": (
-        st.sampled_from([SparseType.FP32, SparseType.FP16, SparseType.BF16])
-        if gpu_available
-        else st.sampled_from([SparseType.FP32, SparseType.FP16])
+    "output_dtype": st.sampled_from(
+        [SparseType.FP32, SparseType.FP16, SparseType.BF16]
     ),
 }
 

fbgemm_gpu/test/tbe/training/backward_adagrad_common.py

@@ -190,7 +190,7 @@ def apply_gwd(
             apply_gwd_per_table(
                 prev_iter_values,
                 weights_values,
-                emb.optimizer_args.learning_rate_tensor.item(),
+                emb.optimizer_args.learning_rate,
                 emb.optimizer_args.weight_decay,
                 step,
                 emb.current_device,

fbgemm_gpu/test/tbe/training/backward_adagrad_global_weight_decay_test.py

@@ -76,9 +76,6 @@
         "test_faketensor__test_forward_gpu_uvm_cache_int8": [
             unittest.skip("Operator not implemented for Meta tensors"),
         ],
-        "test_faketensor__test_forward_cpu_fp32": [
-            unittest.skip("Operator not implemented for Meta tensors"),
-        ],
         # TODO: Make it compatible with opcheck tests
         "test_faketensor__test_forward_gpu_uvm_cache_fp16": [
             unittest.skip(

fbgemm_gpu/test/tbe/training/forward_test.py

@@ -594,7 +594,7 @@ def test_update_hyper_parameters(self) -> None:
         } | {"lr": 1.0, "lower_bound": 2.0}
         cc.update_hyper_parameters(updated_parameters)
         self.assertAlmostEqual(
-            cc.optimizer_args.learning_rate_tensor.item(), updated_parameters["lr"]
+            cc.optimizer_args.learning_rate, updated_parameters["lr"]
         )
         self.assertAlmostEqual(cc.optimizer_args.eps, updated_parameters["eps"])
         self.assertAlmostEqual(cc.optimizer_args.beta1, updated_parameters["beta1"])