Add uvm to allowed compute kernel for zch (pytorch#1695)

Summary: Allow zch to use uvm. Differential Revision: D53598269
henrylhtsang · Feb 9, 2024 · 7c44f3d · 7c44f3d
1 parent 47441cf
commit 7c44f3d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/torchrec/distributed/mc_embedding_modules.py b/torchrec/distributed/mc_embedding_modules.py
@@ -267,6 +267,7 @@ def compute_kernels(
         return [
             EmbeddingComputeKernel.FUSED.value,
             EmbeddingComputeKernel.FUSED_UVM_CACHING.value,
+            EmbeddingComputeKernel.FUSED_UVM.value,
         ]
 
     def sharding_types(self, compute_device_type: str) -> List[str]:

diff --git a/torchrec/distributed/planner/tests/test_enumerators.py b/torchrec/distributed/planner/tests/test_enumerators.py
@@ -782,7 +782,10 @@ def test_filter_compute_kernels_mch_ebc(self) -> None:
 
         self.assertEqual(
             set(allowed_compute_kernels),
-            {EmbeddingComputeKernel.FUSED.value},
+            {
+                EmbeddingComputeKernel.FUSED.value,
+                EmbeddingComputeKernel.FUSED_UVM.value,
+            },
         )
 
     def test_filter_compute_kernels_mch_ebc_no_available(self) -> None: