Repository navigation
[OMNIML-6078] [19/21] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters - #2748
shengliangxu wants to merge 18 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 | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 73.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 30 functions across 11 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
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 clearly and concisely identifies the main change: adding expert-parallel adapters for Kimi-K2.5 and Kimi-VL. 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.
-
✨ Finishing Touches: 💡 2
-
📝 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
Comment @coderabbitai help to get the list of available commands.
f73bc4f to
0098f3a
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-11a-onfly-decompress #2748 +/- ##
=========================================================================
+ Coverage 62.31% 62.36% +0.04%
=========================================================================
Files 658 665 +7
Lines 73926 74040 +114
=========================================================================
+ Hits 46065 46172 +107
- Misses 27861 27868 +7
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 |
0098f3a to
3d6fb9d
Compare
|
/ok to test 3d6fb9d |
shengliangxu
changed the title
1c16b42 to
ecf9d71
Compare
|
/ok to test ecf9d71 |
shengliangxu
changed the title
ecf9d71 to
d525e64
Compare
shengliangxu
changed the title
|
/ok to test d525e64 |
d525e64 to
0c39a27
Compare
|
/ok to test 0c39a27 |
|
/ok to test 3b07654 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for kept_suffix, kept_patch in _KEPT_PATCHES.items(): | ||
| if cls.__module__.endswith(kept_suffix) and kept_patch(sys.modules[cls.__module__]): |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The get_class_in_module wrapper is missing the remote-module guard the eager loop below (L110-114) applies, and it indexes sys.modules unguarded.
get_class_in_module ends in return getattr(module, class_name), so the returned class's __module__ is not necessarily the module that was just loaded — it is whatever module defined the class. A remote modeling_kimi_k25.py that re-exports transformers' port (from transformers.models.kimi_k25.modeling_kimi_k25 import KimiK25ForConditionalGeneration) — exactly the forward-compat shape 5.13 invites a checkpoint to adopt — yields cls.__module__ == "transformers.models.kimi_k25.modeling_kimi_k25", which also ends with .modeling_kimi_k25. The kept patch then runs against transformers' own port module: to_native finds no module-level DeepseekV3ForCausalLM and raises "the checkpoint's modeling code changed" mid-from_pretrained, or (for a suffix whose port module does carry the name) mutates transformers' own module. test_only_remote_modules_are_patched covers the loop for precisely this hazard; the wrapper is the path that runs during the real load and is unguarded.
Separately, sys.modules[cls.__module__] is a KeyError for any class whose __module__ is not a live module entry — this wrapper stays installed process-wide and sees every dynamic-class resolution, so a failure here breaks unrelated remote-code loads.
| for kept_suffix, kept_patch in _KEPT_PATCHES.items(): | |
| if cls.__module__.endswith(kept_suffix) and kept_patch(sys.modules[cls.__module__]): | |
| for kept_suffix, kept_patch in _KEPT_PATCHES.items(): | |
| module = sys.modules.get(cls.__module__) | |
| if ( | |
| module is not None | |
| and cls.__module__.startswith(f"{dmu.TRANSFORMERS_DYNAMIC_MODULE_NAME}.") | |
| and cls.__module__.endswith(kept_suffix) | |
| and kept_patch(module) | |
| ): |
Cleaner still: factor the startswith(remote) and endswith(suffix) test into one helper used by both the wrapper and the loop, so the two can't drift again.
There was a problem hiding this comment.
Fixed in 81ec044: the wrapper takes the module defining the class with sys.modules.get, and re-patches it only when it is remote code with the kept suffix. That test is now one helper, _is_remote(name, suffix), which the scan of loaded modules uses too, so the two cannot drift. A class defined in transformers' port (a remote module importing its class from transformers.models.kimi_k25.modeling_kimi_k25) leaves the port alone, and a class whose module sys.modules does not hold resolves without a KeyError. test_only_remote_modules_are_patched now also resolves a class through the wrapper, once defined in the port and once in a module that is not loaded. On the old wrapper the first case patches the port and the second raises KeyError.
| 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) | ||
|
|
||
| # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep. Install the | ||
| # full-closure staging shim BEFORE any remote-code load (this plugin's get_class_from_dynamic_module | ||
| # below, and from_pretrained later in load_and_prepare_ep). | ||
| _kimi.install_compat_import_shims() | ||
| stage_full_remote_closure() | ||
|
|
||
| # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields. | ||
| # Kimi's __init__ passes config.text_config to DeepseekV3ForCausalLM; the native class rejects | ||
| # Kimi's custom config subclass. to_native_backbone_config also drops auto_map + | ||
| # quantization_config so transformers builds the native modeling instead of re-routing back | ||
| # to Kimi's remote code via auto_map. | ||
| hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] Nothing verifies the top-level quantization_config copy this route depends on, so a non-compressed-tensors scheme carried only on text_config is silently dropped rather than refused.
The comment on L73-74 states that block-FP8 experts are "the expert-parallel load's to read ... by the copy of quantization_config Kimi's config class keeps on the top-level config". That is load-bearing and unchecked:
to_native_backbone_configpopsquantization_configoff the dict it builds the nativeDeepseekV3Configfrom (_onfly_decompress.py:64), so after L90text_confighas none.load_and_prepare_epthen computesquant_method = _ckpt_quant_method(hf_config)(model_load_utils.py:2336) from the prepared config, reading root →text_config→llm_config. With thetext_configcopy gone and no root copy, it returnsNone, the_EP_SOURCE_FORMATSrefusal at L2340-2348 does not fire, andFineGrainedFP8Config(dequantize=True)is never installed — the fp8 blocks are read as stored into the bf16 fused params, withweight_scale_invnever applied. Wrong expert numerics for the whole calibration, with no error.
Today's Kimi-K2.5 config class does keep the copy, so this is latent — but the branch at L76 already has the scheme in hand, and CONTRIBUTING/CLAUDE guidance is to fail loudly rather than silently ignore a route. A few lines in the shared _kimi helper would close it for both adapters, e.g.:
# in _kimi.py, called from both prepares before to_native_backbone_config
def check_scheme_survives_the_swap(label: str, hf_config, handled: bool) -> None:
"""The swap drops text_config's quantization_config; the load reads the root copy instead."""
if handled or _onfly_decompress._find_quantization_config(hf_config) is None:
return
if getattr(hf_config, "quantization_config", None) is None:
raise ValueError(
f"{label}: this checkpoint declares quantization_config only on text_config, which "
"the swap to a native deepseek_v3 config drops, so the expert-parallel load would "
"read its experts as unquantized. Start from an unquantized copy of the model."
)The same gap applies to kimi_vl/ep.py, which considers quantization not at all.
Minor, in the same block: _onfly_decompress._find_quantization_config is reached into privately from a sibling package; exposing the compressed-tensors predicate (or an install_onfly_decompress_if_compressed) in _onfly_decompress would keep the condition next to the thing it gates.
There was a problem hiding this comment.
Fixed in 81ec044: a refusal both adapters share, _kimi.refuse_quantization_only_on_text_config, stops a config whose quantization_config sits on text_config with no top-level copy, before anything changes. Kimi-K2.5 runs it when the scheme is not compressed-tensors (a compressed-tensors one goes to the decompressor, which marks the config for the load), and Kimi-VL runs it always, after the refusals the two share. The message: "kimi_vl: the checkpoint declares quantization_config on text_config only, which the native deepseek_v3 config replacing it drops, so the expert-parallel load would read its experts as stored. Declare it on the top-level config too, or start from an unquantized copy of the model." The published checkpoints are unaffected: Kimi-K2.5's config class copies the scheme to the top level, and the Kimi-VL checkpoints are bf16. The refusal table gives both adapters a block-FP8 scheme on text_config alone; both cases fail without the check. The Kimi-K2.5 rebuild test's block-FP8 case now carries the top-level copy the real config class adds. On the minor point, the compressed-tensors condition stays in prepare: _onfly_decompress is a private module of the same modelopt.torch.models package, and moving the condition into it would change [18/21] without changing behavior.
|
|
||
| Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM`` with | ||
| transformers' native class, so ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 | ||
| for ``self.language_model`` instead of the per-expert Linears version. | ||
| """ | ||
| from transformers.dynamic_module_utils import get_class_from_dynamic_module | ||
| from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM | ||
|
|
||
| # Force the dynamic-module loader to stage + import the module into sys.modules. | ||
| get_class_from_dynamic_module( | ||
| f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code | ||
| ) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] The wrapper drops every kwarg, so its safety argument has to hold for arguments transformers has not added yet.
Two consequences today: recompute_mapping=False from the load becomes an unconditional recompute (correct, just wasted work), and missing_keys is discarded on the docstring's premise that tie_word_embeddings=False. Nothing enforces that premise — a Kimi variant shipping tie_word_embeddings=True would tie correctly (the remote tie_weights still runs) but leave the tied key in the load's missing list, i.e. a spurious "newly initialized" warning on a weight that is fine. Any future kwarg that does affect tying would be dropped just as silently.
Forwarding what the original accepts, instead of dropping all of it, keeps the shim correct as the signature evolves:
def _drop_tie_weights_kwargs(orig_tie):
"""Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it."""
import inspect
accepted = set(inspect.signature(orig_tie).parameters) - {"self"}
@functools.wraps(orig_tie)
def _tie_weights(self, *args, **kwargs):
return orig_tie(self, *args, **{k: v for k, v in kwargs.items() if k in accepted})
_tie_weights._kimi_tie_weights_shim = True
return _tie_weights(functools.wraps also keeps the wrapped function introspectable, which matters when the next person debugs a tie_weights TypeError.)
There was a problem hiding this comment.
Pushing back: for the checkpoints this wraps, forwarding by signature would forward nothing. Both published remote tie_weights (Kimi-K2.5's modeling_kimi_k25.py, Kimi-VL's modeling_kimi_vl.py) are def tie_weights(self): return self.language_model.tie_weights(), so a signature filter drops missing_keys and recompute_mapping just as the wrapper does now, and a tie_word_embeddings=True variant would keep its tied key in the missing list either way. A remote tie_weights that does take keyword arguments needs no wrapper at all. A filter written for one would also need a **kwargs case: inspect.signature lists that as a single parameter named kwargs, so the suggested set would forward nothing to it. The wrapper stays specific to the signature both checkpoints ship.
There was a problem hiding this comment.
Claude review — [19/21] Kimi-K2.5 / Kimi-VL expert-parallel adapters
Findings: CRITICAL 0 · IMPORTANT 2 · SUGGESTION 1
Full review of all 14 changed files (10 source + 4 test); scope was the whole diff, prioritizing modelopt/. I also read models/specs.py, utils/moe.py (apply_distributed_support / _spec_ep_plan / _routed_expert_keys), models/_onfly_decompress.py and the load_and_prepare_ep quantization path for the composition questions below.
What checks out
- Plan/experts agreement.
prepared_ep_plan=BASE_MODEL_EP_PLANnameslayers.*.mlp.experts, which_routed_expert_modulesmatches against thenn.ModuleListthe checkpoint's own code builds atlanguage_model.model.layers.<i>.mlp.experts— so the EP-degree-1 refusal fires without runningprepare, exactly asprepared_ep_plan's contract requires. The newtest_resolve_distribution_plancase pins it. - Plan precedence.
apply_distributed_support("deepseek_v3", text_config)attaches only when the target carries no plan, so from transformers 5.13 the nativedeepseek_v3plan (withep_router) wins and before itBASE_MODEL_EP_PLANapplies. Hoisting the constant out ofEPSpeckeeps all three specs on one definition, and bothapply_distributed_supportand_spec_ep_plancopy it before use, so the shared dict cannot be mutated through a spec. - Refuse-before-mutate.
checkpoint_to_patchruns first in both prepares and covers the transformers port, missingtext_config, missingckpt_pathand missingtrust_remote_code; forkimi_k25the decompressor install precedes thetext_configswap, so a refusal leaves the config and the loader staging untouched. All four refusals plus the decompressor-refusal ordering are tested, andstartswith("transformers.")correctly does not catchtransformers_modules.*. - Spec registration / plugin laziness.
min_transformers_version/modeling_sourcestay consistent for both specs undertest_version_claim_matches_modeling_source; nothing outsidemodels/kimi_*learns about Kimi; no new hard imports of transformers at module load beyond what the package already pulls. - Distributed symmetry.
prepareruns underdist.all_reduce("model spec", ...)on every rank, and the patch paths issue no collectives, so the rank-local re-patchprintcannot desynchronize anything. - Compat shims target the modules the copied-out HF modeling actually imports from (
transformers.utils.import_utils,transformers.pytorch_utils,transformers.modeling_attn_mask_utils), are idempotent, and the unreachable-mask stub raises rather than silently returning a wrong mask.
Most impactful findings
- IMPORTANT — the kept-patch wrapper lacks the remote-module guard the eager loop has (
_remote_code.py:96-97).get_class_in_modulereturnsgetattr(module, class_name), socls.__module__is the defining module, not the one loaded. A remotemodeling_kimi_k25.pythat re-exports transformers' 5.13 port makes the wrapper patchtransformers.models.kimi_k25.modeling_kimi_k25— the exact hazardtest_only_remote_modules_are_patchedguards the loop against — failing the load with the misleading "the checkpoint's modeling code changed".sys.modules[cls.__module__]also wants to be a.get, since this wrapper sees every dynamic-class resolution process-wide. - IMPORTANT — the top-level
quantization_configcopy the block-FP8 route depends on is never verified (kimi_k25/ep.py:75-90, and absent entirely inkimi_vl/ep.py).to_native_backbone_configpopsquantization_config, andload_and_prepare_epreads_ckpt_quant_methodfrom the prepared config: with no root copy it returnsNone, the_EP_SOURCE_FORMATSrefusal does not fire, and fp8 blocks load as stored withweight_scale_invnever applied — wrong expert numerics for the whole calibration, silently. Latent for today's checkpoints; the branch already has the scheme in hand, so a loud refusal is a few shared lines.
Plus one SUGGESTION on _drop_tie_weights_kwargs forwarding by signature rather than dropping all kwargs.
Risk
Moderate and well contained: the whole change is additive and reachable only through the expert-parallel path of two model types, the test suite is unusually thorough for the patching mechanics (reload, second path, port isolation, per-adapter regressions), and the EP32 runs on a real Kimi-K2.7-Code checkpoint back the happy path. Both IMPORTANT findings are about the unhappy paths the design already reasons about in comments but does not yet enforce in code — finding 1 would surface as a confusing hard failure, finding 2 as silent numerics. Neither needs a redesign.
🤖 Generated with Claude Code
Both checkpoints ship their own modeling, which builds the language model's MoE from an nn.ModuleList of per-expert Linears that transformers' expert-parallel machinery cannot shard. Their text backbone is DeepSeek-V3, so the adapters swap in transformers' native deepseek_v3 language model -- fused gate_up_proj / down_proj experts and the plan to shard them -- and keep the bundled vision tower and projector from the checkpoint's own code. Kimi-K2.5's experts are compressed-tensors INT4, decompressed as they load. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The module skipped below 5.6, but the code under test requires 5.8 and raises below it, so a 5.6 or 5.7 run failed instead of skipping. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…s remote code transformers 5.13 added kimi_k25, so its spec can no longer claim remote code; test_specs_vs_transformers fails on the latest transformers until it names the release. A load without trust_remote_code now builds transformers' port, which prepare's patch never reaches, so prepare refuses that config instead of running unpatched. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Both specs' prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. An EP degree of 1 skips prepare, so their specs supplied no plan there, and the remote-code checkpoints, whose configs carry none, loaded whole on every rank with a replicated dense part instead of being refused. Their specs now state deepseek_v3's plan, hoisted to a constant both import. It finds the experts the checkpoint's own code builds, at the same paths. Where transformers ships a deepseek_v3 plan of its own, prepare keeps that one, which names the same experts. The comment that moves with the plan no longer says transformers ships none, which stopped being true in 5.18. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Without prepare's own text_config check, to_native_backbone_config raises the same ValueError, so the refusal tests passed either way. The check exists to refuse before prepare installs the remote-code staging and, for Kimi-VL, the import shims, which are process-wide. The tests now hand the loader an unwrapped stand-in and require it to be left as it was, which fails once the check is gone. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The swap-fails-loudly, missing-text_config refusal and spec-routing tests were written once per adapter. Each is now one test in test_kimi_vl_ep.py, parametrized over Kimi-VL and Kimi-K2.5, and each case still fails on its own adapter's regression: a silent swap, staging before the text_config check, a spec wired to the other adapter's prepare. The VL routing test's modeling_source assert goes, since test_specs_vs_transformers checks that claim. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_prepare_rebuilds_the_language_model_config_without_a_decompressor imported transformers.utils.import_utils inside the test. It is core transformers, which the module already imports after its importorskip and version guards. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The real Kimi-K2.5 remote code (moonshotai/Kimi-K2.5) failed to load through the adapter on every supported transformers release. Its modeling_deepseek.py imports is_torch_fx_available, which transformers 5.x dropped, and its KimiK25ForConditionalGeneration.tie_weights takes no arguments while transformers' post_init and load call it with recompute_mapping=False. Kimi-VL's adapter already handles both; give Kimi-K2.5's the same import shims and tie_weights wrapper. The Kimi-VL tests' stand-in remote code now takes the modeling module and VLM class to stand in for, and the swap-and-shim test runs for both adapters; the Kimi-K2.5 case fails without this change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With a checkpoint path, the Kimi-K2.5 and Kimi-VL adapters load the checkpoint's own modeling through transformers' get_class_from_dynamic_module, which runs it without checking trust_remote_code: the flag lands in its kwargs unused. Refuse in prepare unless the caller passed trust_remote_code=True, before anything is staged or patched. load_pretrained, resolve_distribution_plan and hf_ptq forward the caller's flag, so a trusted load is unchanged. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…-K2.5's transformers floor Both Kimi adapters refuse, before they change anything, a prepare with no checkpoint path: they patch the checkpoint's own modeling code, and without it the config was left half-prepared, a native language-model config and a plan for fused experts under Kimi's own unpatched modeling. The trust_remote_code check now applies to every prepare. Kimi-K2.5's spec declares the transformers 5.8 its on-the-fly decompression needs, so an older release is refused by check_distributed_support, with its way around, before prepare runs. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Kimi-K2.5, like Kimi-K2.6 and Kimi-K2.7-Code, carries its compressed-tensors quantization_config
only on text_config. prepare replaced text_config with a native DeepseekV3Config, which drops it,
before it installed the on-the-fly decompressor, so the decompressor found no scheme and refused the
compressed checkpoint ("config has no compressed-tensors quantization_config"). prepare now installs
it first, ahead of the compat shims and the remote-code staging too, so a checkpoint the
decompressor refuses leaves the config and the loader as they were. The adapter test's stand-in
config carries the scheme only on its text config, as the checkpoint does.
The kimi_k25 and kimi_vl specs say that only their expert-parallel section is registered: the MoE
and export sections stay deepseek_v3's. The kimi_k25 docstrings name Kimi-K2.5, whose model type
Kimi-K2.6 and Kimi-K2.7-Code share, and Kimi-VL's version note states the 5.6 floor it is held to.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
prepare runs on every rank (resolve_distribution_plan applies the model spec on each), so the Kimi-K2.5 and Kimi-VL patch notes and the compat-shim warnings printed once per rank; they now print on rank 0, as the decompressor's install note does. Both adapters use the shims, so their warnings and the attention-mask stub's error say kimi rather than kimi_vl. The deepseek_v3 plan's comment says which releases it serves: transformers ships no deepseek_v3 plan before 5.13, and from 5.13 its own wins, with the same experts and an ep_router gate entry the load swaps for torch_a2a_router. The native-port refusal test passes trust_remote_code=True and matches the part of the message only that refusal has, so it fails without the guard instead of matching the trust_remote_code refusal. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Each Kimi adapter repointed its remote module's DeepseekV3ForCausalLM at transformers' native class once, in prepare. transformers re-executes a remote module in place when the files it hashes change, and loads it afresh for another path to the same checkpoint; either rebinds the module's globals, so from_pretrained built the checkpoint's own per-expert model on a config already swapped to the native one. An expert-parallel Kimi-K2.7-Code load over 32 GPUs failed so, on some ranks, reading rope_theta, which the native config keeps under rope_parameters. keep_patched applies a patch to every loaded module with a name suffix, and again whenever transformers resolves a class from one, so the module from_pretrained builds from carries it. Both Kimi adapters use it for the native class and the tie_weights shim. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
keep_patched's patch now returns whether it changed the module. A re-application that did means transformers loaded the module again after it was patched, so the rank prints which module, and a run shows whether the reload happened rather than only that the load succeeded. The test also covers a second path to the same checkpoint, which transformers loads under another module name. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
prepare installed the on-the-fly decompressor for every kimi_k25 checkpoint, and the decompressor refuses one without a compressed-tensors quantization_config: a bf16 copy stopped on "config has no compressed-tensors quantization_config", a block-fp8 one on pydantic's ValidationError. The expert-parallel load reads both itself: bf16 experts as stored, block-fp8 ones dequantized through the quantization_config Kimi's config class copies to the top level. prepare now installs the decompressor only for compressed-tensors experts, which Kimi-K2.5, Kimi-K2.6 and Kimi-K2.7-Code all ship. Kimi-VL's tie_weights shim also says why dropping missing_keys is safe: the remote tie_weights only calls its language model's, and the Kimi checkpoints set tie_word_embeddings=False. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Since this slice installs the decompressor only for a compressed-tensors checkpoint, the 5.8 floor is the one the spec declares; the module docstring said the decompressor enforces it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Both Kimi adapters read the remote module's DeepseekV3ForCausalLM with the native class as the default, so a module that no longer defines that name (renamed, or imported inside the VLM class's __init__) looked already patched: the swap and the tie_weights wrapper were skipped, the adapter reported success, and the load went on to build the per-expert MoE under a plan naming fused experts. The swap now raises for such a module. keep_patched scans remote-code modules only, so transformers' own port, which shares the module name from 5.13 (transformers.models.kimi_k25.modeling_kimi_k25) and defines no such class, is left alone. The two adapters' patchers were one function three names apart, and Kimi-K2.5's imported Kimi-VL's private shims. Both now live in modelopt/torch/models/_kimi.py with the refusals the adapters share, so Kimi-VL also refuses transformers' own config, and that refusal names the port's module and the auto_map the checkpoint needs. The Kimi test fixtures restore the get_class_in_module wrapper and the kept patches; the Kimi-K2.5 file leaves out the compat shims, which the Kimi-VL file covers for both adapters. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
keep_patched's get_class_in_module wrapper re-applied a kept patch to the module defining the resolved class, matched by name suffix alone. That module need not be the remote one transformers loaded: a remote modeling_kimi_k25.py importing its class from transformers' own 5.13 port made the wrapper patch transformers.models.kimi_k25.modeling_kimi_k25 and fail the load with "the checkpoint's modeling code changed", and a class whose module is not in sys.modules raised KeyError. The wrapper now looks the module up with sys.modules.get and patches remote-code modules only, through the check the scan of loaded modules uses (_is_remote). Both Kimi adapters replace text_config with a native DeepseekV3Config that drops its quantization_config, and the expert-parallel load then reads the checkpoint's quantization off the top-level copy Kimi-K2.5's config class keeps. A scheme declared on text_config alone was dropped silently: block-fp8 experts would load as stored, never scaled by weight_scale_inv. prepare now refuses such a config before changing anything. Kimi-K2.5's compressed-tensors experts, which the on-the-fly decompressor takes over, are unaffected. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 81ec044 |
|
/claude review |
| raise ValueError( | ||
| f"{label}: the checkpoint declares quantization_config on text_config only, which the " | ||
| "native deepseek_v3 config replacing it drops, so the expert-parallel load would read " | ||
| "its experts as stored. Declare it on the top-level config too, or start from an " | ||
| "unquantized copy of the model." | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] The remedy this refusal offers is only right for the formats the expert-parallel load can dequantize, and the function is reached with others.
kimi_vl.prepare calls this unconditionally, and kimi_k25.prepare calls it on the whole non-compressed-tensors branch — so the scheme in hand can be anything other than compressed-tensors on the k25 path, and anything at all on the VL path. For a compressed-tensors scheme that reaches here (a Kimi-VL checkpoint compressed, or a k25 one whose quant_method spells it differently), "Declare it on the top-level config too" sends the user to a config that load_and_prepare_ep then refuses anyway: _EP_SOURCE_FORMATS is {None, "fp8", "mxfp4"} and _decompressed_as_loaded is false because kimi_vl installs no decompressor. The user follows the advice, edits config.json, and gets a second refusal from a different place.
Impact is low — it is advice text on an unhappy path, and today's Kimi-VL checkpoints are bf16 — but the first refusal is where the adapter knows which scheme it is looking at, so it can say the right thing. Either name the scheme and gate the "declare it on the top level too" half on it being one _EP_SOURCE_FORMATS accepts, or drop that half and keep "start from an unquantized copy of the model", which is correct for every case that reaches here.
| model_type="kimi_k25", | ||
| # transformers ships Kimi-K2.5 itself from 5.13. Expert parallelism still runs on the | ||
| # checkpoint's own code (``trust_remote_code=True``), which ``prepare`` patches. | ||
| min_transformers_version="5.13", |
There was a problem hiding this comment.
[SUGGESTION] kimi_k25 takes modeling_source at its "transformers" default so it can record 5.13, but nothing checks that claim and the field now says the opposite of what the adapter does.
Two halves to this:
-
The claim is unverified.
test_moe_block_resolves_from_its_declared_versionis parametrized overVERSIONED_MOE_MODEL_TYPES, which isMOE_MODEL_TYPES(specs with amoe_spec) filtered bymin_transformers_version is not None. This spec deliberately carries nomoe_spec, so it is in neither list, andtest_remote_code_models_are_really_absentskips it too sincemodeling_source != "remote_code". So"5.13"is checked bytest_version_claim_matches_modeling_sourceonly — i.e. only for internal consistency withmodeling_source, never against transformers. If the port'smodel_typeis not literallykimi_k25, or lands in a different release, nothing fails. Contrastkimi_vl, whose"remote_code"claim is checked and will fail loudly the day transformers ships a Kimi-VL port. -
The field reads backwards.
ModelSpec.modeling_source's docstring says it records where the modeling code lives, "somodel_typeis a remote-code spelling rather than a transformers one, and loading the model needs the flag." For this spec, loading does need the flag —checkpoint_to_patchrefusestrust_remote_code=Falseoutright, and refuses the 5.13 port's config as well, so the only config this spec's EP path ever accepts is remote code. A reader who trusts the field concludes the opposite. No runtime consequence today (I grepped:modeling_sourcehas no consumer outsidespecs.py, the spec files andtest_specs_vs_transformers.py), so this is bookkeeping — but it is exactly the kind of claim that goes stale unnoticed.
The comment on lines 35-36 already states the real situation in prose. Consider making the data say it too: either modeling_source="remote_code" with min_transformers_version=None (the 5.13 fact stays in the comment, and the port's arrival is then covered by test_remote_code_models_are_really_absent — which would fail once 5.13 is in the matrix, so this only works if that is acceptable), or keep the current pair and note in the comment that nothing verifies "5.13" until this spec gains a moe_spec in the planned follow-up.
|
Claude review — Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 2 Scope: full review — the trigger was a bare Both IMPORTANT findings from the previous round are fixed — verified against the code, not the threads
What I traced and found correct
Suggestions (non-blocking, both posted inline)
Not re-raised (already on the PR and unchanged in intent): the Risk: low. Additive and reachable only through the expert-parallel 🤖 Generated with Claude Code |
What does this PR do?
Type of change: New feature
Kimi-K2.5 and Kimi-VL ship their own modeling, which builds the language model's MoE from an
nn.ModuleListof per-expert Linears that transformers' expert-parallel machinery cannot shard. Their text backbone is DeepSeek-V3, so the adapters swap in transformers' nativedeepseek_v3language model (fusedgate_up_proj/down_projexperts, and the plan to shard them) and keep the bundled vision tower and projector from the checkpoint's own code. Kimi-K2.5's experts are compressed-tensors INT4 and are decompressed as they load ([18/21]). Its checkpoint'sconfig.jsonkeeps that compressed-tensors config undertext_config, and its config class also copies it to the top level (which the expert-parallel load in [18/21] drops);prepareinstalls the decompressor first of all, so a checkpoint the decompressor refuses leaves the config and the loader untouched. It installs it only for a checkpoint whosequantization_configdeclares compressed-tensors, as Kimi-K2.5, K2.6 and K2.7-Code all do: a bf16 checkpoint of the model type loads as Kimi-VL's does, and a block-FP8 one is dequantized by the expert-parallel load from the copy ofquantization_configKimi's config class keeps at the top level. Both adapters refuse, before changing anything, a checkpoint whose scheme sits ontext_configalone, with no top-level copy, unless it goes to the decompressor: the native text config drops it, and the load would read the experts as stored.Both checkpoints' modeling also predates transformers 5: it imports
is_torch_fx_available, which 5.x dropped, and its VLM class'stie_weights(self)rejects therecompute_mappingthatpost_initand the load pass. Each adapter installs import shims and wrapstie_weightsbefore the remote code loads; without them the real Kimi-K2.5 code fails on every supported release. Both adapters share one patch, inmodelopt/torch/models/_kimi.py, with the compat shims, thetie_weightswrapper and the refusals both prepares run first. The adapters' patch notes and the shims' warnings print on rank 0 only, and the shims both adapters share label their messageskimi. Both adapters keep the checkpoint's modeling module patched (the nativeDeepseekV3ForCausalLMand thetie_weightswrapper) whenever transformers resolves a class from it, so a module transformers executes again, or loads again under another name for another path to the checkpoint, keeps the patch; a patch re-applied that way is reported on the rank that re-applied it. A remote module that no longer defines a module-levelDeepseekV3ForCausalLM(renamed, or imported inside the VLM class) stops the load instead of passing for already patched, andkeep_patchedpatches remote-code modules only, leaving transformers' own same-named port module alone, also when transformers resolves a class defined there (a remote module can import its class from the port).Each adapter is an
EPSpec.preparehook on its model spec, so nothing outsidemodelopt/torch/models/kimi_*learns about Kimi.kimi_k25covers Kimi-K2.5 and the checkpoints sharing its model type (Kimi-K2.6, Kimi-K2.7-Code). The Kimi specs register only the expert-parallel section: the MoE and export sections staydeepseek_v3's, and sharing them, with a Kimi export test, is left to a follow-up. Each spec also states the plan itsprepareattaches (EPSpec.prepared_ep_plan, [13/21]): DeepSeek-V3's, now a sharedBASE_MODEL_EP_PLANconstant. So a run at an EP degree of 1, which skipsprepare, still finds the experts the checkpoint's own modeling builds, and refuses to hold them whole on every rank. That plan serves transformers before 5.13; from 5.13 transformers' owndeepseek_v3plan (the same experts plus anep_routergate entry) wins.transformers 5.13 added its own Kimi-K2.5 port, which a load without
trust_remote_codenow builds. The adapter's patch never reaches that port, so it refuses the port's config: the refusal names the port's module and asks fortrust_remote_code=Truefrom a checkpoint whoseauto_mapnames its ownAutoConfig(withtrust_remote_code=True, a checkpoint without one still gets the port). Kimi-VL's prepare runs the same check. The spec records the 5.13 release.Both adapters refuse a prepare without the checkpoint's path before they change anything, and Kimi-K2.5's spec declares the transformers 5.8 its decompression needs.
Usage
Picked automatically by
apply_distributed_supportforkimi_k25andkimi_vlconfigs undertorchrun; see [21/21].Testing
New
tests/unit/torch/models/test_kimi_k25_ep.pyandtest_kimi_vl_ep.py, against stand-in remote modules:DeepseekV3Config, with the EP plan and no staleauto_maporquantization_config; for Kimi-K2.5 with compressed-tensors, bf16 and block-FP8 experts, the decompressor installed for compressed-tensors ones only (the bf16 and block-FP8 cases fail with an unconditional install);config.jsondoes: the decompressor installed first, while that config is still present, and a decompressor refusal leavingtext_configand the remote-code staging untouched;tie_weightswrapper, against a stand-in that mirrors both checkpoints' remote code; the wrapper bound per class;prepare, when the remote modeling did not load or has noDeepseekV3ForCausalLM(that case fails without its fix), the config has no text config, or the config is transformers' own Kimi-K2.5 port, also withtrust_remote_code=True(that case fails without its guard); the refusal table gives both adapters a transformers config, and a block-FP8 scheme ontext_configalone (both cases fail without the refusal); the Kimi-K2.5 rebuild test's block-FP8 case carries the top-level copy the config class adds. A config without a text config is refused before anything is patched: the remote-code staging, and Kimi-VL's import shims, stay uninstalled;prepareattaches naming the same routed experts as the plan the spec states;get_class_in_modulewrapper, the kept patches and the compat shims.tests/unit/torch/models/test_remote_code.py: the patch put back on the module transformers loads again, after a forced re-execution and through a second path to the same checkpoint, and the re-patch reported. These tests show the reload handling. Another showskeep_patchedleaving transformers' own same-named module alone, also when a class defined there resolves through its wrapper, and a class whose module is not loaded resolving without a KeyError (both resolution cases fail without the fix).tests/unit/torch/utils/test_resolve_distribution_plan.py: a remote-code Kimi-K2.5 or Kimi-VL, as its own code builds it, left whole on every rank at an EP degree of 1 is refused like a native MoE, withpreparenever run.Mutation-checked: each test fails on at least one deliberate break of the code.
The tests written once per adapter are one test each, parametrized over Kimi-VL and Kimi-K2.5; each case still fails on its own adapter's regression.
A prepare without a checkpoint is refused before anything is patched, and transformers 5.7 is refused for Kimi-K2.5 by its spec.
Locally on this revision, the stack rebased on main 0638a27: pre-commit over the whole PR against main, at every PR of the stack; each fix's own tests at its PR; and at the stack's tip the unit tests of the changed areas (
tests/unit/torch/utils,export,models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.GPU tests were not re-run on this revision. Earlier real-model runs, with and without the reload fix (the patch kept across a module reload, above): Kimi-K2.7-Code (Kimi-K2.5 architecture, compressed-tensors pack-quantized experts under
text_config), expert-parallel over 32 GPUs (8 nodes, EP32), transformers 5.13.0, compressed-tensors 0.17.1, recipenvfp4_experts_only_mse-kv_fp8_cast, 4096 calibration samples of 4096 tokens, batch 1:DeepseekV3ForCausalLMon the native config and failed readingrope_thetaThe failure was seen once and not reproduced: the four later runs include two from an empty module cache, and the runs with the fix logged no re-patch. Those runs show load, calibration and export with the fix in place; the reload handling itself is shown by the unit tests above.
Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 19 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [18/21].
Summary by CodeRabbit