megatron_patch/fixes/optimizer_offloading/fix_optimizer_offloading.patch

diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
index 969a8f80..a15d3421 100644
--- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
+++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
@@ -463,3 +463,7 @@ class HybridDeviceOptimizer(torch.optim.Optimizer):
         if self.gpu_optimizer is not None:
             return self.cpu_optimizers + [self.gpu_optimizer]
         return self.cpu_optimizers
+
+    def _update_fp32_param_by_new_param(self):
+        for param, fp32_param in self.param_to_fp32_param.items():
+            fp32_param.data.copy_(param)
\ No newline at end of file
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index 30954645..bfbe248b 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -2025,7 +2025,7 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
         an intermediary.
         """
         if isinstance(self.optimizer, HybridDeviceOptimizer):
-            return
+            return self.optimizer._update_fp32_param_by_new_param()
 
         if self.ddp_config.use_custom_fsdp:
             for model_chunk in self.model_chunks: