fix for when use_precision_aware_optimizer=false

Signed-off-by: Lifu Zhang <[email protected]>

Author: Lifu Zhang

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2473,6 +2473,10 @@ def update_main_grads(self):
             optimizer_grad = group.main_grad_buffer.get_item(
                 item_id, only_shard=sharded_optimizer_state
             )
+            if group.main_weight_buffer is not None:
+                if getattr(self, "use_precision_aware_optimizer", False):
+                    # Convert the gradient to the main weight buffer dtype.
+                    optimizer_grad = optimizer_grad.to(param.dtype)
 
             if name not in self.dist_main_grad:
                 # Register the gradient as a distributed tensor.
@@ -2497,8 +2501,11 @@ def update_main_grads(self):
 
             # The presence of main_grad_buffer but no main_weight_buffer may imply
             # that a precision-aware optimizer is used.
-            if getattr(self, "use_precision_aware_optimizer", False):
+            if getattr(self, "use_precision_aware_optimizer", True):
                 setattr(param, "decoupled_grad", grad)
+            else:
+                # Attach the gradient to the optimizer parameter.
+                setattr(param, "grad", grad.to(param.dtype) if grad is not None else None)
 
     @property
     def num_buckets(self):