From 1cd7838516e7d2c128037e7cb69cc8b51ffa7e9f Mon Sep 17 00:00:00 2001 From: jianbinc Date: Tue, 7 Apr 2026 23:45:42 +0800 Subject: [PATCH] add --megatron-fsdp-enable-fine-grained-param-gather option --- .../distributed_data_parallel_config.py | 11 +++++++++++ .../distributed/fsdp/mcore_fsdp_adapter.py | 3 --- .../distributed_data_parallel_config.py | 11 +++++++++++ .../fsdp/src/megatron_fsdp/fully_shard.py | 1 - .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 5 +++-- megatron/training/arguments.py | 18 +++++++++++++++++- 6 files changed, 42 insertions(+), 7 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index b938447827c..2b2bd68a1cf 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -202,6 +202,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 diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 8993620c779..8d08673a7be 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -141,9 +141,6 @@ def __init__( dist_index=self.megatron_fsdp_dist_index, calculate_per_token_loss=config.calculate_per_token_loss, init_model_with_meta_device=config.init_model_with_meta_device, - enable_fine_grained_param_gather_hook=( - config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather - ), ), ) self.param_and_grad_buffer = self.module.param_and_grad_buffer diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index a2feb99cb23..158a4e54974 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -151,6 +151,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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index 0f5873d6085..6c85db3a56d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py @@ -695,7 +695,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, ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 7536a1ba9d2..5c76cf121ed 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -213,7 +213,6 @@ def __init__( fsdp_double_buffer: bool = False, fsdp_db_use_persist_buf_on_alloc_fail: bool = False, disable_symmetric_registration: bool = False, - enable_fine_grained_param_gather_hook: bool = False, report_nan_in_param_grad: bool = False, ): super().__init__() @@ -265,7 +264,9 @@ def __init__( self.mp_policy = mixed_precision_policy self.calculate_per_token_loss = calculate_per_token_loss self.init_model_with_meta_device = init_model_with_meta_device - self.enable_fine_grained_param_gather_hook = enable_fine_grained_param_gather_hook + self.enable_fine_grained_param_gather_hook = ( + self.ddp_config.megatron_fsdp_enable_fine_grained_param_gather + ) self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0bfe6142d01..ef5ba629b0c 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3267,7 +3267,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