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.

LayerPrimary responsibilityWhat it is not
NemotronA 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-BridgeArchitecture detection, Hugging Face ↔ Megatron checkpoint conversion, model providers, recipes, and PyTorch-native training entry points.Not primarily an inference server.
Megatron-CoreComposable Transformer building blocks, distributed parallelism, schedules, distributed checkpointing, and optimizers.Not a model zoo that automatically understands every HF architecture.
Transformer EngineOptimized 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]

StepWhat Bridge doesWhy it matters
1. Detect and translate configAutoBridge 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 shardsThe 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 planBridge 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 tensorsOnly 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 distributeThe 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 trainThe 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 componentHF representationMegatron representation / mapping
Q, K, V projectionsSeparate q_proj, k_proj, v_proj tensorsQKVMapping fuses or splits them using head-aware, model-specific interleaving; the result is typically column-parallel under TP.
Attention output projectiono_projRowParallelMapping; split along the input dimension under tensor parallelism, then aggregate partial results at runtime.
Gated MLP inputSeparate gate_proj and up_projGatedMLPMapping combines or splits them into/from Megatron linear_fc1, normally column-parallel.
MLP outputdown_projUsually RowParallelMapping to Megatron linear_fc2.
Norms and small global tensorsOne tensor per layerReplicatedMapping when each TP rank must own an identical copy.
MoE expertsHF expert tensors named by global expert IDMappings 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

CheckpointExpected purposeVerification question
HF basePortable source of pretrained knowledgeDoes the tokenizer/config exactly match the weight revision?
Megatron baseParallel-aware initialization for trainingDo rank-local shards load with the intended TP/PP/EP topology?
Megatron run checkpointResumption artifact during SFT/pretrainingCan the same software and topology resume safely, with optimizer state if required?
HF exportPortable evaluation/deployment artifactDo 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 goalRecommended starting pointWhy
Run inference onlyUse 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 scaleHF checkpoint → Bridge import → Megatron training → Bridge export.You preserve pretrained knowledge while gaining Megatron parallelism.
Pretrain a new architecture from random initializationBuild 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 architectureImplement 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

  1. NVIDIA Megatron-Core — Overview
  2. NVIDIA Transformer Engine — Documentation
  3. NVIDIA Megatron Bridge — Conversion Technical Details
  4. NVIDIA-NeMo/Megatron-Bridge GitHub repository and README
  5. NVIDIA — Nemotron 3 Nano in Megatron Bridge
  6. NVIDIA Technical Blog — Inside NVIDIA Nemotron 3
  7. NVIDIA — Nemotron AI Stack
← Previous Post
Next Post →

Leave a Comment