Repository navigation
[OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter - #2748
shengliangxu wants to merge 17 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| Linked Issues check |  |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
| Security Anti-Patterns | |
No listed security anti-pattern was introduced. The Kimi adapters define trust_remote_code: bool = False, reject checkpoint loading when it is false, and pass the caller-controlled value to the dyna… |
| Title check | |
The title accurately identifies the Kimi-K2.5 expert-parallel adapter, which is a major part of the changeset. It does not mention the additional Kimi-VL adapter, but the title remains clear and relat… |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The Kimi adapters define
trust_remote_code: bool = False, reject checkpoint loading when it is false, and pass the caller-controlled value to the dynamic loader. No addedtorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval,exec, or# nosecusage appears in the Python diff. No dependency-file changes were made.
Full details: Title check
-
Explanation
The title accurately identifies the Kimi-K2.5 expert-parallel adapter, which is a major part of the changeset. It does not mention the additional Kimi-VL adapter, but the title remains clear and related.
Finishing Touches
-
Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
Fix failing CI checks
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
95832dd to
e243f04
Compare
f73bc4f to
0098f3a
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-11a-onfly-decompress #2748 +/- ##
=========================================================================
+ Coverage 62.46% 62.52% +0.05%
=========================================================================
Files 658 665 +7
Lines 74106 74227 +121
=========================================================================
+ Hits 46294 46407 +113
- Misses 27812 27820 +8
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/ok to test 0098f3a |
e243f04 to
d6afd32
Compare
0098f3a to
3d6fb9d
Compare
|
/ok to test 3d6fb9d |
shengliangxu
changed the title
d6afd32 to
67355f4
Compare
1c16b42 to
ecf9d71
Compare
67355f4 to
b162532
Compare
|
/ok to test ecf9d71 |
shengliangxu
changed the title
ecf9d71 to
d525e64
Compare
shengliangxu
changed the title
|
/ok to test d525e64 |
269d3d2 to
dddc2ee
Compare
d525e64 to
0c39a27
Compare
|
/ok to test 0c39a27 |
dddc2ee to
c668ec2
Compare
|
/ok to test 27961b6 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| make_importable() | ||
| if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None: | ||
| return | ||
| if not trust_remote_code: | ||
| return |
There was a problem hiding this comment.
[SUGGESTION] make_importable() runs ahead of all three guards, so the two process-wide installs it performs (install_compat_import_shims, and the get_cached_module_file wrap in stage_full_remote_closure) happen even for a config of transformers' own port, a call with no ckpt_path, and a load without trust_remote_code — the three cases the docstring describes as having "nothing of the checkpoint's to build", and the three where prepare refuses before touching anything.
Nothing breaks today (the shims only re-add names transformers dropped, and the staging wrap is a no-op for a non-remote-code load), but it is the one place in these adapters where a mutation precedes the check that says it is not needed, and this PR's whole premise is that the pre-decision step leaves a process that will not run expert-parallel untouched. Moving the call below the guards keeps the code matching the docstring:
if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None:
return
if not trust_remote_code:
return
make_importable()
get_class_from_dynamic_module(...)All three early-return cases are ones where the checkpoint's own code is never imported, so nothing downstream needs the shims there.
There was a problem hiding this comment.
Taken in fab6ddf. make_buildable now returns before installing anything when there is nothing of the checkpoint's to build: a config of transformers' own port, no checkpoint path, or code the load may not run. New test test_the_build_step_installs_nothing_when_there_is_nothing_to_build[no-checkpoint|untrusted] fails with the shims first.
| quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} | ||
| if quantization_config.get("quant_method") == "compressed-tensors": | ||
| _onfly_decompress.install_onfly_decompress(hf_config) |
There was a problem hiding this comment.
[SUGGESTION] "compressed-tensors" is now stated twice and the two must agree: here, as the condition that actually installs the decompressor, and in specs.py as decompressed_formats=("compressed-tensors",), which _ep_load_refusal reads before prepare runs to decide the checkpoint is one the expert-parallel load takes. If they ever drift, the failure is silent in one direction (the decision plans expert-parallel for a format prepare then leaves compressed, and the load fails later) and invisible in the other (a format prepare would decompress is planned without expert parallelism).
The new test already compensates by deriving its expectation from get_spec("kimi_k25").ep_spec.decompressed_formats rather than from the literal — worth making the code share the same single definition, since specs.py already imports from this module:
# ep.py
DECOMPRESSED_FORMATS = ("compressed-tensors",)
...
if quantization_config.get("quant_method") in DECOMPRESSED_FORMATS:
_onfly_decompress.install_onfly_decompress(hf_config)
# specs.py
from .ep import DECOMPRESSED_FORMATS, make_buildable, prepare
...
decompressed_formats=DECOMPRESSED_FORMATS,There was a problem hiding this comment.
Taken in 777ac2a. kimi_k25.ep.DECOMPRESSED_FORMATS is the one definition: prepare installs the decompressor for those formats, and the spec's decompressed_formats is that tuple.
|
/ok to test 777ac2a |
|
/claude review |
| for name, module in list(sys.modules.items()): | ||
| if name.endswith(f".{modeling}") and module is not None: | ||
| _shim_tie_weights(module, vlm_class) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This hand-rolled sys.modules scan drops the transformers_modules. prefix guard that _remote_code._apply exists to apply — and the hazard it guards is named for this model type.
_remote_code.py:113-128:
def _apply(suffix: str, patch: Callable[[ModuleType], bool]) -> dict[str, bool]:
"""...
Only remote code: transformers' own port of a model can share the remote module's name, as
transformers.models.kimi_k25.modeling_kimi_k25 does from 5.13.
"""
remote = f"{dmu.TRANSFORMERS_DYNAMIC_MODULE_NAME}."
... if module is not None and name.startswith(remote) and name.endswith(suffix)transformers.models.kimi_k25.modeling_kimi_k25 ends with .modeling_kimi_k25, so on transformers ≥ 5.13 this loop reaches transformers' own module whenever it happens to be in sys.modules — which the early type(hf_config).__module__.startswith("transformers.") return does not prevent, since that only checks the config this load was handed, not what else the process has imported (a prior load, or anything that touched transformers.models.kimi_k25).
Why it matters: _shim_tie_weights is a no-op only as long as the port's KimiK25ForConditionalGeneration has no tie_weights in its own vars(). If it does — now or in a later release — make_buildable rebinds it to a wrapper that silently swallows recompute_mapping=False/missing_keys, permanently, for every later load of transformers' native Kimi-K2.5 in that process. That is exactly the "changes nothing the model computes or loads" guarantee this function's docstring makes, broken outside the adapter's own model. patch_language_model_to_native doesn't have the exposure because it goes through keep_patched → _apply.
Fix — reuse the guarded helper instead of re-deriving the loop (one import plus one line, and it drops the function-local import sys):
from ._remote_code import _apply, keep_patched, stage_full_remote_closure
...
_apply(f".{modeling}", lambda module: _shim_tie_weights(module, vlm_class))A regression test is cheap here: register a transformers.models.kimi_k25.modeling_kimi_k25 stand-in in sys.modules carrying a KimiK25ForConditionalGeneration with its own zero-arg tie_weights, run kimi_k25.make_buildable(...), and assert the attribute is unchanged — the mirror of test_the_swap_leaves_a_sibling_remote_module_alone, which already pins the same scoping property for the swap.
There was a problem hiding this comment.
Taken in f3df820. make_buildable now applies the tie_weights shim through _remote_code._apply, which reaches remote modules only, as the swap does. New test test_the_build_step_leaves_transformers_own_port_alone registers a transformers.models.kimi_k25.modeling_kimi_k25 stand-in with its own zero-argument tie_weights and checks that the step leaves it alone while it adapts the remote one. It fails with the unscoped scan. This also covers the summary comment's one action item.
|
Claude review — Recommendation: comment. One scoping issue in the new build step; the Needs action:
No action needed:
🤖 Generated with Claude Code |
|
/ok to test f3df820 |
|
/claude review |
| prepared_ep_plan: dict[str, str] | None = None | ||
| """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` | ||
| supplies. | ||
|
|
||
| Stated, not applied: ``prepare`` still attaches it. A run that rules expert parallelism out | ||
| skips ``prepare``, which patches the model's code, yet must still tell an expert-parallel MoE | ||
| it would hold whole on every rank from a dense model. It finds the routed experts through this | ||
| plan, in the model as the checkpoint's own code builds it: ``prepare`` replaces the experts in | ||
| place, so the plan names their containers there too.""" | ||
| Deciding the plan attaches it with no ``prepare`` run, so the decision counts the routed experts | ||
| in the model the checkpoint's own code builds. A run that rules expert parallelism out finds |
There was a problem hiding this comment.
[SUGGESTION] This change makes prepared_ep_plan load-bearing in a way nothing enforces.
Before this PR, the plan decision ran prepare, so a spec that declared prepare got its plan whether or not it also declared prepared_ep_plan — the field only mattered for the narrower, loud EP-degree-1 refusal path. Now resolve_distribution_plan calls apply_distributed_support(..., prepare=False), and moe.py:290 attaches ep.base_model_ep_plan or ep.prepared_ep_plan only. A future spec that declares prepare and neither plan degrades silently: _decide_distribution reads base_model_ep_plan as {}, _routed_expert_modules finds no experts, num_experts is 0, _choose_ep_degree returns 1 — and the _routed_expert_keys/_spec_ep_plan fallback at model_load_utils.py:1704 is gated on ep_size == 1, so with the degree left at auto it is never reached. The MoE is then planned FSDP2 or whole with no refusal at all, which is the opposite of the "fail loudly rather than silently ignore a route" rule in AGENTS.md.
Both current specs declare it, so there is no bug today — but the invariant is only stated in prose. tests/unit/torch/models/test_ep_specs.py:82 already walks every EP spec for exactly this kind of claim; one extra assertion there would pin it:
@pytest.mark.parametrize("model_type", EP_MODEL_TYPES)
def test_a_spec_whose_prepare_supplies_its_plan_states_that_plan(model_type):
"""Deciding the plan never runs prepare, so a plan only prepare attaches must be stated:
unstated, the decision reads the MoE as dense and plans it without expert parallelism, with
no refusal."""
ep = get_spec(model_type).ep_spec
assert not ep.prepare or ep.base_model_ep_plan or ep.prepared_ep_plan, (
f"{model_type} supplies its expert-parallel plan from prepare alone, which the plan "
"decision does not run: state it as prepared_ep_plan"
)Alternatively raise it from register() so the failure lands at import rather than in a test.
There was a problem hiding this comment.
Taken in 7a1f1cf: test_a_spec_with_a_prepare_states_its_plan in tests/unit/torch/models/test_ep_specs.py asserts that every EP spec with a prepare states base_model_ep_plan or prepared_ep_plan. It fails for both Kimi specs with prepared_ep_plan dropped. I kept it a test, beside the other spec invariants that file pins, rather than a check in register(), which would add runtime behaviour no current spec needs.
|
Claude review — Recommendation: approve. The last round's only action item is fixed with a test, and nothing new blocks: the No action needed.
🤖 Generated with Claude Code |
Kimi-K2.5 (and Kimi-K2.6 and Kimi-K2.7-Code, which share its model type) is Kimi-VL's shape with compressed-tensors INT4 experts. Its adapter swaps in transformers' native deepseek_v3 language model as Kimi-VL's does, through the patch they share, and installs the on-the-fly decompressor first of all, so a checkpoint the decompressor refuses leaves the config and the loader untouched. It installs it only for a checkpoint whose quantization_config declares compressed-tensors: bf16 experts load as they are, and any other scheme is left to the expert-parallel load's format handling, which reads the copy of quantization_config Kimi's config class keeps on the top level. A scheme on text_config alone, with no top-level copy, is refused before anything changes. The spec declares the transformers 5.8 the decompression needs, and records that transformers ships its own Kimi-K2.5 port from 5.13, which the adapter refuses: expert parallelism runs on the checkpoint's own code (trust_remote_code=True). The shared adapter tests and the remote-Kimi refusal test run on Kimi-K2.5 too. Moved out of the Kimi adapters change, whose history is archived. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… what its missing MoE section skips transformers ships Kimi-K2.5 from 5.13, and the registry pairs a transformers modeling_source with that release, so a remote_code claim would be stale from 5.13 on, though prepare runs the checkpoint's own code. As for Kimi-VL, is_moe does not recognize the native DeepseekV3MoE blocks prepare installs until the sections are shared. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The shared adapter tests, run for Kimi-K2.5 too, already check the load and the swap. What is left is Kimi-K2.5's own: the swap, scoped to modeling_kimi_k25, leaves modeling_deepseek's DeepseekV3ForCausalLM alone. The test fails when the scope is dropped. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The plan decision ran the model spec's prepare on a copy of the config before the plan was known, wherever it could still be expert-parallel. For the Kimi models prepare patches the checkpoint's own code, the loader and the config for the whole process, so a run whose plan came out FSDP2-sharded, or whole and refused, kept those patches. The decision now reads the plan the spec states (prepared_ep_plan), on the config Kimi's experts are counted from, and finds the experts through it in the model the checkpoint's own code builds; only an expert-parallel plan then runs prepare, on every rank. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…epare The plan decision now builds the model the checkpoint's own code builds, before the spec's prepare runs, and Kimi-VL's and Kimi-K2.5's code only imports on current transformers with the names it lost shimmed and its two-deep imports staged. EPSpec.make_importable runs just that, on every rank in the decision's spec step, whatever plan comes out: it patches no class, config, converter or loader hook. Both Kimi specs name the shared step, which their prepare runs too. EPSpec.prepare, prepared_ep_plan and resolve_distribution_plan now say that prepare runs only for an expert-parallel plan, and the end-to-end test runs the real Kimi prepare through the decision: FSDP2-sharded, or left whole and refused, the checkpoint's own language model class, the loader and the config as read are untouched; an expert-parallel plan swaps in the native language model. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From ep-routing's review of the move: prepare's doc names the collective step it runs in, on what config, and the load that prepares a handed plan; make_importable's names its gate and why it must be idempotent; prepared_ep_plan's says deciding attaches it with no prepare; and resolve_distribution_plan names both collective steps. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…oad takes decision and the expert-parallel load would refuse its checkpoint's format. Kimi-K2.5's compressed-tensors checkpoint reads as such a format on the config the decision sees, before prepare installs the decompressor, so a decided Kimi-K2.5 run planned without expert parallelism. EPSpec.decompressed_formats states the formats a spec's prepare decompresses as the experts load; the decision counts them among those the load takes, and Kimi-K2.5's spec names compressed- tensors. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
As CONTRIBUTING asks, and as the Kimi-VL adapter's tests now do: the Kimi-K2.5 adapter's tests move to tests/unit/torch/models/kimi_k25/ test_ep.py, mirroring kimi_k25/ep.py. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ision Measured on Kimi-K2.5's own code: since the plan decision builds the checkpoint's own model on the meta device before prepare, the build fails on transformers 5.18, as post_init calls the VLM class's tie_weights with recompute_mapping, which the remote tie_weights does not take. prepare's swap used to adapt it first. EPSpec.make_importable becomes make_buildable(hf_config, ckpt_path=..., trust_remote_code=...): what the checkpoint's own code needs to import and build on the installed transformers, which may adapt a method to the arguments transformers now passes but changes nothing the model computes or loads. Kimi-VL's and Kimi-K2.5's step also loads the code and adapts its VLM class's tie_weights; the contract now names what it may not patch, as Claude's review asked: an MoE or expert class, a config, a weight converter or a weight-loading hook. New test test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_ transformers builds the tiny Kimi-VL remote code through transformers' own loader after the step, and fails with that TypeError without it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: the test expected the decompressor for compressed-tensors by a literal, beside the spec's decompressed_formats that the plan decision reads. It now expects an install exactly for the formats the spec states, so the spec and prepare cannot drift apart without failing it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…epared From Claude's review: the step that prepares a decided expert-parallel plan is a collective of its own, which carries every prepare refusal, and no test reached it with a rank failing alone. The rank-failure test gains a prepared case: rank 1 states the spec as rank 0 does and fails only where prepare runs, and every rank must raise its error; it fails when that step is not a collective. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test imports the tiny Kimi-VL checkpoint's code through transformers' loader, which keeps it in sys.modules; a later test then found it as a loaded modeling_kimi_vl. It now drops the remote modules it loaded. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…o build From Claude's review: make_buildable installed the import shims and the staging before the checks that say nothing of the checkpoint's will be built, a config of transformers' own port, no checkpoint path, or code the load may not run. It now returns before them in those cases, so the step leaves the process as it found it. New test test_the_build_step_installs_nothing_when_there_is_nothing_to_build fails with the shims first. The build-step test also checks that the VLM's language model stays the checkpoint's own, per-expert class. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: compressed-tensors was written both in prepare's condition for installing the decompressor and in the spec's decompressed_formats, which the plan decision reads before prepare runs. Both now read kimi_k25.ep.DECOMPRESSED_FORMATS. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: make_buildable adapted tie_weights on every loaded module whose name ends with the modeling module's, which from transformers 5.13 includes transformers' own Kimi-K2.5 port, transformers.models.kimi_k25.modeling_kimi_k25, whose class a later load of the port would then use with its keyword arguments dropped. It now goes through _remote_code._apply, which reaches remote modules only, as the swap does. test_the_build_step_leaves_transformers_own_port_alone fails with the unscoped scan. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Deciding the plan no longer runs prepare, so a spec whose plan only prepare attaches would have its MoE read as dense and planned without expert parallelism, with no refusal. Both Kimi specs state theirs; the test pins it for every EP spec. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…mpresses Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> #2746 no longer keeps a loadable flag in the decision: it holds the expert-parallel load's refusal against an MoE, and the formats a spec states it decompresses are not among the reasons.
|
/ok to test b88d349 |
Note
Split at the reviewers' request (2026-10-30): this PR is now [2/2], the Kimi-K2.5 adapter.
_kimi.py) and the shared deepseek_v3 plan.keep_patchedmoved to [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130.Base: #3138. The diff below is this PR's alone, and it merges after #3138.
Series layout
main([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 have merged)require_transformers, and a VLM built without its LM headprepareon the expert-parallel route only (this PR)What does this PR do?
Type of change: New feature
Kimi-K2.5 (and Kimi-K2.6 and Kimi-K2.7-Code, which share its model type) is Kimi-VL's shape with compressed-tensors INT4 experts. Its adapter swaps in transformers' native
deepseek_v3language model as Kimi-VL's does (#3138), through the patch they share, and decompresses the experts as they load (#3137, #2747).config.jsonkeeps the compressed-tensors config undertext_config, and its config class also copies it to the top level, which the expert-parallel load drops.prepareinstalls the decompressor first of all, so a checkpoint the decompressor refuses leaves the config and the loader untouched.quantization_configdeclares compressed-tensors, as Kimi-K2.5, K2.6 and K2.7-Code all do.quantization_configthe config class keeps on the top level.text_configalone, with no top-level copy, is refused before anything changes.preparenow runs only on the expert-parallel route. Before, the plan decision ran a model spec'sprepareon a copy of the config before the plan was known, and the Kimi patches it made (the swapped language-model class, the decompressor's loader hooks) stayed in the process when the plan came out FSDP2-sharded, or whole and refused. Now:prepared_ep_plan, ontext_configfor both Kimi specs) and finds the experts through it in the model the checkpoint's own code builds;EPSpec.make_buildable, what the checkpoint's own code needs to import and build on the installed transformers: the import shims, the staging of the code's two-deep imports, and the VLM class'stie_weightsadapted to the keyword argumentspost_initpasses it. It changes nothing the model computes or loads, and patches no MoE or expert class, config, weight converter or weight-loading hook;prepare, on every rank, and returns that config; any other plan returns the config as read.EPSpec.prepare's,prepared_ep_plan's andresolve_distribution_plan's docs say so, replacing #2746's interim contract. Kimi-VL's spec (#3138) takes it here too, and Kimi-K3 (#2749), now on this PR, uses it.Why transformers' own Kimi-K2.5 port is refused. transformers 5.13 added a Kimi-K2.5 port, which a load without
trust_remote_codebuilds. Its language model is a nativedeepseek_v3with fused experts, so expert parallelism could run on it without these patches. This adapter does not target it, for three reasons:wqkvinto q/k/v. The export would have to map all of them back to the checkpoint's names for vLLM.So the adapter refuses the port's config before it changes anything, naming
trust_remote_code=Truefrom a checkpoint whoseauto_mapnames its ownAutoConfigas the way in. Supporting the port is a follow-up that needs its own run on the real checkpoints. The spec records the 5.13 release.Why it patches the checkpoint's module, and again on every class resolution: see #3130, where
keep_patchedlives.Usage
Picked automatically by
apply_distributed_supportforkimi_k25configs undertorchrun; see #2754.Testing
tests/unit/torch/models/test_kimi_k25_ep.py, against a stand-in remote module, covers:DeepseekV3Configfor compressed-tensors, bf16 and block-FP8 experts, with the decompressor installed for compressed-tensors ones only. The bf16 and block-FP8 cases fail with an unconditional install.text_configis still present, and a decompressor refusal leavingtext_configand the remote-code staging untouched;The shared adapter tests ([OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138's
test_kimi_vl_ep.py) and the remote-Kimi refusal test intest_model_load_utils.pygain their Kimi-K2.5 cases.Locally at this head:
tests/unit/torch/modelsandtest_model_load_utils.pypass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.GPU tests were not re-run on this revision. Earlier real-model runs (with the module-reload fix, through the decompressor): Kimi-K2.7-Code, expert-parallel over 32 GPUs (8 nodes, EP32), transformers 5.13.0, compressed-tensors 0.17.1, recipe
nvfp4_experts_only_mse-kv_fp8_cast, 4096 calibration samples of 4096 tokens, batch 1. Load, calibration and export completed:DeepseekV3ForCausalLMon the native configRe-cut on 2026-10-29 onto [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747, and split on 2026-10-30; the history before is archived. Review fixes then:
2026-10-30: rebased onto main 99c3834, through the new [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 (which now refuses compressed tensors the model leaves unplaced), [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132, with no conflicts and the own diff unchanged. Pre-commit against main passes, mypy included;
tests/unit/torch/models,test_model_load_utils.pyandtest_ep_loader.pypass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-30, review round: the spec's comment says why
modeling_sourcestaystransformers, and its docstring names theis_moegap as Kimi-VL's does; the Kimi-K2.5 swap test keeps only its sibling-module claim (it fails when the patch's suffix scope is dropped), since the shared adapter tests check the swap. Rebased onto [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138's new head. Pre-commit against main passes, mypy included;tests/unit/torch/models,test_model_load_utils.pyandtest_ep_loader.pypass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-30: rebased onto the rebuilt base: main 7168107, where [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 have merged, joined with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 5d51cd6 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 4494480. There were no conflicts in this PR and its own diff is unchanged. Pre-commit against main passes, mypy included. On CPU, torch 2.14, with transformers 5.18 and 5.5.4, these pass:
tests/unit/torch/models,test_model_load_utils.py,test_ep_loader.py,test_unplaced_source_keys.py.2026-10-30: rebased onto [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138 9e2cd9a, whose Kimi swap now builds the native language model only from the config
preparemade. Without that, a load kept off expert parallelism after deciding had prepared a copy of the config failed. The own diff is unchanged. [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138's new test runs for Kimi-K2.5 too. Pre-commit against main passes.tests/unit/torch/models,test_model_load_utils.pyandtest_ep_loader.pypass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-30: rebased onto the rebuilt base: [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 ffb58cf joined with [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 0f0570b, whose decompressor ([OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137) now reverses its converters as the ones it rewrote, so a decompressed checkpoint's export keeps the checkpoint's names. No conflicts in this PR; the own diff is unchanged. Pre-commit against main passes.
tests/unit/torch/models,test_model_load_utils.py,test_ep_loader.pyandtest_unplaced_source_keys.pypass on CPU, torch 2.14: 593 on transformers 5.18, and 451 passed, 57 skipped on 5.5.4.2026-10-30: rebased onto [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138 9ef474a, on the rebuilt base: [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 d3a058e joined with [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 dc0a618. The own diff is unchanged. Pre-commit against main passes, mypy included. On CPU, torch 2.14,
tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4.2026-10-30: rebased onto [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138 b4847b8, on the rebuilt base: main 059588b, with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 0eadf4c joined with [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 23eb10d. The own diff is unchanged. Pre-commit against main passes, mypy included. On CPU, torch 2.14,
tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4.2026-10-31, on the user's decision:
prepareruns only on the expert-parallel route (53fc43c, 1df64fd, 7691431), and a checkpoint a spec decompresses counts as one the expert-parallel load takes (57be464). Rebased onto main d464d30, with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 9267a0b and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 7d78ca1:test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code[fsdp2|whole|expert-parallel]runs the real Kimi-VL and Kimi-K2.5preparethrough the decision. FSDP2-sharded, or left whole and refused, the checkpoint's own language-model class, the loader hooks and the config as read are untouched; an expert-parallel plan swaps in the native language model. The fsdp2 and whole cases fail on the old decision;test_a_checkpoint_its_spec_decompresses_is_decided_expert_parallel: a compressed-tensors Kimi-K2.5 checkpoint left to the decision is planned expert-parallel, and a Kimi-VL one is not; the Kimi-K2.5 case fails without the spec's formats;test_a_load_prepares_its_config_oncenow counts real preparations only;tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4;load_pretrainedintests/gpu/torch/utils/test_expert_parallel_load.pytake the new decision: at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 63c8267, which contains this head, the whole file passed with transformers 5.18.0, on 2 GPUs (4 passed,[ep-half]skipped as it needs 4) and on 4 GPUs (all 5). No GPU case exercises the decision's fallback to no expert parallelism, which the CPU rows above cover.2026-10-31, Claude's review of the move, at 27961b6, rebased onto main d464d30 with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 77c4117 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 7d78ca1:
the decision's meta build of real Kimi-K2.5 code failed on transformers 5.18, because the remote
tie_weightstakes no keyword arguments; the build step (make_buildable, 385d420) adapts it.test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformersbuilds the tiny Kimi-VL remote code through transformers' own loader after the step, and fails with that TypeError without it;the decision's cost, measured on the real Kimi-K2.5 code and config (61 layers, 384 routed experts) on the meta device, transformers 5.18, one process; only rank 0 builds:
a rank failing alone in the step that prepares a decided expert-parallel plan fails every rank (
test_a_rank_that_cannot_apply_the_spec_fails_every_rank[prepared]), and the Kimi-K2.5 decompressor test derives its expectation from the spec'sdecompressed_formats;pre-commit against main passes, mypy included; on CPU, torch 2.14,
tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4. At [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 183bf95, which contains this PR, the whole GPU file passed with transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed and[ep-half]skipped, as it needs 4; on 4 GPUs, all 6 passed. Its Kimi cases load through the spec and the loader directly, not through the plan decision, so the build step's coverage is the CPU test above and the measurement on Kimi-K2.5's code.2026-10-31, Claude's review, at 777ac2a: the build step installs nothing when there is nothing of the checkpoint's to build (fab6ddf;
test_the_build_step_installs_nothing_when_there_is_nothing_to_buildfails with the shims first), and Kimi-K2.5's decompressed formats are stated once (777ac2a). Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-31, at f3df820: at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 4389363, which contains this PR at 777ac2a, the whole GPU file passed, with transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed and
[ep-half]skipped, as it needs 4; on 4 GPUs, all 6 passed. Its Kimi cases load through the spec and the loader directly, not through the plan decision. Since then, from Claude's review, the build step'stie_weightsshim goes through_remote_code._apply, so it reaches remote modules only and leaves transformers' own Kimi-K2.5 port alone (test_the_build_step_leaves_transformers_own_port_alonefails without it). No GPU case runs the build step. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-31, at b88d349: rebased onto [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138 d2321c1, on main e352cac with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 d8059b3 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 6b87b69. [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 now attaches a spec's plan behind its own existing-plan check, and its
_ep_load_refusalreturns the error and no longer reads the spec; this PR's stated-plan selection keeps that check, and_ep_load_refusalreads the spec only for itsdecompressed_formats. From Claude's review,test_a_spec_with_a_prepare_states_its_plan(7a1f1cf) asserts that every EP spec with apreparestates its plan; it fails for both Kimi specs withprepared_ep_plandropped. At [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 a60c01c, which contains this head, the whole GPU file passed, with transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed and[ep-half]skipped, as it needs 4; on 4 GPUs, all 6 passed. Its Kimi cases load through the spec and the loader directly, not through the plan decision. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models, [2/2] of #2748's split. It merges after #3138.