M-FSDP: Make fine_grained_param_gather configurable for MXFP8 to enable performance–memory trade-offs

megatron/core/distributed/distributed_data_parallel_config.py

@@ -213,6 +213,17 @@ class DistributedDataParallelConfig:
       main gradients to parameter dtype for `.grad`.
     """
 
+    megatron_fsdp_enable_fine_grained_param_gather: bool = False
+    """If set to True, enables fine-grained parameter gathering for Megatron-FSDP.
+      This feature increases the overlap between parameter all-gather and forward computation,
+      at the cost of more frequent communication calls.
+      For MXFP8, this approach helps save memory during fine-grained activation
+      recomputation, because MXFP8 forward and backward passes use different
+      parameter representations (rowwise data for forward, colwise data for backward).
+      In this mode, only the rowwise parameters of modules involved in recomputation
+      will be unsharded.
+    """
+
     def __post_init__(self):
         import os
 

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -199,6 +199,7 @@ def __init__(
                 enable_fine_grained_param_gather_hook=(
                     (config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather)
                     or config.overlap_moe_expert_parallel_comm
+                    or self.ddp_config.megatron_fsdp_enable_fine_grained_param_gather
                 ),
                 enable_fine_grained_param_gather_backward_hook=(
                     config.overlap_moe_expert_parallel_comm

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

@@ -153,6 +153,17 @@ class DistributedDataParallelConfig:
       main gradients to parameter dtype for `.grad`.
     """
 
+    megatron_fsdp_enable_fine_grained_param_gather: bool = False
+    """If set to True, enables fine-grained parameter gathering for Megatron-FSDP.
+      This feature increases the overlap between parameter all-gather and forward computation,
+      at the cost of more frequent communication calls.
+      For MXFP8, this approach helps save memory during fine-grained activation
+      recomputation, because MXFP8 forward and backward passes use different
+      parameter representations (rowwise data for forward, colwise data for backward).
+      In this mode, only the rowwise parameters of modules involved in recomputation
+      will be unsharded.
+    """
+
     def __post_init__(self):
         import os
 

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

@@ -715,7 +715,6 @@ def fully_shard(
         fsdp_double_buffer=fsdp_double_buffer,
         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,
     )
 

megatron/training/arguments.py

@@ -3284,7 +3284,23 @@ def _add_experimental_args(parser):
                              "model gradients during FSDP. If 'auto', then the main gradient data-type will "
                              "be used for the gradient communication / reduction data-type. When using NCCL "
                              "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.")
-
+    group.add_argument(
+            '--megatron-fsdp-enable-fine-grained-param-gather',
+            action='store_true',
+            default=False,
+            dest='megatron_fsdp_enable_fine_grained_param_gather',
+            help=(
+                'If set, enables fine-grained parameter gathering for Megatron-FSDP. '
+                'This allows greater overlap between parameter all-gather operations and '
+                'forward computation, at the cost of additional communication calls. '
+                'For MXFP8, this helps save memory during fine-grained activation '
+                'recomputation, because MXFP8 forward and backward passes use different '
+                'parameter representations (rowwise data for forward, colwise data for '
+                'backward). Only the rowwise parameters of modules involved in '
+                'recomputation will be unsharded.'
+            ),
+        )
+
     return parser