File size: 1,301 Bytes
0143a83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | 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:
|