support VBE -- reland (pytorch#2256)

Summary: Pull Request resolved: pytorch#2256 Reland D52517415 after deprecating a few unused FBGEMM kernels to reduce the binary size Reviewed By: sryap Differential Revision: D52590085 fbshipit-source-id: 4c8ce861930dfedb4356d65a4448ae91dff7a7dc
q10 · Jan 8, 2024 · 0a23158 · 0a23158
1 parent 638e833
commit 0a23158
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/fbgemm_gpu/CMakeLists.txt b/fbgemm_gpu/CMakeLists.txt
@@ -204,6 +204,7 @@ set(GPU_OPTIMIZERS ${COMMON_OPTIMIZERS} ${GPU_ONLY_OPTIMIZERS})
 # Optimizers with the VBE support
 set(VBE_OPTIMIZERS
     rowwise_adagrad
+    rowwise_adagrad_with_counter
     sgd)
 
 # Individual optimizers (not fused with SplitTBE backward)

diff --git a/fbgemm_gpu/codegen/embedding_common_code_generator.py b/fbgemm_gpu/codegen/embedding_common_code_generator.py
@@ -913,7 +913,7 @@ def rowwise_adagrad_with_counter() -> Dict[str, Any]:
         "split_weight_update_cpu": split_weight_update_cpu,
         "has_cpu_support": True,
         "has_gpu_support": True,
-        "has_vbe_support": False,
+        "has_vbe_support": True,
     }