How nvidia/megatron-lm works
Purpose and Entry Points
Megatron-LM combines a reusable GPU-scale deep-learning library (Megatron Core) with reference training applications. It targets distributed training and inference for GPT, BERT, T5, hybrid/Mamba, MoE, and multimodal models. Primary executable entry points are pretrain_gpt.py, pretrain_hybrid.py, pretrain_mamba.py, pretrain_vlm.py, and train_rl.py; examples provide runnable model-specific launch scripts.
Training Orchestration
The training layer parses arguments/configuration, initializes distributed execution, builds models/data/optimizers, and runs forward-backward-update/evaluation/checkpoint loops.
• megatron/training/initialize.py
• megatron/training/training.py
• megatron/training/checkpointing.py
It is the boundary between application scripts/model providers and Core primitives. Training config supports YAML-driven object instantiation under megatron/training/config/, while model-specific construction is supplied through root builders and megatron/training/models/.
Model and Compute Layer
Core defines composable model implementations and transformer internals. GPT, BERT, T5, hybrid Transformer–SSM, vision/audio, LLaVA, and MiMo models assemble common embeddings, attention, MLPs, layer specifications, and transformer blocks.
• megatron/core/models/gpt/gpt_model.py
• megatron/core/transformer/transformer_block.py
• megatron/core/models/hybrid/hybrid_model.py
Mixture-of-Experts is a first-class transformer variant: routers select experts and token dispatchers execute cross-rank expert exchange.
• megatron/core/transformer/moe/moe_layer.py
• megatron/core/transformer/moe/router.py
• megatron/core/transformer/moe/token_dispatcher.py
CUDA/Triton/Transformer Engine-backed fusions and mixed-precision extensions sit beneath this API, isolating optimized kernels from model definitions.
Distributed Execution and Memory
parallel_state and process-group configuration establish TP, PP, DP, context parallelism, and expert parallelism. Tensor-parallel layers shard model computation; pipeline schedules send activations/gradients between stages; DDP/FSDP and distributed optimizers manage parameter and gradient sharding.
• megatron/core/parallel_state.py
• megatron/core/tensor_parallel/layers.py
• megatron/core/pipeline_parallel/schedules.py
• megatron/core/distributed/distributed_data_parallel.py
• megatron/core/optimizer/distrib_optimizer.py
This layer depends on PyTorch distributed and GPU collective infrastructure, principally NCCL; optional resharding backends include Gloo, NIXL, and NVSHMEM.
Data, State, and Checkpoints
Datasets convert indexed, blended, text, and multimodal corpora into rank-aware training samples and schedules.
• megatron/core/datasets/blended_megatron_dataset_builder.py
• megatron/core/datasets/gpt_dataset.py
• megatron/core/tokenizers/utils/build_tokenizer.py
Distributed checkpointing serializes sharded model/optimizer state, supports async and fully parallel strategies, and enables resume or topology changes. Resharding is a related runtime path for moving/refitting sharded state.
• megatron/core/dist_checkpointing/core.py
• megatron/core/dist_checkpointing/strategies/fully_parallel.py
• megatron/core/resharding/planner.py
Dataset object-storage support includes S3-oriented utilities; checkpoints otherwise form a filesystem/storage boundary.
Inference and Serving
Inference is separate from training orchestration: high-level LLM APIs submit requests to static or dynamic engines, whose scheduler batches prefill/decode work and manages KV-cache contexts. Model wrappers adapt Core models to generation controllers and sampling.
• megatron/core/inference/apis/llm.py
• megatron/core/inference/engines/dynamic_engine.py
• megatron/core/inference/scheduler.py
• megatron/core/inference/contexts/dynamic_context.py
The dynamic text-generation server exposes completions/chat/health endpoints with OpenAI-style streaming.
• megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py
• megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py
ZMQ can coordinate asynchronous multi-replica inference; NCCL/NIXL can transfer disaggregated KV/SSM state.
Export, Post-Training, and RL
Conversion/export is deliberately outside the core training loop: TensorRT-LLM export maps Core model state and builds engines.
• megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py
• megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py
Post-training integrates ModelOpt-oriented model specs/state hooks, while RL separates agents, rollout inference, and optional FastAPI environment/inference servers.
• megatron/core/post_training/modelopt/gpt/model_specs.py
• megatron/rl/agent/api.py
• megatron/rl/server/agent/fastapi_env_server.py
Main entry points: training.py Training loop · parallel_state.py Parallel topology · gpt_model.py Core model families · transformer_block.py Transformer blocks
Components 24
Applications
- Training entry points · GPT, hybrid, Mamba, VLM, and RL launch training.
pretrain_gpt.py - RL agent & environment API · Separates agents, rollout inference, and optional FastAPI servers.
megatron/rl/agent/api.py
Training
- Distributed initialization
megatron/training/initialize.py - Training loop · Builds components and runs update, evaluation, and save loops.
megatron/training/training.py - Training checkpoint bridge
megatron/training/checkpointing.py - Model providers
Core Modeling
- Core model families · GPT, BERT, T5, hybrid, multimodal, and MiMo models.
megatron/core/models/gpt/gpt_model.py - Transformer blocks
megatron/core/transformer/transformer_block.py - MoE routing & dispatch · Routes tokens and exchanges them across expert ranks.
megatron/core/transformer/moe/moe_layer.py - Fused GPU kernels · Optimized kernels and mixed-precision extensions.
- Rank-aware datasets · Builds blended, indexed, text, and multimodal samples.
megatron/core/datasets/blended_megatron_dataset_builder.py - Tokenizer builder
megatron/core/tokenizers/utils/build_tokenizer.py
Distributed Runtime
- Parallel topology · Establishes TP, PP, DP, context, and expert parallel groups.
megatron/core/parallel_state.py - Tensor-parallel layers
megatron/core/tensor_parallel/layers.py - Pipeline schedules · Moves activations and gradients between pipeline stages.
megatron/core/pipeline_parallel/schedules.py - Distributed optimizer
megatron/core/optimizer/distrib_optimizer.py - PyTorch distributed & collectives · NCCL is primary; Gloo, NIXL, and NVSHMEM are optional.
- Distributed checkpoints · Supports async, fully parallel save/load, resume, and topology changes.
megatron/core/dist_checkpointing/core.py - Resharding planner
megatron/core/resharding/planner.py
Inference & Delivery
- LLM API
megatron/core/inference/apis/llm.py - Dynamic inference engine
megatron/core/inference/engines/dynamic_engine.py - Batch scheduler & KV context · Batches requests and manages dynamic KV-cache contexts.
megatron/core/inference/scheduler.py - OpenAI-style generation server · Exposes completions, chat, and health endpoints.
megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py - TensorRT-LLM export · Converts distributed Core weights and builds inference engines.
megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py
Connections 26
- Training entry points → Distributed initialization: bootstraps
- Training entry points → Training loop: invokes
- Distributed initialization → Parallel topology: creates groups
- Training loop → Model providers: requests model
- Model providers → Core model families: builds
- Training loop → Rank-aware datasets: consumes batches
- Rank-aware datasets → Tokenizer builder: tokenizes with
- Training loop → Pipeline schedules: runs schedules
- Training loop → Distributed optimizer: updates through
- Training loop → Training checkpoint bridge: saves/loads
- Training checkpoint bridge → Distributed checkpoints: delegates state I/O
- Distributed checkpoints → Resharding planner: changes shard layout
- Core model families → Transformer blocks: composes
- Transformer blocks → MoE routing & dispatch: optionally uses
- Transformer blocks → Fused GPU kernels: accelerated by
- Core model families → Tensor-parallel layers: uses sharded layers
- Parallel topology → Tensor-parallel layers: configures
- Parallel topology → Pipeline schedules: configures
- Parallel topology → Distributed optimizer: configures
- Parallel topology → PyTorch distributed & collectives: uses
- LLM API → Dynamic inference engine: submits requests
- Dynamic inference engine → Batch scheduler & KV context: schedules generation
- Dynamic inference engine → Core model families: wraps
- OpenAI-style generation server → LLM API: serves
- TensorRT-LLM export → Core model families: exports weights from
- RL agent & environment API → LLM API: uses rollout inference