# Data Parallelism in MegaDLMs: DDP vs FSDP2 Implementation Guide

> Explore Data Parallelism in MegaDLMs with our DDP vs FSDP2 implementation guide. Learn to optimize distributed training for standard or memory-constrained extreme-scale models.

- Repository: [Jinjie Ni/megadlms](https://github.com/jinjieni/megadlms)
- Tags: how-to-guide
- Published: 2026-03-04

---

**MegaDLMs supports two data parallelism strategies: DistributedDataParallel (DDP) for standard distributed training and Torch Fully Sharded Data Parallel (FSDP2) for memory-constrained extreme-scale models.**

Data parallelism in the `jinjieni/megadlms` repository determines how model parameters and gradients are distributed across GPUs during training. The framework provides distinct implementations for classic replication-based training and modern parameter-sharding approaches, each with specific compatibility constraints and memory trade-offs.

## DistributedDataParallel (DDP): The Default Strategy

The classic **DistributedDataParallel** implementation serves as the default data parallelism mechanism in MegaDLMs. Located in [`megatron/core/distributed/distributed_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/distributed/distributed_data_parallel.py), this class implements the standard NCCL-based all-reduce pattern for gradient synchronization.

### Implementation Details

DDP maintains a **full replica of model parameters** on every data-parallel rank. During the backward pass, gradients are accumulated into buckets and asynchronously all-reduced across the process group. The `_allocate_buffers_for_parameters` method handles the gradient buffer allocation and bucketing logic to overlap communication with computation.

### When to Use DDP

- Training configurations that require **pipeline parallelism** or **expert parallelism** (MoE)
- Scenarios using the **distributed optimizer** or **gradient accumulation fusion**
- Mixed precision training with **fp16** (FSDP2 does not support fp16)
- Standard large-scale pre-training where per-GPU memory is sufficient to hold the full model replica

## Torch Fully Sharded Data Parallel (FSDP2): Sharded Parameters

**Torch FSDP2** provides a memory-efficient alternative that shards model parameters across data-parallel ranks. Implemented in [`megatron/core/distributed/torch_fully_sharded_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/distributed/torch_fully_sharded_data_parallel.py), this strategy uses PyTorch 2.4+'s composable FSDP API to reduce per-GPU memory footprint dramatically.

### Memory Architecture

Unlike DDP, FSDP2 **shards parameters, gradients, and optimizer states** across the data-parallel group. Each rank maintains only a subset of the model weights, with parameters all-gathered just-in-time during the forward and backward passes. The wrapper applies `fully_shard` to specific sub-modules including `TransformerLayer`, `LanguageModelEmbedding`, `RotaryEmbedding`, and the final `ColumnParallelLinear` layers.

### Compatibility Constraints

FSDP2 imposes strict limitations enforced in [`megatron/training/arguments.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/arguments.py) (lines 375-392):

- **Incompatible with pipeline parallelism** (`--pipeline-model-parallel-size` must be 1)
- **Incompatible with expert parallelism** (MoE)
- **Incompatible with the distributed optimizer**
- **Incompatible with gradient accumulation fusion**
- **Incompatible with fp16** (requires bf16 or fp32)

### Activation Requirements

To enable FSDP2, you must explicitly pass `--use-torch-fsdp2` on the command line. The framework requires PyTorch 2.4 or newer, verified at runtime in [`torch_fully_sharded_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/torch_fully_sharded_data_parallel.py) (line 56). Additional required flags include `--ckpt-format=torch_dist` and `--untie-embeddings-and-output-weights`.

## How MegaDLMs Selects the Data Parallel Wrapper

The training initialization logic in [`megatron/training/utils.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/utils.py) determines which wrapper class to instantiate based on the command-line configuration.

### Wrapper Selection Logic

The `ALL_MODULE_WRAPPER_CLASSNAMES` tuple dynamically includes `TorchFullyShardedDataParallel` when available:

```python

# megatron/training/utils.py

try:
    from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP
    ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, torch_FSDP, Float16Module)
except ImportError:
    ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, Float16Module)

```

When `--use-torch-fsdp2` is detected, the training scripts ([`training.py`](https://github.com/jinjieni/megadlms/blob/main/training.py) and [`training_nan_skipping.py`](https://github.com/jinjieni/megadlms/blob/main/training_nan_skipping.py)) set `DP = torch_FSDP` instead of the default `DDP`.

## Configuration Examples

### Standard DDP Training

Run distributed training with the default DDP wrapper, compatible with pipeline and tensor parallelism:

```bash
torchrun --nnodes=1 --nproc_per_node=8 \
    pretrain_difflm.py \
    --model-type=gpt \
    --micro-batch-size=4 \
    --global-batch-size=32 \
    --tensor-model-parallel-size=2 \
    --pipeline-model-parallel-size=1 \
    --data-parallel-size=4

