Add Megatron-Bridge Kimi-K2 + DeepSeek-V3 test cases: full-parameter SFT + UCCL-EP-over-EFA dispatcher A/B (A/B validated on 256× B300)

[KeitaW]

What this PR does

Adds a reproducible test case for full-parameter SFT of Kimi K2 (Moonshot AI's 1.04T-param DeepSeek-V3-family MoE) on Amazon EKS with NVIDIA Megatron-Bridge, using UCCL-EP's EFA-native deep_ep drop-in for the expert-parallel all-to-all. This replaces NVSHMEM/IB-verbs-bound NVIDIA DeepEP (which cannot run on EFA) without patching Megatron-Core — UCCL ships a top-level deep_ep shadow module so Megatron-Core's flex/deepep dispatcher transparently routes over EFA via UCCL + GDRCopy.

It also lands the first end-to-end DeepEP-on-EFA-vs-NCCL-all-to-all training A/B on B300 (see results below) — no such measurement existed in the literature — measured on two substrates: the DeepSeek-V3 256-expert recipe (2026-06-01) and the literal 384-expert Kimi-K2 built via AutoBridge (2026-06-04), with per-iteration loss-curve work-equivalence on both.

Structure

Laid out framework/library/model so future Megatron-Bridge models can reuse the environment:

• 3.test_cases/megatron/megatron-bridge/ — model-agnostic shared env (Dockerfile, build + single-node sanity scripts, shared convert-checkpoint.sh, CI build test).
• 3.test_cases/megatron/megatron-bridge/kimi-k2/ — Kimi K2 (384-expert): full-parameter SFT recipe (conf + K8s manifests) and the dispatcher A/B (benchmarks/).
• 3.test_cases/megatron/megatron-bridge/dsv3/ — DeepSeek-V3 (256-expert): the same two workloads in a structurally identical layout — a full-parameter SFT recipe and the measured dispatcher A/B below.

The image bakes no model config; each model mounts its conf/ at /workspace/conf at runtime via a ConfigMap, so one image serves every model under the library.

Components

• Framework: NVIDIA Megatron-Bridge — NGC 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. (The older 25.11.01 base shipped Bridge 0.2.0, whose flex/deepep GPU allowlist rejects p6-b300; 0.4.0+ fixes it.)
• Models/Workloads: Kimi-K2 (1.04T MoE, 384 experts, MLA) and DeepSeek-V3 (671B MoE, 256 experts, MLA), each with full-parameter SFT and the UCCL-EP-over-EFA dispatcher A/B, in a structurally identical layout
• Hardware target: 32× p6-b300.48xlarge (256× B300, sm_103), 16× EFA/node
• Orchestration: Kubernetes / EKS (Kubeflow PyTorchJob, etcd rendezvous, optional KAI gang scheduling)
• Transport: UCCL-EP over AWS EFA (deep_ep shadow, pinned commit 0dc87eb; IB verbs/NVSHMEM removed)

Benchmark result — UCCL-EP vs NCCL all-to-all (256× B300, measured 2026-06-01)

A live 32× p6-b300.48xlarge (256× B300) A/B that swaps only the Megatron-Core MoE token dispatcher — everything else (model, data, parallelism TP8/PP8/EP32/DP4, seq 4096, GBS 256, bf16, balanced routing, image) held byte-identical. Substrate is 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 for the substrate caveat).

The dispatcher winner crosses over with micro-batch granularity:

micro-batch	overlap	NCCL all-to-all	UCCL deep_ep	dispatcher delta
1	off	12.54 s	14.12 s	NCCL +12.6% faster
4	off	9.77 s	6.26 s	UCCL −36.0% faster
4	on	5.98 s	3.84 s	UCCL −35.8% faster

Mean training-iteration time (lower = better) over 16 steady-state iters after warmup, 0 stalls, EFA active on every rank. At the throughput-efficient operating point (micro-batch ≥ 4) UCCL deep_ep is ~36% faster, and the advantage holds under deployment-realistic 1F1B overlap (contrary to the usual expectation that overlap compresses the dispatcher delta toward parity). 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.

Work-equivalence verified (no token dropping): both arms' config dumps are byte-identical except the dispatcher selector (no moe_expert_capacity_factor/moe_token_drop_policy set → drop-free by construction), and an iteration-1 loss match (deepep 11.897349 vs alltoall 11.897517, rel diff 1.4e-5 — bf16 round-off, identical num_tokens). So the −36% is a genuine equal-work speedup.

