Merge branch 'fix_chained_opt_with_dist_opt_overlap' into 'main'

Add disable/enable_pre_hook attributes for ChainedOptimizer.

See merge request ADLR/megatron-lm!1289

megatron/core/optimizer/optimizer.py

@@ -753,6 +753,24 @@ def load_state_dict(self, state_dict):
         for optimizer in self.chained_optimizers:
             self.param_groups += optimizer.param_groups
 
+    def disable_pre_hook(self):
+        if not self.config.use_distributed_optimizer or not self.config.overlap_param_gather:
+            raise ValueError(
+                "disable_pre_hook should only be called with 'use_distributed_optimizer' "
+                "and 'overlap_param_gather' are both enabled."
+            )
+        for optimizer in self.chained_optimizers:
+            optimizer.disable_pre_hook()
+
+    def enable_pre_hook(self):
+        if not self.config.use_distributed_optimizer or not self.config.overlap_param_gather:
+            raise ValueError(
+                "enable_pre_hook should only be called with 'use_distributed_optimizer' "
+                "and 'overlap_param_gather' are both enabled."
+            )
+        for optimizer in self.chained_optimizers:
+            optimizer.enable_pre_hook()
+
     def step(self):
         """ChainedOptimizer will step all optimizers one by one.
         """