diff --git a/docs/discussions/README.md b/docs/discussions/README.md index e791ed57cd8..aab65fc65ca 100644 --- a/docs/discussions/README.md +++ b/docs/discussions/README.md @@ -19,13 +19,9 @@ This directory contains in-depth guides, tutorials, and discussions about optimi ### Training Guides -- **[Megatron-FSDP User Guide](megatron-fsdp-user-guide/megatron-fsdp-user-guide.md)** - - A practical guide to enable Megatron-FSDP training, including a quick-start example for DeepSeek-V3, required and recommended configurations, and instructions for checkpoint conversion from torch_dist to fsdp_dtensor. - ## Contributing -If you'd like to contribute a guide or tutorial, please follow this structure: +To contribute a guide or tutorial, follow this structure: 1. Create a new directory: `docs/discussions/your-guide-name/` 2. Add your main guide: `docs/discussions/your-guide-name/your-guide-name.md` diff --git a/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md b/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md deleted file mode 100644 index b5de090ab46..00000000000 --- a/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md +++ /dev/null @@ -1,129 +0,0 @@ ---- -orphan: true ---- - - - -# Megatron-FSDP User Guide - -## Table of Contents - -- [Megatron-FSDP Quick Start](#megatron-fsdp-quick-start) -- [Checkpoint Conversion from 3D-Parallel to Megatron-FSDP](#checkpoint-conversion-from-3d-parallel-to-megatron-fsdp) - -## Megatron-FSDP Quick Start - -We recommend using the latest [NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags), which provides a tested software stack and optimized performance. - -For your reference, we provide an example launch script for DeepSeek-V3: [`sbatch_mfsdp_deepseek_v3.sh`](./example-scripts/sbatch_mfsdp_deepseek_v3.sh). - -### Required Configurations - -To enable Megatron-FSDP, add the following required flags to your training script: - -```bash ---use-megatron-fsdp ---data-parallel-sharding-strategy optim_grads_params ---no-gradient-accumulation-fusion ---use-distributed-optimizer ---ckpt-format fsdp_dtensor -``` - -### Recommended Configurations - -We also recommend adding the following configurations to further improve performance: - -```bash -unset CUDA_DEVICE_MAX_CONNECTIONS -``` -```bash ---calculate-per-token-loss ---init-model-with-meta-device ---grad-reduce-in-bf16 ---fsdp-double-buffer ---use-nccl-ub -``` - -💡 **Detailed explanations of these configurations are provided below.** - -#### 1. Disable `CUDA_DEVICE_MAX_CONNECTIONS` - -To ensure full parallelization of FSDP communication and computation, disable the CUDA_DEVICE_MAX_CONNECTIONS environment variable. This step avoids potential bubbles in the CUDA stream. (But it may slow down TP and CP to some extent.) - -#### 2. Add `--calculate-per-token-loss` - -For gradients sharding mode optimization, include the `--calculate-per-token-loss` flag in your training script. This improves performance by reducing the frequency of gradient scaling, which is also a sizable drain on SM resources. - -#### 3. Add `--init-model-with-meta-device` - -Allows model initialization using meta device, followed by layer-by-layer initialization of distributed model weight buffers via the `Module.reset_parameters` API, facilitating the initialization of extremely large models. - -#### 4. Add `--grad-reduce-in-bf16` - -Enables gradient reduction in BF16 precision instead of FP32, reducing communication volume and accelerating the backward pass. - -#### 5. Add `--fsdp-double-buffer` - -Uses persistently allocated double buffers for temporarily-defined memory needed in `MegatronFSDP` communications. While having persistent double buffers may increase peak VRAM utilization, it is necessary to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`. Currently, this is supported only for simple repetitive model structures such as GPT. - -- **Only effective when using Megatron-LM.** -- Defaults to `False`. Automatically overridden to `True` when `nccl_ub` is enabled. - -#### 6. Add `--use-nccl-ub` - -Allocates and [registers NCCL user buffers](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#) 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](https://docs.nvidia.com/networking/display/sharpv3130) 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 Megatron-LM.** -- Defaults to `False`. -- By default we try to use NCCL window (symmetric) registration if it is available. If not it falls back to conventional local registration. -- **Incompatible with PyTorch's segmentable allocator:** Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when using `--use-nccl-ub`, as this will cause a runtime error due to compatibility issues with the `torch.cuda.MemPool` API. - -## Checkpoint Conversion from 3D-Parallel to Megatron-FSDP - -Megatron-FSDP introduces `fsdp_dtensor`, a DTensor-based distributed checkpoint format that serves as its standard. To help you smoothly transition from 3D-Parallel to Megatron-FSDP, we provide a script for converting checkpoints from the `torch_dist` format to the `fsdp_dtensor` format. Using DeepSeek-V3 as an example, the detailed conversion process is described below. - -### Step 1: Generate 3D-Parallel Checkpoint with `param_to_param_group_map` - -Run your 3D-parallel + EP training script to generate a `torch_dist` checkpoint along with a directory containing `param_to_param_group_map` files. Add the following flag to your training script: - -```bash ---dump-param-to-param-group-map /path/to/param_to_param_group_map -``` - -If you already have a `torch_dist` checkpoint, simply specify the `--dump-param-to-param-group-map /path/to/param_to_param_group_map` flag and run a very short experiment-this will create the `param_to_param_group_map` you need without full pretraining. - -### Step 2: Export `param_to_param_group_map` to a JSON File - -Convert the `param_to_param_group_map` into a JSON file for easier processing by running: - -```bash -python tools/checkpoint/checkpoint_inspector.py print-torch-dcp-in-json /path/to/param_to_param_group_map -``` - -This will create a `param_to_param_group_map.json` file in the `/path/to/param_to_param_group_map` directory. - -### Step 3: Convert Checkpoint from `torch_dist` to `fsdp_dtensor` - -Convert your `torch_dist` checkpoint to the `fsdp_dtensor` format using the parameter to `param_to_param_group_map` JSON file: - -```bash -torchrun --nproc_per_node=8 --nnodes=1 \ - tools/checkpoint/checkpoint_inspector.py \ - convert-torch-dist-to-fsdp-dtensor --swiglu \ - /path/to/input_torch_dist_checkpoint \ - /path/to/output_fsdp_dtensor_checkpoint \ - --param-to-param-group-map-json /path/to/param_to_param_group_map.json -``` - -**Note:** For multi-node conversion tasks, please refer to the example script: [`sbatch_checkpoint_convert.sh`](./example-scripts/sbatch_checkpoint_convert.sh). - -### Step 4: Launch Megatron-FSDP Training - -Start your Megatron-FSDP training job using the converted `fsdp_dtensor` checkpoint. \ No newline at end of file diff --git a/docs/images/custom_fsdp/FSDP_workflow.png b/docs/images/custom_fsdp/FSDP_workflow.png deleted file mode 100644 index 588b6f220a3..00000000000 Binary files a/docs/images/custom_fsdp/FSDP_workflow.png and /dev/null differ diff --git a/docs/images/custom_fsdp/MCore_Custom_FSDP_Class_Diagram.png b/docs/images/custom_fsdp/MCore_Custom_FSDP_Class_Diagram.png deleted file mode 100644 index f9603079b92..00000000000 Binary files a/docs/images/custom_fsdp/MCore_Custom_FSDP_Class_Diagram.png and /dev/null differ diff --git a/docs/images/megatron_fsdp/DDP_vs_FSDP.png b/docs/images/megatron_fsdp/DDP_vs_FSDP.png new file mode 100644 index 00000000000..627821439e2 Binary files /dev/null and b/docs/images/megatron_fsdp/DDP_vs_FSDP.png differ diff --git a/docs/images/custom_fsdp/FSDP_Allreduce.png b/docs/images/megatron_fsdp/FSDP_Allreduce.png similarity index 100% rename from docs/images/custom_fsdp/FSDP_Allreduce.png rename to docs/images/megatron_fsdp/FSDP_Allreduce.png diff --git a/docs/images/megatron_fsdp/fsdp_double_buffer.png b/docs/images/megatron_fsdp/fsdp_double_buffer.png new file mode 100644 index 00000000000..fbfbcef9b28 Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_double_buffer.png differ diff --git a/docs/images/megatron_fsdp/fsdp_streams.png b/docs/images/megatron_fsdp/fsdp_streams.png new file mode 100644 index 00000000000..6b8840783c8 Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_streams.png differ diff --git a/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png b/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png new file mode 100644 index 00000000000..6f6e61dfb21 Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png differ diff --git a/docs/images/megatron_fsdp/hfsdp.png b/docs/images/megatron_fsdp/hfsdp.png new file mode 100644 index 00000000000..3c056d20689 Binary files /dev/null and b/docs/images/megatron_fsdp/hfsdp.png differ diff --git a/docs/images/megatron_fsdp/lcm_dim0_shard.png b/docs/images/megatron_fsdp/lcm_dim0_shard.png new file mode 100644 index 00000000000..910add676f1 Binary files /dev/null and b/docs/images/megatron_fsdp/lcm_dim0_shard.png differ diff --git a/docs/images/megatron_fsdp/mixed_sharding.png b/docs/images/megatron_fsdp/mixed_sharding.png new file mode 100644 index 00000000000..81cbc153f8a Binary files /dev/null and b/docs/images/megatron_fsdp/mixed_sharding.png differ diff --git a/docs/images/megatron_fsdp/quantized_param_gather.png b/docs/images/megatron_fsdp/quantized_param_gather.png new file mode 100644 index 00000000000..e1908e66ad7 Binary files /dev/null and b/docs/images/megatron_fsdp/quantized_param_gather.png differ diff --git a/docs/images/megatron_fsdp/sharded_quantization.png b/docs/images/megatron_fsdp/sharded_quantization.png new file mode 100644 index 00000000000..c65bab5305a Binary files /dev/null and b/docs/images/megatron_fsdp/sharded_quantization.png differ diff --git a/docs/images/megatron_fsdp/uneven_sharding.png b/docs/images/megatron_fsdp/uneven_sharding.png new file mode 100644 index 00000000000..0c34a51b026 Binary files /dev/null and b/docs/images/megatron_fsdp/uneven_sharding.png differ diff --git a/docs/images/megatron_fsdp/zero3_model_state.png b/docs/images/megatron_fsdp/zero3_model_state.png new file mode 100644 index 00000000000..84ad33ff779 Binary files /dev/null and b/docs/images/megatron_fsdp/zero3_model_state.png differ diff --git a/docs/index.md b/docs/index.md index 0dbf7d2e3b7..11337315588 100644 --- a/docs/index.md +++ b/docs/index.md @@ -67,7 +67,7 @@ models/index user-guide/features/moe user-guide/features/context_parallel -user-guide/features/custom_fsdp +user-guide/features/megatron_fsdp user-guide/features/dist_optimizer user-guide/features/optimizer_cpu_offload user-guide/features/pipeline_parallel_layout diff --git a/docs/user-guide/features/custom_fsdp.md b/docs/user-guide/features/custom_fsdp.md deleted file mode 100644 index 3c774e5f493..00000000000 --- a/docs/user-guide/features/custom_fsdp.md +++ /dev/null @@ -1,196 +0,0 @@ - - -# Megatron FSDP - -**Note: In M-Core 0.14, the custom FSDP refactored its checkpoint implementation to use DTensor-based PyTorch distributed checkpointing. The custom FSDP was also renamed Megatron FSDP. The relevant sections of this document are no longer applicable.** - -## How to Use Megatron FSDP - -Add these flags to enable MCore custom FSDP. - -```bash ---use-megatron-fsdp ---data-parallel-sharding-strategy optim_grads_params ---no-gradient-accumulation-fusion ---use-distributed-optimizer -``` - -For a practical guide covering required configurations, checkpoint conversion, and example scripts, refer to the [Megatron-FSDP User Guide](../../discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md). - -## Key Features - -- **Sharding Strategy**: Shards optimizer states, gradients, and parameters to reduce memory consumption. -- **Communication and Computation Overlap**: Overlaps communication with computation where possible during training. -- **Supports automatic mixed precision training**: Compatible with BF16 O1/O2/O3 recipes, as well as FP8 compute with FP32 parameters and FP8 parameter training, with several precision configuration options. -- **Tensor Parallelism (TP), Expert Parallelism (EP), and Context Parallelism (CP)**: Compatible with TP, EP, and CP configurations for scaling large language models. -- **Distributed Model Initialization with Meta Device**: Allows model initialization using the meta device, then layer-by-layer initialization of distributed model weight buffers through the `Module.reset_parameters` API, which supports very large models. - -## Configuration Recommendations - -### Disable `CUDA_DEVICE_MAX_CONNECTIONS` - -To ensure full parallelization of FSDP communication and computation, disable the CUDA_DEVICE_MAX_CONNECTIONS environment variable. This step avoids potential bubble in CUDA stream. (But it may slow down TP and CP to some extent.) - -```bash -unset CUDA_DEVICE_MAX_CONNECTIONS -``` - -### Add `--calculate-per-token-loss` - -For gradients sharding mode optimization, include the `--calculate-per-token-loss` flag in your training script. This improves performance by reducing the frequency of gradient scaling, which is also a sizable drain on SM resources. - -## Design of Custom FSDP - -### Overview - -The custom Fully Sharded Data Parallelism (FSDP) implementation in Megatron Core targets memory use and throughput for large language models. The core design principles include: - - - **Optimized for Large Language Models**: This custom FSDP implementation scales with models containing billions of parameters and supports training at that scale. - - **Efficient Memory Consumption**: By sharding optimizer states, gradients, and model parameters, the custom FSDP cuts memory use so models that would not fit on device with plain DDP can train. - - **Efficient Workflow and Overlapping Communication and Computation**: The implementation reduces communication steps during training and overlaps communication with computation where possible to reduce idle time. - - **Support for MCore's Efficient Training Methods**: The custom FSDP integrates with advanced parallelism in Megatron Core, including tensor parallelism, expert parallelism, and context parallelism. It also supports automatic mixed precision training. - -The design of Custom FSDP draws inspiration from PyTorch FSDP [Zhao, Yanli, et al.](https://arxiv.org/pdf/2304.11277) and the MCore distributed optimizer. The following background on PyTorch FSDP clarifies the concepts behind the custom FSDP design. - -> In DistributedDataParallel, (DDP) training, each process/ worker owns a replica of the model and processes a batch of data, finally it uses all-reduce to sum up gradients over different workers. In DDP the model weights and optimizer states are replicated across all workers. FSDP is a type of data parallelism that shards model parameters, optimizer states and gradients across DDP ranks. - -> When training with FSDP, the GPU memory footprint is smaller than when training with DDP across all workers. This makes the training of some very large models feasible by allowing larger models or batch sizes to fit on device. This comes with the cost of increased communication volume. The communication overhead is reduced by internal optimizations like overlapping communication and computation. - -![Diagram of FSDP workflow showing all-gather, forward, discard, backward, and reduce-scatter across ranks](../../images/custom_fsdp/FSDP_workflow.png) - -The unit processed in the workflow here is the "FSDP instance 1: N layers", where an FSDP instance is the smallest FSDP processing unit (also a PyTorch module). You can release this module's weights after its forward or backward because no other computation depends on those weights. That behavior supports FSDP's layer-by-layer execution and memory-saving strategy. An FSDP instance is also called an **FSDP Unit**. - -An FSDP instance can correspond to multiple FSDP parameter groups. These groups are separated by Data Parallel (DP) communication groups and the data type of the parameter or gradient. Consequently, an FSDP instance may require several parameter-gather tasks before execution (forward or backward). Each **FSDP parameter group** corresponds to one **Data Parallel Buffer** in custom FSDP. - -At a high level, FSDP works as follows: - -In constructor: - - Shard model parameters and each rank only keeps its own shard - -In forward path: - - Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit - - Run forward computation - - Discard parameter shards it has just collected - -In backward path: - - Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit - - Run backward computation - - Run reduce_scatter to sync gradients - - Discard parameters - -One way to view FSDP’s sharding is to decompose the DDP gradient all-reduce into reduce-scatter and all-gather. Specifically, during the backward pass, FSDP reduces and scatters gradients, ensuring that each rank possesses a shard of the gradients. Then it updates the corresponding shard of the parameters in the optimizer step. Finally, in the subsequent forward pass, it performs an all-gather operation to collect and combine the updated parameter shards. - -![Diagram comparing DDP all-reduce with FSDP reduce-scatter and all-gather](../../images/custom_fsdp/FSDP_Allreduce.png) - -### Custom FSDP Underlying Data Structure - -To implement the FSDP functionality described above, the custom FSDP is designed with the following Python classes and data structure: - -![Class diagram of Megatron Core custom FSDP Python types](../../images/custom_fsdp/MCore_Custom_FSDP_Class_Diagram.png) - -### The Custom FSDP Interface: FullyShardedDataParallel - -The custom FSDP provides the same programming interface as PyTorch's DistributedDataParallel (DDP) as FullyShardedDataParallel (FSDP). For example, you can apply FSDP to models as follows: - -```python -# Initialize model and optimizer -ddp_config.use_megatron_fsdp = True -ddp_config.data_parallel_sharding_strategy = "optim_grads_params" -model = GPTModel(transformer_config) -model = FullyShardedDataParallel( - transformer_config, - model, - ddp_config, - fsdp_unit_modules = [TransformerLayer, LanguageModelEmbedding], -) -optimizer = torch.optim.AdamW(model.parameters(), lr=lr) -optimizer = DistributedOptimizer(optimizer, [model], [model.param_and_grad_buffer]) - -# Training loop -def train_step(inputs, labels): - optimizer.zero_grad() - for mbs_input, mbs_label in zip(inputs, labels): - outputs = model(mbs_input) - loss = loss_fn(outputs, mbs_label) - loss.backward() - optimizer.step() - -# Save and load model and optimizer state dict -def model_and_optimizer_state_dict(): - state_dict = { - "model": model.sharded_state_dict(), - "optimizer": optimizer.sharded_state_dict(), - } - return state_dict - -def load_model_and_optimizer_state_dict(state_dict): - model.load_state_dict(state_dict["model"]) - optimizer.load_state_dict(state_dict["optimizer"]) -``` - -Key notes: - - - You can configure which modules should be treated as FSDP units through the `fsdp_unit_modules` argument. This configuration is mandatory. - - The custom FSDP must be used with a distributed optimizer since it provides distributed checkpointing. - - The data-parallel communication group for parameters is not explicitly shown. Custom FSDP configures these groups as either DP (data-parallel) or EDP (expert data-parallel) based on parameter markings. - -#### Initializing Models on the Meta Device - -For training particularly large models with FSDP, you can initialize the model on the meta device. Using PyTorch's `reset_parameters` API, you can initialize model weights layer by layer during the construction of the `ParamAndGradBuffer`. Most PyTorch native modules and TransformerEngine modules support this API (for example, [PyTorch Linear](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/nn/modules/linear.py#L114), [TE LayerNormLinear](https://github.com/NVIDIA/TransformerEngine/blob/release_v2.0/transformer_engine/pytorch/module/layernorm_linear.py#L1107)). - -```python -# Initialize model on meta device -with torch.device("meta"): - model = GPTModel(config) - -model = FullyShardedDataParallel( - transformer_config, - model, - ddp_config, - fsdp_unit_modules=[TransformerLayer, LanguageModelEmbedding], -) -``` - -**Important Considerations:** -- *Custom Modules*: If your model contains custom modules, ensure they implement the `reset_parameters` API. Otherwise, you may need to force parameter initialization on a CUDA or CPU device. -- *Tensor Initialization*: Be cautious of tensors created during model initialization without a specified device; they default to the meta device. To avoid issues, explicitly specify the device for these tensors to ensure compatibility with this function. - -### Interaction Between Custom FSDP and Model Forward/Backward Propagation - -Custom FSDP implements Fully Sharded Data Parallelism (FSDP) through a series of module hooks, gradient hooks, or by adding functions between modules. This involves inserting communications and manipulating parameters and gradients during PyTorch's module forward or backward propagation. - -Module hooks summary: -- Module pre-forward hook(`module.register_forward_pre_hook`): This hook unshards model weights before the forward pass. In the case of an FSDP Unit Module, add a RegisterFSDPBackwardFunction function that will reshard model weights and reduce gradients after module backward propagation. -- Module post-forward hook(`module.register_forward_hook`): This hook is used to reshard model weights after the forward pass. -- Root module pre-backward hook(`root_module.register_full_backward_pre_hook`): This hook checks that all model parameters are resharded, in order to avoid unnecessary memory spikes. It also marks all modules as being in the `TrainingState.PRE_BACKWARD` state. -- Module pre-backward hook(`module.register_full_backward_pre_hook`): This hook is used to unshard the model weights before the backward pass. -- Root module post-backward hook(`torch.autograd.Variable._execution_engine.queue_callback`): This hook is used to make sure all gradients in the backprop are properly handled / available. - -The gradient reduction pipeline maintains a map of gradients to FSDP parameter groups. If all gradients in an FSDP parameter group are ready, it launches a gradient reduction. Note that this assumes that the model's gradients are always generated in a certain order (reverse of `module.parameters()`), as otherwise, FSDP would maintain too many parameter group grad buffers, leading to excessive memory usage. - -#### Optimized for Activation Recompute - -Using activation recomputation runs the same module forward first and then its backward during backprop, which can unshard and reshard model weights twice. If the runtime can treat that as one forward-plus-backward region, it can unshard once and reshard once. - -To make this determination, the implementation tracks the model state with `training_state`: `FORWARD`, `PRE_BACKWARD`, `POST_BACKWARD`, `IDLE`. It is worth noting that the pre-backward hook runs before the pre-forward hook: the pre-backward hook performs the model weight unshard, then marks the model as `PRE_BACKWARD`, and when the pre-forward hook observes that mark it skips unshard. Similarly, for duplicate reshard logic, the post-forward hook runs before the post-backward path, and checking for the `PRE_BACKWARD` flag in the post-forward hook can cancel unshard. - -### Memory Mechanisms and Features of Custom FSDP - -FSDP can distribute model parameters, gradients, optimizer states, and (for mixed-precision training) high-precision main weights. That covers most memory outside activations, but FSDP can still hit allocator and spike issues. - -FSDP frequently unshards and reshards model weights, which can lead to busy memory allocation and deallocation. This results in untimely tensor releases, causing memory spikes (or even out-of-memory errors), crashes of the PyTorch memory allocator cache, and many `cudaMalloc` and `cudaFree` calls. These issues can slow the system noticeably. - -You can often address untimely tensor release with the `tensor._typed_storage()._resize_(0)` API, which deallocates storage immediately. Custom FSDP exposes hooks in `AllGatherPipeline` and `GradReducePipeline` to swap the temporary buffer allocator used for parameter gathering and gradient reduction with `StorageResizeBasedBucketAllocator`, using that `_resize_(0)` path for releases. - -The PyTorch memory allocator cache can fail when real usage nears GPU capacity, which hurts performance. Mitigation is limited; avoiding repeated pressure at the memory limit helps. A self-managed allocator such as `RotaryBucketAllocator` is another option, though it is not yet mature. - -## References - -- [Getting Started with Fully Sharded Data Parallel (FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html) diff --git a/docs/user-guide/features/index.md b/docs/user-guide/features/index.md index d080fe3f256..9dea6bd34a4 100644 --- a/docs/user-guide/features/index.md +++ b/docs/user-guide/features/index.md @@ -17,7 +17,7 @@ Guides for Megatron Core training features. fine_grained_activation_offloading moe context_parallel -custom_fsdp +megatron_fsdp dist_optimizer optimizer_cpu_offload pipeline_parallel_layout diff --git a/docs/user-guide/features/megatron_fsdp.md b/docs/user-guide/features/megatron_fsdp.md new file mode 100644 index 00000000000..6b55bf9fcce --- /dev/null +++ b/docs/user-guide/features/megatron_fsdp.md @@ -0,0 +1,608 @@ + + +# Megatron-FSDP + +## ✨ Overview + +**Megatron-FSDP** is an NVIDIA-developed distributed parallelism library written in native PyTorch that provides a high-performance implementation of **Fully Sharded Data Parallelism (FSDP)**. It offers seamless cross-compatibility with various deep learning frameworks and parallelism libraries such as Megatron-Core, and is performance-optimized to support training and inference of extremely large PyTorch models at data-center scale on NVIDIA GPUs. + +- PyPI: https://pypi.org/project/megatron-fsdp/ +- Source Code: https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/distributed/fsdp/src + +### 🧩 Compatibility + +- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/stable/distributed.html#devicemesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)** +- **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)** +- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)** +- **[NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)** + +### 💡 Features + +- **Performant & Scalable**: Optimized for NVIDIA CUDA with efficient memory management and performance. Sports near-linear scaling up from single compute nodes to entire data-centers. +- **Multiple Algorithms in One**: Supports sharding your choice of optimizer states, gradients, and model parameters (FSDP), including hierarchical data parallelism strategies such as **Hybrid-Sharded Data Parallelism (HSDP)** and **Hybrid-FSDP (HFSDP / Fully-Sharded Optimizer State)** for optimizing intra-node and inter-node memory, communication, and performance. +- **"Bring Your Own Parallelism"**: Works seamlessly with PyTorch, Megatron-LM, Megatron-Bridge, and TransformerEngine, and can be plugged into other frameworks such as HuggingFace Transformers and TorchTitan. +- **Simple & Powerful**: Similar to PyTorch FSDP, the `fully_shard` API doesn't depend on any complex training framework or distributed environment. + +### ⏱️ Optimizations + +- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) Mixed-Precision & Fused Kernels**: Native performance- and memory-optimal _compatibility with MXFP8, NVFP4, and various other quantization recipes and fused kernels provided by TransformerEngine_. +- **Advanced Bucketing**: `dtype`-customizable and precision-aware bucketing system to _tune the memory overhead, numerical accuracy, and latency of collectives_. Avoids redundant `COPY` operations before and after collectives, while remaining compatible with **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)** features such as **[Torch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**. +- **Buffer Management**: Efficient use of storage and [NCCL User Buffer Registration](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#user-buffer-registration) enable _direct communication into NCCL-managed memory_, achieving true zero-`COPY` data movement. Introduced in NCCL `v2.27`, **NCCL Symmetric Memory** communications employ _symmetric kernels_ that drastically reduce SM utilization and include networking optimizations such as high-precision (`FP32`) reduction over-the-wire. +- **Optimized Communication & SM Utilization via SHARP**: Leverages [**SHARP** (Scalable Hierarchical Aggregation and Reduction Protocol)](https://docs.nvidia.com/networking/display/sharpv3130) to _offload FSDP collectives to network switches (InfiniBand or NVLink-Switch)_ and significantly reduce utilization of GPU streaming multi-processors (SM) from 16-32 to 1-6 for **Multi-Node NVLink (MNNVL)** systems (Grace-Blackwell, Vera-Rubin, etc.), which lowers communication latency in large scaled-out workloads and frees up GPU-hosted processors for overlapped compute (GEMM) kernels. When FSDP sharding domains span both NVLink and InfiniBand, **hierarchical SHARP collectives** (NVL-SHARP and IB-SHARP) _optimize communication paths across the entire system topology_. +- [**Hybrid-FSDP (HFSDP)**](#understanding-hybrid-fsdp-hfsdp), a variation of _Hybrid-Sharded Data Parallelism (HSDP)_ that further shards the optimizer state across intra- and inter-node data-parallel ranks, _bridges the memory-communication trade-off between HSDP and FSDP_, unlocking memory efficiency at minimal cost to performance. + +## 🚀 Quick Start + +### 📦 Installation + +#### NeMo Framework Container + +Megatron-FSDP is pre-installed with Megatron-Core in the [NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags). + +#### Megatron-Core + +Megatron-FSDP is bundled with Megatron-Core, which can be installed via `pip`: + +``` +# Install via PyPI +pip install --no-build-isolation megatron-core[mlm,dev] + +# Install from Source +git clone https://github.com/NVIDIA/Megatron-LM.git +cd Megatron-LM +pip install --no-build-isolation .[mlm,dev] +``` + +To import Megatron-FSDP in Python: +```python +import megatron.core.distributed.fsdp.src.megatron_fsdp +``` + +#### PyPI + +To install Megatron-FSDP as a standalone package to use the `fully_shard` API: + +``` +pip install megatron-fsdp +``` + +To import Megatron-FSDP in Python: + +```python +import megatron_fsdp +``` + +### 🎛️ Megatron-FSDP `fully_shard` + +Megatron-FSDP supports a simple `fully_shard` API that seamlessly enables FSDP with very few lines of code. + +```python +import torch +from megatron_fsdp import ( + fully_shard_model, + fully_shard_optimizer, +) + +# Initialize Torch Distributed. +torch.distributed.init_process_group() +torch.cuda.set_device(torch.distributed.get_rank()) + +# Fully-shard the model. +model = torch.nn.Transformer() +fsdp_model = fully_shard_model( + module=model, + fsdp_unit_modules=[ + torch.nn.TransformerEncoder, + torch.nn.TransformerDecoder + ] +) + +# Fully-shard the optimizer. +toy_adam = torch.optim.AdamW(params=fsdp_model.parameters(), lr=0.01) +optimizer = fully_shard_optimizer(optimizer=toy_adam) + +# Forward pass. +inp = torch.randn(1, 512, 512).to("cuda") +tgt = torch.randn(1, 512, 512).to("cuda") +output = fsdp_model(inp, inp) + +# Backward pass. +torch.nn.functional.mse_loss(output, tgt).backward() + +# Optimizer step. +optimizer.step() +optimizer.zero_grad() + +# Checkpoint the model and optimizer. +torch.distributed.checkpoint.save({ + "model": fsdp_model.state_dict(), + "optimizer": optimizer.state_dict(), +}, checkpoint_id="ckpt/") + +# Load the saved checkpoint. +ckpt = { + "model": fsdp_model.state_dict(), + "optimizer": optimizer.state_dict(), +} +torch.distributed.checkpoint.load(state_dict=ckpt, checkpoint_id="ckpt/") +fsdp_model.load_state_dict(ckpt["model"], strict=False) +optimizer.load_state_dict(ckpt["optimizer"]) +``` + +> ℹ️ `fully_shard` is an _**experimental**_ API. Please check back for updates as we fine-tune our user experience! For more examples using `fully_shard` for Megatron-FSDP, refer to our suite of unit tests: [`tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py`](../../../tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py) + +### 🤖 Megatron-LM + +Megatron-FSDP is deeply integrated into Megatron-Core. To enable FSDP (where optimizer states, gradients, and compute parameters are sharded) in Megatron, use the following arguments: + +``` +# Train models in Megatron-LM using Megatron-FSDP. +--use-megatron-fsdp +--data-parallel-sharding-strategy {no_shard, optim, optim_grads, optim_grads_params} +--ckpt-format fsdp_dtensor +``` + +Complete Llama-8B and DeepSeek-V3 training scripts using Megatron-FSDP with recommended settings can be found in [Megatron-LM/examples/megatron_fsdp](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/megatron_fsdp). + +#### Recommended Configuration for Megatron-LM + +Frequently-used options use with Megatron-FSDP include: + +```bash +# Un-set CUDA_DEVICE_MAX_CONNECTIONS to ensure stream independence / full-parallelization of FSDP computation and communication. May slightly affect TP and CP performance though. +unset CUDA_DEVICE_MAX_CONNECTIONS + +# Meta-Device Initialization - Load large model onto CUDA devices in shards to avoid OOM. +--init-model-with-meta-device + +# Per-Token Loss / No Gradient Scaling - Deactivate DP scaling during gradient reduction, which can be a drain on SM resources. +--calculate-per-token-loss + +# Decrease gradient reduction and accumulation precision to recommended data-types based on the precision of the model parameters, usually BF16. Reduces communication volume during the backwards pass. Can be further customized with `--megatron-fsdp-main-grads-dtype` and `--megatron-fsdp-grad-comm-dtype`, which are enabled by this argument. +--grad-reduce-in-bf16 + +# Register NCCL user buffers and Megatron-FSDP double buffers to enable zero-copy symmetric kernels and low-SM utilization via SHARP. Improves overall performance but increases memory overhead due to double-buffering and is NOT compatible with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. +--use-nccl-ub +--fsdp-double-buffer +--fsdp-manual-registration +``` + +### 🤖 Megatron-Core + +Megatron-FSDP has a lower-level `FullyShardedDataParallel` class API that can be used with a simplified version of Megatron-LM's training loop. + +```python +# Initialize model and optimizer. +ddp_config.use_megatron_fsdp = True +# Megatron-FSDP Base Sharding Strategies: +# no_shard, optim, optim_grads, optim_grads_params +ddp_config.data_parallel_sharding_strategy = "optim_grads_params" +model = GPTModel(transformer_config) +model = FullyShardedDataParallel( + transformer_config, + model, + ddp_config, + fsdp_unit_modules = [TransformerLayer, LanguageModelEmbedding], +) +optimizer = torch.optim.AdamW(model.parameters(), lr=lr) +optimizer = DistributedOptimizer(optimizer, [model], [model.param_and_grad_buffer]) + +# Training loop +def train_step(inputs, labels): + optimizer.zero_grad() + for mbs_input, mbs_label in zip(inputs, labels): + outputs = model(mbs_input) + loss = loss_fn(outputs, mbs_label) + loss.backward() + optimizer.step() + +# Save and load model and optimizer state dict +def model_and_optimizer_state_dict(): + state_dict = { + "model": model.sharded_state_dict(), + "optimizer": optimizer.sharded_state_dict(), + } + return state_dict + +def load_model_and_optimizer_state_dict(state_dict): + model.load_state_dict(state_dict["model"]) + optimizer.load_state_dict(state_dict["optimizer"]) +``` + +### 🔁 Checkpoint Conversion + +Megatron-FSDP checkpointing supports [PyTorch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html). In Megatron-LM, this is the `--ckpt-format fsdp_dtensor` checkpointing format. + +#### Converting Torch DCP to Torch Save (Non-Distributed) Checkpoints + +PyTorch has utilities to convert Torch DCP checkpoints to and from regular Torch checkpoints: +```shell +python -m torch.distributed.checkpoint.format_utils --help +usage: format_utils.py [-h] {torch_to_dcp,dcp_to_torch} src dst + +positional arguments: + {torch_to_dcp,dcp_to_torch} + Conversion mode + src Path to the source model + dst Path to the destination model + +options: + -h, --help show this help message and exit +``` +For example: +```shell +python -m torch.distributed.checkpoint.format_utils dcp_to_torch dcp_ckpt/ torch_ckpt.pt +``` +or: +```python +from torch.distributed.checkpoint.format_utils import ( + dcp_to_torch_save, + torch_save_to_dcp, +) + +# Convert DCP model checkpoint to torch.save format. +dcp_to_torch_save(CHECKPOINT_DIR, TORCH_SAVE_CHECKPOINT_PATH) + +# Convert torch.save model checkpoint back to DCP format. +torch_save_to_dcp(TORCH_SAVE_CHECKPOINT_PATH, f"{CHECKPOINT_DIR}_new") +``` +Torch Save checkpoints can then be converted into HuggingFace SafeTensors or other checkpoint formats for distribution. + +> ℹ️ Megatron-FSDP checkpoints have a `module.` prefix pre-pended to all model parameter names in the state dictionary, and converting a Torch Save checkpoint to a Megatron-FSDP Torch DCP checkpoint requires testing. Work-in-progress! + +#### Converting N-D Parallel (`torch_dist`) to Megatron-FSDP (`fsdp_dtensor`) Checkpoints + +As a pre-requisite for checkpoint conversion, dump the parameter group mapping when training with 3D-parallel (DDP, TP, PP) and/or EP: + +```bash +--dump-param-to-param-group-map /path/to/param_to_param_group_map +``` + +and convert the map to a `param_to_param_group_map.json` JSON file in the `/path/to/param_to_param_group_map` directory: + +```bash +python tools/checkpoint/checkpoint_inspector.py print-torch-dcp-in-json /path/to/param_to_param_group_map +``` + +> ℹ️ If you already have a `torch_dist` checkpoint, simply specify the `--dump-param-to-param-group-map /path/to/param_to_param_group_map` flag and run a trivial training or checkpointing experiment to create the `param_to_param_group_map` you need without full pretraining. + +Finally, convert your `torch_dist` checkpoint to the `fsdp_dtensor` format using the `param_to_param_group_map.json`: + +```bash +torchrun --nproc_per_node=8 --nnodes=1 \ + tools/checkpoint/checkpoint_inspector.py \ + convert-torch-dist-to-fsdp-dtensor (--swiglu) \ # --swiglu for specific models. + /path/to/input_torch_dist_checkpoint/ \ + /path/to/output_fsdp_dtensor_checkpoint/ \ + --param-to-param-group-map-json /path/to/param_to_param_group_map.json +``` + +> ℹ️ For multi-node conversion tasks, please refer to the DeepSeek-V3 example script (`sbatch_checkpoint_convert.sh`) in [Megatron-LM/examples/megatron_fsdp](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/megatron_fsdp). + +## Megatron-FSDP Feature Guide & API + +| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config | +|--------------|-------------|----------------------|----------------------| +| **Megatron-FSDP** | Use Megatron-FSDP in Megatron-LM. | `--use-megatron-fsdp` | `fully_shard_model(module)` | +| **Megatron-FSDP Checkpointing** | Save and load un-even DTensor checkpoints using [Torch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html). | `--ckpt-format fsdp_dtensor` | `preproc_state_dict_for_dcp_ckpt=True` | +| **Meta Device Initialization** | Megatron-FSDP initializes a meta-initialized model to the CUDA device in shards to avoid OOM on large models. Requires implementation of `Module.reset_parameters()` for per-Module sharded initialization. | `--init-model-with-meta-device` | `init_model_with_meta_device=True` | +| **Distributed Optimizer** | Megatron-FSDP uses Megatron-Core's `DistributedOptimizer`. Automatically set when using Megatron-FSDP. | `--use-distributed-optimizer` | `fully_shard_optimizer(optimizer)` | + +### FSDP Fundamentals + +```{figure} ../../images/megatron_fsdp/DDP_vs_FSDP.png +:alt: FSDP Pipeline +:align: center + +Comparison between Distributed Data Parallelism (DDP) and Fully-Sharded Data Parallelism (FSDP). While gradients are all-reduced in DDP, they are sharded and reduce-scattered with FSDP. + +Source: Meta AI, Ott, Myle, et al. “Fully Sharded Data Parallel: Faster AI Training with Fewer GPUs.” _Facebook Engineering_, 15 July 2021, https://engineering.fb.com/2021/07/15/open-source/fsdp/. +``` + +**Fully Sharded Data Parallelism (FSDP)** is a type of distributed data parallelism (DDP) that shards optimizer state, weight gradients (`wgrad`), and model weights across devices that ingest data-parallel samples for data-parallel training or inference. Activations (`fprop`) and data gradients (`dgrad`) are not sharded or distributed, and are preserved for the backward pass, but can be recomputed during the backward pass, offloaded to CPU, or sharded / routed using other parallelisms such as tensor parallelism (TP), context parallelism (CP), or expert parallelism (EP). + +```{figure} ../../images/megatron_fsdp/zero3_model_state.png +:alt: ZeRO-3 Model State +:align: center + +Sharded memory profiles for ZeRO-1 (optimizer state), ZeRO-2 (optimizer state and gradients), and ZeRO-3 (optimizer state, gradients, and parameters). + +Source: Zero-Redundancy Optimizer Model State Partition Diagram. From _The Ultra-Scale Playbook: Training LLMs on GPU Clusters_ by Tazi, Nouamane, et al. HuggingFace, 2025, https://huggingface.co/spaces/nanotron/ultrascale-playbook. +``` + +The core principles of FSDP are: + +- Only a small depth-wise fraction of the model state can exist un-sharded at any point in time. +- Communication should overlap computation. + +From these core principles, software requirements can be derived: + +0. Model states sharded by FSDP are directly initialized across devices in shards. +1. Model parameters are all-gathered (AG) in pre-designated groups or modules pre-forward and pre-backward to un-shard a small fraction of the model state at any point in time during training or inference. After `fprop` and `dgrad` computation, the un-sharded weights are immediately de-allocated. +2. `wgrad` are reduce-scattered (RS) and accumulated in pre-designated groups or modules immediately post-backward to limit the amount of un-sharded gradients at any point in time during training or inference. +3. Distributed optimizers, optimizers that are initialized with respect to a sharded model state and support distributed mechanics, update the sharded model state using the reduced gradient shard to implement data parallelism (DP). +4. Computation and communication are overlapped across multiple CUDA streams, expending multiple streaming multi-processors (SM). Weights from subsequent groups or modules are pre-fetched, which ideally hides the communication latency required for FSDP behind model computation kernels (GEMM). + +FSDP can also be visualized as a decomposition of the all-reduce collective used in DDP into a gradient reduce-scatter, distributed optimization step, and parameter all-gather. + +```{figure} ../../images/megatron_fsdp/FSDP_Allreduce.png +:alt: FSDP RS & AG +:align: center + +Source: Feng, Wei, Will Constable, and Yifan Mao. “Getting Started with Fully Sharded Data Parallel (FSDP2).” _PyTorch Tutorials_, 17 Mar. 2022, https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html. +``` + +### FSDP Unit Modules + +| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config | +|--------------|-------------|----------------------|----------------------| +| **FSDP Unit Modules** | A list of `str` or `class` import paths for `torch.nn.Module`(s) that are considered FSDP unit modules and sharded by Megatron-FSDP. Parameters and sub-modules that are not members of an FSDP unit are not sharded. | Defaults to supported Megatron-Core modules (`TransformerLayer`, etc.) in Megatron-LM. | `fsdp_unit_modules=[...]` | +| **FSDP Double Buffer Allocator** | Megatron-FSDP uses the double-buffer allocator, which persistently allocates a buffer pair assigned to alternating FSDP units that temporarily stores parameters and gradients. Automatically used with NCCL user buffer registration. | `--fsdp-double-buffer` | `fsdp_double_buffer=True` | +| **Param All-Gather Overlap** | Whether to overlap parameter all-gather with compute. Automatically activated for the ZeRO-3 sharding strategy. | `--overlap-param-gather` | `overlap_param_gather=True` | +| **Gradient Reduce-Scatter Overlap** | Whether to overlap gradient reduce-scatter or all-reduce with compute. Automatically activated for ZeRO-2 and ZeRO-3 sharding strategies. | `--overlap-grad-reduce` | `overlap_grad_reduce=True` | +| **FSDP Communication Size** | Customize the size (in `numel()` elements) of AG and RS communications in Megatron-FSDP, by limiting how many elements are concurrently pre-fetched or reduced for AG and RS. Effectively suggests how many FSDP units are processed concurrently, which may launch collectives earlier and improve performance. Optionally, tune this value depending on system memory and performance requirements. | `--suggested-communication-unit-size ` | N/A (Megatron-Core Only) | + +> Only a small depth-wise fraction of the model state can exist un-sharded at any point in time. + +**FSDP Unit Modules** represent fractions of the model state that are computed and communicated as a (coalesced) group, un-sharded when needed for computation, and re-sharded after computation to release memory for subsequent model states. Implicitly, an FSDP unit module is also a **_modeling contract_**, requiring that FSDP-managed unit module parameters are not accessed or modified beyond the scope of the forward pass, backward pass, or optimization step. + +Megatron-FSDP accepts a list of `str` or `class` paths representing FSDP unit modules via the `fsdp_unit_modules` argument, which is currently hard-coded to supported model classes (like `TransformerLayer`) in Megatron-Core. It performs a depth-first traversal of the model (via `torch.nn.Module.named_modules()`) and groups the parameters of each matching module for sharding and coalesced communication. Nested units are resolved by precedence: if a module matches an FSDP unit class but is already a sub-module of a previously registered FSDP unit, it is skipped, so the outermost (and necessarily largest) FSDP unit class in any module sub-tree becomes the effective FSDP unit module. + +> Communication should overlap computation. + +Once a model is partitioned into unit modules, computation is overlapped with communication based on the granularity of the FSDP unit module. Depending on the size of the compute and communication kernels, fine-tuning the unit module size and grouping configuration can impact performance and elicit trade-offs between overlap and memory when using FSDP. + +```{figure} ../../images/megatron_fsdp/fsdp_streams.png +:alt: FSDP Streams +:align: center + +Each color-coded block in the compute and communication streams, merged and categorized in the simplified (and worst-case) scenario where SM resources are under contention, correspond to a _single_ FSDP unit module. +``` + +Compute-communication overlaps are orchestrated using **CUDA streams** that capture and parallelize serial operations. All collectives associated with all combinations of `{DP-Inner, DP-Outer}` and `{AG, RS}` are scheduled and tracked with separate streams and communicators / `ProcessGroup`(s). + +- Parameters are un-sharded prior to `fprop` and `dgrad` computation. To overlap the pre-fetch all-gather with computation, at least two FSDP units worth of un-sharded weight memory is required at any point in time. +- Gradients are reduced and sharded after `wgrad` computation. To overlap gradient reduce-scatter with `wgrad` computation, at least two FSDP units worth of un-sharded gradient memory is required at any point in time. + +#### FSDP Module Hooks + +To implement these "unit-periodic" mechanics, Megatron-FSDP uses `Module` hooks to install a variety of (pre- and post-) forward and backward operations: + +- **Pre-Forward** + - Un-shards the model parameters of the current and (via pre-fetching) forward-subsequent FSDP unit modules. + - When `MegatronFSDP.forward()` is invoked, Megatron-FSDP will swap all parameter references to point to the un-sharded `Tensor` compute weights for the forward and backward pass. +- **Post-Forward** + - Re-shards model weights after the forward pass, if the module is an FSDP unit. Non-unit modules remain persistently un-sharded. + - When using activation recomputation during the backwards pass, computing both `fprop` and `dgrad` requires these parameters, so parameters are resharded during **Post-Backward**. + - Releases the transpose cache of quantized parameters (in FSDP / ZeRO-3) for specific quantization recipes in `TransformerEngine`. +- **Pre-Backward** + - Un-shards the model parameters of the current and (via pre-fetching) backward-subsequent FSDP unit modules. + - Implemented as a `torch.autograd.graph.register_multi_grad_hook` triggered by the output `dgrad`, and installed via a `Module` _post-forward_ hook. +- **Post-Backward** + - Re-shards model weights after the backward pass, if the module is an FSDP unit. Non-unit modules remain persistently un-sharded. + - Implemented by injecting an Autograd function (`RegisterFSDPBackwardFunction`) that is installed during a `Module` _pre-forward_ hook. + - Reduces gradients after the backward pass. + - Implemented using a `Tensor.register_post_accumulate_grad_hook` triggered by `param.grad`, as well as a root-level post-backward hook installed during **Pre-Backward** (`torch.autograd.Variable._execution_engine.queue_callback`). +- **State Dictionary** + - When `module.state_dict()` (for any module managed by Megatron-FSDP) is invoked, Megatron-FSDP will swap all parameter references to point to sharded `DTensor` main weights for distributed optimization and checkpointing. + - When `MegatronFSDP.load_state_dict()` is invoked, both the main and compute weights are updated. When using quantized model compute, the main weights are quantized and sharded. + +#### Double Buffering + +Megatron-FSDP uses a `Tensor._typed_storage()._resize_(bytes)`-based allocator to instantly allocate and de-allocate memory without depending on the `CUDACachingAllocator` for un-sharded parameters and gradients by default. (Cache fragmentation and garbage collection can procrastinate large quantities of `cudaMalloc` and `cudaFree` operations that can block programs and spike memory, particularly when memory utilization is maxed out.) However, modifying the underlying storage of a buffer is not compatible with NCCL symmetric registration or CUDA graphability, which require a persistent state during runtime. + +To support these optimizations, Megatron-FSDP uses **double-buffering**, which assigns 2 persistently-allocated buffers to FSDP units in an alternating pattern, hard-limiting the memory overhead for parameter and gradient buffer allocation and ensuring that no more than 2 FSDP units are computed or communicated concurrently. + +```{figure} ../../images/megatron_fsdp/fsdp_double_buffer.png +:alt: FSDP Double Buffering +:align: center + +Visualization of double buffering in Megatron-FSDP. Even- and odd-indexed FSDP units share the same un-sharded parameter and gradient buffers, overwriting incumbent data as needed during runtime. Megatron-FSDP ensures that no more than two FSDP units are un-sharded at any point during runtime. +``` + +With double-buffering, Megatron-FSDP does not need to allocate memory after initialization, which can reduce memory fragmentation and improve performance. However, double-buffering requires _depth-wise model symmetry_, where even- and odd-indexed FSDP units have identical size during runtime. If double-buffering is utilized, Megatron-FSDP computes the **_mode_** of FSDP unit sizes as the symmetrical double-buffer size, and any FSDP units not symmetrical to the computed size will default to the `_resize_(bytes)`-based allocator (or persistently allocated for extremely large and asymmetrical layers that affect performance significantly like `torch.nn.Embedding` when the low-level argument `fsdp_db_use_persist_buf_on_alloc_fail` is set). + +### Data-Parallel Sharding Strategies + +| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config | +|--------------|-------------|----------------------|----------------------| +| **Data Parallel Sharding Strategy** | Primary data-parallel sharding strategy for FSDP, which supports DDP, ZeRO-1 (optimizer), ZeRO-2 (optimizer and gradients), and ZeRO-3 (optimizer, gradients, and parameters). Typically uses intra-node communications, i.e. "inner" or "intra" DP. | `--data-parallel-sharding-strategy {no_shard, optim, optim_grads, optim_grads_params}` | `zero_dp_strategy={no_shard, optim, optim_grads, optim_grads_params, 0, 1, 2, 3}` | +| **DP-Outer Sharding Strategy** | Secondary data-parallel sharding strategy for HSDP, which supports Hybrid-Sharded Data Parallel (HSDP / `no_shard`) and Hybrid-FSDP (HFSDP / `optim`). Typically uses inter-node communications, i.e. "outer" or "inter" DP. | `--outer-dp-sharding-strategy {no_shard, optim}` | `outer_dp_sharding_strategy={no_shard, optim, 0, 1}` | +| **Hybrid Data Parallelism Size** | Specify the DP-Outer / Inter-DP parallel size. DP-Inner / Intra-DP sizes will be deduced from the sizes of other parallelisms and `torch.distributed.get_world_size()`. | `--num-distributed-optimizer-instances ` | `dp_outer_dim=` (Cumulative DP groups `hybrid_fsdp_group` / `hybrid_fsdp_expt_group` are required for HFSDP.) | + +Megatron-FSDP supports a variety of sharding strategies over a variety of distributed topologies: + +- **Distributed Data Parallelism (DDP)** + - Model state is replicated across DP ranks. + - Gradient all-reduce is overlapped with backward compute and launched during the last backward pass before the optimization step. +- **ZeRO-1** + - Optimizer state is sharded across DP ranks. + - Gradient reduce-scatter is overlapped with backward compute and launched during the last backward pass before the optimization step. (Reduce-scatter is used in lieu of all-reduce for performance, because only a shard of the gradient is needed for optimization.) +- **ZeRO-2** + - Optimizer state and gradients are sharded across DP ranks. + - Gradient reduce-scatter is overlapped with backward compute and accumulated during every backward pass. +- **Fully-Sharded Data Parallelism (FSDP / ZeRO-3)** + - Optimizer state, gradients, and parameters are sharded across DP ranks. + - Gradient reduce-scatter is overlapped with backward compute and accumulated during every backward pass. +- **Hybrid-Sharded Data Parallelism (HSDP)** + - Optimizer state, gradients, and parameters are sharded across the "inner" or "intra" DP ranks. + - Model state is replicated across "outer" / "inter" DP ranks, and outer data-parallel gradients are all-reduced during the last backward pass before the optimization step. +- **Hybrid-FSDP (HFSDP)** + - Optimizer state, gradients, and parameters are sharded across the "inner" or "intra" DP ranks. + - Optimizer state is _further_ sharded across "outer" / "inter" DP ranks. + - Outer data-parallel gradients are reduce-scattered after during the last backward pass before the optimization step. + - Outer data-parallel parameters are all-gathered during the first forward pass after the optimization step. + - FSDP primary sharding (`optim_grads_params`) is required for HFSDP secondary sharding (`optim`). + - Requires passing cumulative data-parallel groups (`hybrid_fsdp_group` / `hybrid_fsdp_expt_group`), which include ALL data-parallel ranks, to Megatron-FSDP. + - To create these using `DeviceMesh`, create a data-parallel `DeviceMesh` for the cumulative DP group and use `DeviceMesh._unflatten(dp_dim, mesh_sizes=(dp_outer_size, dp_inner_size), mesh_dim_names=("dp_outer_dim", "dp_shard_dim"))` to construct a `DeviceMesh` with DP-Inner and DP-Outer mesh dimensions for Hybrid-FSDP. + +#### Understanding Hybrid-FSDP (HFSDP) + +```{figure} ../../images/megatron_fsdp/hfsdp.png +:alt: Hybrid-FSDP Topology +:align: center + +Hybrid-FSDP (HFSDP) is a variation of HSDP where the optimizer state in particular is sharded across both DP-Inner and DP-Outer, i.e. all data-parallel ranks, which further reduces memory utilization. In other words, intra-node sharding and communication uses ZeRO-3, while inter-node sharding and communication uses ZeRO-1. Parameters and gradients are converted from and to the fully-sharded optimizer state during optimization steps only, reducing the frequency of inter-node communications. + +Inspired by the artistry in the DHEN (Zhang, Luo, Liu, Meta, et al., 2022) paper: https://arxiv.org/abs/2203.11014 +``` + +**Hybrid-Fully Sharded Data Parallelism (HFSDP)** is a slight modification to HSDP that fully-shards the optimizer state across all data-parallel ranks and introduces outer-level all-gather and reduce-scatter collectives to map fully-sharded parameters and gradients into partially-sharded parameters and gradients. + +The memory profile of HFSDP is a "hybrid" of FSDP (optimizer state) and HSDP (gradients and model weights). Another elegant way to understand HFSDP functionality is ZeRO-1 composed with ZeRO-3. + +$$\text{Hybrid-FSDP Memory Profile} = \frac{\text{Optimizer State}}{\text{DP-Inner} \ \times \ \text{DP-Outer}} + \frac{\text{Gradient} + \text{Weight}}{\text{DP-Inner}}$$ + +The modified algorithm has the following characteristics: + +- Megatron-FSDP maintains a view of the model parameters sharded across all data-parallel ranks. + - Distributed checkpoints save and load the fully-sharded model parameters. + - Distributed optimizer state is initialized on the fully-sharded model parameters. +- During the first forward pass after checkpointing or optimization, fully-sharded model weights are all-gathered into partially-sharded model weights. +- During the last backward pass before optimization, partially-sharded model gradients are reduce-scattered into fully-sharded model gradients. +- Otherwise, FSDP is performed on the partially-sharded model weights and accumulated gradients. Because model weights and gradients are only updated and ingested once per optimization cycle, we can skip or postpone all expensive inter-node / DP-outer collectives until an optimization step.​ + +In addition to improved memory utilization, HFSDP communications are split in communication size (bytes communicated), communication topology (DP-Inner and DP-Outer groups), and communication domain (NVLink and InfiniBand) across two sharding stages. + +```{figure} ../../images/megatron_fsdp/fsdp_v_hfsdp_streams.png +:alt: Hybrid-FSDP Streams +:align: center + +Inter-node communications can also be parallelized with intra-node communications using separate CUDA streams. +``` + +#### Mixing FSDP & Model Parallelism + +Megatron-FSDP is also compatible with a variety of model parallelisms that shard the model state, such as **Tensor Parallelism (TP)** and **Expert Parallelism (EP)**. When sharding model states across multiple dimensions in the device topology, _**FSDP sharding is always performed last**_, because FSDP collectives un-shard and re-shard parameters and gradients immediately before and after computation. Thus, FSDP sharding mechanics are implemented over tensor and expert parallel (strided) shards. + +```{figure} ../../images/megatron_fsdp/mixed_sharding.png +:alt: Mixed Model Parallelism +:align: center + +Wheneveer FSDP is composed with other model parallelisms, FSDP sharding is always exercised last to seamlessly integrate with existing model shards. +``` + +Megatron-FSDP uses `torch.distributed.DeviceMesh` to describe and configure communications across devices in data-parallel group(s). Because heterogeneous models that have mixed layers, such as [Hybrid Mamba-Transformer](https://arxiv.org/abs/2504.03624) or [Mixture-of-Experts (MoE)](https://arxiv.org/abs/1701.06538) models, require different parallelism configurations, multiple `DeviceMesh`(s) may be required for specific layers that require distinct distributed topologies for optimal memory efficiency and performance. + +Currently, Megatron-FSDP supports two `DeviceMesh`(s), one for dense / non-expert `Module`(s) and another for Megatron-Core MoE sparse / expert `Module`(s). (Expert modules and parameters in Megatron-Core are automatically detected.) + +- Dense modules typically have a `DeviceMesh` with data parallel, tensor parallel, and context parallel dimensions, where the data parallel dimension is used for FSDP. Typically, both data-parallel and context-parallel ranks are used for sharding in FSDP. +- Mixture-of-experts modules typically have a `DeviceMesh` with data parallel, tensor parallel, and expert parallel dimensions, where the data parallel dimension is used for FSDP. + +For more information about Mixture-of-Experts in Megatron-Core, refer to the [Megatron-Core User Guide - MoE](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/moe.html). + +#### Non-Uniform / Un-Even Model Sharding + +While `torch.distributed.tensor.DTensor` defaults to per-parameter sharding, where Tensors are split evenly on `dim=0` across the data-parallel domain, Megatron-FSDP uses **non-uniform or un-even `DTensor` shards** of a (flattened) group of parameters associated with an FSDP unit. + +```{figure} ../../images/megatron_fsdp/uneven_sharding.png +:alt: Non-Uniform Sharding +:align: center + +Comparison of FSDP2 per-parameter sharding and Megatron-FSDP per-unit or per-module sharding. FSDP2 requires `COPY` operations to move parameters and gradients in and out of communication buffers to reduce the frequency of NCCL collective calls, while Megatron-FSDP assigns sliced views of contiguous communication buffers to parameters associated with an FSDP unit. +``` + +While complex and less user-intuitive, an un-evenly sharded data structure enables a few performance benefits without introducing expensive `COPY` operations to set up communication and computation buffers: + +- **Fewer NCCL calls**, reducing kernel launch and synchronization overhead. Only parameters in FSDP units that have different communication-related properties, such as their `dtype` or distributed topology, are coalesced into separate NCCL calls. +- Flat communication and computation buffers are **contiguous-by-design**, supporting optimized CUDA kernels that require buffers backed by contiguous memory, such as grouped GEMMs used in MoE. + +Effectively, this implies that the same `DTensor`-sharded model parameters may have completely different shapes on different ranks, and if entire parameters are assigned to other ranks, the local `Tensor` will be empty. + +> ℹ️ Megatron-FSDP has a handy library ([`megatron_fsdp.uneven_dtensor`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py)) for manipulating un-evenly sharded `DTensors`, focused on per-parameter operations like un-sharding or reducing parameters that have different shapes across ranks. While the parameter group is evenly-sharded for FSDP collectives, per-parameter collectives (that assume a symmetrical amount of bytes are communicated between devices) will hang waiting on bytes that will never arrive for un-evenly sharded `DTensors`. + +In particular, contiguous memory is only half the requirement for high-performance CUDA kernels. The other requirement is **locality**, which FSDP can violate, that introduces compatibility issues when combining FSDP with present and future optimizations. For example, block-wise quantization (scaling factor / `absmax` calculations for MXFP8, NVFP4, etc.) requires DP communication and custom max-reduce kernels if the block is sharded by FSDP. + +Megatron-FSDP supports `dim=0` sharding, which computes the _**least-common multiple (LCM) of `p.shape[1:]` for all parameters `p` in an FSDP unit**_ and _**pads the un-sharded buffer to the closest multiple of `DP x LCM(p.shape[1:])`**_, forming a "DP-LCM" partition with `LCM`-length parts to ensure that DP-sharding boundaries do not violate chunks of data for coordinates of `dim=0`. + +```{figure} ../../images/megatron_fsdp/lcm_dim0_shard.png +:alt: Flat Buffer Sharding Algorithm +:align: center + +Visualization of how parameters are assigned un-evenly to the flat per-unit buffer sharded across DP ranks. With the LCM algorithm, every slice of `dim=0` is never bisected by FSDP. Algorithms and compute kernels can leverage this locality and contiguity. +``` + +1. When a parameter is _divisble by the LCM_, it can be inserted at any index multiple of the LCM in the buffer that is free. `p[i]` chunks of this parameter by definition divide the LCM, and thus align with the DP-LCM sharding grid. +2. When a parameter _is larger than but not divisible by the LCM_, the remainder `r` populates a fraction of another LCM part, so a "conjugate" parameter that also exceeds the LCM with a "conjugate" remainder `r'` that is less than or equal to `LCM - r` is installed to fill the remaining space and align with the DP-LCM sharding grid. +3. When a parameter _is smaller than but not divisible by the LCM_, a post-assignment sweep on the leftover space in the flat buffer is run, and all gaps that are multiples of the LCM that are large enough to support the entire parameter are utilized. Once all gaps are filled, the final parameters are assigned to the tail of the buffer respecting the DP-LCM sharding grid. + +> ℹ️ Generalized support for contiguity and locality in Megatron-FSDP is a **_work-in-progress_** and will evolve with contribution from the OSS community and PyTorch. For more information about how kernel buffer requirements affect the design of FSDP data structures, refer to the [veScale: Consistent and Efficient Tensor Programming with Eager-Mode SPMD (Li, Youjie, ByteDance Seed, et al.)](https://arxiv.org/abs/2509.07003) paper that comprehensively analyzes these requirements. + +### Mixed-Precision & Quantization + +| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config | +|--------------|-------------|----------------------|----------------------| +| **Quantized Parameters** | Megatron-FSDP will shard and all-gather TransformerEngine-quantized parameters for computation. Quantized parameters are updated every optimization step, and both row-wise (FWD) and column-wise (BWD) data are managed for non-transposable 1-D quantization recipes like MXFP8. Otherwise, only activations are quantized. | `--fp8-param-gather` | TransformerEngine `quantized_model_init()` | +| **Main Parameter (Optimization / Checkpoint) Data-Type** | Data-type for optimization and checkpointing parameters. If set to `auto`, model compute weights are utilized instead. Required for `--fp8-param-gather`. Defaults to FP32. | `--megatron-fsdp-main-params-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(main_params_dtype=...)` | +| **Main Gradient (Accumulation) Data-Type** | Data-type for gradient accumulation. If set to `auto`, main gradient precision will be derived from model parameter precision. Defaults to `auto`. | `--megatron-fsdp-main-grads-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(main_grads_dtype=...)` | +| **Gradient Communication (Reduction) Data-Type** | Data-type for gradient communication and reduction. If set to `auto`, the main gradient precision will be used for communication. (When using NCCL symmetric registration, low-precision gradients are reduced in FP32 over-the-wire.) Defaults to `auto`. | `--megatron-fsdp-grad-comm-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(grad_comm_dtype=...)` | +| **Weight Gradient Accumulation Fusion** | When using TransformerEngine modules, Megatron-FSDP implements `get_main_grad` to allocate un-sharded gradient buffers called by TransformerEngine, to avoid `COPY`-ing the gradient to Megatron-FSDP communication buffers. Used by default and can be deactivated with `--no-gradient-accumulation-fusion`. | `--no-gradient-accumulation-fusion` | N/A (Megatron-Core Only) | +| **Precision-Aware Optimizer** | Use the TransformerEngine `FusedAdam` optimizer, and Megatron-FSDP will install the gradient in a temporary attribute `Parameter.decoupled_grad` which is consumed by `FusedAdam`. Megatron-FSDP manages the main parameters, but the optimizer state precision can be customized with `--exp-avg-dtype` and `--exp-avg-sq-dtype`, which both support `fp8` optimization state. | `--use-precision-aware-optimizer` | `use_decoupled_grad=True` | + +#### Quantization + +Quantization is an extremely important feature for Megatron-FSDP as it reduces memory utilization and communication size for both activations and parameters, which directly affects the viability and performance of FSDP. + +```{figure} ../../images/megatron_fsdp/quantized_param_gather.png +:alt: Quantized Model Parameters & FSDP +:align: center + +Visualization of Megatron-FSDP's training loop when using quantized weights from TransformerEngine. Every optimization step updates the quantized representation of sharded model weights, which have reduced communication size. +``` + +While TransformerEngine handles activation quantization, Megatron-FSDP shards quantized weights for AG. + +0. _**Quantized Model Initialization**_ - Model is initialized with quantized weights, e.g. MXFP8 or NVFP4. If using `meta` device initialization, Megatron-FSDP will call `reset_parameters()` to initialize quantized weights layer-by-layer. If row-wise and column-wise data are not transposable, Megatron-FSDP will shard and buffer both. Additionally, high-precision main weights are retrieved and sharded for distributed optimization, checkpointing, and quantization. +0. _**Forward / Backward Pass**_ - Quantized weights are un-sharded for both the forward and backward pass. If row-wise and column-wise data aren't transposable, the row-wise weights are gathered for forward, and the column-wise weights are gathered for backward. +0. _**Distributed Optimization Step**_ - Non-quantized accumulated gradient shards from quantized GEMMs are applied to high-precision main weight shards. +0. _**Sharded Quantization**_ - Sharded main weights are quantized to update the quantized compute weights for subsequent training steps. + +```{figure} ../../images/megatron_fsdp/sharded_quantization.png +:alt: Sharded Quantization +:align: center + +Sharded quantization involves reducing maxima to compute a global set of scaling factors for local / sharded quantization. +``` + +In particular, _sharded quantization_ minimizes communication size and memory utilization by communicating scaling factors instead of main weights. + +1. _**Local Abs-Max**_ - For a group of parameters in an FSDP unit, compute local tensor-wise or block-wise maxima across the global un-sharded shape, with zero padding for non-local data. +1. _**Global Abs-Max**_ - Globally all-reduce maxima and derive scaling factors from maxima. +1. _**Local Quantization**_ - Locally quantize sharded main weights and install into compute weight buffers. + +#### Mixed-Precision + +Megatron-FSDP sharding and communication buffers support mixed-precision, such that users can customize the `dtype` used for main weights, gradient communication (reduction), and gradient accumulation in addition to the native or quantized `dtype` used for model computation. These options are wrapped in a `MixedPrecisionPolicy` dataclass. + +- _**Main Weight Precision**_ - Controls the data-type for parameters responsible for distributed optimization, distributed checkpointing, and quantization. If set to `auto` (`None`), the native model compute parameter data-type will be utilized. Required for parameter quantization with `--fp8-param-gather`. Defaults to `torch.float32`. +- _**Main Gradient Precision**_ - Controls the data-type for `wgrad` accumulation and distributed optimization. Defaults to `auto` (`None`), the model native gradient data-type will be utilized. While `torch.float32` (or higher) is recommended for accuracy at scale, as `main_grads_dtype` controls the data-type for gradient accumulation, `auto` is more flexible and uses pre-determined parameter gradient logic in mixed-precision scenarios, such as `BF16` for `FP8`/`FP4` parameters quantized via TransformerEngine. +- _**Gradient Communication Precision**_ - Controls the data-type for gradient communications when reducing gradients. Lower precision improves (communication) performance. Defaults to `auto` (`None`), in which the main gradient data-type will be utilized. If using `no_shard`, `optim`, HSDP, or HFSDP, allocating `dtype`-custom gradient communication buffers may increase per-unit memory overhead, so users should consider the performance-memory trade-off when using this feature. + - If using NCCL symmetric registration `v2.27+`, gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced and accumulated in communication and accumulation precision as usual. + +### NCCL + +| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config | +|--------------|-------------|----------------------|----------------------| +| **NCCL User Buffers** | Allocate and register Megatron-FSDP communication buffers with NCCL, which enables zero-`COPY`, high-precision reduction, copy-engine collectives, and symmetric kernels. Uses double buffering. | `--use-nccl-ub` | `nccl_ub=True` | +| **NCCL Manual Registration** | Instead of registering NCCL user buffers on first allocation, batch registration of all communication buffers at the end of the initial training step. Reduces registration latency. | `--fsdp-manual-registration` | N/A (Megatron-Core Only) | +| **Disable Symmetric Registration** | Disable symmetric registration with NCCL. Optional, as symmetric registration failure defaults to normal registration. | `--disable-symmetric-registration` | `disable_symmetric_registration=True` | + +[NVIDIA Collective Communications Library (NCCL)](https://developer.nvidia.com/nccl) implements multi-device and multi-node communication primitives optimized for CUDA devices and networking from NVIDIA. Megatron-FSDP communications are registered and deeply integrated with NCCL, which enables a variety of hardware-level networking optimizations such as copy-engine AG, high-precision RS, SHARP reduction offloading, and symmetric kernels. + +To leverage NCCL networking optimizations, **NCCL user buffer registration (UBR)** is required to inform NCCL of PyTorch Tensors ("user buffers") that act directly as the input and target of NCCL collectives for PyTorch `ProcessGroup`(s). Because registered communication buffers are known to NCCL, `COPY` operations that send collective inputs to NCCL buffers and collective outputs to PyTorch buffers are no longer required, which enables Megatron-FSDP to be zero-`COPY` end-to-end. + +NCCL (`v2.27+`) supports symmetric allocation or registration for communicators over the NVLink domain, which allow buffers that share identical virtual addresses across devices to benefit from optimized collectives: + +- **Symmetric Kernels** - On the NVLink domain, symmetric kernels operating on symmetric memory reduces the SM utilization for a single communication kernel to 1. +- **NVSwitch SHARP Offloading** - To further minimize SM utilization for AG and RS collectives, NCCL SHARP offloads reduction and aggregation work to NVLink and IB Switch hardware that uses 1-6 SM depending on the domain: NVL, IB, or NVL + IB. +- **Copy-Engine (CE) Collectives**: Instead of using SMs (or CTAs) for common non-computational collectives like AG in Megatron-FSDP, copy engines are instead used to perform all-gather collectives, dedicating SM resources to compute and reduction during FSDP. Requires NCCL `v2.28+`. +- **High-Precision Reduction**: When training large models, high-precision gradient reduction and accumulation is desired for accuracy and convergence, but communicating FP32 gradients is expensive. With symmetric registration, FP32 accumulators enable gradients to be reduced in FP32 but communicated in BF16, which decreases gradient RS communication latency while maintaining high accuracy during training. Megatron-FSDP supports FP32 main gradient accumulation but BF16 gradient communication, customizable through `megatron_fsdp.MixedPrecisionPolicy`. + +These optimizations significantly reduce SM resource contention for overlapped compute and communication kernels in FSDP. Symmetric registration, allocation, and pooling is also supported in PyTorch: [`torch.distributed._symmetric_memory`](https://docs.pytorch.org/docs/stable/symmetric_memory.html). \ No newline at end of file diff --git a/docs/user-guide/parallelism-guide.md b/docs/user-guide/parallelism-guide.md index e09848f8800..d78a7939c7b 100644 --- a/docs/user-guide/parallelism-guide.md +++ b/docs/user-guide/parallelism-guide.md @@ -15,19 +15,20 @@ Megatron Core supports multiple parallelism strategies that can be combined to e The following table summarizes supported parallelism strategies. -| Strategy | What it parallelizes | Best for | +| Strategy | Parallelism Objective | Best For | |----------|---------------------|----------| -| **Data Parallelism (DP)** | Batch dimension | Standard training, most common | -| **Tensor Parallelism (TP)** | Individual layers | Large layers, GPU memory constraints | -| **Pipeline Parallelism (PP)** | Model depth | Very deep models | -| **Context Parallelism (CP)** | Sequence length | Long sequences (8K+ tokens) | -| **Expert Parallelism (EP)** | MoE experts | Mixture-of-Experts models | +| **Data Parallelism (DP)** | Batch Dimension | Data Scalability, Standard Training | +| **Tensor Parallelism (TP)** | Individual Layers | Large Layers & Activation, GPU Memory Constraints | +| **Pipeline Parallelism (PP)** | Model Depth | Very Deep Models | +| **Context Parallelism (CP)** | Sequence Length | Long Sequences (8K+ Tokens) | +| **Expert Parallelism (EP)** | MoE Experts | Mixture-of-Experts Models | +| **Fully-Sharded Data Parallelism (Megatron-FSDP)** | Model State | Extremely Large Models & DP Interchangeability | ## Data Parallelism (DP) -Replicate the model across GPUs and split the batch. +### Standard Distributed Data Parallel (DDP) -### Standard Data Parallel (DDP) +Replicate the model across GPUs and split the batch. ```bash torchrun --nproc_per_node=8 pretrain_gpt.py \ @@ -36,22 +37,40 @@ torchrun --nproc_per_node=8 pretrain_gpt.py \ Each GPU has a full copy of the model and processes a portion of the batch. -### Fully Sharded Data Parallel (FSDP) +### Megatron Fully-Sharded Data Parallel (Megatron-FSDP) -Shard model parameters, gradients, and optimizer states to reduce memory: +Shard model parameters, gradients, and optimizer states across GPUs to reduce memory utilization. -```bash -# Megatron FSDP (~15% faster than PyTorch FSDP2) ---use-megatron-fsdp \ +``` +--use-megatron-fsdp --data-parallel-sharding-strategy optim_grads_params +--ckpt-format fsdp_dtensor +--init-model-with-meta-device ``` -**Sharding strategies** +**Sharding Strategies** + +`--data-parallel-sharding-strategy` supports the following options: - `optim` - Shard optimizer states only (ZeRO-1) - `optim_grads` - Shard gradients + optimizer (ZeRO-2) - `optim_grads_params` - Shard parameters + gradients + optimizer (ZeRO-3) +If `--num-distributed-optimizer-instances` is > 1, then hierarchical data parallelism is enabled. + +`--outer-dp-sharding-strategy` supports the following options: + +- `no_shard` (**Hybrid-Sharded Data Parallelism**) - Replicate the model state across outer data parallel ranks. +- `optim` (**Hybrid-FSDP**) - Shard the optimizer state across the outer data parallel ranks. + - Requires `--data-parallel-sharding-strategy optim_grads_params`. + +**When to Use** + +- Large models with large or fused compute kernels to hide communications under. +- Integrated with TP, CP, EP, and easily composable with heterogeneous parallelisms. +- With SM-reducing optimizations from NCCL and activation offloading from TransformerEngine. +- Using `fully_shard` without depending on Megatron-LM. + ## Tensor Parallelism (TP) Split individual model layers across GPUs. Recommended for large hidden dimensions. @@ -61,7 +80,7 @@ Split individual model layers across GPUs. Recommended for large hidden dimensio --sequence-parallel # Enable sequence parallelism (recommended) ``` -**When to use** +**When to Use** - Model layers do not fit on a single GPU - Large hidden dimensions (4096+) @@ -76,7 +95,7 @@ Split model layers across GPUs vertically (by depth). --num-layers-per-virtual-pipeline-stage 4 # Virtual pipeline for load balancing ``` -**When to use** +**When to Use** - Very deep models (50+ layers) - Combine with TP for large models @@ -91,7 +110,7 @@ Split long sequences across GPUs for efficient long-context training. --cp-comm-type p2p # Communication type ``` -**When to use** +**When to Use** - Long sequences (8K+ tokens) - Reduces activation memory @@ -208,24 +227,24 @@ Reduces activation memory by sharding sequence dimension in LayerNorm and Dropou ## Choosing the Right Strategy ### Start Simple -1. Begin with **Data Parallelism** (DP) only -2. Add **Tensor Parallelism** (TP) if the model does not fit -3. Add **Pipeline Parallelism** (PP) for very large models -4. Add **Context Parallelism** (CP) for long sequences +1. Begin with **Data Parallelism** (DP) only. +2. Add **Tensor Parallelism** (TP) if the model does not fit. +3. Add **Pipeline Parallelism** (PP) for very large models. +4. Add **Context Parallelism** (CP) for long sequences. ### Memory Constraints -- Use **FSDP** to reduce memory per GPU -- Use **TP** to split large layers -- Use **PP** to split model depth -- Enable **activation checkpointing** for extreme cases +- Use **FSDP** to split model state per GPU. +- Use **TP** to split large layers. +- Use **PP** to split model depth. +- Enable **activation checkpointing or offloading** for extreme cases. ### Communication Bottlenecks -- Reduce **TP** degree (increases memory per GPU) -- Increase **PP** degree (may reduce efficiency) -- Use **CP** instead of larger TP for long sequences +- Reduce **TP** degree (increases memory per GPU). +- Increase **PP** degree (may reduce efficiency). +- Use **CP** instead of larger TP for long sequences. ## Next Steps - **API Reference**: Refer to [Tensor Parallel](../api-guide/core/tensor_parallel.md) and [Pipeline Parallel](../api-guide/core/pipeline_parallel.md) in the API documentation -- **Advanced Features**: Refer to [Megatron FSDP](features/custom_fsdp.md) and [Distributed Optimizer](features/dist_optimizer.md) +- **Advanced Features**: Refer to [Megatron-FSDP](features/megatron_fsdp.md), [MoE](features/moe.md), and [Distributed Optimizer](features/dist_optimizer.md) - **Performance Tuning**: Refer to the [NVIDIA NeMo Performance Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance-guide.html) diff --git a/examples/megatron_fsdp/README.md b/examples/megatron_fsdp/README.md new file mode 100644 index 00000000000..eaf5eca1364 --- /dev/null +++ b/examples/megatron_fsdp/README.md @@ -0,0 +1,157 @@ +# Megatron-FSDP Examples + +Example scripts for training and checkpoint conversion using [Megatron-FSDP](../../docs/user-guide/features/megatron_fsdp.md). These demonstrate recommended configurations for Llama 3 8B and DeepSeek-V3 671B models, as well as checkpoint format conversion between `torch_dist` (N-D parallel) and `fsdp_dtensor` formats. + +## Scripts + +### `train_llama3_8b_fsdp_h100_fp8.sh` + +Single-node training script for **Llama 3 8B** using Megatron-FSDP with FP8 precision on H100 GPUs. Uses `torchrun` for local distributed training and supports both mock data (for benchmarking) and real datasets. + +#### Usage + +Run from the root of the Megatron-LM repository: + +```bash +# With mock data (default, for benchmarking) +bash examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh + +# With real data +bash examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh \ + checkpoints/llama3_8b_fsdp_fp8 \ + tensorboard_logs/llama3_8b_fsdp_fp8 \ + /path/to/tokenizer \ + /path/to/data_prefix +``` + +| Positional Argument | Default | Description | +|---------------------|---------|-------------| +| `$1` — Checkpoint path | `checkpoints/llama3_8b_fsdp_fp8` | Directory for saving and loading checkpoints. | +| `$2` — TensorBoard path | `tensorboard_logs/llama3_8b_fsdp_fp8` | Directory for TensorBoard logs. | +| `$3` — Tokenizer | `MOCK` | Path to a tokenizer model, or `MOCK` for `NullTokenizer`. | +| `$4` — Data path | `MOCK` | Data prefix for training data, or `MOCK` for mock data. | + +#### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `USE_MEGATRON_FSDP` | `1` | Set to `1` to enable Megatron-FSDP. Set to `0` to train with standard DDP. | +| `SHARDING_STRATEGY` | `optim_grads_params` | FSDP sharding strategy (ZeRO-3). Options: `no_shard`, `optim`, `optim_grads`, `optim_grads_params`. | +| `OUTER_SHARDING_STRATEGY` | `no_shard` | DP-Outer sharding strategy for HSDP/HFSDP. Options: `no_shard`, `optim`. | +| `MASTER_ADDR` | `localhost` | Master node address for distributed training. | +| `MASTER_PORT` | `6000` | Master node port. | +| `NODE_RANK` | `0` | Rank of the current node. | + +#### Configuration Summary + +- **Model**: Llama 3 8B (GQA with 32 heads / 8 KV groups, RoPE, SwiGLU, RMSNorm) +- **Parallelism**: TP=1, CP=1, PP=1, 8 GPUs per node, FSDP ZeRO-3 +- **Precision**: FP8 (hybrid format) with BF16 training and BF16 gradient reduction +- **Batch size**: micro-batch=1, global-batch=128, sequence length=8192 +- **Optimizations**: NCCL user buffers, FSDP double buffering, manual registration, meta-device initialization, per-token loss, overlapped grad-reduce and param-gather + +--- + +### `sbatch_mfsdp_deepseek_v3.sh` + +Multi-node SLURM training script for **DeepSeek-V3** (671B MoE) using Megatron-FSDP. Submits an `sbatch` job with containerized execution via `srun`. + +#### Usage + +Set the required configuration variables and submit: + +```bash +export MEGATRON_PATH=/path/to/Megatron-LM +export CONTAINER_IMAGE=/path/to/container.sqsh # or docker image URL +export OUTPUT_PATH=/path/to/output +export DATA_PATH=/path/to/training/data + +bash examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh +``` + +Before running, update the `#SBATCH` directives and `--container-mounts` in the script to match your cluster configuration. + +#### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `MEGATRON_PATH` | *(required)* | Path to the Megatron-LM repository. | +| `CONTAINER_IMAGE` | *(required)* | Container image (`.sqsh` file or Docker URL). | +| `OUTPUT_PATH` | *(required)* | Base directory for checkpoints, TensorBoard logs, SLURM logs, and Nsight profiles. | +| `DATA_PATH` | *(required)* | Training data prefix path. | +| `USE_MEGATRON_FSDP` | `1` | Enable Megatron-FSDP. Set to `0` for standard DDP. | +| `SHARDING_STRATEGY` | `optim_grads_params` | FSDP sharding strategy (ZeRO-3). | +| `TP` | `1` | Tensor parallel size. | +| `EP` | `8` | Expert parallel size. | +| `MBS` | `4` | Micro-batch size. | +| `GBS` | `2048` | Global batch size. | +| `PROFILE` | `0` | Set to `1` to enable Nsight Systems profiling (steps 10–12). | +| `WANDB` | `1` | Set to `1` to enable Weights & Biases logging. Requires `WANDB_API_KEY`. | +| `COMMENT` | N/A | Tag appended to W&B experiment names and Nsight profile filenames. | + +#### Configuration Summary + +- **Model**: DeepSeek-V3 (61 layers, 256 routed experts, top-8 routing, Multi-Latent Attention, MTP) +- **Parallelism**: TP=1, EP=8, CP=1, FSDP ZeRO-3 +- **Precision**: BF16 +- **MoE**: Flex dispatcher with HybridEP backend, grouped GEMM, sigmoid routing with expert bias, auxiliary sequence loss +- **Recomputation**: Selective recomputation of `mlp`, `moe`, `mla_up_proj`, and `layernorm` modules +- **Optimizations**: NCCL user buffers, FSDP double buffering, meta-device initialization, per-token loss, overlapped grad-reduce and param-gather +- **Tokenizer**: `deepseek-ai/DeepSeek-V3` via HuggingFace + +--- + +### `sbatch_checkpoint_convert.sh` + +SLURM batch script for converting checkpoints from **`torch_dist`** (N-D parallel) format to **`fsdp_dtensor`** (Megatron-FSDP) format. This enables resuming training under Megatron-FSDP from checkpoints originally saved with tensor/pipeline/expert parallelism. + +#### Prerequisites + +Before converting, you need a `param_to_param_group_map.json` file. Generate it by running a `torch_dist` training job with the `--dump-param-to-param-group-map` flag, then converting the output: + +```bash +# 1. Run a training job (or trivial experiment) with the dump flag +--dump-param-to-param-group-map /path/to/param_to_param_group_map + +# 2. Convert the dumped map to JSON +python tools/checkpoint/checkpoint_inspector.py \ + print-torch-dcp-in-json /path/to/param_to_param_group_map +``` + +See the [Checkpoint Conversion](../../docs/user-guide/features/megatron_fsdp.md#checkpoint-conversion) section in the Megatron-FSDP docs for details. + +#### Usage + +Set the required configuration variables, update the checkpoint paths in `RUN_CMD`, and submit: + +```bash +export MEGATRON_PATH=/path/to/Megatron-LM +export CONTAINER_IMAGE=/path/to/container.sqsh +export OUTPUT_PATH=/path/to/output + +bash examples/megatron_fsdp/sbatch_checkpoint_convert.sh +``` + +Before running, you must edit the script to fill in: +- The input `torch_dist` checkpoint path +- The output `fsdp_dtensor` checkpoint path +- The path to `param_to_param_group_map.json` +- The `#SBATCH` directives and `--container-mounts` for your cluster + +#### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `MEGATRON_PATH` | *(required)* | Path to the Megatron-LM repository. | +| `CONTAINER_IMAGE` | *(required)* | Container image (`.sqsh` file or Docker URL). | +| `OUTPUT_PATH` | *(required)* | Base directory for SLURM logs. | + +#### Conversion Command + +The script runs `checkpoint_inspector.py convert-torch-dist-to-fsdp-dtensor` with the `--swiglu` flag (for models using SwiGLU activations). Remove `--swiglu` if converting a non-SwiGLU model. + +## Further Reading + +- [Megatron-FSDP User Guide](../../docs/user-guide/features/megatron_fsdp.md) — full feature guide, API reference, and sharding strategy documentation. +- [Megatron-FSDP on PyPI](https://pypi.org/project/megatron-fsdp/) — standalone `fully_shard` API. +- [Megatron-FSDP Source](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/distributed/fsdp/src) — implementation source code. diff --git a/docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_checkpoint_convert.sh b/examples/megatron_fsdp/sbatch_checkpoint_convert.sh similarity index 100% rename from docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_checkpoint_convert.sh rename to examples/megatron_fsdp/sbatch_checkpoint_convert.sh diff --git a/docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_mfsdp_deepseek_v3.sh b/examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh similarity index 99% rename from docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_mfsdp_deepseek_v3.sh rename to examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh index 7b93d25d943..22a8f22f68c 100644 --- a/docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_mfsdp_deepseek_v3.sh +++ b/examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh @@ -23,7 +23,7 @@ TP=${TP:-1} EP=${EP:-8} MBS=${MBS:-4} GBS=${GBS:-2048} -COMMENT=${COMMENT:-"hybridep-selective-recompute"} +COMMENT=${COMMENT:-""} PRETRAIN_ARGS=( --distributed-timeout-minutes 60 diff --git a/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh b/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh new file mode 100644 index 00000000000..ddd3f160fa7 --- /dev/null +++ b/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh @@ -0,0 +1,212 @@ +#!/bin/bash + +CHECKPOINT_PATH=${1:-"checkpoints/llama3_8b_fsdp_fp8"} +TENSORBOARD_LOGS_PATH=${2:-"tensorboard_logs/llama3_8b_fsdp_fp8"} +TOKENIZER_ARG=${3:-"MOCK"} # Path to tokenizer model, or "MOCK" +DATA_ARG=${4:-"MOCK"} # Data prefix, or "MOCK" + +# Create directories if they don't exist +mkdir -p "$(dirname "$CHECKPOINT_PATH")" +mkdir -p "$(dirname "$TENSORBOARD_LOGS_PATH")" + +# Distributed training setup +GPUS_PER_NODE=8 +NUM_NODES=1 +MASTER_ADDR=${MASTER_ADDR:-localhost} +MASTER_PORT=${MASTER_PORT:-6000} +NODE_RANK=${NODE_RANK:-0} +WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES)) + +# Path to the pretrain_gpt.py script, assuming this script +# is run from the root of the Megatron-LM repository. +PRETRAIN_SCRIPT_PATH="pretrain_gpt.py" + +# Model & Training Parameters +USE_MEGATRON_FSDP=${USE_MEGATRON_FSDP:-1} +SHARDING_STRATEGY=${SHARDING_STRATEGY:-"optim_grads_params"} +OUTER_SHARDING_STRATEGY=${OUTER_SHARDING_STRATEGY:-"no_shard"} +TP_SIZE=1 +CP_SIZE=1 +PP_SIZE=1 +MICRO_BATCH_SIZE=1 +GLOBAL_BATCH_SIZE=128 +NUM_LAYERS=32 +DTYPE="fp8" +SEQ_LENGTH=8192 +MAX_POSITION_EMBEDDINGS=8192 + +# Data cache path (useful for both mock and real data) +DATA_CACHE_PATH="${PWD}/benchmark_cache_llama3_8b_fsdp_fp8" +mkdir -p "$DATA_CACHE_PATH" + +DISTRIBUTED_ARGS=( + --nproc_per_node $GPUS_PER_NODE + --nnodes $NUM_NODES + --node_rank $NODE_RANK + --master_addr $MASTER_ADDR + --master_port $MASTER_PORT +) + +MODEL_ARGS=( + --use-mcore-models + --num-layers $NUM_LAYERS + --hidden-size 4096 + --ffn-hidden-size 14336 + --num-attention-heads 32 + --group-query-attention + --num-query-groups 8 + --kv-channels 128 + --seq-length $SEQ_LENGTH + --max-position-embeddings $MAX_POSITION_EMBEDDINGS + --position-embedding-type rope + --rotary-base 1000000 + --rotary-percent 1.0 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --swiglu + --normalization RMSNorm + --init-method-std 0.0134 + --attention-backend fused + --apply-layernorm-1p + --untie-embeddings-and-output-weights + --disable-bias-linear +) + +TRAINING_ARGS=( + --micro-batch-size $MICRO_BATCH_SIZE + --global-batch-size $GLOBAL_BATCH_SIZE + --train-samples 1953125000 + --lr-decay-samples 1949218748 + --lr-warmup-samples 3906252 + --lr 0.00015 + --min-lr 0.00001 + --decoupled-lr 5.0e-4 + --decoupled-min-lr 4.5e-5 + --lr-decay-style cosine + --clip-grad 1.0 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.95 + --bf16 + --cross-entropy-loss-fusion + --manual-gc + --empty-unused-memory-level 1 + --exit-duration-in-mins 235 +) + +if [ "${USE_MEGATRON_FSDP}" = 1 ]; then + unset CUDA_DEVICE_MAX_CONNECTIONS + TRAINING_ARGS=( + "${TRAINING_ARGS[@]}" + --use-megatron-fsdp + --data-parallel-sharding-strategy ${SHARDING_STRATEGY} + --no-gradient-accumulation-fusion + --calculate-per-token-loss + --init-model-with-meta-device + --ckpt-format fsdp_dtensor + --grad-reduce-in-bf16 + --use-nccl-ub + --fsdp-double-buffer + --fsdp-manual-registration + # To enable HFSDP, DP full-sharding of the optimizer state with + # hierarchical data parallelism (DP-Outer=2, DP-Inner=DP//2)... + # --num-distributed-optimizer-instances 2 + # --outer-dp-sharding-strategy ${OUTER_SHARDING_STRATEGY} + # To further customize Megatron-FSDP data precision... + # --megatron-fsdp-main-params-dtype fp32 + # --megatron-fsdp-main-grads-dtype auto + # --megatron-fsdp-grad-comm-dtype auto + ) +fi + +# Conditional arguments based on DTYPE (FP8) +DTYPE_ARGS=() +if [[ "$DTYPE" == "fp8" ]]; then + DTYPE_ARGS+=( + "--fp8-format hybrid" + "--fp8-amax-history-len 1024" + "--fp8-amax-compute-algo max" + "--fp8-param-gather" + ) +fi + +# Model parallelism arguments +MODEL_PARALLEL_ARGS=( + --tensor-model-parallel-size $TP_SIZE + --context-parallel-size $CP_SIZE + --sequence-parallel +) + +# Distributed Data Parallel (DDP) arguments +# From original script's ddp_args +DDP_ARGS=( + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather +) +TRAINING_ARGS+=("${DDP_ARGS[@]}") + + +# Data arguments (conditional for mock vs real data) +DATA_ARGS_LIST=() +if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then + DATA_ARGS_LIST+=( + "--mock-data" + "--tokenizer-type NullTokenizer" + "--vocab-size 128256" + "--data-cache-path ${DATA_CACHE_PATH}" + "--tiktoken-pattern v2" + "--split '99,1,0'" + "--no-create-attention-mask-in-dataloader" + "--no-mmap-bin-files" + "--num-workers 1" + ) +else + # Settings for real data + DATA_ARGS_LIST+=( + "--data-path $DATA_ARG" + "--tokenizer-type HuggingFaceTokenizer" + "--tokenizer-model $TOKENIZER_ARG" + "--data-cache-path ${DATA_CACHE_PATH}" + "--split '99,1,0'" + "--no-create-attention-mask-in-dataloader" + "--no-mmap-bin-files" + "--num-workers 1" + # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit. + "--vocab-size 128256" + ) +fi + +EVAL_AND_LOGGING_ARGS=( + --log-interval 1 + --eval-iters 32 + --eval-interval 100 + --save-interval 1000 + --log-throughput + --profile + --profile-step-start 4 + --profile-step-end 6 + --distributed-timeout-minutes 60 + --save "$CHECKPOINT_PATH" + --load "$CHECKPOINT_PATH" + --tensorboard-dir "$TENSORBOARD_LOGS_PATH" +) + +# Ensure pretrain_gpt.py is found +if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then + echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH" + echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present." + exit 1 +fi + +# Run the training command +torchrun ${DISTRIBUTED_ARGS[@]} \ + "$PRETRAIN_SCRIPT_PATH" \ + ${MODEL_ARGS[@]} \ + ${TRAINING_ARGS[@]} \ + ${DTYPE_ARGS[@]} \ + ${MODEL_PARALLEL_ARGS[@]} \ + ${DATA_ARGS_LIST[@]} \ + ${EVAL_AND_LOGGING_ARGS[@]} + +set +x \ No newline at end of file diff --git a/megatron/core/distributed/README.md b/megatron/core/distributed/README.md index c4a75284414..489e381f9e0 100644 --- a/megatron/core/distributed/README.md +++ b/megatron/core/distributed/README.md @@ -1,11 +1,27 @@ -## How to use pytorch FSDP2? +# Distributed Data Parallelism -Add these flag to enable Torch FSDP2. +This module contains algorithms, data structures, and utilities used for different types of distributed data parallelism, such as DDP and FSDP. + +## Distributed Data Parallelism + +This is the default data parallelism used with all parallelism topologies in Megatron-LM. + +## Megatron-FSDP + +To use Megatron-FSDP in Megatron-LM, enable the following arguments: + +``` +--use-megatron-fsdp +--ckpt-format fsdp_dtensor +--init-model-with-meta-device +``` + +## FSDP2 + +To use FSDP2 in Megatron-LM, enable the following arguments: ``` --use-torch-fsdp2 --no-gradient-accumulation-fusion --ckpt-format torch_dist ``` - -It is worth noting that CUDA_MAX_CONNECTIONS=1 should not be enabled to ensure that the communication of FSDP and the computation on the primary stream can be fully parallelized. diff --git a/megatron/core/distributed/fsdp/src/README.md b/megatron/core/distributed/fsdp/src/README.md index 98c01a759eb..d3422d03abb 100644 --- a/megatron/core/distributed/fsdp/src/README.md +++ b/megatron/core/distributed/fsdp/src/README.md @@ -1,6 +1,6 @@
-# 🚀 Megatron-FSDP +# Megatron-FSDP
@@ -12,38 +12,16 @@ ## ✨ What is Megatron-FSDP? -**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP). It offers seamless cross-compatibility with major deep learning frameworks and parallelism libraries, making it easy to scale your PyTorch models across multiple GPUs and nodes. +**Megatron-FSDP** is an NVIDIA-developed distributed parallelism library written in native PyTorch that provides a high-performance implementation of **Fully Sharded Data Parallelism (FSDP)**. It offers seamless cross-compatibility with various deep learning frameworks and parallelism libraries such as Megatron-Core, and is performance-optimized to support training and inference of extremely large PyTorch models at data-center scale on NVIDIA GPUs. -Megatron-FSDP can provide up to 25% speed up and 23% memory savings compared to FSDP2. +For comprehensive information about Megatron-FSDP, refer to: [Megatron-FSDP | Megatron-Core Developer Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/) -### Compatibility +### 🧩 Compatibility -- **[PyTorch DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)** +- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/stable/distributed.html#devicemesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)** - **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)** - **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)** - -## ✨ Features - -- **Easy Integration**: Simple `fully_shard` function for quick model parallelization -- **High Performance**: Optimized for NVIDIA GPUs with efficient memory management -- **Cross-Framework**: Works seamlessly with PyTorch, Huggingface Transformers, Megatron-LM, Megatron Bridge and TransformerEngine -- **Scalable**: Supports both single-node multi-GPU and multi-node distributed training -- **Flexible Configuration**: Configurable sharding strategies and process groups - -## ⚡ Optimizations - -- **Advanced Bucketing**: Data-type aware bucketing system to minimize the overhead of collective operations -- **Buffer Management**: Zero copy communication is achieved by reorganizing the storage of parameters and main grad with `ParamAndGradBuffer` class -- **Communication Overlapping**: Improved communication overlap of paramter all-gather and gradient reduce-scatter -- **FP8 Mixed Precision with Transformer Engine**: Compatibility with Transformer Engine enables efficient FP8 mixed precision training -- **Gradient accumulate fusion support with Transformer Engine**: Remove the explicit gradient copy to the communication buffer in backwards pass - -### Advanced Collective Communication -- **SM Usage Reduction with SHARP**: FSDP's `All-Gather` (AG) and `Reduce-Scatter` (RS) collectives are designed to overlap with compute kernels. However, standard NCCL communication kernels can consume a significant number of GPU SMs (e.g., 16-32 SMs), "stealing" resources from compute (GEMM) kernels and reducing overall TFLOPS. -- **In-Switch Processing**: We leverage **SHARP** (Scalable Hierarchical Aggregation and Reduction Protocol) to offload these collective operations. SHARP performs aggregation and reduction computations directly on the network switches (InfiniBand or NVLink Switch) instead of on the GPU SMs. This dramatically reduces the SM consumption for communication to **1-6 SM** freeing up GPU resources for compute. It also provides lower communication latency, especially in large, scaled-out workloads. -- **Symmetric Optimizations for MNNVL**: We support **symmetric-based optimizations**, introduced in NCCL v2.27, which enable switch offloading for **Multi-Node NVLink (MNNVL)** systems such as GB200/GB300. This allows the same SM-saving benefits over the high-bandwidth NVLink fabric itself. -- **Hierarchical Collectives**: When an FSDP sharding domain spans both NVLink and InfiniBand, the library utilizes **hierarchical SHARP collectives** (e.g., NVL-SHARP + IB-SHARP) to optimize the communication path across the entire system topology. - +- **[NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)** ## 📦 Installation @@ -56,226 +34,57 @@ pip install megatron-fsdp ## 🚀 Quick Start -### Basic Usage - -Transform your PyTorch model to use Fully Sharded Data Parallelism with just a few lines: - -```python -import torch -from megatron_fsdp import ( - fully_shard_model, - fully_shard_optimizer, -) - -""" -Enable FSDP with Megatron-FSDP via the `fully_shard_*` API. -""" -# Shard your model. -model = fully_shard_model( - model, - fsdp_unit_modules=[ - YourModelLayerClass, - "import.path.to.model.class.YourModelLayerClass", - ], - ... -) -# Shard your optimizer. -optimizer = fully_shard_optimizer( - torch.optim.Adam(model.parameters(), lr=1e-3) -) - -# Your model is now ready for distributed training! -``` - -### Comparison with FSDP-2 - -`fully_shard` / `fully_shard_model` / `fully_shard_optimizer` are simple entrypoints into `MegatronFSDP`. - -- No need to call `fully_shard` on all the sub-modules, just pass your sub-module classes or import paths to `fully_shard`! -- Seamlessly preserves the identity of your training loop with only a few lines of code and multiple options for initialization: - - `fully_shard_*` is a two-line change when sharding the model and optimizer separately. - - `fully_shard` is a one-line change for previously-initialized models and optimizers. - -Compare this with FSDP2: - -```python -import torch -from torch.distributed.fsdp import fully_shard - -# Your existing model and optimizer. -model = YourModel() -optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) - -# Enable FSDP with FSDP2. -for module in model.modules(): - # Sub-Modules to shard. - if isinstance(module, YourModelLayerClass): - fully_shard(module) -fully_shard(model) - -# Your model is now ready for distributed training! -``` - -### `torch.compile` Compatibility - -Megatron-FSDP is compatible with `torch.compile`, but this feature is still experimental and may introduce performance regressions in some workloads. - -## 📖 Megatron-FSDP Comprehensive Walkthrough - -### Import `megatron_fsdp`. - ```python import torch from megatron_fsdp import ( fully_shard_model, fully_shard_optimizer, - MixedPrecisionPolicy, ) -``` - -### Set up a distributed environment using `DeviceMesh`. - -`DeviceMesh` simplifies the construction of complex arrangements of devices -to support various parallelisms. - -```python -from torch.distributed.device_mesh import DeviceMesh - -# Initialize DeviceMesh. -device_mesh = torch.distributed.device_mesh.init_device_mesh( - "cuda", - mesh_shape=(dp_outer_size, dp_shard_size, cp_size, tp_size), - mesh_dim_names=("dp_outer", "dp_shard", "cp", "tp"), -) -# Only relevant when using HSDP, where we also need the full DP group for data parallelism, -# This sub-mesh can be provided to distributed samplers or dataloaders. -device_mesh[("dp_outer", "dp_shard")]._flatten("dp") -# Only required if using CP. Otherwise, just pass dp_shard to FSDP. -device_mesh[("dp_shard", "cp")]._flatten("dp_shard_cp") -# Only required if using HSDP. Otherwise, don't pass hybrid_fsdp_group. -device_mesh[("dp_outer", "dp_shard", "cp")]._flatten("hsdp") -hsdp_group = device_mesh["hsdp"].get_group() - -# Initialize DeviceMesh for expert parallel (EP) modules when using FSDP + EP. -expert_device_mesh = torch.distributed.device_mesh.init_device_mesh( - "cuda", - mesh_shape=(dp_outer_size, expt_dp_shard_size, expt_tp_size), - mesh_dim_names=("dp_outer", "dp_shard_cp", "tp"), -) -expert_device_mesh[("dp_outer", "dp_shard_cp")].flatten("hsdp") -hsdp_expt_group = expert_device_mesh["hsdp"].get_group() -``` - -### Convert models into fully-sharded `MegatronFSDP` models with `fully_shard_model`. - -This wraps the model in a MegatronFSDP class that schedules the sharding -lifecycle of the model parameters and gradients during training and inference. - -```python -model = fully_shard_model( - # PyTorch (Root) Module - model, - # Sharded Modules - fsdp_unit_modules=[...], - # Device Mesh - device_mesh=device_mesh - # Always required for FSDP or HSDP. - dp_shard_dim="dp_shard_cp", - # Set this required argument to use HSDP instead of FSDP. Otherwise, set this to None. - dp_outer_dim="dp_outer", - # Only required for TP-sensitive models (i.e. Megatron-LM / TransformerEngine) - # or when using DTensor-based TP. Otherwise, set this to None. - tp_dim="tp", - # Only required when using HSDP. Otherwise, set this to None. - hybrid_fsdp_group=hsdp_group, - # Only required when using HSDP + EP. Otherwise, set this to None. - hybrid_fsdp_expt_group=hsdp_expt_group, - # Only required for FSDP + EP. Otherwise, set this to None. - expt_device_mesh=expt_device_mesh, - # FSDP Sharding Strategy: no_shard (0) / optim (1) / optim_grads (2) / optim_grads_params (3) - zero_dp_strategy=3, - outer_dp_sharding_strategy=1, - # Initialize the model on devices in shards to avoid OOM. Requires device("meta")-init for model. - init_model_with_meta_device=True, - # Mixed-Precision Policy for controlling compute and communication precision in Megatron-FSDP. - mixed_precision_policy=MixedPrecisionPolicy(), - # Sync parameters and gradients each step. Allows for gradient transformations after backward pass, - # and synchronizes parameters and gradients across HSDP groups, but deactivates compute-communication - # overlap going into the subsequent training step. - sync_model_each_microbatch=True, - # Preprocess state dict for DCP checkpointing. Required for Torch Distributed Checkpoint. - preproc_state_dict_for_dcp_ckpt=True, -) -``` - -The original `torch.nn.Module` can be accessed at `MegatronFSDP.module`. - -### Initialize and fully-shard your optimizer on the `MegatronFSDP` model. -Initialize your optimizer on the Megatron-FSDP model distributed `Parameter`(s). -If your optimizer has already been initialized, either use the `fully_shard` -entrypoint, or use `optimizer.add_param_group({"params": model.parameters()})` -after resetting your optimizer state via `optimizer.param_groups.clear()` -and `optimizer.state.clear()`. +# Initialize Torch Distributed. +torch.distributed.init_process_group() +torch.cuda.set_device(torch.distributed.get_rank()) -```python -optimizer = torch.optim.Optimizer(model.parameters()) -``` - -`fully_shard_optimizer` modifies your `optimizer.step()`, `optimizer.zero_grad()`, -and distributed optimizer parameters to punctually trigger scheduled FSDP operations -for Megatron-FSDP. - -```python -fully_shard_optimizer( - # PyTorch Optimizer - optimizer, - # Preprocess state dict for DCP checkpointing. - # Required for Torch Distributed Checkpoint. - preproc_state_dict_for_dcp_ckpt=True, -) -``` - -Extended arguments to `step()` and `zero_grad()` control these FSDP operations: - -```python - optimizer.step( - ..., - # Sync all gradients before the optimizer step. Alternatively enabled using - # `sync_model_each_microbatch=True` in MegatronFSDP. - sync_grad_before_optimizer_step=True, - # After `optimizer.step()`, install optimized weights into MegatronFSDP's buffers. - install_optimized_model_weights=True, - ) - - optimizer.zero_grad( - ..., - # Also zero out MegatronFSDP's gradient accumulation buffers. - zero_grad_buffer=True - ) -``` - -### `MegatronFSDP` Distributed Checkpointing - -Distributed checkpoints can be saved and loaded using Torch DCP. Alternatively, -you can load non-distributed checkpoints before fully-sharding your model with -any existing checkpoint utility compatible with PyTorch Modules. - -```python -# Save model and optimizer state. -torch.distributed.checkpoint.save( - {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, - checkpoint_id=str(CKPT_DIR) +# Fully-shard the model. +model = torch.nn.Transformer() +fsdp_model = fully_shard_model( + module=model, + fsdp_unit_modules=[ + torch.nn.TransformerEncoder, + torch.nn.TransformerDecoder + ] ) -# Load model and optimizer state. -ckpt_state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()} -torch.distributed.checkpoint.load(state_dict=ckpt_state_dict, checkpoint_id=str(CKPT_DIR)) -# `model.load_state_dict(strict=False)` is only necessary to ignore TE FP8 extra state -# that is missing from the DCP checkpoint but present in TEBaseModule. -# Megatron-FSDP does not support TE FP8 extra state checkpointing with DCP. -model.load_state_dict(ckpt_state_dict["model"], strict=False) -optimizer.load_state_dict(ckpt_state_dict["optimizer"]) +# Fully-shard the optimizer. +toy_adam = torch.optim.AdamW(params=fsdp_model.parameters(), lr=0.01) +optimizer = fully_shard_optimizer(optimizer=toy_adam) + +# Forward pass. +inp = torch.randn(1, 512, 512).to("cuda") +tgt = torch.randn(1, 512, 512).to("cuda") +output = fsdp_model(inp, inp) + +# Backward pass. +torch.nn.functional.mse_loss(output, tgt).backward() + +# Optimizer step. +optimizer.step() +optimizer.zero_grad() + +# Checkpoint the model and optimizer. +torch.distributed.checkpoint.save({ + "model": fsdp_model.state_dict(), + "optimizer": optimizer.state_dict(), +}, checkpoint_id="ckpt/") + +# Load the saved checkpoint. +ckpt = { + "model": fsdp_model.state_dict(), + "optimizer": optimizer.state_dict(), +} +torch.distributed.checkpoint.load(state_dict=ckpt, checkpoint_id="ckpt/") +fsdp_model.load_state_dict(ckpt["model"], strict=False) +optimizer.load_state_dict(ckpt["optimizer"]) ``` ## ⚙️ `fully_shard` / `MegatronFSDP` API - Advanced Features @@ -305,17 +114,17 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin - Defaults to `False`. - Note that the `device` argument which installs your model on a specific device or rank will be deactivated when `init_model_with_meta_device=True`. - `mixed_precision_policy` takes a `megatron_fsdp.MixedPrecisionPolicy` that configures mixed-precision compute and communication for Megatron-FSDP. Configuration options include: - - `main_params_dtype` controls the data-type for parameters used in distributed optimization or quantization. + - `main_params_dtype` controls the data-type for parameters responsible for distributed checkpointing, distributed optimization, and quantization. - Defaults to `torch.float32`. - If set to `None`, the native model compute parameter data-type will be utilized. - - Requires specification (cannot be `None`) when using `FP8` parameters with Megatron-FSDP. + - Requires specification (cannot be `None`) when using quantized parameters with Megatron-FSDP. - `main_grads_dtype` controls the data-type for gradients used in distributed optimization. - - Defaults to `None`, the model native gradient data-type will be utilized. + - Defaults to `None`, in which the model native gradient data-type will be utilized. - While `torch.float32` (or higher) is recommended for accuracy at scale, as `main_grads_dtype` controls the data-type for gradient accumulation, `None` is more flexible and uses pre-determined parameter gradient logic in mixed-precision scenarios, such as `BF16` for `FP8`/`FP4` parameters quantized via TransformerEngine. - - `grad_comm_dtype` controls the data-type for gradient communications (RS / AR) when reducing gradients. Lower precision `grad_comm_dtype` improves (communication) performance, but may increase memory utilization or sacrifice gradient precision in certain cases. - - Defaults to `None`, the `main_grads_dtype` data-type will be utilized, and no additional memory is allocated when `grad_comm_dtype == main_grads_dtype`. - - If using HSDP (either DP-Replicate or DP-Outer in `outer_dp_sharding_strategy`), `no_shard`, `optim`, or a `FixedPoolAllocator` (`fsdp_double_buffer`), allocating `dtype`-custom gradient communication buffers (per FSDP group) adds memory overhead of up to 10% or more, and users should consider the performance-memory trade-off when using this feature. - - If using NCCL UBR v2.27+ (`nccl_ub=True`), gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced in `grad_comm_dtype` (and accumulated in `main_grads_dtype`) as usual. + - `grad_comm_dtype` controls the data-type for gradient communications when reducing gradients. Lower precision `grad_comm_dtype` improves (communication) performance, but may increase memory utilization or sacrifice gradient precision in certain cases. + - Defaults to `None`, in which the `main_grads_dtype` data-type will be utilized. No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`. + - If using HSDP (either DP-Replicate or DP-Outer in `outer_dp_sharding_strategy`), `no_shard`, or `optim`, allocating `dtype`-custom gradient communication buffers may increase per-unit memory overhead, so users should consider the performance-memory trade-off when using this feature. + - If using NCCL user buffer registration `v2.27+`, gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced in `grad_comm_dtype` (and accumulated in `main_grads_dtype`) as usual. - `overlap_grad_reduce` and `overlap_param_gather` will overlap gradient [`reduce-scatter`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#reducescatter) and parameter [`all-gather`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather) group communications with backward and forward compute with asynchronous calls and pre-fetching. (In the case of `no_shard`, parameters are not gathered but gradient [`all-reduce`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce) is overlapped.) - Both default to `True`. - `sync_model_each_microbatch` will trigger a `wait` (`MegatronFSDP.finish_grad_sync()`) on gradient reduction, parameter de-allocation, and optimizer parameter / gradient installation (in preparation for `optimizer.step()`) after every forward-backward pass. When using HSDP, parameters and gradients will be all-gathered and reduced respectively on the "outer" DP group each training step instead of each optimization cycle. This behavior is desirable for a transparent and user-friendly sharded training loop where post-backward transformations on the gradient and a clean compute / memory state are necessary within and between training iterations, but damages performance in situations where optimization is delayed (e.g. gradient accumulation) when the communications of the previous training iteration can be overlapped with the compute of the next training iteration. Will also override `is_last_microbatch` / `microbatch_count` logic in `MegatronFSDP`. @@ -362,7 +171,7 @@ Megatron-FSDP natively supports mixed-precision activations and parameter shardi - Within the [`transformer_engine.pytorch.autocast(recipe: transformer_engine.common.recipe.Recipe)`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.autocast) context, model activations are converted based on the recipe. - Within the [`transformer_engine.pytorch.quantized_model_init(recipe: transformer_engine.common.recipe.Recipe)`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.quantized_model_init) context, TransformerEngine native modules (e.g. [`transformer_engine.pytorch.TransformerLayer`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.TransformerLayer)) have their parameters converted based on the recipe. - - Requires FP8 model activations, i.e. `transformer_engine.pytorch.autocast`. + - Requires quantized model activations, i.e. `transformer_engine.pytorch.autocast`. ```python # FP8 Recipe @@ -397,4 +206,4 @@ with transformer_engine.pytorch.autocast(recipe=fp8_recipe): mfsdp_model(x).sum().backward() ``` -ℹ️ `TransformerEngine` kernels have a fair bit of configuration constraints when using FP8-quantized parameters, such as using fused QKV parameters or defining activations and parameters with shapes compatible to FP8 CuBLAS kernels on supported hardware from NVIDIA. To properly initialize `TransformerLayer`, you can refer to the toy model used in our FP8 unit tests: `Megatron-LM/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py::TestMegatronFsdpFullyShard::test_fully_shard_te_quantized`. \ No newline at end of file +ℹ️ `TransformerEngine` kernels have various constraints related to quantized Tensors, such as using fused QKV parameters or defining activations and parameters with shapes compatible to CuBLAS kernels on supported hardware from NVIDIA. To properly initialize `TransformerLayer`, you can refer to the example model used in our unit tests: `Megatron-LM/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py::TestMegatronFsdpFullyShard::test_fully_shard_te_quantized`. \ No newline at end of file