make lazy init tunable

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -150,6 +150,7 @@ def __init__(
         enable_async_update: bool = True,  # whether enable L2/rocksdb write to async background thread
         # if > 0, insert all kv pairs to rocksdb at init time, in chunks of *bulk_init_chunk_size* rows
         bulk_init_chunk_size: int = 0,
+        lazy_bulk_init_enabled: bool = False,
     ) -> None:
         super(SSDTableBatchedEmbeddingBags, self).__init__()
 
@@ -437,7 +438,7 @@ def __init__(
                 f"passed_in_path={ssd_directory}, num_shards={ssd_rocksdb_shards},num_threads={ssd_rocksdb_shards},"
                 f"memtable_flush_period={ssd_memtable_flush_period},memtable_flush_offset={ssd_memtable_flush_offset},"
                 f"l0_files_per_compact={ssd_l0_files_per_compact},max_D={self.max_D},rate_limit_mbps={ssd_rate_limit_mbps},"
-                f"size_ratio={ssd_size_ratio},compaction_trigger={ssd_compaction_trigger},"
+                f"size_ratio={ssd_size_ratio},compaction_trigger={ssd_compaction_trigger}, lazy_bulk_init_enabled={lazy_bulk_init_enabled},"
                 f"write_buffer_size_per_tbe={ssd_rocksdb_write_buffer_size},max_write_buffer_num_per_db_shard={ssd_max_write_buffer_num},"
                 f"uniform_init_lower={ssd_uniform_init_lower},uniform_init_upper={ssd_uniform_init_upper},"
                 f"row_storage_bitwidth={weights_precision.bit_rate()},block_cache_size_per_tbe={ssd_block_cache_size_per_tbe},"
@@ -470,7 +471,10 @@ def __init__(
             if self.bulk_init_chunk_size > 0:
                 self.ssd_uniform_init_lower: float = ssd_uniform_init_lower
                 self.ssd_uniform_init_upper: float = ssd_uniform_init_upper
-                self._lazy_initialize_ssd_tbe()
+                if lazy_bulk_init_enabled:
+                    self._lazy_initialize_ssd_tbe()
+                else:
+                    self._insert_all_kv()
         else:
             # pyre-fixme[4]: Attribute must be annotated.
             # pyre-ignore[16]

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -218,6 +218,7 @@ def generate_ssd_tbes(
         share_table: bool = False,
         prefetch_pipeline: bool = False,
         bulk_init_chunk_size: int = 0,
+        lazy_bulk_init_enabled: bool = False,
     ) -> Tuple[SSDTableBatchedEmbeddingBags, List[torch.nn.EmbeddingBag]]:
         """
         Generate embedding modules (i,e., SSDTableBatchedEmbeddingBags and
@@ -295,9 +296,10 @@ def generate_ssd_tbes(
             bounds_check_mode=BoundsCheckMode.WARNING,
             l2_cache_size=8,
             bulk_init_chunk_size=bulk_init_chunk_size,
+            lazy_bulk_init_enabled=lazy_bulk_init_enabled,
         ).cuda()
 
-        if bulk_init_chunk_size > 0:
+        if bulk_init_chunk_size > 0 and lazy_bulk_init_enabled:
             self.assertIsNotNone(
                 emb.lazy_init_thread,
                 "if bulk_init_chunk_size > 0, lazy_init_thread must be set and it should not be force-synchronized yet",
@@ -696,9 +698,12 @@ def test_ssd_backward_adagrad(
 
     @given(
         bulk_init_chunk_size=st.sampled_from([0, 100]),
+        lazy_bulk_init_enabled=st.booleans(),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
-    def test_ssd_emb_state_dict(self, bulk_init_chunk_size: int) -> None:
+    def test_ssd_emb_state_dict(
+        self, bulk_init_chunk_size: int, lazy_bulk_init_enabled: bool
+    ) -> None:
         # Constants
         lr = 0.5
         eps = 0.2
@@ -732,6 +737,7 @@ def test_ssd_emb_state_dict(self, bulk_init_chunk_size: int) -> None:
             output_dtype=output_dtype,
             share_table=True,
             bulk_init_chunk_size=bulk_init_chunk_size,
+            lazy_bulk_init_enabled=lazy_bulk_init_enabled,
         )
 
         Es = [emb.embedding_specs[t][0] for t in range(T)]