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

Summary: Pull Request resolved: #1695 Allow zch to use uvm. Reviewed By: dstaay-fb Differential Revision: D53598269 fbshipit-source-id: 53951d8ef0998202026343590d096cc0ea4e415b
pytorch · Feb 9, 2024 · eaf7f79 · eaf7f79
1 parent 1c47e60
commit eaf7f79
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: