# nvidia/megatron-lm architecture

> Architecture diagram of the GitHub repository nvidia/megatron-lm, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-08). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/nvidia/megatron-lm
- Repository: https://github.com/nvidia/megatron-lm

## Overview

#### 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: [`megatron/training/training.py`](https://github.com/nvidia/megatron-lm/blob/main/megatron/training/training.py) (Training loop), [`megatron/core/parallel_state.py`](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/parallel_state.py) (Parallel topology), [`megatron/core/models/gpt/gpt_model.py`](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/models/gpt/gpt_model.py) (Core model families), [`megatron/core/transformer/transformer_block.py`](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/transformer/transformer_block.py) (Transformer blocks)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_apps["Applications"]
  node_entrypoints["Training entry points<br/>Python executables<br/>[pretrain_gpt.py]"]
  node_rl_agent["RL agent &amp; environment API<br/>post-training runtime<br/>[api.py]"]
end

subgraph group_training["Training"]
  node_training_init["Distributed initialization<br/>training bootstrap<br/>[initialize.py]"]
  node_training_loop["Training loop<br/>orchestrator<br/>[training.py]"]
  node_training_checkpointing["Training checkpoint bridge<br/>checkpoint integration<br/>[checkpointing.py]"]
  node_model_providers["Model providers<br/>model construction"]
end

subgraph group_core["Core Modeling"]
  node_core_models["Core model families<br/>model implementations<br/>[gpt_model.py]"]
  node_transformer["Transformer blocks<br/>compute composition"]
  node_moe["MoE routing &amp; dispatch<br/>expert parallel compute<br/>[moe_layer.py]"]
  node_kernel_extensions["Fused GPU kernels<br/>CUDA/Triton/TE backend"]
  node_datasets["Rank-aware datasets<br/>data builder"]
  node_tokenizer["Tokenizer builder<br/>tokenization<br/>[build_tokenizer.py]"]
end

subgraph group_runtime["Distributed Runtime"]
  node_parallel_state["Parallel topology<br/>process-group state<br/>[parallel_state.py]"]
  node_tensor_parallel["Tensor-parallel layers<br/>sharded compute<br/>[layers.py]"]
  node_pipeline["Pipeline schedules<br/>activation scheduling<br/>[schedules.py]"]
  node_distributed_optimizer["Distributed optimizer<br/>parameter/gradient sharding"]
  node_collectives["PyTorch distributed &amp; collectives<br/>external GPU runtime"]
  node_dist_checkpointing[("Distributed checkpoints<br/>sharded state I/O<br/>[core.py]")]
  node_resharding["Resharding planner<br/>state migration<br/>[planner.py]"]
end

subgraph group_serving["Inference &amp; Delivery"]
  node_llm_api["LLM API<br/>inference API<br/>[llm.py]"]
  node_dynamic_engine["Dynamic inference engine<br/>generation engine<br/>[dynamic_engine.py]"]
  node_scheduler_context["Batch scheduler &amp; KV context<br/>prefill/decode runtime<br/>[scheduler.py]"]
  node_text_server["OpenAI-style generation server<br/>HTTP streaming server"]
  node_trt_export["TensorRT-LLM export<br/>engine export"]
end

node_entrypoints -->|"bootstraps"| node_training_init
node_entrypoints -->|"invokes"| node_training_loop
node_training_init -->|"creates groups"| node_parallel_state
node_training_loop -->|"requests model"| node_model_providers
node_model_providers -->|"builds"| node_core_models
node_training_loop -->|"consumes batches"| node_datasets
node_datasets -->|"tokenizes with"| node_tokenizer
node_training_loop -->|"runs schedules"| node_pipeline
node_training_loop -->|"updates through"| node_distributed_optimizer
node_training_loop -->|"saves/loads"| node_training_checkpointing
node_training_checkpointing -->|"delegates state I/O"| node_dist_checkpointing
node_dist_checkpointing -.->|"changes shard layout"| node_resharding
node_core_models -->|"composes"| node_transformer
node_transformer -.->|"optionally uses"| node_moe
node_transformer -->|"accelerated by"| node_kernel_extensions
node_core_models -->|"uses sharded layers"| node_tensor_parallel
node_parallel_state -->|"configures"| node_tensor_parallel
node_parallel_state -->|"configures"| node_pipeline
node_parallel_state -->|"configures"| node_distributed_optimizer
node_parallel_state -->|"uses"| node_collectives
node_llm_api -->|"submits requests"| node_dynamic_engine
node_dynamic_engine -->|"schedules generation"| node_scheduler_context
node_dynamic_engine -->|"wraps"| node_core_models
node_text_server -->|"serves"| node_llm_api
node_trt_export -->|"exports weights from"| node_core_models
node_rl_agent -->|"uses rollout inference"| node_llm_api

