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

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). The exact combination of tools depends on the tested software versions and the model architecture.

• Megatron-LM / Megatron Core: 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. It also provides native state-sharding options like the distributed optimizer.
• DeepSpeed ZeRO: 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 can 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, the foundation is often set using the official NVIDIA NGC PyTorch container.

# Example 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. The cluster size dictates the split based on the formula: Total GPUs = Tensor Parallelism (TP) × Pipeline Parallelism (PP) × Data Parallelism (DP).

• Tensor Parallelism (TP): Shards the math 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 may be appropriate for some configurations but 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 can reduce the frequency of inter-node communication compared with some tensor-parallel configurations, but its efficiency depends on activation size, pipeline scheduling, 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 additionally shards gradients, and Stage 3 also shards model parameters.

4. DeepSpeed ZeRO-2 vs ZeRO-3

When evaluating DeepSpeed Stage 2 vs Stage 3, the primary consideration is memory pressure versus communication overhead.

If a model's parameters do not fit onto a single GPU, ZeRO-3 dynamically shards and gathers parameters during forward/backward passes. However, ZeRO-2 generally retains higher throughput due to lower communication overhead. Benchmark both configurations against your specific memory constraints, batch size, and network topology to find the optimal trade-off.

5. Container and NCCL Networking Configuration

When running distributed workloads, optimizing the NCCL environment variables ensures your collective operations leverage high-speed RDMA fabrics.

# Enable InfiniBand/RoCEv2 transport
export NCCL_IB_DISABLE=0

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

6. Multi-Node Training with Slurm

A production-grade #SBATCH script initializes the distributed environment properly.

Topology Math Validation: In a 4-node, 32-GPU setup with TP=8 and PP=2, Data Parallelism calculates as: 32 / (8 × 2) = 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. Rank 0 aggregating the entire State Dict creates massive memory and I/O bottlenecks. Use PyTorch Distributed Checkpoint (DCP) when supported by your framework and checkpoint format to allow ranks to write their shards to a high-throughput Parallel File System.