Enable Distributed Shampoo (v2) in HPC TorchRec (#1459)

Summary: Pull Request resolved: #1459 See title. Pair programmed with Tsung-Hsien Lee (zong). NOTE: Checkpointing in HPC TorchRec currently does not work. Next step is to enhance TorchRec checkpointing (cc: YLGH). Differential Revision: D50479625 fbshipit-source-id: 04fafa41a4aa1f5f0a779b3f7e34f9f146022c8c
pytorch · Dec 9, 2023 · 6099753 · 6099753
1 parent fba7f3f
commit 6099753
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchrec/distributed/embedding_types.py b/torchrec/distributed/embedding_types.py
@@ -44,6 +44,7 @@ class OptimType(Enum):
     ADAGRAD = "ADAGRAD"
     ROWWISE_ADAGRAD = "ROWWISE_ADAGRAD"
     SHAMPOO = "SHAMPOO"
+    SHAMPOO_V2 = "SHAMPOO_V2"
     LION = "LION"