How tanzeyy/megatron-lm works
Megatron-LM is a GPU-optimized system for training transformer models at scale, with Megatron Core providing composable model, parallelism, optimization, and data components. The main workflow is training entry point → dataset construction → model computation distributed across parallel strategies → optimization and checkpointing. The repository also supports inference serving and TensorRT-LLM export. The README documents these capabilities, but most implementation excerpts were not sampled; the graph therefore shows documented subsystem relationships at a high level and omits uncertain internal wiring.
Main entry points: training.py Training loop · gpt_model.py GPT model · pretrain_gpt.py Pretraining entry · engines Inference engines
Components 24
Training orchestration
- Pretraining entry
pretrain_gpt.py - Training loop
megatron/training/training.py - Checkpointing
megatron/training/checkpointing.py
Data pipelines
- Dataset builders
megatron/core/datasets - Data preparation
tools/preprocess_data.py
Model computation
- GPT model
megatron/core/models/gpt/gpt_model.py - Transformer layers
megatron/core/transformer/transformer_block.py - MoE layers
megatron/core/transformer/moe/moe_layer.py - Mamba models
megatron/core/models/mamba/mamba_model.py - Multimodal models
megatron/core/models/multimodal
Distributed runtime
- Optimizers
megatron/core/optimizer/optimizer.py - Data parallelism
megatron/core/distributed - Tensor parallelism
megatron/core/tensor_parallel - Pipeline parallelism
megatron/core/pipeline_parallel
Inference and export
- Generation API
megatron/inference/text_generation/api.py - Inference engines
megatron/core/inference/engines - Generation controller
megatron/core/inference/text_generation_controllers/text_generation_controller.py - Inference wrappers
megatron/core/inference/model_inference_wrappers - HTTP endpoints
megatron/inference/endpoints/completions.py - TensorRT-LLM export
megatron/core/export/trtllm
Other components
- Trainer
- Transformer Engine
- TensorRT-LLM
- Checkpoint files
Connections 26
- Trainer → Pretraining entry: starts
- Pretraining entry → Training loop: runs
- Pretraining entry → Dataset builders: configures
- Data preparation → Dataset builders: prepares data
- Dataset builders → Training loop: supplies batches
- Training loop → GPT model: trains
- GPT model → Transformer layers: uses layers
- GPT model → MoE layers: supports
- Training loop → Mamba models: trains
- Training loop → Multimodal models: trains
- Training loop → Optimizers: updates weights
- Training loop → Data parallelism: uses
- Training loop → Tensor parallelism: uses
- Training loop → Pipeline parallelism: uses
- Transformer layers → Transformer Engine: optional backend
- Training loop → Checkpointing: saves state
- Checkpointing → Checkpoint files: writes
- Trainer → HTTP endpoints: requests generation
- HTTP endpoints → Generation API: dispatches
- Generation API → Inference engines: submits requests
- Inference engines → Generation controller: steps generation
- Generation controller → Inference wrappers: runs forward pass
- Inference wrappers → GPT model: wraps model
- Inference engines → HTTP endpoints: returns results
- Trainer → TensorRT-LLM export: exports model
- TensorRT-LLM export → TensorRT-LLM: builds engine