Remove debug_split_optimizer_states (#4397)

Summary:
Pull Request resolved: #4397

- Remove `debug_split_optimizer_states` from training.py, since it has been
  superseded by `split_optimizer_states`

Reviewed By: duduyi2013

Differential Revision: D77256897

fbshipit-source-id: 2897a7dba9b0477be9ee00feacf1273e6158027e

Author: q10

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -940,7 +940,6 @@ def __init__(
             )
             # pyre-ignore
             self.stats_reporter.register_stats(self.l2_num_cache_misses_stats_name)
-            # pyre-ignore
             self.stats_reporter.register_stats(self.l2_num_cache_lookups_stats_name)
             self.stats_reporter.register_stats(self.l2_num_cache_evictions_stats_name)
             self.stats_reporter.register_stats(self.l2_cache_free_mem_stats_name)
@@ -1083,7 +1082,7 @@ def _report_duration(
         """
         recorded_itr, stream_cnt, report_val = self.prefetch_duration_us
         duration = dur_ms
-        if time_unit == "us":  # pyre-ignore
+        if time_unit == "us":
             duration = dur_ms * 1000
         if it_step == recorded_itr:
             report_val = max(report_val, duration)
@@ -1124,7 +1123,6 @@ def record_function_via_dummy_profile_factory(
 
             def func(
                 name: str,
-                # pyre-ignore[2]
                 fn: Callable[..., Any],
                 *args: Any,
                 **kwargs: Any,
@@ -2168,64 +2166,10 @@ def forward(
         )
 
         if self.optimizer == OptimType.EXACT_ROWWISE_ADAGRAD:
-            # pyre-fixme[7]: Expected `Tensor` but got implicit return value of `None`.
             return invokers.lookup_rowwise_adagrad_ssd.invoke(
                 common_args, self.optimizer_args, momentum1
             )
 
-    @torch.jit.ignore
-    def debug_split_optimizer_states(self) -> List[Tuple[torch.Tensor, int, int]]:
-        """
-        Returns a list of optimizer states, table_input_id_start, table_input_id_end, split by table
-        Testing only
-        """
-        (rows, _) = zip(*self.embedding_specs)
-
-        rows_cumsum = [0] + list(itertools.accumulate(rows))
-        if self.kv_zch_params:
-            opt_list = []
-            table_offset = 0
-            for t, row in enumerate(rows):
-                # pyre-ignore
-                bucket_id_start, bucket_id_end = self.kv_zch_params.bucket_offsets[t]
-                # pyre-ignore
-                bucket_size = self.kv_zch_params.bucket_sizes[t]
-                table_input_id_start = (
-                    min(bucket_id_start * bucket_size, row) + table_offset
-                )
-                table_input_id_end = (
-                    min(bucket_id_end * bucket_size, row) + table_offset
-                )
-
-                # TODO: this is a hack for preallocated optimizer, update this part once we have optimizer offloading
-                unlinearized_id_tensor = self._ssd_db.get_keys_in_range_by_snapshot(
-                    table_input_id_start,
-                    table_input_id_end,
-                    0,  # no need for table offest, as optimizer is preallocated using table offset
-                    None,
-                )
-                sorted_offsets, _ = torch.sort(unlinearized_id_tensor.view(-1))
-                opt_list.append(
-                    (
-                        self.momentum1_dev.detach()[sorted_offsets],
-                        table_input_id_start - table_offset,
-                        table_input_id_end - table_offset,
-                    )
-                )
-                table_offset += row
-            return opt_list
-        else:
-            return [
-                (
-                    self.momentum1_dev.detach()[
-                        rows_cumsum[t] : rows_cumsum[t + 1]
-                    ].view(row),
-                    -1,
-                    -1,
-                )
-                for t, row in enumerate(rows)
-            ]
-
     @torch.jit.ignore
     def _split_optimizer_states_non_kv_zch(
         self,
@@ -2344,6 +2288,7 @@ def split_optimizer_states(
             table_offset += emb_height
         logging.info(
             f"KV ZCH tables split_optimizer_states query latency: {(time.time() - start_time) * 1000} ms, "
+            # pyre-ignore [16]
             f"num ids list: {[ids.numel() for ids in sorted_id_tensor]}"
         )
         return opt_list
@@ -2623,6 +2568,7 @@ def split_embedding_weights(
         )
         if self.kv_zch_params is not None:
             logging.info(
+                # pyre-ignore [16]
                 f"num ids list: {[ids.numel() for ids in bucket_sorted_id_splits]}"
             )
 
@@ -2946,7 +2892,7 @@ def _report_ssd_l1_cache_stats(self) -> None:
                     / passed_steps
                 ),
             )
-            # pyre-ignore
+
             self.stats_reporter.report_data_amount(
                 iteration_step=self.step,
                 event_name=f"ssd_tbe.prefetch.cache_stats.{stat_index.name.lower()}",
@@ -2973,35 +2919,35 @@ def _report_ssd_io_stats(self) -> None:
         bwd_l1_cnflct_miss_write_back_dur = ssd_io_duration[3]
         flush_write_dur = ssd_io_duration[4]
 
-        # pyre-ignore
+        # pyre-ignore [16]
         self.stats_reporter.report_duration(
             iteration_step=self.step,
             event_name="ssd.io_duration.read_us",
             duration_ms=ssd_read_dur_us,
             time_unit="us",
         )
-        # pyre-ignore
+
         self.stats_reporter.report_duration(
             iteration_step=self.step,
             event_name="ssd.io_duration.write.fwd_rocksdb_read_us",
             duration_ms=fwd_rocksdb_read_dur,
             time_unit="us",
         )
-        # pyre-ignore
+
         self.stats_reporter.report_duration(
             iteration_step=self.step,
             event_name="ssd.io_duration.write.fwd_l1_eviction_us",
             duration_ms=fwd_l1_eviction_dur,
             time_unit="us",
         )
-        # pyre-ignore
+
         self.stats_reporter.report_duration(
             iteration_step=self.step,
             event_name="ssd.io_duration.write.bwd_l1_cnflct_miss_write_back_us",
             duration_ms=bwd_l1_cnflct_miss_write_back_dur,
             time_unit="us",
         )
-        # pyre-ignore
+
         self.stats_reporter.report_duration(
             iteration_step=self.step,
             event_name="ssd.io_duration.write.flush_write_us",
@@ -3023,25 +2969,25 @@ def _report_ssd_mem_usage(
         memtable_usage = mem_usage_list[2]
         block_cache_pinned_usage = mem_usage_list[3]
 
-        # pyre-ignore
+        # pyre-ignore [16]
         self.stats_reporter.report_data_amount(
             iteration_step=self.step,
             event_name="ssd.mem_usage.block_cache",
             data_bytes=block_cache_usage,
         )
-        # pyre-ignore
+
         self.stats_reporter.report_data_amount(
             iteration_step=self.step,
             event_name="ssd.mem_usage.estimate_table_reader",
             data_bytes=estimate_table_reader_usage,
         )
-        # pyre-ignore
+
         self.stats_reporter.report_data_amount(
             iteration_step=self.step,
             event_name="ssd.mem_usage.memtable",
             data_bytes=memtable_usage,
         )
-        # pyre-ignore
+
         self.stats_reporter.report_data_amount(
             iteration_step=self.step,
             event_name="ssd.mem_usage.block_cache_pinned",

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -526,8 +526,6 @@ def generate_ssd_tbes(
         if share_table:
             # autograd with shared embedding only works for exact
             table_to_replicate = T // 2
-            # pyre-fixme[6]: For 2nd param expected `Embedding` but got
-            #  `Union[Embedding, EmbeddingBag]`.
             feature_table_map.insert(table_to_replicate, table_to_replicate)
             emb_ref.insert(table_to_replicate, emb_ref[table_to_replicate])
 
@@ -746,6 +744,17 @@ def execute_ssd_forward_(
         )
         return output_ref_list, output
 
+    def split_optimizer_states_(
+        self, emb: SSDTableBatchedEmbeddingBags
+    ) -> List[torch.Tensor]:
+        _, bucket_asc_ids_list, _ = emb.split_embedding_weights(
+            no_snapshot=False, should_flush=True
+        )
+
+        return emb.split_optimizer_states(
+            bucket_asc_ids_list, no_snapshot=False, should_flush=True
+        )
+
     @given(
         **default_st, backend_type=st.sampled_from([BackendType.SSD, BackendType.DRAM])
     )
@@ -937,7 +946,7 @@ def test_ssd_backward_adagrad(
         )
 
         # Compare optimizer states
-        split_optimizer_states = [s for (s, _, _) in emb.debug_split_optimizer_states()]
+        split_optimizer_states = self.split_optimizer_states_(emb)
         for f, t in self.get_physical_table_arg_indices_(emb.feature_table_map):
             # pyre-fixme[16]: Optional type has no attribute `float`.
             ref_optimizer_state = emb_ref[f].weight.grad.float().to_dense().pow(2)
@@ -1079,7 +1088,7 @@ def test_ssd_emb_state_dict(
             else 1.0e-2
         )
 
-        split_optimizer_states = [s for (s, _, _) in emb.debug_split_optimizer_states()]
+        split_optimizer_states = self.split_optimizer_states_(emb)
         emb.flush()
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
@@ -1168,7 +1177,7 @@ def execute_ssd_cache_pipeline_(  # noqa C901
         )
 
         optimizer_states_ref = [
-            s.clone().float() for (s, _, _) in emb.debug_split_optimizer_states()
+            s.clone().float() for s in self.split_optimizer_states_(emb)
         ]
 
         Es = [emb.embedding_specs[t][0] for t in range(T)]
@@ -1312,15 +1321,12 @@ def _prefetch(b_it: int) -> int:
                 emb.flush()
 
             # Compare optimizer states
-            split_optimizer_states = [
-                s for (s, _, _) in emb.debug_split_optimizer_states()
-            ]
+            split_optimizer_states = self.split_optimizer_states_(emb)
             for f, t in self.get_physical_table_arg_indices_(emb.feature_table_map):
                 optim_state_r = optimizer_states_ref[t]
                 optim_state_t = split_optimizer_states[t]
                 emb_r = emb_ref[f]
 
-                # pyre-fixme[16]: Optional type has no attribute `float`.
                 optim_state_r.add_(
                     # pyre-fixme[16]: `Optional` has no attribute `float`.
                     emb_r.weight.grad.float()
@@ -2252,7 +2258,9 @@ def test_apply_kv_state_dict(
             )
 
             torch.testing.assert_close(
+                # pyre-ignore [16]
                 emb_state_dict_list[t].full_tensor()[sorted_ids.indices],
+                # pyre-ignore [16]
                 emb_state_dict_list2[t].full_tensor()[sorted_ids2.indices],
                 atol=tolerance,
                 rtol=tolerance,