Megatron Bridge

Megatron Bridge distributed training on AWS

Megatron-Bridge + UCCL-EP over EFA

This directory is the library-level, model-agnostic home for NVIDIA Megatron-Bridge test cases that run Mixture-of-Experts (MoE) training on Amazon EKS with UCCL-EP carrying the expert-parallel all-to-all over AWS EFA.

The crux is carrying the Megatron-Core flex/deepep all-to-all over AWS EFA — without patching Megatron-Core. NVIDIA DeepEP is built on NVSHMEM + InfiniBand verbs and does not run on EFA out of the box. Two EFA-native providers plug into the same flex/deepep path by shipping a top-level deep_ep module (installed into site-packages, so import deep_ep resolves to it):

1. UCCL (uccl-project/uccl) — an EFA-native deep_ep drop-in (UCCL-EP + GDRCopy). This is the default image (EP_BACKEND=uccl).
2. NVIDIA DeepEP, EFA-patched — DeepEP itself, rebuilt to run NVSHMEM over libfabric (host-proxy, IBGDA off) on EFA, via the patches in ../../micro-benchmarks/expert-parallelism/deepep-benchmark (vendored at deepep/). Built with EP_BACKEND=nvshmem.

So “stock DeepEP can’t run on EFA” is true only of the IB/IBGDA build — patched, it does. The qwen3-235b/ case uses this to compare NCCL all-to-all vs DeepEP+UCCL vs DeepEP+NVSHMEM head-to-head. Select the provider with the EP_BACKEND Docker build arg (one Dockerfile, two image tags); see Shared environment workflow.

Layout

The container environment (Dockerfile + its build/validation scripts) lives here at the library level and is shared by every model under it. Per-model recipes (checkpoint conversion, the SFT conf, deployment manifests, benchmarks) live in a model subdirectory.

megatron-bridge/                  # <library> — model-agnostic environment
├── Dockerfile                    # NGC NeMo base + EFA/GDRCopy + UCCL + deep_ep shadow
├── 1.build-and-push.sh           # build the shared env image and push the pinned tag to ECR
├── 2.sanity-singlenode.sh        # single-node 8-GPU deep_ep/EFA/EP smoke gate (run in the image)
├── convert-checkpoint.sh         # shared HF -> BF16 -> Megatron-Core conversion (parameterized per model)
├── test_megatron_bridge_uccl.py  # CI build smoke test for the shared image
├── kimi-k2/                      # <model> — Kimi K2 (384-expert): SFT + dispatcher A/B
│   ├── README.md
│   ├── conf/                     # SFT ConfigContainer (mounted into the image at runtime)
│   ├── kubernetes/               # PyTorchJob manifest + deploy guide
│   └── benchmarks/               # dispatcher A/B entrypoint + measured RESULTS.md
└── dsv3/                         # <model> — DeepSeek-V3 (256-expert): SFT + dispatcher A/B
    ├── README.md
    ├── conf/                     # same layout as kimi-k2/ (both models are structural siblings)
    ├── kubernetes/
    └── benchmarks/

The image is model-agnostic: SFT configs are not baked in. Each model mounts its own conf/ at /workspace/conf at runtime (e.g. via a ConfigMap — see the model’s kubernetes/README.md), so one image serves every model under this library.

Shared environment workflow

These two steps build and validate the shared image and apply to all models. Run them from this directory, then continue in the model subdirectory.

1. Build the environment image and push to ECR

1.build-and-push.sh builds Dockerfile, creates the ECR repository if needed (megatron-bridge-uccl), logs in, and pushes the pinned tag nemo-26.04.01-uccl-0dc87eb (no latest). The Dockerfile starts from nvcr.io/nvidia/nemo:26.04.01 (CUDA 13.1, PyTorch 2.11), which already ships Megatron-Bridge v0.4.2, Megatron-Core 0.17.1 (with the flex/deepep dispatcher), and TransformerEngine — these are not reinstalled. (The older 25.11.01 base shipped Megatron-Bridge 0.2.0, whose flex/deepep GPU allowlist rejects p6-b300; 0.4.0+ fixes it — see the Dockerfile header.) It then strips the IB fabric, lays down GDRCopy v2.5.2 + the EFA installer 1.48.0, and builds UCCL (pinned to commit 0dc87eb) / UCCL-EP for sm_103 (B300) plus the deep_ep shadow.

