Scaling LLM Pretraining with DeepSpeed and Megatron-LM: 3D Parallelism, ZeRO-3, and Multi-Node Systems

Distributed Memory Bottlenecks in Frontier Model Pretraining

Training large foundation models (13B to 70B+ parameters) exceeds the physical High Bandwidth Memory (HBM) of any single accelerator:

Model Parameters ( $Φ$ ): Stored in 16-bit precision ( $BF16 / FP16$ ), requiring $2Φ$ bytes ( $140 GB$ for a 70B model).
Gradients: Stored in 16-bit precision, requiring $2Φ$ bytes ( $140 GB$ ).
AdamW Optimizer States: Consists of 32-bit master weights ( $4Φ$ bytes), first-order momentum ( $4Φ$ bytes), and second-order uncentered variance ( $4Φ$ bytes), totaling $12Φ$ bytes ( $840 GB$ ).
Activations & Temporary Buffers: Scales quadratically with sequence length $S$ and batch size $B$ , consuming hundreds of additional gigabytes unless managed via activation recomputation.

Total static state requirements alone scale to $16Φ$ bytes ( $1.12 TB$ for 70B), far exceeding an 80 GB NVIDIA H100 GPU. Scaling requires orthogonal decomposition across 3D Parallelism (Tensor, Pipeline, Data/ZeRO) and high-throughput inter-node collective communication.

Architectural Comparison

Dimension	Pure DDP (PyTorch)	DeepSpeed ZeRO-3	Megatron-LM (3D Parallelism)
Memory Sharding Target	Replicated Parameters & Gradients	Parameters, Gradients, Optimizer States ( $\frac{16Φ}{N _{DP}}$ )	Intra-Layer Matmul ( $TP$ ) + Layer Staging ( $PP$ )
Communication Pattern	All-Reduce on Gradients	All-Gather (Forward/Backward) + Reduce-Scatter	All-Reduce in Intra-Node NVLink ( $TP$ ) + P2P ( $PP$ )
Scale Limit	$\leq 7 B$ per 80GB GPU	Up to $\approx 30 B$ single-node, $70 B$ multi-node	$70 B \sim 1 T +$ across multi-node clusters
Model Code Adaptation	Standard PyTorch Module	Standard Hugging Face Trainer	Sharding-Aware Tensor & Sequence Custom Blocks
FP8 Integration	External Plugins	ZeRO++ FP8	Native TransformerEngine Fused Kernels

Mathematical Formulation

flowchart TD MODEL["Foundation Model 70B (80 Layers, Hidden=8192)"] subgraph NODE["Intra-Node NVLink Domain (900 GB/s)"] TP["Tensor Parallelism (TP = 8)\nColumnParallelLinear -> RowParallelLinear\n+ Sequence Parallelism (SP)"] end subgraph INTER["Inter-Node InfiniBand Domain (400 Gb/s NDR)"] PP["Pipeline Parallelism (PP = 2)\n1F1B Schedule across 2 Stage-Groups"] DP["Data Parallelism (DP = 4)\nZeRO-1 Distributed Optimizer Sharding"] end MODEL --> TP TP --> PP PP --> DP DP --> ALLREDUCE["All-Reduce / Reduce-Scatter Synchronizations"]

Figure 1: 3D Parallelism topology for a 70B parameter model across 8 nodes (64× H100 GPUs). Tensor and sequence parallelism operate within intra-node NVLink domains ( $T P = 8$ ), while pipeline ( $P P = 2$ ) and distributed optimizer data parallelism ( $D P = 4$ ) span inter-node InfiniBand networks.

1. ZeRO Memory Partitioning Hierarchy

The Zero Redundancy Optimizer partitions training state across $N_{DP}$ data-parallel ranks:

$Memory_{ZeRO-1} = 2Φ + 2Φ + \frac{12Φ}{N _{DP}} = 4Φ + \frac{12Φ}{N _{DP}}$

$Memory_{ZeRO-2} = 2Φ + \frac{2Φ + 12Φ}{N _{DP}} = 2Φ + \frac{14Φ}{N _{DP}}$

$Memory_{ZeRO-3} = \frac{2Φ + 2Φ + 12Φ}{N _{DP}} = \frac{16Φ}{N _{DP}}$

Under ZeRO-3 with $N_{DP} = 64$ , the static state per GPU drops from $1.12 TB$ to $17.5 GB$ , leaving ample headroom for activations.

2. Megatron Tensor & Sequence Parallelism