```

*Omitting `--use-torch-fsdp2` automatically selects `DistributedDataParallel`.*

### FSDP2 Training for Memory-Constrained Regimes

Enable parameter sharding with FSDP2 (requires PyTorch 2.4+):

```bash
torchrun --nnodes=1 --nproc_per_node=8 \
    pretrain_difflm.py \
    --model-type=gpt \
    --micro-batch-size=4 \
    --global-batch-size=32 \
    --tensor-model-parallel-size=1 \
    --pipeline-model-parallel-size=1 \
    --data-parallel-size=8 \
    --use-torch-fsdp2 \
    --ckpt-format=torch_dist \
    --untie-embeddings-and-output-weights

```

**Critical constraints:** Tensor and pipeline parallelism must be disabled (`size=1`), and you must use `--ckpt-format=torch_dist` with untied embeddings.

### Debug Helper: Verify Active Wrapper

Inspect which data parallelism class is currently wrapping your model:

```python
from megatron.core import DistributedDataParallel as DDP
from megatron.training.utils import ALL_MODULE_WRAPPER_CLASSNAMES

print("Available wrappers:", ALL_MODULE_WRAPPER_CLASSNAMES)

# After model construction:

print("Active wrapper:", type(model).__name__)

```

This outputs either `DistributedDataParallel` or `TorchFullyShardedDataParallel` depending on whether `--use-torch-fsdp2` was passed.

## Summary

- **DistributedDataParallel (DDP)** is the default data parallelism strategy in MegaDLMs, maintaining full model replicas on each GPU and supporting all parallelism combinations including pipeline and expert parallelism.
- **Torch FSDP2** provides memory-efficient parameter sharding for extreme-scale models but requires PyTorch 2.4+ and disables pipeline parallelism, expert parallelism, and fp16 training.
- The wrapper selection occurs in [`megatron/training/utils.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/utils.py) based on the `--use-torch-fsdp2` command-line flag, with validation logic enforced in [`megatron/training/arguments.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/arguments.py).
- DDP is implemented in [`megatron/core/distributed/distributed_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/distributed/distributed_data_parallel.py), while FSDP2 resides in [`megatron/core/distributed/torch_fully_sharded_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/core/distributed/torch_fully_sharded_data_parallel.py).

## Frequently Asked Questions

### What is the minimum PyTorch version required for FSDP2 in MegaDLMs?

MegaDLMs requires **PyTorch 2.4 or newer** to use the Torch Fully Sharded Data Parallel implementation. This is verified at runtime in [`torch_fully_sharded_data_parallel.py`](https://github.com/jinjieni/megadlms/blob/main/torch_fully_sharded_data_parallel.py) (line 56). Attempting to use FSDP2 with earlier PyTorch versions will raise an import or version assertion error.

### Can I use pipeline parallelism with FSDP2?

No, **FSDP2 is incompatible with pipeline parallelism** in MegaDLMs. The argument parser in [`megatron/training/arguments.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/arguments.py) (lines 375-392) enforces that `--pipeline-model-parallel-size` must be 1 when `--use-torch-fsdp2` is enabled. This restriction exists because FSDP2's parameter sharding conflicts with the layer partitioning and activation management required by pipeline parallelism.

### Why does DDP remain the default instead of FSDP2?

DDP remains the default because it offers **broader compatibility** with MegaDLMs' advanced parallelism features. Unlike FSDP2, DDP supports pipeline parallelism, expert parallelism (MoE), the distributed optimizer, and fp16 training. DDP is also more mature within the Megatron codebase, making it the safer choice for standard large-scale training jobs where per-GPU memory is sufficient to hold full model replicas.

### How do I verify which data parallelism wrapper is active during training?

You can inspect the model's wrapper class at runtime using the `ALL_MODULE_WRAPPER_CLASSNAMES` tuple from [`megatron/training/utils.py`](https://github.com/jinjieni/megadlms/blob/main/megatron/training/utils.py). After model construction, check `type(model).__name__` to see whether it returns `DistributedDataParallel` or `TorchFullyShardedDataParallel`. Alternatively, examine the training logs for the wrapper initialization messages, or check whether the `--use-torch-fsdp2` flag was passed in your launch command.