bigscience-workshop/megatron-deepspeed

Ongoing research training transformer language models at scale, including: BERT & GPT-2

Structural Verdict

How Data Flows Through the System

Text data flows from indexed binary files through tokenization and batching to transformer models with distributed training

1. Data Loading — Load preprocessed text from memory-mapped indexed datasets (config: data_path, data_impl)
2. Tokenization — Convert text to token sequences using configured tokenizer (config: tokenizer_type, vocab_file)
3. Batching — Create training batches with attention masks and position embeddings (config: micro_batch_size, seq_length)
4. Model Forward — Process tokens through transformer layers with parallel computation (config: tensor_model_parallel_size, pipeline_model_parallel_size)
5. Loss Computation — Calculate language modeling loss with optional label smoothing (config: label_smoothing)
6. Backward Pass — Distributed gradient computation with DeepSpeed optimization (config: deepspeed_config, zero_stage)
7. Checkpointing — Periodic saving of model state and optimizer state (config: save_interval, save)

System Behavior

How the system actually operates at runtime — where data accumulates, what loops, what waits, and what controls what.

Data Pools

Feedback Loops

• Training Loop (training-loop, balancing) — Trigger: Training step completion. Action: Forward pass, loss computation, backward pass. Exit: Maximum iterations reached.
• Gradient Accumulation (convergence, reinforcing) — Trigger: Microbatch completion. Action: Accumulate gradients across microbatches. Exit: Global batch size reached.
• Learning Rate Schedule (convergence, balancing) — Trigger: Training step. Action: Adjust learning rate based on schedule. Exit: Training completion.

Delays & Async Processing

• Checkpoint Intervals (scheduled-job, ~save_interval steps) — Training pauses to save model state to disk
• Data Loading (async-processing) — Background data loading while training continues
• DeepSpeed Initialization (async-processing) — ZeRO stage 3 parameter distribution across devices

Control Points

• DeepSpeed ZeRO Stage (env-var) — Controls: Memory optimization level (0-3). Default: zero_stage argument
• Tensor Parallel Size (env-var) — Controls: Number of GPUs for tensor parallelism. Default: tensor_model_parallel_size
• Pipeline Parallel Size (env-var) — Controls: Number of pipeline stages. Default: pipeline_model_parallel_size
• Sequence Length (env-var) — Controls: Maximum input sequence length. Default: seq_length

Technology Stack

Key Components

• initialize_megatron (function) — Initializes distributed training environment, DeepSpeed integration, and global configuration megatron/initialize.py
• GPTModelPipe (class) — GPT model implementation with pipeline parallelism support for DeepSpeed megatron/model/gpt_model.py
• pretrain (function) — Main training loop with distributed optimization and checkpointing megatron/training.py
• build_train_valid_test_datasets (function) — Creates training, validation, and test datasets from indexed data files megatron/data/gpt_dataset.py
• BlendableDataset (class) — Combines multiple datasets with configurable mixing weights for training megatron/data/blendable_dataset.py
• parse_args (function) — Parses command-line arguments for model architecture, training, and DeepSpeed configuration megatron/arguments.py
• mpu (module) — Model parallel utilities for distributed tensor and pipeline parallelism megatron/mpu/__init__.py
• MegatronPretrainingSampler (class) — Data sampler that handles resumption from checkpoints and distributed data loading megatron/data/data_samplers.py
• save_checkpoint (function) — Saves model state, optimizer state, and training progress to disk megatron/checkpointing.py
• DecoderPackedMTFDataset (class) — Multi-task finetuning dataset with sequence packing for efficient training megatron/data/decoder_packed_mtf_dataset.py

Science Pipeline

1. Load Binary Dataset — Memory-map indexed dataset files with document boundaries [(num_documents, variable_length) → (total_tokens,)] megatron/data/indexed_dataset.py
2. Sample Generation — Create random samples mapping documents to training sequences [(total_tokens,) → (num_samples, seq_length)] megatron/data/dataset_utils.py
3. Tokenization & Batching — Group samples into microbatches with attention masks [(num_samples, seq_length) → (batch_size, seq_length)] megatron/data/data_samplers.py
4. Embedding Layer — Token and position embeddings with tensor parallelism [(batch_size, seq_length) → (batch_size, seq_length, hidden_size)] megatron/model/gpt_model.py
5. Transformer Layers — Multi-head attention and MLP blocks with pipeline parallelism [(batch_size, seq_length, hidden_size) → (batch_size, seq_length, hidden_size)] megatron/model/transformer.py
6. Output Projection — Final linear layer to vocabulary logits [(batch_size, seq_length, hidden_size) → (batch_size, seq_length, vocab_size)] megatron/model/gpt_model.py

Assumptions & Constraints

• [warning] Assumes input token tensors are (batch_size, seq_length) but shape validation is minimal (shape)
• [info] Token IDs assumed to be within vocabulary size but no explicit bounds checking (value-range)
• [critical] Assumes all tensors are on same CUDA device for parallel operations (device)
• [info] Dataset mixing weights assumed to sum to 1.0 but normalization is applied silently (dependency)

Related Ml Training Repositories

Frequently Asked Questions

Analyzed on March 31, 2026 by CodeSea. Written by Karolina Sarna.