For Multi-Head Attention and Multi-Layer Perceptrons (MLP), Megatron splits weight matrices across the column and row dimensions:

Column-Parallel MLP: $Y_{1} = GeLU (X W_{1, i})$ , where $W_{1} = [W_{1, 1}, \dots, W_{1, k}]$ .
Row-Parallel MLP: $Y_{2} = \sum_{i = 1}^{k} Y_{1, i} W_{2, i}$ , followed by an All-Reduce operation.

Sequence Parallelism (SP) shards non-tensor-parallel operations (LayerNorm and Dropout) along the sequence length dimension $S$ , reducing activation memory from $O (B \cdot S \cdot h)$ to $O (\frac{B \cdot S \cdot h}{N _{TP}})$ .

3. Model FLOPs Utilization (MFU) Calculation

The theoretical computation required for one forward and backward pass per token is $6Φ$ . Given global batch size $B_{global}$ , sequence length $S$ , step time $Δ t$ , and theoretical GPU peak throughput $T_{peak}$ :

$MFU = \frac{6 \cdot Φ \cdot B _{global} \cdot S}{Δ t \cdot N _{GPUs} \cdot T _{peak}}$

For an H100 SXM ( $T_{peak} = 1979 TFLOPs$ in BF16 dense), well-optimized 3D parallelism targets $MFU \geq 55% \sim 65%$ .

Implementation: Multi-Node Pretraining Engine

Environment Setup

FROM nvcr.io/nvidia/pytorch:25.04-py3

ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_NO_CACHE_DIR=1