bash 1.build-and-push.sh
# Image: <account>.dkr.ecr.us-west-2.amazonaws.com/megatron-bridge-uccl:nemo-26.04.01-uccl-0dc87eb

For the DeepEP+NVSHMEM provider (used by the qwen3-235b/ 3-way comparison), build the same Dockerfile with EP_BACKEND=nvshmem — it skips UCCL and instead builds NVIDIA DeepEP (567632d) over NVSHMEM v3.7 (libfabric/EFA, IBGDA off) into /opt/venv:

EP_BACKEND=nvshmem bash 1.build-and-push.sh
# Image: <account>.dkr.ecr.us-west-2.amazonaws.com/megatron-bridge-uccl:nemo-26.04.01-deepep-nvshmem-567632d-cu13

2. Single-node sanity gate

Do not skip this. It is far cheaper to fail on 1 node than to burn 32 capacity-block nodes. 2.sanity-singlenode.sh runs a single-node, 8-GPU smoke test inside the image that confirms the UCCL deep_ep wrapper is active, EFA is present, and the Megatron-Core flex/deepep dispatcher config is wired.

# inside the container on one p6-b300.48xlarge node:
bash 2.sanity-singlenode.sh

Note (nemo:26.04.01 / Megatron-Core 0.17.1): Gates 1–4 (EFA device present, UCCL deep_ep active with Buffer, MCore flex/deepep config fields, NCCL all-reduce over EFA) pass. Gate 5 — a hand-rolled MoEFlexTokenDispatcher micro-step — is stale on Core 0.17.1: its standalone setup predates the ProcessGroupCollection API and raises a pg_collection error. This is not a UCCL/image fault — the real pretrain() path builds the process groups internally and dispatches correctly (the multi-node benchmark runs clean through MoEFlexTokenDispatcher(backend="deepep")). Treat the multi-node run / benchmarks as the authoritative end-to-end dispatch check until Gate 5 is ported to the 0.17.1 API.

Models

Both models provide the same two workloads — full-parameter SFT (conf/ + kubernetes/) and the UCCL-EP vs NCCL all-to-all dispatcher A/B (benchmarks/) — in a structurally identical directory layout.

To add a model: create megatron-bridge/<model>/ with its conf/, deployment manifests, and a model README (and a benchmarks/ entrypoint if you want the dispatcher A/B). Reuse the shared image from step 1 (mount the model’s conf at runtime) and the shared convert-checkpoint.sh — do not add a second Dockerfile.

Benchmark results & dispatcher recommendation

DeepSeek-V3, 32× p6-b300 / 256× B300 — UCCL-EP vs NCCL all-to-all

The headline measurement this environment was built for: swapping only the Megatron-Core MoE token dispatcher — NCCL all-to-all (baseline) vs UCCL’s EFA-native deep_ep (treatment) — on a live 32× p6-b300.48xlarge (256× B300) block. Model: the deepseek_v3 recipe — DeepSeek-V3 256-expert MoE (the architecture family Kimi-K2 belongs to, but not the literal 384-expert Kimi-K2 — see RESULTS.md), TP8/PP8/EP32/DP4, seq 4096, bf16, balanced routing. Everything else is held byte-identical across arms. Full methodology, caveats, and raw numbers: dsv3/benchmarks/RESULTS.md.

