Megatron-DeepSpeed 3D Parallelism: Scaling 100B+ LLM Training

Originally published at Servermo

Scaling models beyond 100B parameters with conventional DistributedDataParallel (DDP) becomes increasingly difficult because each data-parallel replica must maintain substantial model and training state. For perspective, a 400B parameter model requires roughly 800 GB of VRAM just to hold FP16 weights. When factoring in gradients, optimizer states, and forward activations, the total memory requirement can balloon to several terabytes.

To overcome this memory wall, enterprise AI teams rely on a composite architecture, deploying the Megatron-DeepSpeed fusion framework.

This blueprint provides Site Reliability Engineers (SREs) and MLOps teams with a production-grade guide to configuring 3D Parallelism (TP + PP + DP), evaluating ZeRO optimizations, and diagnosing NCCL RoCEv2 networks on GPU clusters.

Reference Hardware and Software Configuration

To ensure reproducibility and prevent version-mismatch failures, the configurations in this guide align with the following production stack:

• GPU Topology: 4 Nodes × 8x NVIDIA H100 (80GB) SXM5 with NVLink 4.0
• Network Fabric: 400 Gbps RoCEv2 (ConnectX-7 HCAs)
• Base Container: nvcr.io/nvidia/pytorch:24.08-py3
• Software Engine: PyTorch 2.5.0a0, CUDA 12.6, NCCL 2.22.3
• Framework: Megatron Core v0.15+ integrated with DeepSpeed v0.14+

1. Megatron-LM vs Megatron-DeepSpeed: The Compute & Memory Synergy

When scaling models beyond 100B parameters, engineering teams face two distinct bottlenecks: Compute (interconnect bandwidth) and Memory (VRAM capacity):

• Megatron-LM / Megatron Core (Compute Optimization): Specializes in Tensor Parallelism (TP) and Pipeline Parallelism (PP). It provides the building blocks to slice layers across GPUs to maximize raw compute speed over high-bandwidth interconnects (NVLink). It also provides native state-sharding options like the distributed optimizer.
• DeepSpeed ZeRO (Memory Sharding): A robust memory-sharding engine. It eliminates redundancy by partitioning optimizer states, gradients, and parameters across stages, specifically targeting Out of Memory (OOM) constraints.

By integrating them into a Megatron-DeepSpeed configuration, SREs leverage Megatron's optimized intra-node compute execution alongside DeepSpeed's aggressive memory distribution across Data Parallel (DP) replicas.

2. Enterprise Container & Environment Setup

Before launching multi-node jobs via Slurm or MPI, establish the baseline environment using the official NVIDIA NGC PyTorch container.

⚠️ Alert: Docker IPC and Network Considerations
Container IPC and networking settings heavily affect NCCL initialization and communication performance. While omitting --ipc=host does not guarantee an automatic fallback to Ethernet, it can cause shared-memory issues that throttle high-speed interconnects. Validate shared-memory availability, GPU peer access, RDMA device visibility, and the selected NCCL transport inside the container.

# SRE-approved container launch command
docker run --gpus all --ipc=host --network=host --ulimit memlock=-1 \
  -v /mnt/shared_pfs:/workspace \
  -it nvcr.io/nvidia/pytorch:24.08-py3 /bin/bash

3. TP, PP and DP: Topology and Sizing

Megatron-LM (compute) and DeepSpeed (memory) must be mathematically aligned to your hardware topology. Cluster sizing follows the fundamental formula:

$$\text{Total GPUs} = \text{Tensor Parallelism (TP)} \times \text{Pipeline Parallelism (PP)} \times \text{Data Parallelism (DP)}$$

• Tensor Parallelism (TP): Shards matrix multiplications inside a layer (e.g., attention heads) across GPUs. Keep tensor-parallel groups within a high-bandwidth intra-node GPU domain when practical (e.g., TP=8 over NVLink). Extending TP across nodes requires careful evaluation of interconnect bandwidth and collective latency.
• Pipeline Parallelism (PP): Divides model layers into sequential stages. Adjacent stages exchange activations during the forward pass and gradients during the backward pass. PP reduces the frequency of inter-node communication compared to cross-node TP, but its efficiency depends on activation size, pipeline scheduling (e.g., 1F1B schedule), micro-batch count, and network latency.
• ZeRO Data Parallelism (DP): DeepSpeed ZeRO partitions training states across the data-parallel group:
  • Stage 1: Shards optimizer states.
  • Stage 2: Shards optimizer states and gradients.
  • Stage 3: Shards optimizer states, gradients, and model parameters.

4. DeepSpeed ZeRO-2 vs ZeRO-3

When evaluating DeepSpeed Stage 2 versus Stage 3, the primary trade-off is memory pressure versus communication overhead:

• ZeRO-2: Shards optimizer states and gradients. Retains full model parameters on each worker during forward and backward passes. It offers higher training throughput due to lower collective communication overhead.
• ZeRO-3: Dynamically shards and gathers parameters on the fly during forward and backward passes. Necessary when full parameters cannot fit onto individual GPUs, but incurs additional communication overhead.

The CPU Offload Trade-Off

DeepSpeed ZeRO-Offload (CPU/NVMe) allows training large models on constrained hardware by offloading optimizer states to host RAM. However, transferring hundreds of gigabytes over PCIe buses introduces severe throughput penalties. While it reduces GPU VRAM requirements, it trades massive throughput for memory, substantially slowing down synchronous training. Benchmark this trade-off using Model FLOPs Utilization (MFU) and tokens/sec metrics.

5. Container and NCCL Networking Configuration

Optimizing NCCL environment variables ensures collective communication leverages high-speed RDMA fabrics:

# Enable InfiniBand/RoCEv2 transport
export NCCL_IB_DISABLE=0

# Pin the exact Host Channel Adapters (HCAs) with exact-match prefix
# (Verify available interfaces on your nodes via `ibv_devinfo`)
export NCCL_IB_HCA="=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1"

⚠️ Alert: GPUDirect RDMA Verification
NCCL_NET_GDR_LEVEL controls the GPU-to-NIC topology distance at which NCCL may use GPUDirect RDMA. Prefer topology-based values such as PIX, PXB, PHB, or SYS when explicit tuning is necessary. Leave automatic selection enabled unless measurements justify a change. Validate the transport using NCCL_DEBUG=INFO and check for GDR memory type: CUDA to confirm active GPUDirect RDMA.

6. Multi-Node Training with Slurm

A production-grade #SBATCH script initializes the distributed environment across worker nodes:

⚠️ Important Note: Slurm torchrun and Variable Escaping
When using srun to trigger torchrun, do not evaluate $SLURM_NODEID beforehand. By passing the torchrun command in single quotes ('...'), $SLURM_NODEID evaluates correctly inside the task environment on each specific worker node. Ensure srun launches exactly 1 task per node, allowing torchrun to handle the 8 local GPU worker processes.

Topology Math Validation: In a 4-node, 32-GPU setup with TP=8 and PP=2, Data Parallelism calculates as:
$$32 / (8 \times 2) = \text{DP } 2$$

#!/bin/bash
#SBATCH --job-name=megatron-deepspeed-train
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1   # 1 torchrun launcher per node
#SBATCH --gres=gpu:8
#SBATCH --exclusive

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=29500

srun --ntasks=4 --ntasks-per-node=1 bash -c 'torchrun \
  --nproc_per_node=8 \
  --nnodes=$SLURM_NNODES \
  --node_rank=$SLURM_NODEID \
  --master_addr=$MASTER_ADDR \
  --master_port=$MASTER_PORT \
  pretrain_gpt.py \
  --tensor-model-parallel-size 8 \
  --pipeline-model-parallel-size 2 \
  --deepspeed \
  --deepspeed_config ds_config.json'

7. Distributed Checkpointing and Recovery

Saving a 100B+ model state dumps terabytes of data simultaneously. Having Rank 0 aggregate the entire State Dict creates catastrophic memory and I/O bottlenecks. Use PyTorch Distributed Checkpoint (DCP) when supported by your framework and checkpoint format to allow all ranks to write their shards in parallel to a high-throughput Parallel File System (Lustre/GPFS).