Benchmark result — literal Kimi-K2, 384 experts (256× B300, measured 2026-06-04)

The same A/B re-run on the literal Kimi-K2 (61 layers, 384 routed experts, top-8, MLA, no MTP), with the provider built by AutoBridge.from_hf_pretrained("moonshotai/Kimi-K2-Base", trust_remote_code=True).to_megatron_provider() — same image, same TP8/PP8/EP32/DP4 layout, mock data, 24 iters/run with log_interval=1 (first 4 dropped). See kimi-k2/benchmarks/RESULTS.md.

micro-batch	overlap	NCCL all-to-all	UCCL deep_ep	dispatcher delta
1	off	12.88 s	14.98 s	NCCL +16.3% faster
4	off	9.30 s	6.04 s	UCCL −35.1% faster
4	on	5.79 s	3.83 s	UCCL −33.8% faster

The DSV3 findings transfer to the real K2 architecture: UCCL deepep −34/−35% at mb=4 in both overlap regimes; NCCL faster only at mb=1 no-overlap. Work-equivalence held across the full run: per-iteration lm loss curves for deepep vs alltoall match within ≤4.1e-4 relative over iters 1–10 in all three cell pairs (bf16 round-off accumulation, no systematic offset). The mb=1+overlap pair was not measured — the Capacity Block expired mid-campaign after 6 of 8 cells (partial logs preserved); a 16-node TP8/PP4/EP32/DP2 appendix table is in the results doc (within-cell deltas only, not comparable across layouts).

Testing performed

• Local CI-equivalent lint: shell bash -n + copyright + set -e, Python py_compile + copyright, Dockerfile version pins (EFA 1.48.0, GDRCopy v2.5.2, CUDA 13.1) and pinned FROM, no :latest/binaries/experimental.
• Both PyTorchJob manifests validated as multi-doc YAML; ConfigMap volume/mount placement checked.
• test_megatron_bridge_uccl.py builds the image and asserts import deep_ep resolves to the UCCL wrapper (deep_ep.Buffer present).
• Single-node 8-GPU sanity gate (2.sanity-singlenode.sh): Gates 1–4 pass (EFA device present, UCCL deep_ep active with Buffer, MCore flex/deepep config fields, NCCL all-reduce over EFA). Gate 5's hand-rolled MoEFlexTokenDispatcher micro-step is stale on Core 0.17.1 (predates the ProcessGroupCollection API) — not a UCCL/image fault; the real pretrain() path builds the process groups internally and the multi-node benchmark runs clean through MoEFlexTokenDispatcher(backend="deepep"), which is the authoritative end-to-end dispatch check.
• Multi-node validation: the 256× B300 dispatcher A/B above ran end-to-end on both arms; every rank logged NET/OFI Selected provider is efa, fabric is efa-direct (found 16 nics) (true EFA RDMA, no socket fallback) and the UCCL arm logged the UCCL-EP EFA proxy active.

Scope note. Both the DeepSeek-V3 256-expert and the literal 384-expert Kimi-K2 dispatcher A/Bs are validated at 256× B300 (tables above). The K2 provider comes from AutoBridge (no hand-derived routing overrides needed — K2's HF config carries n_group=1/topk_group=1 and num_nextn_predict_layers=0, and the AutoBridge-built provider preserves them). Benchmarks use random init + mock data and measure throughput + dispatcher work-equivalence; real-data SFT convergence is not benchmarked here. The K2 HF→MCore checkpoint conversion path was exercised end-to-end (1.6 TiB MCore checkpoint produced). Both models now also ship a full-parameter SFT recipe (conf/ + K8s manifests) and share one model-agnostic convert-checkpoint.sh; these SFT recipes — including the new DeepSeek-V3 one — are scaffolded and unrun (TODOs flagged inline; DeepSeek-V3 keeps its 1 MTP layer where K2 has none), distinct from the measured dispatcher A/Bs above. The shared convert-checkpoint.sh runs the AutoBridge import_ckpt weight-conversion path, which is exercised for K2 but flagged unverified for DeepSeek-V3 on this image.

Checklist

• Dockerfile pins versions (EFA, GDRCopy, CUDA, UCCL commit) and pinned FROM
• Scripts have copyright headers
• README documents prerequisites and usage
• No large binaries committed