From 60997538b58c6c40f18a88cdf434002013ae981b Mon Sep 17 00:00:00 2001 From: Michael Shi Date: Fri, 8 Dec 2023 22:12:38 -0800 Subject: [PATCH] Enable Distributed Shampoo (v2) in HPC TorchRec (#1459) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/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 --- torchrec/distributed/embedding_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchrec/distributed/embedding_types.py b/torchrec/distributed/embedding_types.py index 2ea4241a4..eeb30703b 100644 --- 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"