⚠️ Alert: Asynchronous Validation
Distributed Checkpointing is not automatically asynchronous. If non-blocking checkpointing is required, integrate and benchmark the appropriate asynchronous API (e.g., torch.distributed.checkpoint.async_save). Validate staging overhead, storage throughput, checkpoint integrity, and restore behavior under realistic training conditions.

8. Bare Metal vs Cloud for Distributed Training

While public cloud environments offer instant elasticity, virtualization layers and SR-IOV network stacks can introduce overhead and jitter during multi-node all-reduce operations. Dedicated bare-metal infrastructure provides the uninterrupted RoCEv2 fabrics and non-blocking storage required for synchronous 3D parallelism workloads.

MLOps & Distributed Training FAQ

What is the difference between DeepSpeed ZeRO-3 and Megatron-LM Pipeline Parallelism?

ZeRO-3 is a memory-sharding technique that distributes parameters across GPUs and gathers them dynamically, increasing communication overhead but saving memory. Pipeline Parallelism (Megatron-LM) physically splits layers across GPUs sequentially, exchanging activations and gradients at boundaries. In production, they are combined: PP spans across nodes, while ZeRO optimizes memory within Data Parallel replicas.

How do I troubleshoot OOM errors during large-model initialization?

First, identify whether the failure occurs during model construction, checkpoint loading, optimizer initialization, or the first forward/backward pass. For compatible DeepSpeed ZeRO-3 workflows, wrapping model loading inside deepspeed.zero.Init() can reduce initialization memory pressure by partitioning parameters during construction.

Can I run Tensor Parallelism across multiple nodes?

While technically possible, extending Tensor Parallelism across nodes is rarely recommended for large models because TP requires an all-reduce operation per transformer block. This demands the massive bandwidth of intra-node interconnects (like NVLink). Pushing TP across an inter-node InfiniBand or Ethernet link typically causes severe pipeline stalling.

Why is Sequence Parallelism required when using Tensor Parallelism?

Tensor Parallelism (TP) shards weight matrices, but intermediate operations like LayerNorm and Dropout duplicate full sequence activations across all GPUs in the TP group, wasting critical VRAM. Sequence Parallelism (SP) shards these activations along the sequence dimension, utilizing All-Gather and Reduce-Scatter operations instead of standard All-Reduce. Enabling SP alongside TP (e.g., --sequence-parallel) reduces activation memory footprints by a factor equal to your TP degree.

How do I diagnose silent NCCL hangs during distributed training?

Silent hangs often occur due to asymmetric collective calls or an Out of Memory (OOM) failure on a single rank, which leaves all other ranks waiting indefinitely. Configure PyTorch-specific variables like TORCH_NCCL_ASYNC_ERROR_HANDLING=1 and TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=1800 to surface these silent timeouts as catchable exceptions. Deploy NCCL_DEBUG=WARN in production to capture hardware drops without flooding your console with per-collective initialization logs.

How much GPU VRAM is actually required to fine-tune a 70B LLM?

A 70B parameter model requires approximately 140 GB of VRAM strictly to load BF16 weights. Full fine-tuning—which must store weights, gradients, and FP32 AdamW optimizer states—demands between 860 GB and 1.2 TB of aggregate VRAM. This requires a cluster of 11 to 16 H100 (80GB) GPUs utilizing DeepSpeed ZeRO-3. Parameter-Efficient Fine-Tuning (PEFT) methods like QLoRA quantize the base model to 4-bit, dropping total VRAM to roughly 52 GB, which fits comfortably on a single 80GB GPU.

Can Gradient Checkpointing cause deadlocks in ZeRO-3 or FSDP?

Yes. Gradient checkpointing recomputes forward activations during the backward pass to conserve VRAM. If your distributed unit wrapping boundaries do not perfectly align with the checkpointed regions, it forces overlapping collective communication calls. One rank may call an all-gather for a forward recomputation while another rank executes a backward reduce-scatter for the same parameter, triggering an immediate deadlock. Always align your ZeRO-3 partitioning boundaries precisely with your gradient checkpointing wrappers.