add functions to support checkpointing in megatron ddp

481f5c4f · Rick Ho · 79ccb7b6 · 481f5c4f
Commit 481f5c4f authored Feb 05, 2021 by Rick Ho
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 0 deletions

fmoe/megatron.py fmoe/megatron.py +9 -0

No files found.
--- a/fmoe/megatron.py
+++ b/fmoe/megatron.py
@@ -49,3 +49,12 @@ class DistributedDataParallel(DistributedGroupedDataParallel):
            mp_group=mpu.get_model_parallel_group(),
            dp_group=mpu.get_data_parallel_group()
        )
+
+    def state_dict(self, *args, **kwargs):
+        return self.module.state_dict(*args, **kwargs)
+
+    def state_dict_for_save_checkpoint(self, *args, **kwargs):
+        return self.module.state_dict_for_save_checkpoint(*args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        return self.module.load_state_dict(*args, **kwargs)