add unit test for offloading (metaheader included) (#3202)

Summary:
Pull Request resolved: #3202

X-link: facebookresearch/FBGEMM#1560

Added unit test for the following conditions, and fixed related bugs:
1. ZCH fused optimizer with offloading
2. Added DRAM kernel for stat dict loading (with metaheader)and numerical accuracy
3. Applied return whole row for DRAM kernel.

Reviewed By: emlin, bobbyliujb

Differential Revision: D77474510

fbshipit-source-id: e307077acce9f40cf733d8382bc45c6a931ae1e4

Author: Kathy Xu

torchrec/distributed/batched_embedding_kernel.py

@@ -430,7 +430,7 @@ def __init__(  # noqa C901
         )
 
         all_optimizer_states = emb_module.get_optimizer_state(
-            sorted_id_tensor=sorted_id_tensors
+            sorted_id_tensor=sorted_id_tensors,
         )
         opt_param_list = [param["momentum1"] for param in all_optimizer_states]
         emb_table_config_copy = copy.deepcopy(self._config.embedding_tables)

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -421,20 +421,21 @@ def _compare_models(
                             src_wid = sd1[wid_key].local_shards()[local_shard_id].tensor
                             dst_wid = sd2[wid_key].local_shards()[local_shard_id].tensor
 
-                            sorted_src_wid, _ = torch.sort(src_wid.view(-1))
-                            sorted_dst_wid, _ = torch.sort(dst_wid.view(-1))
+                            sorted_src_wid = torch.sort(src_wid.view(-1))[0]
+                            sorted_dst_wid = torch.sort(dst_wid.view(-1))[0]
                             assert torch.equal(sorted_src_wid, sorted_dst_wid)
-                            src_tensor = src.tensor.get_weights_by_ids(src_wid)
-                            dst_tensor = dst.tensor.get_weights_by_ids(dst_wid)
+                            # kvz zch emb table comparison, id is non-continuous
+                            src_tensor = src.tensor.get_weights_by_ids(sorted_src_wid)
+                            dst_tensor = dst.tensor.get_weights_by_ids(sorted_dst_wid)
                         else:
                             # normal ssd offloading emb table comparison
                             src_tensor = src.tensor.full_tensor()
                             dst_tensor = dst.tensor.full_tensor()
                     else:
-                        src_tensor = src.tensor
-                        dst_tensor = dst.tensor
+                        src_tensor = torch.sort(src.tensor.flatten()).values
+                        dst_tensor = torch.sort(dst.tensor.flatten()).values
                     if is_deterministic:
-                        self.assertTrue(torch.equal(src_tensor, dst_tensor))
+                        self.assertTrue(torch.allclose(src_tensor, dst_tensor))
                     else:
                         rtol, atol = _get_default_rtol_and_atol(src_tensor, dst_tensor)
                         torch.testing.assert_close(

torchrec/distributed/tests/test_model_parallel_nccl_ssd_single_gpu.py

@@ -985,6 +985,7 @@ def _copy_fused_modules_into_ssd_emb_modules(
         kernel_type=st.sampled_from(
             [
                 EmbeddingComputeKernel.SSD_VIRTUAL_TABLE.value,
+                EmbeddingComputeKernel.DRAM_VIRTUAL_TABLE.value,
             ]
         ),
         sharding_type=st.sampled_from(
@@ -1387,6 +1388,10 @@ def _copy_ssd_emb_modules(
                         pmt2 = sharded_t2.local_shards()[0].tensor
                         pmt2.wrapped.set_weights_and_ids(w1, w1_id.view(-1))
 
+                    # Remove the cache to force state dict read from backend again
+                    emb_module1._split_weights_res = None
+                    emb_module2._split_weights_res = None
+
                     # purge after loading. This is needed, since we pass a batch
                     # through dmp when instantiating them.
                     emb_module1.purge()