feat(megatron-bridge): add NeMo 26.08 Kimi K2 EP backend qualification

[KeitaW]

Purpose

Update the Megatron Bridge Kimi K2 test case to NeMo release 26.08 and add a reproducible four-arm expert-parallel qualification for NCCL all-to-all, UCCL, DeepEP v1 over NVSHMEM/EFA, and DeepEP v2 over NCCL GIN/GDA.

Relates to #1116 and #1189.

Changes

• Pin the NeMo release 26.08 software and communication stack, build one immutable image per arm, and record image and source manifests.
• Add the Megatron Core deepep_v2 backend using ElasticBuffer, including autograd support, eager communicator initialization, pipeline overlap lifetime handling, and the validated DeepEP v2 patches.
• Add named-node Amazon EKS campaign orchestration with capacity and ownership checks, fail-fast handling, route tracing, runtime identity gates, durable artifact harvest, checksum custody, and scoped teardown.
• Add qualification, headline, and optimized benchmark matrices for the four EP arms.
• Add NCCL-derived correctness gates for loss, full-precision gradient norm, sampled optimizer update L2 norm, route identity, finite metrics, and dropped selections.
• Add deterministic analysis tooling and checked-in JSON, Markdown, PNG, and SVG results for the accepted NeMo release 26.08 campaigns.

Test Plan

Environment:

• AWS service: Amazon EKS in ap-south-1
• Instance type: p6-b200.48xlarge
• Scale: 32 nodes with 8 B200 GPUs per node, 256 B200 GPUs total
• Parallelism: TP 8 ranks, PP 8 stages, EP 32 ranks, ETP 1 rank, DP 4 replicas
• Workload: Kimi K2 architecture with random initialization and mock data, BF16 precision, global batch 256 samples, microbatch 4 samples, and sequence length 4,096 tokens per sample
• Scored jobs: 40 optimizer iterations per job, 8 discarded warmup iterations, 32 scored steady iterations, and 3 fresh job starts per arm
• Optimizer control: constant learning rate 0.000005 dimensionless and random seed 1,234 dimensionless

Local test commands:

python3 -m pytest -q 3.test_cases/megatron/megatron-bridge/kimi-k2/tests/test_parse_runs.py
git diff --name-only -z origin/main...HEAD -- '*.py' | xargs -0 -r python3 -m py_compile
git diff --name-only -z origin/main...HEAD -- '*.py' | xargs -0 -r ruff check

Real-GPU qualification:

• Run the four-arm primary scored cell as 12 independent jobs on the same 32-node set.
• Run Gate A for the synchronous and asynchronous ElasticBuffer contract on 16 ranks.
• Run Gate B router replay directly against NCCL on 2 ranks and 16 ranks.
• Run Gate C for pipeline and overlap lifetime on 32 nodes and 256 GPUs.
• Run Gate D short Kimi K2 training for all 4 arms on 32 nodes and 256 GPUs, plus an independent NCCL self-repeat used to derive fixed BF16 output envelopes.
• Verify every harvested artifact against the campaign checksum manifest, then verify scoped Kubernetes teardown after releasing the campaign lease.

Test Results

All 12 primary scored jobs passed without retries. Each value is the median across 3 fresh job-start medians. Each job discarded its first 8 optimizer iterations and scored the remaining 32 optimizer iterations.

Arm	Median iteration time, ms	Median throughput, tokens/s	Median performance, TFLOPS/GPU	Run-to-run CV, %	Paired time reduction vs NCCL, mean [95% CI]
nccl-alltoall	7,518.80	139,461	106.513	0.048	Reference
uccl	6,215.65	168,699	128.228	0.226	17.22% [17.02%, 17.37%]
deepep-v1-nvshmem	6,813.85	153,889	117.306	0.193	9.43% [9.31%, 9.66%]
deepep-v2-gin-gda	5,861.70	178,886	135.666	0.116	21.98% [21.90%, 22.04%]

The interval uses 10,000 paired repeat-level bootstrap resamples with a 1,234 dimensionless seed and a 95% percentile interval. DeepEP v2 reduced mean iteration time by 5.75% relative to UCCL, with a 95% CI from 5.48% to 6.05%, in this B200, microbatch 4 samples, overlap-off cell.

All 12 jobs recorded 40 of 40 loss rows, a constant learning rate, finite loss and gradient norm, 0 skipped iterations, 0 NaN iterations, and 0 dropped token selections. Every non-NCCL arm matched the same-repeat NCCL route hash on all 32 ranks. Maximum absolute loss deltas were 0.000223 dimensionless for UCCL, 0.000125 dimensionless for DeepEP v1, and 0.000154 dimensionless for DeepEP v2, all below the fixed NCCL BF16 self-repeat envelope of 0.001114845276 dimensionless.

The short training gate also passed full-precision gradient-norm and sampled update-L2 checks for all 4 arms. DeepEP v2 runtime manifests prove _DeepepV2Manager, ElasticBuffer, a 32-rank EP group, NCCL_GIN_TYPE=5, the Libfabric_GDAKI GIN plugin version 14, GPU Direct RDMA, and all 8 EFA HCAs per node.

Absolute curves are faceted into one column per EP arm because the trajectories are nearly coincident at this scale. Each metric uses a shared y-axis across arm columns; repeat markers are staggered, but metric values and optimizer iterations are not offset.

The requested 256-GPU B300 headline remains NOT_RUN_INSUFFICIENT_CAPACITY: 0 B300 nodes were available and 32 B300 nodes were required. No B200 measurement is used as a B300 result.

The accepted scored artifact root contains 8,213 checksum-verified files. Post-run verification found the owned namespace absent, 0 owned pods, 0 requested GPUs, and an empty campaign lease holder. Capacity Block, EC2 instances, EKS nodes, and foreign workload objects were preserved.

Full methodology, per-repeat data, output curves, runtime proof, provenance hashes, exclusions, and teardown evidence are in RESULTS-26.08.md.

Directory Structure

3.test_cases/megatron/megatron-bridge/
├── bench/                       # Arm definitions, campaign orchestration, parsing, and telemetry
├── docker/                      # Version-pinned installers, patches, manifests, and image verification
├── kimi-k2/
│   ├── benchmarks/              # Training entrypoint, matrices, analysis, figures, and results
│   └── tests/                   # Distributed contract, replay, scale, and short-training gates
├── Dockerfile                   # Common NeMo release 26.08 image
├── run-ab-rawpods.sh            # Named-node raw-pod runner
└── versions.lock.yaml           # Immutable dependency and image inputs

Checklist

• I have read the contributing guidelines.
• I am working against the latest main branch.
• I have searched existing open and recently merged PRs to confirm this is not a duplicate.
• The contribution is self-contained with documentation and scripts.
• External dependencies are pinned to a specific version or tag, with no latest tag.
• The README includes the NeMo release 26.08 entry point and links to implementation and result details.
• The updated test case follows the expected directory structure.