click node_entrypoints "https://github.com/nvidia/megatron-lm/blob/main/pretrain_gpt.py"
click node_training_init "https://github.com/nvidia/megatron-lm/blob/main/megatron/training/initialize.py"
click node_training_loop "https://github.com/nvidia/megatron-lm/blob/main/megatron/training/training.py"
click node_training_checkpointing "https://github.com/nvidia/megatron-lm/blob/main/megatron/training/checkpointing.py"
click node_core_models "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/models/gpt/gpt_model.py"
click node_transformer "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/transformer/transformer_block.py"
click node_moe "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/transformer/moe/moe_layer.py"
click node_datasets "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/datasets/blended_megatron_dataset_builder.py"
click node_tokenizer "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/tokenizers/utils/build_tokenizer.py"
click node_parallel_state "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/parallel_state.py"
click node_tensor_parallel "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/tensor_parallel/layers.py"
click node_pipeline "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/pipeline_parallel/schedules.py"
click node_distributed_optimizer "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/optimizer/distrib_optimizer.py"
click node_dist_checkpointing "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/dist_checkpointing/core.py"
click node_resharding "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/resharding/planner.py"
click node_llm_api "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/apis/llm.py"
click node_dynamic_engine "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/engines/dynamic_engine.py"
click node_scheduler_context "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/scheduler.py"
click node_text_server "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py"
click node_trt_export "https://github.com/nvidia/megatron-lm/blob/main/megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py"
click node_rl_agent "https://github.com/nvidia/megatron-lm/blob/main/megatron/rl/agent/api.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_entrypoints,node_rl_agent toneBlue
class node_training_init,node_training_loop,node_training_checkpointing,node_model_providers toneAmber
class node_core_models,node_transformer,node_moe,node_kernel_extensions,node_datasets,node_tokenizer toneMint
class node_parallel_state,node_tensor_parallel,node_pipeline,node_distributed_optimizer,node_collectives,node_dist_checkpointing,node_resharding toneRose
class node_llm_api,node_dynamic_engine,node_scheduler_context,node_text_server,node_trt_export toneIndigo
```

## Components

### Applications

- [Training entry points](https://github.com/nvidia/megatron-lm/blob/main/pretrain_gpt.py): Python executables · GPT, hybrid, Mamba, VLM, and RL launch training. · `pretrain_gpt.py`
- [RL agent & environment API](https://github.com/nvidia/megatron-lm/blob/main/megatron/rl/agent/api.py): post-training runtime · Separates agents, rollout inference, and optional FastAPI servers. · `megatron/rl/agent/api.py`

### Training

- [Distributed initialization](https://github.com/nvidia/megatron-lm/blob/main/megatron/training/initialize.py): training bootstrap · `megatron/training/initialize.py`
- [Training loop](https://github.com/nvidia/megatron-lm/blob/main/megatron/training/training.py): orchestrator · Builds components and runs update, evaluation, and save loops. · `megatron/training/training.py`
- [Training checkpoint bridge](https://github.com/nvidia/megatron-lm/blob/main/megatron/training/checkpointing.py): checkpoint integration · `megatron/training/checkpointing.py`
- **Model providers**: model construction

### Core Modeling

- [Core model families](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/models/gpt/gpt_model.py): model implementations · GPT, BERT, T5, hybrid, multimodal, and MiMo models. · `megatron/core/models/gpt/gpt_model.py`
- [Transformer blocks](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/transformer/transformer_block.py): compute composition · `megatron/core/transformer/transformer_block.py`
- [MoE routing & dispatch](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/transformer/moe/moe_layer.py): expert parallel compute · Routes tokens and exchanges them across expert ranks. · `megatron/core/transformer/moe/moe_layer.py`
- **Fused GPU kernels**: CUDA/Triton/TE backend · Optimized kernels and mixed-precision extensions.
- [Rank-aware datasets](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/datasets/blended_megatron_dataset_builder.py): data builder · Builds blended, indexed, text, and multimodal samples. · `megatron/core/datasets/blended_megatron_dataset_builder.py`
- [Tokenizer builder](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/tokenizers/utils/build_tokenizer.py): tokenization · `megatron/core/tokenizers/utils/build_tokenizer.py`

### Distributed Runtime

- [Parallel topology](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/parallel_state.py): process-group state · Establishes TP, PP, DP, context, and expert parallel groups. · `megatron/core/parallel_state.py`
- [Tensor-parallel layers](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/tensor_parallel/layers.py): sharded compute · `megatron/core/tensor_parallel/layers.py`
- [Pipeline schedules](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/pipeline_parallel/schedules.py): activation scheduling · Moves activations and gradients between pipeline stages. · `megatron/core/pipeline_parallel/schedules.py`
- [Distributed optimizer](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/optimizer/distrib_optimizer.py): parameter/gradient sharding · `megatron/core/optimizer/distrib_optimizer.py`
- **PyTorch distributed & collectives**: external GPU runtime · NCCL is primary; Gloo, NIXL, and NVSHMEM are optional.
- [Distributed checkpoints](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/dist_checkpointing/core.py): sharded state I/O · Supports async, fully parallel save/load, resume, and topology changes. · `megatron/core/dist_checkpointing/core.py`
- [Resharding planner](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/resharding/planner.py): state migration · `megatron/core/resharding/planner.py`

### Inference & Delivery

- [LLM API](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/apis/llm.py): inference API · `megatron/core/inference/apis/llm.py`
- [Dynamic inference engine](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/engines/dynamic_engine.py): generation engine · `megatron/core/inference/engines/dynamic_engine.py`
- [Batch scheduler & KV context](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/scheduler.py): prefill/decode runtime · Batches requests and manages dynamic KV-cache contexts. · `megatron/core/inference/scheduler.py`
- [OpenAI-style generation server](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py): HTTP streaming server · Exposes completions, chat, and health endpoints. · `megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py`
- [TensorRT-LLM export](https://github.com/nvidia/megatron-lm/blob/main/megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py): engine export · Converts distributed Core weights and builds inference engines. · `megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py`

## Connections

- 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