m-fsdp: wire use_precision_aware_optimizer from ddp_config to ParamAn…

megatron/core/distributed/distributed_data_parallel_config.py

@@ -196,6 +196,12 @@ class DistributedDataParallelConfig:
       No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
     """
 
+    megatron_fsdp_use_decoupled_grad: bool = False
+    """If true, Megatron-FSDP's ParamAndGradBuffer uses the precision-aware optimizer
+      gradient path (e.g. `decoupled_grad` on optimizer parameters) instead of casting
+      main gradients to parameter dtype for `.grad`.
+    """
+
     def __post_init__(self):
         import os
 

megatron/core/distributed/fsdp/src/README.md

@@ -326,6 +326,9 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
     - Defaults to `False`.
 - `keep_fp8_transpose_cache` will keep the fp8 transpose cache when using `MegatronFSDP`. This option will cause (number of parameter $\times$ 1 Byte) of memory overhead, but can skip the weight transpose operation in the backward propagation. This feature will not give any benefit from the Blackwell architecture.
     - Defaults to `False`.
+- `use_decoupled_grad` installs the reduced gradient into a separate buffer: `Parameter.decoupled_grad`. This buffer is utilized by specific optimizers, such as TransformerEngine's `FusedAdam`, and can be used to temporarily store your gradient for custom `torch.nn.Optimizer`(s).
+    - Defaults to `False`.
+    - Required for `transformer_engine.pytorch.optimizers.FusedAdam`.
 - `nccl_ub` will allocate and register the NCCL userbuffer for param and grad buffers. This option enables an SM-efficient NCCL algorithm that could improve the performance of overlapped computations. This flag will be much more effective when used together with SHARP if the FSDP communication includes both NVL and IB domains. Enabling this option will cause additional memory overhead due to the requirement to enable the `fsdp_double_buffer` option.
     - **Only effective when using with Megatron-Core.**
     - Defaults to `False`.

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

@@ -145,6 +145,12 @@ class DistributedDataParallelConfig:
       No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
     """
 
+    megatron_fsdp_use_decoupled_grad: bool = False
+    """If true, Megatron-FSDP's ParamAndGradBuffer uses the precision-aware optimizer
+      gradient path (e.g. `decoupled_grad` on optimizer parameters) instead of casting
+      main gradients to parameter dtype for `.grad`.
+    """
+
     def __post_init__(self):
         import os
 

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

@@ -101,6 +101,7 @@ def fully_shard_model(
     fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
     disable_symmetric_registration: bool = False,
     enable_fine_grained_param_gather: bool = False,
+    use_decoupled_grad: bool = False,
 ) -> torch.nn.Module:
     """
     Fully-shard the model for Megatron-FSDP. This wraps the model in a MegatronFSDP
@@ -247,6 +248,10 @@ class that schedules the sharding lifecycle of the model parameters and gradient
             unshards parameters per-Module instead of unsharding all sub-modules of an FSDP
             unit module simultaneously. Defaults to False.
 
+        use_decoupled_grad (bool):
+            If true, reduced gradients are installed into `Parameter.decoupled_grad` instead
+            of `Parameter.grad`. Defaults to False.
+
     Returns:
         model (MegatronFSDP): The wrapped Megatron-FSDP model configured for FSDP.
     """
@@ -341,6 +346,7 @@ class that schedules the sharding lifecycle of the model parameters and gradient
         fsdp_double_buffer=fsdp_double_buffer or nccl_ub,
         fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
         disable_symmetric_registration=disable_symmetric_registration,
+        megatron_fsdp_use_decoupled_grad=use_decoupled_grad,
     )
 
     # Create FSDPDistributedIndex.
@@ -641,6 +647,7 @@ def fully_shard(
     fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
     disable_symmetric_registration: bool = False,
     enable_fine_grained_param_gather: bool = False,
+    use_decoupled_grad: bool = False,
 ) -> tuple[MegatronFSDP, torch.optim.Optimizer]:
     """
     Fully shard the model and the optimizer for Megatron-FSDP.
@@ -689,6 +696,7 @@ def fully_shard(
         fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
         disable_symmetric_registration=disable_symmetric_registration,
         enable_fine_grained_param_gather=enable_fine_grained_param_gather,
+        use_decoupled_grad=use_decoupled_grad,
     )
 
     # Extend optimizer methods to support Megatron-FSDP operations.

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

@@ -1642,6 +1642,7 @@ def __init__(
         )
 
         self.ddp_config = ddp_config
+        self.use_decoupled_grad = ddp_config.megatron_fsdp_use_decoupled_grad
         self.module = module
         self.bucketing_policy = bucketing_policy
         self.param_to_name = {p: name for name, p in self.module.named_parameters()}
@@ -2842,8 +2843,9 @@ def update_main_grads(self):
                 item_id, only_shard=sharded_optimizer_state
             )
             if group.main_weight_buffer is not None:
-                if not getattr(self, "use_precision_aware_optimizer", False):
+                if not self.use_decoupled_grad:
                     # Convert the gradient to the main weight buffer dtype.
+                    # TODO(@cspades): Why this is necessary? Casted below.
                     optimizer_grad = optimizer_grad.to(param.dtype)
 
             if name not in self.dist_main_grad:
@@ -2867,9 +2869,9 @@ def update_main_grads(self):
             if optimizer_grad.numel() == 0:
                 grad = None
 
-            # 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 use_decoupled_grad (i.e. for precision-aware optimizers like TE FusedAdam),
+            # install the gradient into param.decoupled_grad.
+            if self.use_decoupled_grad:
                 setattr(param, "decoupled_grad", grad)
             else:
                 # Attach the gradient to the optimizer parameter.

megatron/training/training.py

@@ -1464,6 +1464,7 @@ def build_model():
             kwargs['megatron_fsdp_main_params_dtype'] = args.megatron_fsdp_main_params_dtype
             kwargs['megatron_fsdp_main_grads_dtype'] = args.megatron_fsdp_main_grads_dtype
             kwargs['megatron_fsdp_grad_comm_dtype'] = args.megatron_fsdp_grad_comm_dtype
+            kwargs['megatron_fsdp_use_decoupled_grad'] = args.use_precision_aware_optimizer
 
             # Initialize DDPConfig.
             ddp_config = DistributedDataParallelConfig(**kwargs)