Megatron-Bridge in Practice: A Production 101 Guide to the Nemotron, Megatron-Core, and Transformer Engine Stack
Megatron-Bridge in Practice: A Production 101 Guide to the Nemotron, Megatron-Core, and Transformer Engine Stack
Scope. This tutorial explains how an existing Hugging Face checkpoint becomes a distributed Megatron training model, where Nemotron and Transformer Engine fit, and how to organize a maintainable project around that workflow. It is written as a practical research-and-implementation guide; a Manus 1.6-assisted workflow can help inspect documentation, draft configuration inventories, and review logs, but the actual distributed training still runs in NVIDIA’s software and GPU environment.
1. The one-sentence mental model
Nemotron is the model family; Megatron-Bridge is the compatibility, conversion, and training-workflow layer; Megatron-Core is the distributed-training substrate; Transformer Engine (TE) is the low-precision kernel and numerical-runtime layer. The layers are complementary rather than interchangeable.
| Layer | Primary responsibility | What it is not |
|---|---|---|
| Nemotron | A family of NVIDIA open models, datasets, and recipes. Nemotron 3 Nano, for example, uses a hybrid Mamba-Transformer MoE design. | Not a general conversion library. |
| Megatron-Bridge | Architecture detection, Hugging Face ↔ Megatron checkpoint conversion, model providers, recipes, and PyTorch-native training entry points. | Not primarily an inference server. |
| Megatron-Core | Composable Transformer building blocks, distributed parallelism, schedules, distributed checkpointing, and optimizers. | Not a model zoo that automatically understands every HF architecture. |
| Transformer Engine | Optimized Transformer kernels and mixed-precision support, including FP8 on supported GPUs and MXFP8/NVFP4 on Blackwell. | Not a checkpoint converter. |
Megatron-Core exposes tensor, pipeline, data, expert, context, and sequence parallelism. Megatron-Bridge builds usable workflows on those primitives, while TE supplies low-precision modules and fused operations used by the runtime. NVIDIA documents Megatron-Core as a composable library and Megatron-Bridge as one of the libraries built on it. [1] TE keeps the scaling metadata required for FP8 training inside its modules, which is why mixed precision is a runtime concern rather than simply a different checkpoint-file extension. [2]
2. System architecture: model, workflow, framework, kernels
MODEL & RECIPE LAYER
Nemotron / Llama / Qwen HF config + tokenizer + safetensors
|
v
WORKFLOW & COMPATIBILITY LAYER
AutoBridge -> architecture-specific ModelBridge -> MappingRegistry
| | |
| | QKV / Gated-MLP / row / column mappings
v v
ModelProvider -> distributed Megatron model construction
|
v
DISTRIBUTED TRAINING LAYER
Megatron-Core: TP + PP + DP + EP + CP + SP, schedulers, checkpoints
|
v
NUMERICAL / KERNEL LAYER
Transformer Engine: fused kernels, BF16 / FP8 / MXFP8 / NVFP4 policies
|
v
CUDA, NCCL, NVLink/InfiniBand, NVIDIA GPUs
The diagram exposes a crucial separation of concerns. Bridge changes model representation and organizes training; Core distributes the computation; TE accelerates eligible computation. TE does not “turn an HF checkpoint into FP8” merely by importing it. Its quantization, scaling, autocast, and fused-kernel behavior are activated during the configured model execution and training/inference runtime. [2]
3. What a checkpoint means at each boundary
A Hugging Face checkpoint is usually a portable model package: config.json, tokenizer files, and one or more .safetensors weight files. Its tensor names and layouts follow the HF architecture implementation. A Megatron checkpoint is a training-state artifact designed for a particular distributed model layout. Depending on the run, a rank owns only a local shard of a tensor or a subset of layers and experts. Optimizer and training-state information can be distributed as well.
Therefore, conversion is not a filename conversion. It is a deterministic sequence of name mapping, structure mapping, layout transformation, and distributed scatter/gather. Megatron-Bridge streams the conversion parameter by parameter, so it does not need to load full HF and Megatron models into one GPU’s memory. It is also aware of tensor, pipeline, virtual pipeline, and expert parallelism. [3]
4. An implementation-oriented repository map
The official repository’s top-level layout includes src/megatron/bridge, examples, scripts, tutorials, tests, and a pinned 3rdparty/Megatron-LM submodule. The source tree below is a conceptual map of the important paths; exact files can change by release. [4]
Megatron-Bridge/
├── src/megatron/bridge/
│ ├── models/
│ │ ├── conversion/ # AutoBridge, ModelBridge, registry, mappings
│ │ ├── llama/ # architecture-specific provider + mapping rules
│ │ ├── qwen/ # model-specific differences, e.g. QK norm
│ │ └── nemotron*/ # Nemotron model definitions and recipes
│ ├── training/ # pretrain / finetune entry points and utilities
│ ├── recipes/ # model + hardware-oriented configuration recipes
│ ├── data/ # dataset and SFT data utilities
│ └── peft/ # LoRA, DoRA, and related adapters
├── examples/
│ ├── conversion/ # import, export, round-trip, output comparison
│ └── models/ # supported-model recipes
├── scripts/conversion/ # command-line conversion helpers
├── tutorials/ # data, recipes, precision, and training tutorials
├── tests/ # conversion and training regression tests
└── 3rdparty/Megatron-LM/ # pinned Megatron-Core source
For an application team, keep model data, run configuration, launch scripts, and artifacts separate from the vendor checkout:
my-megatron-project/
├── configs/
│ ├── cluster/ # launcher, topology, and environment choices
│ ├── data/ # dataset blending and tokenization settings
│ └── runs/ # model, optimizer, precision, checkpoint policy
├── data/
│ ├── raw/ ├── processed/ └── manifests/
├── scripts/
│ ├── import_hf.sh ├── train_sft.sh └── export_hf.sh
├── checkpoints/
│ ├── hf-base/ # immutable source checkpoint or its model ID record
│ ├── megatron-base/ # imported distributed base checkpoint
│ └── runs/<run-id>/ # resumable training checkpoints
├── outputs/
│ ├── hf-export/ # portable model for evaluation/deployment
│ └── eval/ # metrics, prompts, parity reports
└── README.md # exact software/container versions and commands
5. How HF → Megatron conversion works internally
The official conversion design centers on AutoBridge, an architecture-specific MegatronModelBridge, a model provider, and a mapping registry. The same high-level sequence applies to Llama, Qwen, Nemotron, and other supported architectures, while individual bridge definitions handle architectural differences. [3]
| Step | What Bridge does | Why it matters |
|---|---|---|
| 1. Detect and translate config | AutoBridge reads the HF configuration, selects a registered bridge, and creates a Megatron-compatible provider. | Hidden size, head counts, GQA/MQA, rotary embedding, MoE, and normalization choices must agree before weights can be loaded. |
| 2. Instantiate local model shards | The provider finalizes TP/PP/EP and creates distributed Megatron model instances. | Each rank can allocate only the tensors it owns. |
| 3. Build a deterministic parameter plan | Bridge enumerates parameters and buffers across pipeline stages, establishes globally sorted names, and resolves a mapping for every destination parameter. | Every rank participates in collectives in the same order; this avoids distributed deadlocks and nondeterministic behavior. |
| 4. Stream source tensors | Only the referenced HF tensors are read from storage for the current mapping task. | Memory usage scales with one or a few tensors, not the entire model. |
| 5. Transform and distribute | The mapping fuses or splits logical matrices, casts when needed, and scatters/gathers or broadcasts according to TP/PP/EP ownership. | The destination local parameter has the exact shape, layout, dtype, and placement expected by Megatron-Core. |
| 6. Persist or train | The populated distributed model can be saved as a Megatron checkpoint or passed directly into a training loop. | There is no mandatory full, intermediate monolithic checkpoint. |
5.1 Per-parameter mapping types
In a common decoder-only Transformer, PyTorch linear weights are conventionally shaped as [out_features, in_features]. Bridge applies different mapping rules based on the logical layer:
| Logical component | HF representation | Megatron representation / mapping |
|---|---|---|
| Q, K, V projections | Separate q_proj, k_proj, v_proj tensors | QKVMapping fuses or splits them using head-aware, model-specific interleaving; the result is typically column-parallel under TP. |
| Attention output projection | o_proj | RowParallelMapping; split along the input dimension under tensor parallelism, then aggregate partial results at runtime. |
| Gated MLP input | Separate gate_proj and up_proj | GatedMLPMapping combines or splits them into/from Megatron linear_fc1, normally column-parallel. |
| MLP output | down_proj | Usually RowParallelMapping to Megatron linear_fc2. |
| Norms and small global tensors | One tensor per layer | ReplicatedMapping when each TP rank must own an identical copy. |
| MoE experts | HF expert tensors named by global expert ID | Mappings place the local subset with the owning EP rank; export gathers and emits tensors using global expert IDs. |
For grouped-query attention, Q, K, and V are not necessarily equal-sized. A safe bridge must understand the query-head and key/value-head counts, then interleave and shard in the order required by the destination implementation. For this reason, blind concatenation of Q, K, and V is not a valid general converter. Bridge’s architecture-specific registry owns that knowledge. Similarly, the MLP fusion format and bias conventions are model dependent. The Qwen3 bridge, for instance, explicitly maps QK layer norms and represents its lack of QKV bias. [3]
5.2 Parallelism semantics during conversion
Assume tensor parallel size TP = 4. A column-parallel matrix W ∈ R^(out × in) is split along out, so each TP rank receives approximately out / 4 rows. A row-parallel matrix is split along in, and each rank computes a partial output that is combined at runtime. Pipeline parallelism assigns complete layer ranges to stages; expert parallelism assigns MoE experts to ranks. Import and export must coordinate these facts in reverse directions.
HF weight file
|
| import: read current tensor only
v
[model-specific transform: QKV fuse / gated MLP fuse / dtype adaptation]
|
+-- TP: split or scatter local tensor shards
+-- PP: owning pipeline stage writes its local parameter
+-- EP: owning rank receives its local experts
v
Distributed Megatron checkpoint / in-memory distributed model
|
| export: PP broadcast -> TP/EP gather -> inverse transform
v
HF safetensors + config + tokenizer
Bridge supports AutoMapping, column-parallel, row-parallel, QKV, gated-MLP, replicated, and custom mappings. The documented export path first makes pipeline-stage data available to collective participants, then gathers TP and EP shards as needed and applies the inverse structural transformation. [3]
6. Where Nemotron changes the design
Nemotron 3 Nano illustrates why the stack needs more than a basic dense-Transformer converter. NVIDIA documents it as a hybrid MoE model with 23 Mamba-2-and-MoE layers plus six attention layers; each MoE layer has 128 experts and one shared expert, five experts are active per token, and the model has 30B total / 3.5B active parameters. [5] The wider Nemotron 3 family combines Mamba sequence modeling, Transformer attention, and MoE routing to target long-context agentic workloads. [6]
That model composition drives the training design. Mamba/attention blocks require architecture-aware model definitions; MoE requires expert parallelism and correct per-expert checkpoints; long contexts may motivate context and sequence parallelism; and sparse activation changes capacity planning. Megatron-Bridge provides the model definition, recipe, conversion, and training integration, while Megatron-Core supplies the distributed mechanisms underneath. NVIDIA’s Nemotron stack documentation places Bridge in pretraining and SFT, and NeMo-RL in the later RL/post-training phase. [7]
7. Practical 101: import, fine-tune, export, and verify
This example explains the lifecycle with a small, supported HF causal-language model. Start with TP=1 for conversion debugging; introduce TP/PP/EP only after a round-trip test is clean. Use an official NeMo container or a release-compatible installation, a licensed model you may download, and a GPU-capable environment. Never treat the illustrative paths or batch sizes below as production defaults.
7.1 Prepare the workspace
mkdir -p my-megatron-project/{configs,scripts,checkpoints,outputs}
cd my-megatron-project
# Recommended environment pattern; choose a release-compatible NeMo image tag.
docker run --rm -it --gpus all -w /workdir \
-v "$PWD:/workdir" --entrypoint bash nvcr.io/nvidia/nemo:<TAG>
# In the container, authenticate only if the model repository requires it.
huggingface-cli login
7.2 Programmatic import for inspection
from megatron.bridge import AutoBridge
model_id = "meta-llama/Llama-3.2-1B" # Substitute a model you are licensed to use.
bridge = AutoBridge.from_hf_pretrained(model_id, trust_remote_code=True)
provider = bridge.to_megatron_provider()
provider.tensor_model_parallel_size = 1
provider.pipeline_model_parallel_size = 1
provider.finalize()
model = provider.provide_distributed_model(wrap_with_ddp=False)
bridge.load_hf_weights(model) # streams HF tensors into Megatron parameters
# After a real Bridge training or finetuning loop has updated `model`:
bridge.save_hf_pretrained(model, "./outputs/hf-export")
Important: this code establishes a conversion-ready distributed model; it does not itself perform gradient updates. A real SFT or pretraining job uses a Bridge recipe/entry point, a tokenizer, data pipeline, optimizer, scheduler, checkpoint policy, and a distributed launcher. NVIDIA’s official README provides recipe examples and distinguishes this pretrained-weight workflow from random-initialized pretraining, where to_megatron_provider(load_weights=False) only reads the source architecture configuration. [4]
7.3 Checkpoint-oriented CLI flow
# 1) Import an HF base model into a distributed Megatron checkpoint.
python examples/conversion/convert_checkpoints.py import \
--hf-model <HF_MODEL_ID> \
--megatron-path ./checkpoints/megatron-base \
--trust-remote-code
# 2) Run your selected recipe with the imported checkpoint as its initialization.
torchrun --nproc-per-node=<GPU_COUNT> <recipe_or_finetune_script.py> \
checkpoint.pretrained_checkpoint=./checkpoints/megatron-base
# 3) Export the trained checkpoint back to a portable HF package.
python examples/conversion/convert_checkpoints.py export \
--hf-model <HF_MODEL_ID> \
--megatron-path ./checkpoints/runs/<run-id> \
--hf-path ./outputs/hf-export
For Nemotron 3 Nano specifically, NVIDIA documents the same import/export structure and a model-specific finetuning recipe that expects a pretrained Megatron checkpoint. It also documents LoRA target modules such as linear_qkv, linear_proj, linear_fc1, linear_fc2, in_proj, and out_proj. [5]
7.4 What “correct” looks like
| Checkpoint | Expected purpose | Verification question |
|---|---|---|
| HF base | Portable source of pretrained knowledge | Does the tokenizer/config exactly match the weight revision? |
| Megatron base | Parallel-aware initialization for training | Do rank-local shards load with the intended TP/PP/EP topology? |
| Megatron run checkpoint | Resumption artifact during SFT/pretraining | Can the same software and topology resume safely, with optimizer state if required? |
| HF export | Portable evaluation/deployment artifact | Do fixed prompts give matching or appropriately close logits after the round trip? |
Use the repository’s conversion examples and model-comparison utilities for parity testing instead of checking only that files exist. Test a fixed token batch before and after conversion; account for configured dtype/precision behavior; and save the exact container, commit, topology, tokenizer revision, and recipe alongside the run. A successful file write is not proof that GQA order, RoPE configuration, expert IDs, or tied weights are correct.
8. Choosing the right flow
| Your goal | Recommended starting point | Why |
|---|---|---|
| Run inference only | Use the existing HF package with an inference stack such as TensorRT-LLM or another compatible engine. | Bridge is not principally a serving runtime. |
| Fine-tune an existing HF model at scale | HF checkpoint → Bridge import → Megatron training → Bridge export. | You preserve pretrained knowledge while gaining Megatron parallelism. |
| Pretrain a new architecture from random initialization | Build the provider/recipe with load_weights=False; no HF weights need be imported. | You need an architecture definition, not existing knowledge. |
| Add support for a new HF architecture | Implement an architecture-specific provider bridge and mapping registry, then write round-trip tests. | The difficult part is semantics, not serialization. |
9. Final takeaways
Megatron-Bridge is valuable because it turns a portable HF model into a parallelism-aware training representation without asking one GPU to hold the complete source and destination models. It maps model semantics, fuses and splits architecture-specific tensors, and distributes them in the topology Megatron-Core expects. Megatron-Core then supplies scalable training primitives, while Transformer Engine provides low-precision execution and optimized kernels. Nemotron is a concrete model family that uses this stack for pretraining and SFT, especially where hybrid Mamba-Transformer-MoE behavior makes distributed design central rather than optional.
References
- NVIDIA Megatron-Core — Overview
- NVIDIA Transformer Engine — Documentation
- NVIDIA Megatron Bridge — Conversion Technical Details
- NVIDIA-NeMo/Megatron-Bridge GitHub repository and README
- NVIDIA — Nemotron 3 Nano in Megatron Bridge
- NVIDIA Technical Blog — Inside NVIDIA Nemotron 3
- NVIDIA — Nemotron AI Stack