Re-organize TBE tests, pt 4 (#2274)

Summary: Pull Request resolved: #2274 - Re-organize TBE tests, pt 4 Reviewed By: spcyppt Differential Revision: D52893922 fbshipit-source-id: 6d916b9ecf0c4b1f212552ca62706bbc1a2d56c5
pytorch · Jan 20, 2024 · f78a12b · f78a12b
1 parent 2cea7db
commit f78a12b
Show file tree

Hide file tree

Showing 9 changed files with 614 additions and 597 deletions.
diff --git a/.github/scripts/fbgemm_gpu_test.bash b/.github/scripts/fbgemm_gpu_test.bash
@@ -84,14 +84,12 @@ run_fbgemm_gpu_tests () {
 
   # These are either non-tests or currently-broken tests in both FBGEMM_GPU and FBGEMM_GPU-CPU
   local files_to_skip=(
-    ./tbe/split_table_batched_embeddings_test.py
     ./ssd_split_table_batched_embeddings_test.py
   )
 
   if [ "$fbgemm_variant" == "cpu" ]; then
     # These are tests that are currently broken in FBGEMM_GPU-CPU
     local ignored_tests=(
-      ./tbe/forward_test.py
       ./uvm_test.py
     )
   elif [ "$fbgemm_variant" == "rocm" ]; then

diff --git a/fbgemm_gpu/test/tbe/backward_adagrad_test.py b/fbgemm_gpu/test/tbe/backward_adagrad_test.py
@@ -70,7 +70,6 @@
 
 
 @optests.generate_opcheck_tests(fast=True)
-@unittest.skipIf(*gpu_unavailable)
 class BackwardAdagradTest(unittest.TestCase):
     def execute_backward_adagrad_(  # noqa C901
         self,
@@ -460,6 +459,7 @@ def execute_backward_adagrad_(  # noqa C901
                     torch.zeros_like(table_indice_weight_grad_mask),
                 )
 
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=128),
@@ -530,6 +530,7 @@ def test_backward_adagrad_fp16_pmSUM(  # noqa C901
             output_dtype,
         )
 
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=128),
@@ -600,6 +601,7 @@ def test_backward_adagrad_fp16_pmMEAN(  # noqa C901
             output_dtype,
         )
 
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=128),
@@ -805,6 +807,7 @@ def test_backward_adagrad_fp32_pmMEAN(  # noqa C901
             output_dtype,
         )
 
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=128),

diff --git a/fbgemm_gpu/test/tbe/backward_none_test.py b/fbgemm_gpu/test/tbe/backward_none_test.py
@@ -74,8 +74,8 @@
 
 
 @optests.generate_opcheck_tests(fast=True, additional_decorators=additional_decorators)
-@unittest.skipIf(*gpu_unavailable)
 class BackwardNoneTest(unittest.TestCase):
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=256),
@@ -103,6 +103,7 @@ class BackwardNoneTest(unittest.TestCase):
     def test_backward_none(self, **kwargs: Any) -> None:
         self.execute_backward_none_(**kwargs)
 
+    @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=5),
         D=st.integers(min_value=2, max_value=256),

diff --git a/fbgemm_gpu/test/tbe/failures_dict_fast.json b/fbgemm_gpu/test/tbe/failures_dict_fast.json
@@ -41,17 +41,17 @@
         "comment": "",
         "status": "xsuccess"
       },
-      "NBitSplitEmbeddingTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xsuccess"
       }
     },
     "fbgemm::int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
-      "NBitSplitEmbeddingTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xfail"
       },
-      "NBitSplitEmbeddingTest.test_schema__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_schema__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xfail"
       }
@@ -125,7 +125,7 @@
         "comment": "",
         "status": "xfail"
       },
-      "NBitSplitEmbeddingTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xfail"
       },
@@ -238,7 +238,7 @@
         "comment": "",
         "status": "xfail"
       },
-      "NBitSplitEmbeddingTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xfail"
       },
@@ -269,7 +269,7 @@
         "comment": "",
         "status": "xfail"
       },
-      "NBitSplitEmbeddingTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
         "comment": "",
         "status": "xfail"
       }
@@ -278,21 +278,13 @@
       "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
         "comment": "",
         "status": "xfail"
-      },
-      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_pruning": {
-        "comment": "",
-        "status": "xfail"
       }
     },
     "fbgemm::pruned_hashmap_insert": {},
     "fbgemm::pruned_hashmap_lookup": {
       "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
         "comment": "",
         "status": "xfail"
-      },
-      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_pruning": {
-        "comment": "",
-        "status": "xfail"
       }
     },
     "fbgemm::reset_weight_momentum": {

diff --git a/fbgemm_gpu/test/tbe/forward_test.py b/fbgemm_gpu/test/tbe/forward_test.py
@@ -268,8 +268,10 @@ def execute_forward_(  # noqa C901
             output_dtype=output_dtype,
             use_experimental_tbe=use_experimental_tbe,
         )
-        # NOTE: test TorchScript-compatible!
-        cc = torch.jit.script(cc)
+
+        if not use_cpu and torch.cuda.is_available():
+            # NOTE: test TorchScript-compatible!
+            cc = torch.jit.script(cc)
 
         for t in range(T):
             cc.split_embedding_weights()[t].data.copy_(

diff --git a/...gpu/test/tbe/nbit_split_embedding_test.py → ...pu/test/tbe/nbit_split_embeddings_test.py b/...gpu/test/tbe/nbit_split_embedding_test.py → ...pu/test/tbe/nbit_split_embeddings_test.py
@@ -44,7 +44,7 @@
 
 
 @optests.generate_opcheck_tests(fast=True)
-class NBitSplitEmbeddingTest(unittest.TestCase):
+class NBitSplitEmbeddingsTest(unittest.TestCase):
     @unittest.skipIf(*gpu_unavailable)
     @given(
         T=st.integers(min_value=1, max_value=10),