BLOOM 176B Training Technology Overview

TL;DR

BLOOM’s 176 B parameter model was trained on 384 × 80 GB NVIDIA A100 GPUs over 3.5 months using a custom Megatron‑DeepSpeed stack that combines ZeRO data parallelism, tensor parallelism, and pipeline parallelism with BF16 mixed‑precision, enabling a 1 M GPU‑hour run on the French Jean Zay supercomputer.

Project Summary

  • Hardware: 384 A100 80 GB GPUs (48 nodes, 8 GPUs per node) with AMD EPYC 7543 CPUs, 512 GB RAM per node, Omni‑Path interconnect, dedicated NCCL subnet, GPFS storage.
  • Software: Megatron‑DeepSpeed (a fork merging Microsoft DeepSpeed and NVIDIA Megatron‑LM).
  • Architecture: GPT‑3‑style transformer with added improvements (ALiBi positional encoding, embedding LayerNorm).
  • Dataset: 350 B tokens from 46 languages (≈1.5 TB cleaned text), vocab size 250 680.
  • Training duration: ~3.5 months (≈1 M compute hours).

Key Contributors

The effort relied on six main groups:

  1. Hugging Face BigScience team (in‑house engineers and funding).
  2. Microsoft DeepSpeed team (provided DeepSpeed library and integration help).
  3. NVIDIA Megatron‑LM team (framework and advice).
  4. IDRIS/GENCI staff managing the Jean Zay supercomputer (donated compute).
  5. PyTorch core team (bug fixes and usability improvements).
  6. Volunteers in the BigScience Engineering workgroup.

Notable individuals include Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari, and Rémi Lacroix.

Megatron‑DeepSpeed Stack

Component Provided by DeepSpeed Provided by Megatron‑LM
ZeRO Data Parallelism ✅
Tensor Parallelism ✅
Pipeline Parallelism ✅
BF16 Optimizer ✅
Fused CUDA Kernels ✅
DataLoader ✅

The stack implements 3‑D parallelism:

  • Data Parallelism (DP) replicates the model across groups of GPUs, each processing a data slice.
  • Tensor Parallelism (TP) shards individual tensors across GPUs, reducing per‑GPU memory.
  • Pipeline Parallelism (PP) splits the model vertically across GPUs, processing micro‑batches in a pipeline to avoid idle GPUs.
  • ZeRO further shards optimizer states, gradients, and optionally weights to fit large models.

Parallelism Details

ZeRO Data Parallelism

Instead of full model replication, each GPU holds only a partition of parameters, gradients, and optimizer states, reconstructing full tensors on‑the‑fly as needed. This reduces memory overhead dramatically.

Tensor Parallelism

Weight matrices are split column‑wise across GPUs; each GPU computes its slice of the matrix multiplication and applies activation locally. This requires a high‑speed interconnect; in BLOOM the TP degree was limited to 4 per node (one GPU per tensor slice).

Pipeline Parallelism

Model layers are divided into stages; micro‑batches flow through the stages, keeping all GPUs busy. The "chunks" (or GAS) hyper‑parameter controls the number of micro‑batches, balancing pipeline bubbles against micro‑batch size. BLOOM used 72 pipeline stages (including two embedding stages) to balance memory across GPUs.

Combined DP + PP + TP (3‑D Parallelism)

The final training configuration employed DP for data distribution, TP for tensor sharding, and PP for layer distribution, achieving efficient utilization of the 384‑GPU cluster.

BF16 Optimizer

Training in FP16 caused divergence on earlier experiments (e.g., a 104 B model). BLOOM switched to BF16 mixed‑precision using a custom BF16Optimizer that:

  • Keeps the exponent range of FP32, avoiding overflow.
  • Performs all accumulations in FP32.
  • Accumulates gradients in FP32 across pipeline micro‑batches. The optimizer enabled stable loss curves for the 176 B model.

Fused CUDA Kernels

To minimize GPU idle time, custom fused kernels from Megatron‑LM were used for:

  • LayerNorm
  • Combined scaling, masking, and softmax
  • Bias‑added GeLU (via PyTorch JIT) These kernels reduce memory traffic by keeping intermediate results in registers.

Dataset Handling

The training data pipeline:

  • Tokenized 1.5 TB of cleaned multilingual text into 350 B tokens.
  • Created per‑sample indices for a fixed sequence length of 2048.
  • Shuffled epoch‑wise ordering to ensure uniform exposure.
  • Saved indices to disk to avoid recomputation on restarts. Multiple datasets were blended with configurable weights.

Architectural Tweaks

  • Embedding LayerNorm: Adding a LayerNorm after the first embedding stabilized training, inspired by the StableEmbedding implementation in bitsandbytes.
  • ALiBi Positional Encoding: Replaced absolute positional embeddings with linear bias attention (ALiBi), allowing extrapolation to sequences longer than the training length (2048).

Engineering Challenges

  • Hardware failures: 1–2 GPU failures per week; checkpoints saved every 3 h limited lost work to ~1.5 h per failure.
  • Software bugs: Required CUDA_LAUNCH_BLOCKING=1, optimizer group splitting, and custom SLURM kill‑switch for multi‑user job control.
  • Downtime: 5–10 h interruptions from deadlocks, disk space exhaustion, and other issues; overall training stayed within the planned 3.5‑month window.
  • On‑call coordination: Distributed across Europe and West‑Coast Canada, enabling 24/7 monitoring without a dedicated pager.

Conclusions

The most intensive phase was the two months of preparation, including the late‑stage development of the BF16 optimizer and debugging of large‑scale parallelism. Once the stack was stable, the 176 B model trained smoothly, demonstrating that open‑source teams can train state‑of‑the‑art multilingual models when provided with sufficient compute and collaborative tooling.

Resources

Key papers

  • Megatron‑LM: Efficient Large‑Scale Language Model Training on GPU Clusters (arXiv:2104.04473)
  • DeepSpeed ZeRO: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv:1910.02054)
  • ALiBi: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (arXiv:2108.12409)

Sources