| |
|
|
| from dataclasses import dataclass |
| from typing import Optional |
|
|
|
|
| @dataclass |
| class DistributedDataParallelConfig: |
| """Configuration for DistributedDataParallel.""" |
|
|
| grad_reduce_in_fp32: bool = False |
| """If true, reduce grads in fp32.""" |
|
|
| overlap_grad_reduce: bool = False |
| """If true, overlap grad all-reduce / reduce-scatter with backward compute.""" |
|
|
| overlap_param_gather: bool = False |
| """If true, overlap param all-gather with forward compute.""" |
|
|
| align_param_gather: bool = False |
| """If true, all PP stages will launch param all-gathers simultaneously. Otherwise, each |
| PP stage will independently launch as needed. |
| """ |
|
|
| use_distributed_optimizer: bool = False |
| """If true, issue reduce-scatter collectives to aggregate gradients and clean up |
| originally allocated model parameters, otherwise issue all-reduce collectives. |
| """ |
|
|
| num_distributed_optimizer_instances: int = 1 |
| """Sets the factor by which the DP domain is sharded to have the partial DistOpt |
| enabled. Defaults to 1, which means DistOpt is across entire DP domain. |
| """ |
|
|
| check_for_nan_in_grad: bool = False |
| """If true, check for NaNs and Infs in gradients _before_ communication collective.""" |
|
|
| check_for_large_grads: bool = False |
| """If true, check for unexpectedly large gradients _before_ communication collective.""" |
|
|
| bucket_size: Optional[int] = None |
| """Maximum number of parameters in each bucket. If unspecified, MCore uses a default |
| value of max(40000000, 1000000 * dp_size) parameters (larger DP sizes need larger |
| buckets to ensure collectives do not become latency-bound).""" |
|
|
| pad_buckets_for_high_nccl_busbw: bool = False |
| """If true, make sure the bucket size is divisible by a large power of 2 (2^16) to |
| ensure NCCL collectives have high bus bandwidth at large DP counts, since NCCL |
| message size (which for ring algorithms is bucket_size / dp_size) apparently needs |
| to be divisible by a power of 2 for high busbw.""" |
|
|
| reduce_scatter_with_fp32_accumulation: bool = False |
| """If true, use a reduce-scatter implementation which sends lower-precision values |
| over the wire (using an all-to-all to keep total communication overhead in line |
| with the standard ring implementation) but performs accumulation locally in FP32.""" |
|
|
| average_in_collective: bool = False |
| """If true, compute average in collective directly, as opposed to dividing by the |
| dp_size first and then computing sum in the collective.""" |
|
|
| fp8_param_gather: bool = False |
| """If true, keep the compute param in fp8 (do not use any other intermediate dtype) and |
| perform the param all-gather in fp8.""" |
|
|
| reuse_grad_buf_for_mxfp8_param_ag: bool = False |
| """If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be |
| set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True.""" |
|
|
| use_megatron_fsdp: bool = False |
| """If true, use the FSDP code path for DDP.""" |
|
|
| use_custom_fsdp: bool = False |
| """ |
| NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. |
| Please use `use_megatron_fsdp` instead, as all functionality will be migrated there. |
| Future updates will drop support for `use_custom_fsdp` to avoid confusion. |
| """ |
|
|
| data_parallel_sharding_strategy: str = 'no_shard' |
| """Sharding strategy for FSDP. Valid values are 'no_shard', 'optim', |
| 'optim_grads', 'optim_grads_params'.""" |
|
|
| gradient_reduce_div_fusion: bool = True |
| """If true, perform gradient reduce and division fusion.""" |
|
|
| suggested_communication_unit_size: int = None |
| """Specifies the number of elements to communicate at once during |
| FSDP (Fully Sharded Data Parallel) operations. |
| This flag also affects FSDP all-gather prefetch behavior. Setting a larger |
| value increases the communication buffer size, while a smaller value |
| disables prefetching and may degrade performance. Adjust this value |
| based on your system's memory and performance requirements.""" |
|
|
| preserve_fp32_weights: bool = True |
| """If true, preserve fp32 weights in the Megatron FSDP ParamAndGradBuffer.""" |
|
|
| keep_fp8_transpose_cache: bool = False |
| """If true, keep the fp8 transpose cache when using Megatron FSDP.""" |
|
|
| nccl_ub: bool = False |
| """If true, allocate and register NCCL userbuffer for param and grad buffer. |
| This flag enables SM efficient nccl algorithm that could improve the performance |
| of FSDP and DP with comm_overlap. This flag will be much more effective when used |
| together with sharp. |
| The follwoing will be the expected number of SM usage for various cases. |
| (Note that this is just a reference number and the number of SM usage could vary |
| on message size, communication domain size and nccl version.) |
| ---------------------------------------------------------- |
| | Communication domain | use_sharp | SM usage of "AG/RS" | |
| |----------------------|-----------|---------------------| |
| | NVL | N/A | 4 / 5 | |
| | NVL+IB | False | 16 / 16 | |
| | NVL+IB | True | 6 / 6 | |
| | IB | False | 1 / 4 | |
| | IB | True | 1 / 1 | |
| ---------------------------------------------------------- |
| """ |
|
|
| fsdp_double_buffer: bool = False |
| """If true, use persistently allocated double buffers for the |
| temporary memory needed in the Megatron FSDP communications. |
| This option will cause additional memory overhead, however, it is necessary for |
| to register user buffer (nccl_ub=True) for the Megatron FSDP. |
| This option will be automatically set to True when nccl_ub=True. |
| """ |
|
|
| outer_dp_sharding_strategy: str = 'no_shard' |
| """ |
| Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. |
| Valid values are 'no_shard', 'optim', 'optim_grads', 'optim_grads_params'. |
| This option is only effective when Hybrid FSDP is enabled. |
| """ |
|
|
| disable_symmetric_registration: bool = False |
| """If true, disable symmetric (window) registration for NCCL userbuffer registration. |
| This option will force to use conventional (local) userbuffer registration |
| when nccl_ub is set. |
| """ |
|
|
| fsdp_manual_registration: bool = False |
| """If true, manually register the FSDP communication buffers to NCCL user buffer. |
| This option is only effective when use_megatron_fsdp and nccl_ub is set. |
| For symmetric registration with large models, the registration itself can take |
| a significant amount of time. This option minimizes the number of registration calls |
| to minimize the registration time. |
| """ |
|
|
| delay_wgrad_compute: bool = False |
| """Delay the weight gradient computation to improve batch-level communication overlapping""" |
|
|
| def __post_init__(self): |
| import os |
|
|
| """Check the validity of the config.""" |
| if self.reuse_grad_buf_for_mxfp8_param_ag: |
| assert self.fp8_param_gather, "Reuse grad buffer only when keeping params in MXFP8." |
|
|
| if self.nccl_ub: |
| if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','): |
| raise ValueError( |
| "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported " |
| "with nccl_ub due to compatibility issue with torch.cuda.MemPool API." |
| ) |
|
|
