build(base): rebase the Baseten patches onto NVIDIA-NeMo main (e1967a9c)

[XiaohanZhangCMU]

Cherry-picked the Baseten patches from trainers-main onto NVIDIA-NeMo/Megatron-Bridge main e1967a9c (2026-08-19), pinned rather than tracking the branch. History is [NVIDIA commits][11 Baseten commits], linear, no merge commits. Delta above the base: 13 files, +315/-20.

The submodule points at a7234ad552fa, the new basetenlabs/Megatron-LM:trainers-main (basetenlabs/Megatron-LM#33, merged).

trainers-main was NVIDIA main at 7ee8b69d (2026-07-09) plus 17 commits; NVIDIA has moved 504 commits since. This PR no longer vendors Kimi-K3 - the new base already contains it, so there is nothing to cherry-pick for K3.

How to review 1600 files

Don't. The upstream range is upstream's. git diff e1967a9cb6d6..HEAD is our entire delta - 13 files - and that is the part worth reading.

Accounting for all 17 original commits

11 applied - Paras's GLM-5.2/HF compat, GLM DSA fused backend and the HybridStack recompute hook; Jack's FP8 context-parallel launches (#19); Shreya's six check_imports cold-cache commits; Ian's async_ckpt_use_cpu_shm.

3 were gitlink-only MCore bumps (f7be4c14, 57d5f291, 20fcf2ea), superseded by the new pin.

2 were merge commits, whose constituents applied individually.

1 skipped as obsolete - 1270ea0e "keep checkpoint save strategy compatible with upstream MCore". It existed as a try/except fallback for MCore versions without get_default_save_sharded_strategy. The new MCore defines it (serialization.py:58), and upstream's bridge now calls neither that helper nor get_default_load_sharded_strategy - 0 call sites, down from 3 on our old branch. Replaying it would add a dead shim plus an unused import that trips ruff F401.

Conflicts

Two, both trivial:

glm_moe_dsa/glm5_bridge.py - upstream added from typing import Any where our patch adds import os. Kept both.

3rdparty/Megatron-LM - the gitlink, twice. Resolved to the new MCore commit.

One defect caught during the replay

Ian's async_ckpt_use_cpu_shm hunk calls inspect.signature, and upstream's checkpointing.py does not import inspect - our old branch did, at module scope. Cherry-picking the hunk alone left the file using a name it never imports, so the first save_checkpoint with the feature enabled would die with NameError. import inspect is folded into that commit.

This is the second time that exact gap appeared in this file, so it is worth stating the rule: after resolving a conflict, check the whole file compiles against its imports, not just that the hunk looks right.

Verifying

git merge-base --is-ancestor e1967a9cb6d6 HEAD    # NVIDIA history untouched
git log --oneline e1967a9cb6d6..HEAD              # exactly our eleven
git log --merges --oneline e1967a9cb6d6..HEAD     # empty
git ls-files -s 3rdparty/Megatron-LM              # a7234ad552fa

Testing

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

[Copilot]

Pull request overview

Vendors NVIDIA’s upstream Kimi-K3 language-backbone implementation into megatron.bridge, including the Megatron↔HF conversion bridge, K3-specific layer/spec wiring, and MXFP4 quantization utilities needed by K3 checkpoints.

Changes:

• Add the Kimi-K3 model provider, layer spec, custom layers/ops, and pipeline stage-boundary packing helpers.
• Add a Kimi-K3 MegatronModelBridge with parameter mappings and MXFP4 expert dequant/requant handling.
• Extend MXFP4 E2M1 packed (de)quantization to support uint8 E8M0 scales, plus a unit test.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit_tests/models/test_quantization_utils.py	Adds coverage for MXFP4 E2M1 packed roundtrip when scales are uint8 E8M0.
src/megatron/bridge/models/kimi/kimi_k3_spec.py	Defines a K3-specific MCore layer spec (heterogeneous KDA/MLA + SiTU activation wiring).
src/megatron/bridge/models/kimi/kimi_k3_provider.py	Introduces KimiK3ModelProvider with K3-specific config knobs.
src/megatron/bridge/models/kimi/kimi_k3_pipeline.py	Implements AttnRes stage-boundary pack/unpack helpers for PP.
src/megatron/bridge/models/kimi/kimi_k3_ops.py	Adds K3 numerical ops (KDA wrapper, SiTU, Kimi RMSNorm, AttnRes aggregation).
src/megatron/bridge/models/kimi/kimi_k3_layers.py	Adds K3 attention/MoE/transformer-layer implementations.
src/megatron/bridge/models/kimi/kimi_k3_bridge.py	Adds HF↔Megatron conversion bridge and mapping registry for K3, including MXFP4 expert handling.
src/megatron/bridge/models/kimi/init.py	Exports K3 bridge/provider from the Kimi package.
src/megatron/bridge/models/conversion/quantization_utils.py	Adds uint8 E8M0 scale decode/encode support for MXFP4 packed weights.
src/megatron/bridge/models/init.py	Exposes K3 bridge/provider at the top-level models package.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[XiaohanZhangCMU]

/code-review

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (3)

src/megatron/bridge/models/kimi/init.py:17

• The new Kimi K3 re-export imports are missing # noqa: F401, which is used consistently in other model __init__.py files to avoid F401 (unused import) lint errors for symbols that are re-exported via __all__ (e.g., src/megatron/bridge/models/llama/__init__.py:15).

from megatron.bridge.models.kimi.kimi_bridge import KimiK2Bridge  # noqa: F401
from megatron.bridge.models.kimi.kimi_k3_bridge import KimiK3Bridge
from megatron.bridge.models.kimi.kimi_k3_provider import KimiK3ModelProvider

src/megatron/bridge/models/kimi/init.py:24

• The PR description says the upstream squash SHA (2f0f8c80) is recorded in models/kimi/__init__.py, but it is not present in this module. If provenance is meant to survive outside the PR description, consider recording it (e.g., as an exported constant) here.

__all__ = [
    "KimiK2Bridge",
    "KimiK3Bridge",
    "KimiK3ModelProvider",
]

src/megatron/bridge/models/kimi/kimi_k3_provider.py:32

• The PR description mentions the upstream squash SHA (2f0f8c80) is recorded on KimiK3ModelProvider, but the class currently has no provenance marker. If this is required for vendoring/auditing, add a ClassVar so it’s available without becoming a dataclass field.

@dataclass
class KimiK3ModelProvider(MLAModelProvider):
    """Megatron configuration and provider for Kimi K3."""

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (1)

src/megatron/bridge/models/kimi/init.py:23

• PR description says the upstream squash SHA 2f0f8c80 is recorded in models/kimi/__init__.py and on KimiK3ModelProvider, but this diff only adds K3 imports/exports here and the provider has no provenance field/comment. Please either add the provenance marker to the code (as described) or update the PR description so it matches what’s actually in this vendored layer.

from megatron.bridge.models.kimi.kimi_bridge import KimiK2Bridge  # noqa: F401
from megatron.bridge.models.kimi.kimi_k3_bridge import KimiK3Bridge
from megatron.bridge.models.kimi.kimi_k3_provider import KimiK3ModelProvider


__all__ = [
    "KimiK2Bridge",
    "KimiK3Bridge",
    "KimiK3ModelProvider",

[JackRao123]

Wait, actually, I think since you're taking from the main branch instead of the dev branch, we should just rebase up to that point and then apply the baseten patches on top of that.

because currently trainers-main history looks like this:

[nvidia main, time 0 to t][our own patches]

so like if we want to take stuff from main, we should rebase and have

[nvidia main, time 0 to t][nvidia main, from t to t'][our own patches]

at least, this is my understanding of what it currently looks like - wdyt?

cause this way we get the base upon which yaoyu-33's kimi implementation might have depended on, and also keeps the lineage clean for when we want to take MORE commits from main because of future models like kimi k4, k5, ... etc

[XiaohanZhangCMU]

@JackRao123 done, and you were more right than either of us thought.

This is now a base bump: trainers-main to NVIDIA main e89eef9d (2026-08-19), our patches replayed on top, zero cherry-picks - because the new base already contains Kimi-K3. The six vendored commits are gone.

Two things fell out that I did not expect:

Upstream had already taken our FP32 patch. NVIDIA main carries _apply_mixed_precision_wrapper in models/model_provider.py character-identical to #27's version, our comment wording included, and MiniMax-M3 uses the same _keep_in_float32_parameter_names marker. So #27 shrank from 7 files to 2 - the fla floor and a test for the helper, since upstream added it without one. #36 went from 13 files to 7 for the same reason (upstream now has iter_finalized, the K3 docs, the E8M0 decode and its own toy-model generator).

The rebase surfaced a real incompatibility the cherry-pick was hiding. Upstream MCore now pins flash-linear-attention==0.5.1, which is below K3's 0.5.2 floor - under 0.5.2 chunk_kda takes A_log/dt_bias through **kwargs and silently discards them, training a different KDA gate. Resolved to >=0.5.2,<0.6 in basetenlabs/Megatron-LM#32. That conflict simply did not exist on the Jul 9 base.

Scope, so you can judge the risk: 496 upstream bridge commits and 280 MCore commits adopted; our MCore was diverged from what the new bridge expects (ahead 10, behind 280), which is why the paired MCore PR exists and has to merge first. Our whole delta above the new base is 13 files, +329/-21 - git diff e89eef9d..HEAD - and that is the only part worth reading in a 1600-file PR.

Four conflicts total across both repos, all unions or supersets, all hand-resolved. Worth flagging: on a first pass I let -X theirs resolve them and it exited zero while silently deleting Ian's async_ckpt_use_cpu_shm from CheckpointConfig and one of two grad_output.to(ctx.input_dtype) casts in MCore. Both are restored and verified by grepping the specific upstream lines. Please do not trust an automated resolution on this one.

Nothing is GPU-tested yet. The matrix needs to cover GLM-5.2 THD-CP, DSv4-Flash, MiMo and Qwen as well as K3, since these patches serve all of them - your pad_between_seqs and DSA fixes among them. Kimi-K3's trainer PRs wait on this, then repin, re-run the 131K validation and regenerate the report before merging.

[XiaohanZhangCMU]

Automated review of the replay: three real defects fixed, the rest triaged

I ran a review over e89eef9d..HEAD (the 13-file delta, not the upstream range) and verified every finding against the code rather than taking it on trust. Fixed in a030aa56:

import inspect was lost. trainers-main imports it at module scope; upstream does not need it. My resolution of the conflicting import block took a "superset" of the hunk and missed the file. Consequence: the first save_checkpoint with async_ckpt_use_cpu_shm=True dies with NameError - the very feature that hunk adds, and nothing in the delta exercises it.

get_default_load_sharded_strategy was imported with no call site. Upstream refactored away both places trainers-main called it (lines 2082 and 3294 there), so keeping the import only trips ruff F401, which this repo selects.

The get_default_save_sharded_strategy compat shim is unreachable. Nothing in this repo or in trainers calls it, and MCore now defines the helper as (backend="torch_dist") while the fallback is (ckpt_format). It covered MCore versions without the helper; the new pin has it. Removed - @paras-stefanopoulos flagging since it was your patch, say the word if you want it kept.

All three come from the same mistake, and it is the same one the -X theirs experiment made: resolving a conflict by looking only at the conflicting lines.

Verified as pre-existing, not from this PR

glm5_bridge.py raw-config read. _load_raw_hf_config(getattr(hf_config, "_name_or_path", "")) with an empty value makes os.path.join("", "config.json") probe the working directory, and otherwise hf_hub_download(repo_id="") raises. The repo's own test_glm5_bridge.py::_provider_from_hf_config builds SimpleNamespace(**config) with no _name_or_path and does reach that line. This is byte-identical on trainers-main (line 99 there, 112 here), so the base bump does not introduce it - but it is very likely broken today and nobody would know, because this fork runs no unit-test job on PRs. Worth a separate fix; I did not touch it here to keep the bump honest.

Same category, lower severity, all pre-existing: qk_pos_emb_head_dim set from raw config then overwritten 30 lines later; qk_head_dim duplicating the generic mapping in model_bridge.py; an unguarded raw_config["qk_nope_head_dim"] KeyError; a dsa_indexer_skip_topk_offset default change that the parametrization ({}, 1, 0) still asserts as 0 while the fixture sets first_k_dense_replace: 3; and dsa_kernel_backend = "cudnn" hard-coded in the bridge where upstream sets it per recipe.

Latent, worth knowing before this fork enables its CI

Upstream's vendored cicd-main.yml:381 runs "Record baseline MCore revision" ungated, and validate_mcore_revision.sh requires the submodule sha to be an ancestor of NVIDIA main/dev or a valid refs/pull/*/merge. Our pin is a basetenlabs fork commit, so it satisfies none of those and the step exits 1. install-test.yml:163 gates the same way, and .gitmodules already points at the fork (also pre-existing on trainers-main).

It is not firing today - cicd-main does not run on this fork's PRs; #34's checks are check, collect, link-check, secrets-detector - so this is a note for whoever turns it on, e.g. after #33 unblocks Lint.

Not fixed, with reasons

uv.lock is not regenerated. It still records flash-linear-attention==0.5.1 from megatron-core's metadata while the new MCore pin declares >=0.5.2,<0.6, so uv lock --check fails. I have not regenerated it because it needs the merged MCore commit to resolve against - it should be regenerated in the same change that refreshes the gitlink after basetenlabs/Megatron-LM#32 lands.

transformers_compat.py:53 - seen.add(stem) before the exists() check means non-resolving stems reach the caller, and get_cached_module_file's local-directory branch has no existence guard, so a transitive import naming a package directory raises FileNotFoundError where upstream only followed direct imports. Real, pre-existing, @ShreyaRavi's area.

quantization_utils.py:84 - the repeat_interleave vectorization materializes a full M×N fp32 scale tensor, taking peak transient from 2 fp32 copies to 3 on a path that converts multi-hundred-GB checkpoints. Pre-existing (Jack's #19), and a real trade worth revisiting, but not in a base bump.

recompute.py:69 - the HybridStack import sits above the try: that the function's own docstring says is there so PEFT degrades to a warning rather than dying. The module exists in the pinned MCore, so it is latent; the submodule is a moving pin, so it should move inside the try.