diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index a15410d75e3..aa3cb53f927 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/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 diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index a852ef226b4..ea6b695988f 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/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 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 5047e652b93..4b9f1297cdb 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 @@ -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 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 a1f7fabd50a..6e2d5bbdec2 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/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, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f3c2ded6907..9b5a4830332 100644 --- a/megatron/training/arguments.py +++ b/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