At the throughput-efficient operating point (micro-batch ≥ 4), UCCL deep_ep is ~36% faster than NCCL all-to-all, and the advantage holds under deployment-realistic 1F1B overlap. NCCL wins only at micro-batch 1 (64 tiny dispatches — UCCL-EP’s per-dispatch overhead unamortized), an operating point no throughput-tuned run uses.

Mean training-iteration time (lower = better) over 16 steady-state iters after warmup, 0 stalls, EFA-active on every rank. Work-equivalence (no token dropping) verified two ways: drop-free config + an iteration-1 loss match (deepep 11.897349 vs alltoall 11.897517). overlap=on is a separate within-regime A/B (VPP=2 + recompute off on both arms) — do not subtract its numbers against the overlap=off rows.

Qwen3 three-way comparison, 8 nodes — NCCL vs DeepEP+UCCL vs DeepEP+NVSHMEM

Adds the NVIDIA DeepEP+NVSHMEM-over-EFA arm and a second fabric (H100). Mean iter time at mb=4, overlap=off; Δ vs NCCL all-to-all (− = DeepEP faster). Full tables: qwen3-235b/benchmarks/RESULTS.md (B300), qwen3-30b/benchmarks/RESULTS.md (H100).

At this 8-node scale NCCL all-to-all is fastest almost everywhere (lone exception: B300 EP16, where DeepEP+NVSHMEM edges it by 5.2%). DeepEP+NVSHMEM is the stronger of the two DeepEP backends in most cells (both H100 points and B300 EP16), though UCCL edges it at B300 EP32 (mb=4); the DeepEP penalty shrinks as EP drops (EP32→EP16) and grows on the lower-bandwidth H100 fabric / smaller (30B) model.

Recommendation (scale-dependent)

The dispatcher choice on EFA depends primarily on scale (EP span / node count), and only secondarily on the instance/fabric:

• Small scale (≲8 nodes, up to EP32): default to NCCL all-to-all. Without IBGDA, DeepEP on EFA runs the NVSHMEM/UCCL host-proxy path and can’t use GPU-initiated RDMA; at small EP span that overhead dominates and NCCL’s dense all-to-all wins (Qwen3 tables above).
• Large scale (≳16–32 nodes, large EP span): prefer DeepEP. DeepEP does sparse dispatch (tokens go only to their routed experts), so its advantage grows with node count and overtakes the host-proxy penalty. Measured: UCCL −36% on DSV3 256× B300, and −34% at 32 nodes / −21% at 16 nodes on Kimi-K2 (kimi-k2/benchmarks/RESULTS.md) — the win itself scales with node count.
• DeepEP+NVSHMEM is validated at K2 scale (32 nodes) and competitive with UCCL. It is the best DeepEP backend at most 8-node points here (both H100 EPs and B300 EP16, where it beats NCCL; UCCL edges it only at B300 EP32 mb=4), and the 2026-07-14 32-node run on literal Kimi-K2 (256× B300, TP8/PP8/EP32) confirmed it at scale: 3.74 s/iter at mb=4+overlap — −35% vs the June NCCL baseline and on par with the June UCCL number (within cross-campaign noise; at mb=4 overlap=off it trails UCCL by ~10% — read the two DeepEP transports as competitive, not ranked; kimi-k2/benchmarks/RESULTS.md). Both DeepEP transports are now proven large-scale choices on EFA; note the NVSHMEM arm may need GDRCOPY_DEV=on (see kimi-k2/README.md).
• Fabric is second-order. DeepEP looked worst on H100 (p5.48xlarge, lower EFA bandwidth exposes the host-proxy cost) and better on B300; p5en/newer should help. But it still lost at B300 8-node/EP32 — so treat this as scale-primary, fabric-secondary.

All numbers use mock data + balanced routing; skewed / real expert routing is not yet tested and could shift the balance.

References

• NVIDIA Megatron-Bridge
• Megatron-Bridge docs
• UCCL project
• Sibling case: ../megatron-lm (EFA/GDRCopy Dockerfile + PyTorchJob template)