# Install InfiniBand, RDMA, and NCCL system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
        libibverbs1 libibverbs-dev librdmacm1 librdmacm-dev \
        libnccl2 libnccl-dev infiniband-diags rdma-core \
        openssh-client openssh-server \
    && rm -rf /var/lib/apt/lists/*

# Install core training frameworks
RUN pip install --no-cache-dir \
        "transformers>=4.46.0" "datasets>=3.1.0" "accelerate>=1.1.0" \
        "flash-attn==2.6.3" "transformer-engine[pytorch]==1.10.0" \
        "megatron-core" "deepspeed>=0.16.0"

Step 1: Pre-tokenization to Megatron Binary Format

Offline memory-mapped tokenization eliminates runtime CPU bottlenecks and ensures deterministic batch sampling across distributed ranks:

# Tokenize raw JSONL text into memory-mapped .bin / .idx files
python Megatron-LM/tools/preprocess_data.py \
    --input /datasets/raw/pretrain_corpus.jsonl \
    --output-prefix /datasets/tokenized/corpus_llama3 \
    --tokenizer-type HuggingFaceTokenizer \
    --tokenizer-model meta-llama/Meta-Llama-3-70B \
    --append-eod \
    --workers 64 \
    --log-interval 10000

Step 2: DeepSpeed ZeRO-3 Single/Multi-Node Configuration

{
  "train_batch_size": 2048,
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 32,
  "steps_per_print": 10,

  "bf16": {
    "enabled": true
  },
  "fp16": {
    "enabled": false
  },

  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true,
    "offload_optimizer": {
      "device": "none"
    },
    "offload_param": {
      "device": "none"
    }
  },

  "activation_checkpointing": {
    "partition_activations": false,
    "contiguous_memory_optimization": true,
    "profile": false
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 1.5e-4,
      "betas": [0.9, 0.95],
      "eps": 1e-8,
      "weight_decay": 0.1
    }
  },

  "scheduler": {
    "type": "WarmupCosineLR",
    "params": {
      "warmup_min_ratio": 0.0,
      "warmup_num_steps": 2000,
      "total_num_steps": 100000,
      "cos_min_ratio": 0.1
    }
  },

  "gradient_clipping": 1.0,
  "wall_clock_breakdown": false
}

Step 3: Production 8-Node SLURM Launch Script (64× H100 GPUs)

#!/bin/bash
#SBATCH --job-name=megatron_70b_pretrain
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=12
#SBATCH --mem=0
#SBATCH --exclusive
#SBATCH --output=/logs/slurm-%j.out
#SBATCH --error=/logs/slurm-%j.err
#SBATCH --gres-flags=enforce-binding

# Network & NCCL topology tuning
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=mlx5
export NCCL_SOCKET_IFNAME=eth0
export NCCL_NET_GDR_LEVEL=5
export CUDA_DEVICE_MAX_CONNECTIONS=1
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# Resolve Master Address to numeric IP to prevent DNS rendezvous timeouts
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29500

DATA_PATH=/datasets/tokenized/corpus_llama3_text_document
CHECKPOINT_PATH=/checkpoints/llama3_70b_pretrain
TOKENIZER_MODEL=/workspace/tokenizers/llama3

# 3D Parallelism execution: TP=8 (Intra-Node NVLink), PP=2 (Inter-Node IB), DP=4 (Distributed Optimizer)
srun --export=ALL torchrun \
    --nproc-per-node=8 \
    --nnodes=8 \
    --node-rank=$SLURM_PROCID \
    --master-addr=$MASTER_ADDR \
    --master-port=$MASTER_PORT \
    pretrain_gpt.py \
    --num-layers 80 \
    --hidden-size 8192 \
    --num-attention-heads 64 \
    --ffn-hidden-size 28672 \
    --seq-length 8192 \
    --max-position-embeddings 8192 \
    --position-embedding-type rope \
    --rotary-base 500000 \
    --norm-rms-norm \
    --swiglu \
    --untied-embeddings-and-output-weights \
    --tensor-model-parallel-size 8 \
    --pipeline-model-parallel-size 2 \
    --sequence-parallel \
    --recompute-activations \
    --recompute-granularity selective \
    --recompute-modules "core_attn,mlp" \
    --use-distributed-optimizer \
    --accumulate-allreduce-grads-in-fp32 \
    --transformer-impl transformer_engine \
    --bf16 \
    --te-ln-scale-bf16 \
    --lr 1.5e-4 \
    --min-lr 1.5e-5 \
    --lr-decay-style cosine \
    --lr-warmup-iters 2000 \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --global-batch-size 4096 \
    --micro-batch-size 1 \
    --train-iters 1500000 \
    --data-path $DATA_PATH \
    --tokenizer-type HuggingFaceTokenizer \
    --tokenizer-model $TOKENIZER_MODEL \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH \
    --save-interval 500 \
    --eval-interval 1000

Empirical Benchmark Evaluation

We evaluated distributed scaling throughput and Model FLOPs Utilization (MFU) across multi-node H100 SXM5 clusters:

Model Scale	Cluster Topology	Parallelism Strategy	Throughput (TFLOPs/GPU)	Effective MFU (%)	Tokens / Day
13B (Llama-3 Architecture)	1 Node (8× H100)	$DP = 8, ZeRO-3$	1,650	66.2%	85.0 Billion
34B (CodeLlama Scale)	2 Nodes (16× H100)	$TP = 4, DP = 4, ZeRO-1$	1,450	58.2%	38.0 Billion
70B (Llama-3 Architecture)	8 Nodes (64× H100)	$TP = 8, PP = 2, DP = 4$	1,310	52.6%	22.5 Billion
70B (FP8 TransformerEngine)	8 Nodes (64× H100)	$TP = 8, PP = 2, DP = 4$	1,700	68.3%	29.2 Billion

Troubleshooting Common Cluster Faults

1. Silent Inter-Node NCCL Hangs During Rendezvous

Symptom: torchrun halts indefinitely at Process group initialized.
Remedy: Ensure MASTER_ADDR is resolved to an explicit IP address rather than a DNS hostname, and verify network connectivity on port 29500 using nc -zv $MASTER_ADDR 29500.

2. Mid-Run Loss Spikes to NaN

Symptom: Loss explodes to inf or nan after tens of thousands of stable steps.
Remedy: Enforce FP32 master gradient accumulation via --accumulate-allreduce-grads-in-fp32, activate gradient norm clipping (--clip-grad 1.0), and inspect tokenized datasets for zero-length sequences.

3. Asymmetric Cross-Node Bandwidth Drops

Symptom: Inter-node All-Reduce latency increases from 50 ms to $> 200 ms$ .
Remedy: Bind each GPU process to its corresponding NUMA domain via --gres-flags=enforce-binding in SLURM, and verify InfiniBand HCA status with ibstat.

Subscribe to Unlock the Rest

This section is exclusive to active subscribers. Support our work and unlock this article immediately.

Subscribe Now

Unlock this post and get unlimited access to all premium articles.

References

Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053.
Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20.
Korthikanti, V. A., et al. (2023). Reducing Activation Recomputation in Large Transformer Models (Sequence Parallelism). MLSys.
NVIDIA. (2024). TransformerEngine: A Library for Accelerating Transformer Models on Hopper & Blackwell Architecture.

Overview