build(base): rebase the Baseten patches onto NVIDIA main (82894c4d)

[XiaohanZhangCMU]

Cherry-picked the ten Baseten patches from trainers-main onto NVIDIA main 82894c4d (2026-08-19), pinned rather than tracking the branch. History is [NVIDIA commits][10 Baseten commits], linear, no merge commits. Delta above the base: 13 files, +358/-25.

trainers-main was NVIDIA main at a2496aa1 (2026-07-09) plus those ten commits; NVIDIA has moved 370 commits since. basetenlabs/Megatron-LM:main is untouched and unused.

Replaces #32, which was based on the single MCore commit the Kimi-K3 bridge pins rather than on NVIDIA main.

Exactly what I ran

# remotes: origin = NVIDIA/Megatron-LM, baseten = basetenlabs/Megatron-LM
git fetch origin main
git fetch baseten trainers-main

# where our fork left NVIDIA main, and what is ours above it
MB=$(git merge-base baseten/trainers-main origin/main)   # a2496aa17562, 2026-07-09
git log --oneline $MB..baseten/trainers-main             # 10 commits, all Baseten-authored
git rev-list --count $MB..origin/main                    # 370 NVIDIA commits since

# replay our ten onto pinned NVIDIA main
git checkout -b mcore/baseten-on-nvidia-main-20260819 82894c4d5c30
git cherry-pick $MB..baseten/trainers-main

That last command stops here:

error: could not apply dcf0a0bc0... [baseten] fix(tensor-parallel): avoid int32 stride overflow
CONFLICT (content): Merge conflict in megatron/core/tensor_parallel/layers.py
CONFLICT (content): Merge conflict in tests/unit_tests/tensor_parallel/test_layers.py

Reproducible from a clean checkout with just those commands.

Why the conflicts exist

Seven of the ten applied automatically. The three that did not are all files upstream edited in those 370 commits:

git log --oneline $MB..origin/main -- megatron/core/tensor_parallel/layers.py            # 6
git log --oneline $MB..origin/main -- megatron/core/transformer/moe/token_dispatcher.py  # 5
git log --oneline $MB..origin/main -- pyproject.toml                                     # 12

And each of the three applies cleanly onto its own merge-base, which is what rules out a wrong base:

for c in dcf0a0bc0 922ebcbdb f14fdf17c; do
  git checkout -q --detach $MB && git cherry-pick -n $c && echo "$c clean onto merge-base"
  git cherry-pick --abort 2>/dev/null; git reset --hard -q $MB
done

So the conflicts are upstream having touched the same lines, not the base being wrong. The only conflict-free target is a2496aa1 - the July base this PR exists to leave.

How each was resolved

megatron/core/tensor_parallel/layers.py + tests/unit_tests/tensor_parallel/test_layers.py - upstream added grad_output = grad_output.to(ctx.input_dtype) ahead of the dim() > 2 dispatch our int32-stride patch replaces. Kept upstream's cast and our dim() == 3 fast path. The test conflict was additive on both sides with no overlapping names: union of upstream's three FP32-output tests and our non-contiguous-grad test.

megatron/core/transformer/moe/token_dispatcher.py - two additive blocks. Union: upstream's get_expert_zero_copy_buffers plus our _clear_forward_state, and our five = None resets in _HybridEPManager.

pyproject.toml - rewritten instead of replayed. The old patch relaxed ~=0.4.0 to >=0.4.0,<0.6 so 0.5.x was admitted; upstream now pins ==0.5.1, which does that but sits below what Kimi-K3 needs. Before 0.5.2, chunk_kda accepts A_log and dt_bias through **kwargs and silently discards them, so the KDA forget gate trains wrong with nothing raised. Now >=0.5.2,<0.6.

Checking the result

git merge-base --is-ancestor 82894c4d5c30 HEAD    # NVIDIA history untouched
git log --oneline 82894c4d5c30..HEAD              # exactly our ten
git log --merges --oneline 82894c4d5c30..HEAD     # empty
git diff --stat 82894c4d5c30..HEAD                # 13 files, +358/-25

Testing

Not run - the suite needs GPUs. Static only: no conflict markers, every changed file parses. The gate is the validation matrix, which needs to cover GLM-5.2 THD-CP, DSv4-Flash, MiMo and Qwen as well as Kimi-K3, since these patches serve all of them.