Repository navigation
[OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter - #3138
shengliangxu wants to merge 9 commits into
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review: 3
Please check the settings in the CodeRabbit UI or the
 |
|
/ok to test 3dbb3e2 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def _drop_tie_weights_kwargs(orig_tie): | ||
| """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. | ||
|
|
||
| Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote | ||
| ``tie_weights`` only calls its language model's, and the Kimi checkpoints set | ||
| ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. | ||
| """ | ||
|
|
||
| def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. | ||
| return orig_tie(self) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] The shim's safety rests on tie_word_embeddings=False, which nothing checks at runtime.
Dropping recompute_mapping is unconditionally fine, but dropping missing_keys is only fine under the docstring's premise that nothing is tied. to_native_backbone_config carries tie_word_embeddings over from the remote text_config verbatim, so a Kimi-VL/K2.5 variant (or a user-edited config) that ties embeddings would silently lose the tied-key bookkeeping: the native LM ties the weight but lm_head.weight is never removed from the load's missing_keys, so a strict load reports a spurious missing key — or, depending on the release's ordering, re-initializes it.
Cheap guard inside the wrapper, where the config is reachable:
def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc.
text_config = getattr(self.config, "text_config", self.config)
if getattr(text_config, "tie_word_embeddings", False):
raise RuntimeError(
"kimi: the tie_weights shim drops missing_keys, which is only safe for an untied "
"checkpoint, but this one sets tie_word_embeddings=True."
)
return orig_tie(self)That turns an assumption documented in a comment into one the code enforces, at no cost on the supported checkpoints.
(2) Why it matters: the failure is silent-to-confusing at load time rather than at patch time, on a path that only triggers for a checkpoint nobody tested. (3) Fix: the guard above, or state the assumption as a refusal in checkpoint_to_patch next to the other pre-patch refusals.
There was a problem hiding this comment.
Declined. I'm not adding the guard, because a tied checkpoint loads correctly through the shim. I documented what the shim loses instead, in d6ca10f.
I checked this with a small VLM whose tie_weights() takes no kwargs, as the Kimi remote code's does, wrapped by this shim:
- Its native DeepSeek-V3 language model has
tie_word_embeddings=True. - I saved it without
lm_head.weightand loaded it withfrom_pretrained. - On transformers 5.18 and 5.5.4,
lm_head.weightcomes back as the same tensor as the embedding, holding the saved values. - The only loss is the bookkeeping:
lm_head.weightstays in the load's missing-keys report. Nothing is re-initialized.
A refusal would therefore block a load that works. No supported checkpoint is tied in any case: Kimi-VL-A3B (Instruct and Thinking-2506) and Kimi-K2.5, K2.6 and K2.7-Code all set tie_word_embeddings=false at the top level and on text_config.
| model_type="kimi_vl", | ||
| modeling_source="remote_code", | ||
| # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. | ||
| ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), |
There was a problem hiding this comment.
[SUGGESTION] With only the EP section registered, is_moe(block, "kimi_vl") returns False for the very blocks prepare installs.
The docstring says the MoE/export sections "stay deepseek_v3's, so lookups by kimi_vl find neither yet", but it's worth being explicit about what that costs, because is_moe has fallbacks that do not rescue this case. After prepare, the language model's MoE blocks are native DeepseekV3MoE, and:
match_moe_block(module, "kimi_vl")→ the spec has nomoe_spec→None;- the name fallback misses (
"deepseekv3moe"ends in neithersparsemoeblocknor containsmoelayer); - the structural fallback misses (the router is
gate, notrouter).
That is exactly what deepseek_v3/specs.py documents as the reason it needs block_names=("DeepseekV3MoE",). So modelopt/torch/export/unified_export_hf.py:565/:949 and layer_utils.py:369 will skip the experts of a kimi_vl model rather than fail loudly. get_expert_linear_names happens to survive via its fused-container fallback, which makes the gap quieter, not smaller.
Nothing in this PR reaches those call sites, and kimi_vl was unregistered before, so there's no regression — hence a suggestion rather than a blocker. But since the follow-up that wires the PTQ entrypoint (#2754) will need it, and BASE_MODEL_EP_PLAN already establishes the precedent of sharing deepseek_v3's constants, reusing its sections here is a two-line change that closes the hole before anything can silently take the non-MoE path:
from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN
from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register
register(
ModelSpec(
model_type="kimi_vl",
modeling_source="remote_code",
# prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan.
ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN),
# prepare's language model is native deepseek_v3: its MoE blocks and export are too.
moe_spec=MoESpec(
block_names=("DeepseekV3MoE",),
expert_linear_names=("gate_proj", "down_proj", "up_proj"),
..., # deepseek_v3's, ideally hoisted to a shared constant as the plan was
),
export_spec=ExportSpec(grouped_expert_export=True),
)
)If you'd rather keep this PR to the EP section, consider making the gap loud instead of silent — e.g. have the follow-up tracked in the docstring name the lookups that miss (is_moe, hence the HF export's expert handling), so the next reader doesn't read "find neither yet" as "falls back to the generic detection".
There was a problem hiding this comment.
Taking the docstring half in d6ca10f; the shared sections stay with the follow-up.
The docstring now names the miss: is_moe's name and structure fallbacks do not recognize the native DeepseekV3MoE blocks that prepare installs, so the HF export's MoE-block handling skips them. #2748 adds the same sentence to Kimi-K2.5's spec, which has the same gap.
I'm keeping the sections out of this PR because registering them changes how the HF export treats these models:
- It enables the MoE pre-processing at the two export call sites you name.
- It enables grouped expert export.
That needs its own real-model export check, for both Kimi-VL and Kimi-K2.5 together. Today, for the fused experts these models load with, the skipped input-preparation handler is a no-op, since the amax fallback for fused experts runs during export. The AWQ pre_quant_scale averaging is the step that is actually skipped.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (the trigger comment carried no extra scoping). All 8 changed files opened (638 added lines): modelopt/torch/models/{_kimi.py, kimi_vl/{__init__,ep,specs}.py, deepseek_v3/specs.py, __init__.py} and both test files. Also read the base-branch machinery the adapter rides on — _remote_code.py (keep_patched, stage_full_remote_closure), _onfly_decompress.to_native_backbone_config, specs.py (EPSpec), utils/moe.py (apply_distributed_support, _ep_config, _routed_expert_keys, is_moe) — since the correctness of this PR is almost entirely in how it composes with them.
Findings
CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2
_kimi.py:118— thetie_weightsshim dropsmissing_keyson the premisetie_word_embeddings=False, which nothing enforces;to_native_backbone_configcarries that flag over from the remote config verbatim. A one-line guard in the wrapper turns the documented assumption into an enforced one.kimi_vl/specs.py:37— registering only the EP section meansis_moe(block, "kimi_vl")isFalsefor the nativeDeepseekV3MoEblocksprepareinstalls: both ofis_moe's fallbacks miss it (name doesn't end insparsemoeblock; router isgate, notrouter), which is precisely whydeepseek_v3needsblock_names. Not a regression —kimi_vlwas unregistered before, and no path here reaches the export call sites — but the follow-up will need it and it's a two-line reuse.
What I checked and found correct
- Plan plumbing.
prepareattaches the plan totext_config, and_ep_configreads the nested text config when the root has none, so the globs land on the tree they're written against.plan_config_attrbeing unset is right:base_model_ep_planisNoneon this spec, so the attr never selects a target. prepared_ep_planvs. whatprepareactually attaches. The stated plan's routed-expert keys reduce to{layers.*.mlp.experts}under_routed_expert_keys— identical whether modelopt'sBASE_MODEL_EP_PLANor transformers' own 5.13 plan wins, since the extraep_routergate entry is in neither_EXPERTS_STYLESnorgrouped_gemm. So the EP-degree-1 refusal finds the ModuleList experts the checkpoint's own code builds (language_model.model.layers.*.mlp.experts, matched byendswith), and the test asserts exactly this equivalence rather than the literal dict.- No aliasing through the shared constant.
BASE_MODEL_EP_PLANis now one dict referenced by two specs, but every consumer copies (dict(ep.base_model_ep_plan)inapply_distributed_support,dict(...)in_spec_ep_plan), so neither spec can mutate the other's plan. - Refuse-before-mutate. All five refusals in
checkpoint_to_patch/refuse_quantization_only_on_text_configrun before the shims, the staging wrapper, the config swap and the remote import — and the test asserts the loader is left unwrapped and the checkpoint's code unimported. Thetrust_remote_coderefusal in particular is load-bearing, sinceget_class_from_dynamic_moduledoesn't check it itself. Thestartswith("transformers.")port test doesn't false-positive ontransformers_modules.*(the trailing dot). - Shim targets match the remote import sites. The three stubs are set on the submodules DeepSeek's remote modeling imports from directly (
transformers.utils.import_utils,transformers.pytorch_utils,transformers.modeling_attn_mask_utils), not on thetransformers.utilsnamespace — which is what makesfrom X import Ysee them.is_torch_fx_available() -> Falsekeeps thetorch.fx.wrapbranch (and hence the_prepare_4d_causal_attention_maskstub) unreached; the stub raises rather than returning something plausible, which is the right call. - Patch durability.
keep_patchedis given the module suffix, so the swap survives transformers re-executing the module or loading it afresh under another hash — and since re-execution rebinds the VLM class too,to_native's earlyreturn Falseon an already-native module can't leave a fresh unshimmedtie_weightsbehind. Failing loudly when there's no module-levelDeepseekV3ForCausalLM, rather than treating it as "already patched", is the correct choice: silence there would build the per-expert layout EP cannot shard. - Version floor.
ep.py's "no floor beyond the EP framework's 5.7" matches_EP_MIN_TRANSFORMERS = "5.7", so the omittedrequires_transformersis right, and the test's module-level skip agrees. - Lazy imports.
modelopt.torch.utils{,.moe}are imported inside functions with the circularity reason named;_remote_codeis stdlib-only at module scope, so the registry import stays clean. - Tests. The refusal table, both loud-failure modes, shim installation, idempotency of the
tie_weightswrap, and the EP-degree-1 refusal against a ModuleList stand-in are all covered, with process-wide patches restored per test. Theinstall_onfly_decompressmonkeypatch in the first test is vacuous for Kimi-VL (bf16, nothing calls it) but reads as a deliberate guard on the shared path, so I left it alone.
Risk
Low. Additive and confined to the expert-parallel path: a new model type, one shared helper module, and a constant hoisted out of deepseek_v3's spec with no behavior change (the comment that moved with it is intact and still accurate). The only process-wide mutations are idempotent and installed after the refusals. The 184-line _kimi.py is shared with #2748, so its two parametrize lists standing at one entry each is expected rather than dead scaffolding.
🤖 Generated with Claude Code
3dbb3e2 to
e41e06d
Compare
6c49640 to
0864f57
Compare
|
/ok to test e41e06d |
|
/ok to test d6ca10f |
0864f57 to
a5d243e
Compare
d6ca10f to
4a459dc
Compare
|
/ok to test 4a459dc |
|
/ok to test 9e2cd9a |
|
/claude review |
| def _native_for_the_prepared_config(own_class, native_class): | ||
| """The ``DeepseekV3ForCausalLM`` the patched module builds its language model with. | ||
|
|
||
| transformers' native class for the native ``DeepseekV3Config`` ``prepare`` puts on | ||
| ``text_config``, and the checkpoint's own class for its own config. Deciding the plan prepares a | ||
| copy of the config before the plan is known, and the patch outlives that copy: a load the plan | ||
| keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own | ||
| language model, as it would with no patch. | ||
| """ | ||
| from transformers import DeepseekV3Config | ||
|
|
||
| def build(config, *args, **kwargs): | ||
| built = native_class if isinstance(config, DeepseekV3Config) else own_class | ||
| return built(config, *args, **kwargs) | ||
|
|
||
| setattr(build, "_kimi_native_swap", True) |
There was a problem hiding this comment.
[SUGGESTION] The swap now binds a plain function where a class used to live, so modeling_kimi_vl.DeepseekV3ForCausalLM stops behaving like a class for anything other than a direct call.
Under the two routes this PR exercises that holds — the LM is only ever built by the VLM's self.language_model = DeepseekV3ForCausalLM(config.text_config) — but any class-shaped access on that name now fails with an opaque AttributeError/TypeError instead of working: .from_pretrained, ._from_config, .config_class, issubclass, or AutoModelForCausalLM.from_config(text_config) resolving a text_config.auto_map["AutoModelForCausalLM"] that names modeling_kimi_vl.DeepseekV3ForCausalLM (a non-EP run keeps the checkpoint's auto_map on text_config, and keep_patched's get_class_in_module wrapper hands the patched binding back to from_pretrained). _BaseAutoModelClass.register would even accept it, since hasattr(build, "config_class") is False, and only fail a line later on _from_config.
A callable object that forwards attribute lookups to the native class keeps the dispatch and the class-like surface, for about the same number of lines:
def _native_for_the_prepared_config(own_class, native_class):
"""..."""
from transformers import DeepseekV3Config
class _Build:
_kimi_native_swap = True
def __call__(self, config, *args, **kwargs):
built = native_class if isinstance(config, DeepseekV3Config) else own_class
return built(config, *args, **kwargs)
def __getattr__(self, name): # from_pretrained, config_class, ... land on the native class
return getattr(native_class, name)
return _Build()getattr(current, "_kimi_native_swap", False) in to_native keeps working unchanged. If you'd rather keep the function, a sentence in the docstring saying the binding is no longer a class, and that only a direct call is supported, would at least make the constraint explicit for the next reader.
There was a problem hiding this comment.
Kept as a function: nothing reaches this name except the direct call.
- The checkpoints' own code. In both Kimi-VL's
modeling_kimi_vl.pyand Kimi-K2.5'smodeling_kimi_k25.py, the module-level name's only use isself.language_model = DeepseekV3ForCausalLM(config.text_config). Nothing takesisinstance/issubclassof it, subclasses it, or reads an attribute from it. - transformers' auto paths.
get_class_in_module, including thekeep_patchedwrapper, is only asked for names anauto_mapdeclares.- The top-level
auto_maps name the VLM classes. - Kimi-VL's
text_confighas noauto_map. - Kimi-K2.5, K2.6 and K2.7-Code name
modeling_deepseek.DeepseekV3ForCausalLM. That module's class is the one the swap leaves alone: the swap rebinds only the namemodeling_kimi_k25imported, andkeep_patchedapplies only to modules named*.modeling_kimi_k25/*.modeling_kimi_vl.
- The top-level
- Our code. ModelOpt refers to
DeepseekV3ForCausalLMonly as an architecture string.
The forwarding object would also make config_class, from_pretrained and similar answer for the native class even when the call would build the checkpoint's own. A docstring line stating the call-only contract is owed for this PR's next push.
| 2. Swap ``hf_config.text_config`` to a transformers-native ``DeepseekV3Config`` (drops the remote | ||
| ``auto_map`` so the native class -- and its ``base_model_ep_plan`` -- is used for the LM). | ||
| 3. Apply the deepseek_v3 EP plan to that native ``text_config`` (post_init prefixes it to the VLM | ||
| root as ``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}``). | ||
| 4. Monkey-patch ``modeling_kimi_vl.DeepseekV3ForCausalLM`` -> the transformers-native class, so | ||
| ``KimiVLForConditionalGeneration.__init__`` (``self.language_model = DeepseekV3ForCausalLM( | ||
| config.text_config)``) builds the EP-capable native LM while MoonViT + the KimiVL wrapper stay | ||
| from the bundled code. Native deepseek_v3's per-expert WeightConverter fuses the checkpoint's | ||
| per-expert ``.weight`` tensors into the 3-D params on load (no decompressor -- weights are bf16). | ||
|
|
||
| Version: no floor beyond the EP framework's 5.7; validated on transformers 5.9.0. The compat shims |
There was a problem hiding this comment.
[SUGGESTION] Three lines of this module docstring now describe the swap as it was before _native_for_the_prepared_config landed, and one records a transformers version the PR no longer tests:
- Step 2 (line 32-33) and the comment on line 67 both say dropping the remote
auto_mapis what makes the native class get used for the LM. It is not, any more: the VLM's__init__calls the module-levelDeepseekV3ForCausalLM, and which class that builds is now decided byisinstance(config, DeepseekV3Config). Droppingauto_mapstill matters (it keepstext_configfrom naming remote classes), but it is no longer the mechanism. - Step 4 (line 36) says the patch points the name at "the transformers-native class". It points it at a factory that picks the native class for the prepared config and the checkpoint's own class for the config as read — which is the whole point of the load-off-EP fix, and the one thing a reader of this file would most want to know.
- Line 42 says "validated on transformers 5.9.0", while the PR reports 5.5.4 and 5.18.
Per CLAUDE.md's "Keep skill/doc edits concise … comments say why, not what", this is a tightening pass rather than new text — e.g. fold steps 2 and 4 into one line saying the native config selects the native LM, and update the version to the pair actually run.
There was a problem hiding this comment.
Agreed; postponed to this PR's next push to save a CI round on an approved PR. The pass will:
- fold steps 2 and 4 into one line saying the native config selects the native language model, and that dropping
auto_maponly keepstext_configfrom naming remote classes; - describe step 4 as the per-config build;
- replace "validated on transformers 5.9.0" with the versions this PR runs.
Kimi-K2.5's module docstring in #2748 says the same about the swap, so it gets the same pass.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (the trigger comment carried no extra scoping). All 8 changed files opened (686 added lines): modelopt/torch/models/{_kimi.py, kimi_vl/{__init__,ep,specs}.py, deepseek_v3/specs.py, __init__.py} and both test files. The checkout is shallow, so I could not diff against the previously-approved head (3dbb3e2) directly and re-reviewed the whole PR diff instead; per the description the only code delta since that approval is _native_for_the_prepared_config, plus docstrings answering last round's two suggestions.
Findings
CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2
_kimi.py:134— the swap now binds a plain function where a class was, so.from_pretrained/._from_config/.config_class/issubclassonmodeling_kimi_vl.DeepseekV3ForCausalLMfail opaquely. Correct for the two routes here (the LM is only ever built by a direct call), but a callable object forwarding__getattr__to the native class keeps the class-like surface for the same line count.kimi_vl/ep.py:32— three docstring lines (steps 2 and 4, and the comment on line 67) still say dropping the remoteauto_mapis what selects the native LM; the factory'sisinstancecheck is. And line 42's "validated on transformers 5.9.0" is behind the 5.5.4/5.18 the PR reports.
What I checked and found correct
- The new fix is right, and necessary.
resolve_distribution_planreturnsprepared[0]only whenplan.is_expert_parallel, and the unpreparedhf_configotherwise (model_load_utils.py:1789) — while the patch is process-wide, installed by theapply_distributed_supportevery rank ran before the plan was known. So a replicate or FSDP2 plan builds the VLM from the config as read, and previously handed the native class the checkpoint's own text config. Dispatching onisinstance(config, DeepseekV3Config)picks the checkpoint's own class there, which is what no patch at all would have done. The FSDP2 route confirms it independently:model_load_utils.py:862re-reads the config fresh and_sharded_load_kwargsprepares only for_EXPERT_PARALLEL. - The EP path keeps the native type end to end. Nothing between
prepareand the build serializes the config:_decide_distributionandload_pretrained:877usecopy.deepcopy(type-preserving), the plan broadcast carries the plan only (each rank prepares its own copy), and_from_config's own deepcopy preserves it too. Ato_dict()/from_dict()round trip would have rebuilttext_configthrough the remoteKimiVLConfig.__init__and silently selectedown_classunder an EP plan; there is no such round trip on these paths. - Idempotency and durability of the swap.
prepareruns twice in a real EP load (deciding, thenload_pretrained:883); the secondto_nativereturnsFalseon the_kimi_native_swapmarker, andkeep_patchedstill reports a non-emptypatchedlist (it lists matched module names, not changed ones), so no spuriousRuntimeError. On re-execution the module-level name is the fresh remote class, unmarked, so both the swap and thetie_weightsshim re-apply together — they are coupled on the same branch, which is what keeps a re-executed module from ending up swapped but unshimmed. - Plan plumbing.
plan_config_attrbeing unset is right:base_model_ep_planisNoneon this spec, so theattrsbranch inapply_distributed_supportnever selects a target, andprepareattaches the plan totext_configitself where_ep_configreads it.apply_distributed_support("deepseek_v3", text_config)cannot recurse (deepseek_v3 has noprepare), and its floor is the framework's 5.7, matchingep.py's claim and the test's module-level skip. prepared_ep_planmatches what the model actually builds._routed_expert_keys(BASE_MODEL_EP_PLAN)reduces to{layers.*.mlp.experts}— identical whether modelopt's plan or transformers' own 5.13 plan wins, since the extraep_routergate entry is in neither_EXPERTS_STYLESnorgrouped_gemm— and_routed_expert_modulesmatcheslanguage_model.model.layers.*.mlp.expertsbyendswith. So the EP-degree-1 refusal finds the ModuleList experts the checkpoint's own code builds, andtest_model_load_utils.py's new test exercises exactly that withpreparestubbed to fail if it runs.- No aliasing through the hoisted constant.
BASE_MODEL_EP_PLANis one dict referenced by two specs, butapply_distributed_supportand_spec_ep_planboth copy it, so neither spec can mutate the other's plan. The comment that moved with it is intact and still accurate. - Refuse-before-mutate. All five refusals run before the shims, the staging wrapper, the config swap and the remote import, and the test asserts the loader is left unwrapped and the checkpoint's code unimported. The
trust_remote_coderefusal is load-bearing, sinceget_class_from_dynamic_moduledoes not check it. Thestartswith("transformers.")port test does not false-positive ontransformers_modules.*(the trailing dot). - Registering
kimi_vlis safe for the existing spec suites.modeling_source="remote_code"withmin_transformers_version=Nonesatisfiestest_specs_vs_transformers.py:119, and transformers ships nokimi_vl, sotest_remote_code_models_are_really_absentholds. TheEP_MODEL_TYPESsweeps intest_ep_specs.pypass: the plan is well formed,preparesatisfies the declares-something check, andplan_config_attris empty.deepseek_v3is imported beforekimi_vlin theisort: skipblock, and the transitive import fromkimi_vl.specscould not double-register either way. - Lazy imports.
modelopt.torch.utils{,.moe}are imported inside functions with the circularity reason named;_remote_codeis stdlib-only at module scope, so the registry import stays clean. - Tests. The
_restorehelper's absent-attribute branch unwinds correctly (monkeypatch undoes in reverse, leaving the attr deleted), so the compat shims do not leak across tests. The two single-entry parametrize lists are scaffolding for #2748's cases, as the description says, not dead code.
Last round's two suggestions were answered by docstrings rather than code — the tie_weights shim now states that a tied checkpoint still loads with lm_head.weight reported missing, and kimi_vl/specs.py states that is_moe misses the native DeepseekV3MoE blocks until the MoE and export sections are shared. Both are the author's call and I am not re-raising them.
Risk
Low. Additive and confined to the expert-parallel path. The one new code path narrows a previously unconditional monkey-patch, and the narrowing is in the safe direction: off expert parallelism the model is now built exactly as it would be with no patch at all, which is what the new test pins.
🤖 Generated with Claude Code
9e2cd9a to
c22b9b2
Compare
a5d243e to
a4cfa16
Compare
|
/ok to test c22b9b2 |
3c641ac to
8bb2e35
Compare
b4847b8 to
1c2910a
Compare
|
/ok to test 1c2910a |
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: the prior loading-coverage gap and unusable quantization remedy remain unresolved.
Needs action:
- Add bounded Kimi-VL loading/logit-parity coverage in
tests/gpu/torch/utils/test_expert_parallel_load.py, including expert slices, preserved dense/vision weights, and a non-EP forward after preparation. - Fix
_kimi.py’s quantization remedy for Kimi-VL; recommend an unquantized/dequantized checkpoint rather than copying the scheme to the root. - Move
test_kimi_vl_ep.pytotests/unit/torch/models/kimi_vl/test_ep.py, followingCONTRIBUTING.md. - 💬 Author postponed the documentation pass — shorten/correct
kimi_vl/ep.pyand fix the PR’s “EP-only” claim to reflect preparation before plan selection. - Document the optional-Transformers rationale for
_kimi.py’s local imports; retain the existing circular-import explanation.
No action needed:
- ✔️ Earlier suggestions about tied-key bookkeeping and the MoE/export detection gap are now explicitly documented.
| with torch.device("meta"): | ||
| assert type(remote.vlm(prepared).language_model) is DeepseekV3ForCausalLM | ||
| assert type(remote.vlm(as_read).language_model) is remote.remote_lm | ||
|
|
There was a problem hiding this comment.
Bot comment.
The previous critical coverage request is still unresolved. The positive fixture replaces get_class_from_dynamic_module with an in-memory module, and this test constructs on meta without loading tensors or running a forward. I rechecked tests/gpu/torch/utils/test_expert_parallel_load.py: its loading/parity cases still use Qwen3-MoE, not this nested DeepSeek-V3 adapter. Extend it with a tiny local remote-code VLM checkpoint that exercises the real adapter and Transformers loader. Verify each rank's expected expert slices, dense/vision/projector values, and logits against an unpatched whole-model reference, with bounded collectives and rank-safe failure reporting. Include a real non-EP load/forward after preparation to protect the process-wide fallback. This is the central behavior, not extra helper coverage; CONTRIBUTING.md requires exercising the actual runtime path.
There was a problem hiding this comment.
Added in 3925d67: test_kimi_vl_loads_expert_parallel_and_matches_the_whole_model in tests/gpu/torch/utils/test_expert_parallel_load.py.
- It writes a tiny Kimi-VL-shaped remote checkpoint (
tests/_test_utils/torch/kimi_vl_tiny_modeling.py), with its experts stored per expert, and loads it expert-parallel through the spec'sprepareand transformers' loader. - Each rank must hold its own experts, equal to the per-expert tensors fused.
- The vision tower and projector must load as stored, and nothing may be left unplaced.
- Each rank's logits must match the whole model's on one GPU.
- Every rank's checks are gathered under bounded collectives.
Its non-distributed parts pass on CPU on transformers 5.18: the swap, the exact fusion, the nested plan at the VLM root as language_model.model.layers.*, and the logits. A run on 2 and 4 GPUs is requested; the description will record it. A non-expert-parallel load after preparation is left out: the tiny code's own language model is transformers' DeepSeek-V3, so it would not exercise the fallback.
| 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 " |
There was a problem hiding this comment.
Bot comment.
The earlier remedy issue remains unchanged. Copying quantization_config to the root gets past this guard, but load_and_prepare_ep then reads it through _ckpt_quant_method and refuses it: _EP_SOURCE_FORMATS contains only None, and kimi_vl.prepare installs no decompressor. Could Kimi-VL recommend starting from an unquantized/dequantized checkpoint instead, reserving root-config advice for an adapter that actually supports decompressing the scheme? Cover the supported remedy through loading rather than only asserting the first refusal.
There was a problem hiding this comment.
Fixed in fed6e2f. The refusal now names the way in, an unquantized copy of the model. It no longer suggests copying the scheme to the top-level config, which only leads to the expert-parallel load's own refusal of a quantized checkpoint.
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| """Expert-parallel adapter for Kimi-VL-A3B (``kimi_vl``), a remote-code VLM on a Moonlight MoE. |
There was a problem hiding this comment.
Bot comment.
The previous filename request is still open. These cases primarily exercise modelopt/torch/models/kimi_vl/ep.py::prepare; CONTRIBUTING.md says test filenames must mirror their production module. Move the adapter tests to tests/unit/torch/models/kimi_vl/test_ep.py. If shared _kimi.py cases become direct helper tests, place those separately in test_kimi.py rather than retaining the adapter-specific name.
There was a problem hiding this comment.
Moved in 6a84c62 to tests/unit/torch/models/kimi_vl/test_ep.py, mirroring kimi_vl/ep.py. The directory is a package, so it can hold a test_ep.py beside the other adapters'.
| decompressor (the checkpoint ships plain bf16 per-expert weights). | ||
|
|
||
| Challenge (identical to kimi_k25): the bundled ``modeling_kimi_vl.py`` builds its language model from | ||
| an INLINE ``DeepseekV3ForCausalLM`` whose ``DeepseekV3MoE`` is an ``nn.ModuleList`` of per-expert |
There was a problem hiding this comment.
Bot comment.
The agreed documentation pass has not landed. Please replace the challenge/solution walkthrough with a short adapter contract and keep the design discussion in the PR description, per CONTRIBUTING.md. The current explanation also remains stale: _native_for_the_prepared_config dispatches on the native config's type, not on dropping auto_map, and the patched binding is a per-config function, not simply the native class. Remove the old validation-version prose or put actual testing information in the PR body.
There was a problem hiding this comment.
Done in 6a84c62. The module docstring is now the adapter's contract, close to your wording, and the walkthrough is in the PR description.
| keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own | ||
| language model, as it would with no patch. | ||
| """ | ||
| from transformers import DeepseekV3Config |
There was a problem hiding this comment.
Bot comment.
The earlier local-import documentation request is still open. The module-scope comment explains the local ModelOpt imports, but not these Transformers imports or those in the shim/patch helpers. pyproject.toml makes Transformers an optional [hf] dependency, and model registration imports this module. Keep these imports lazy and add a brief comment naming that concrete reason, as CONTRIBUTING.md requires.
There was a problem hiding this comment.
Done in fed6e2f. _kimi's module-scope comment now says why both modelopt.torch.utils and transformers are imported inside its functions: the model registry imports this module, and it must import without transformers, the optional [hf] extra.
| modeling_source="remote_code", | ||
| # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. | ||
| ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), | ||
| ) |
There was a problem hiding this comment.
Bot comment.
The PR checklist still says this is additive and 'runs only on the expert-parallel path', but registering this hook makes resolve_distribution_plan run it on every rank before deciding the plan whenever EP remains a candidate. It can therefore install shims/patches or raise during a decision that ultimately selects a non-EP layout. The main description already explains the per-config fallback; please make the checklist consistent with it. The later #2748 change cannot establish an EP-only guarantee for this independently landable PR.
There was a problem hiding this comment.
Fixed in the description. The checklist now says what registering the spec reaches beyond the expert-parallel load in this PR on its own:
- In a multi-rank run where expert parallelism is still a candidate, every rank runs
preparebefore the plan is decided. That shims the code's imports and patches its modeling module, which builds the native class only for the configpreparemade. It can also refuse the checkpoint, failing the decision. - A run that would leave the MoE whole on every rank, with nothing asking for that, is refused.
- A single-process load is unchanged.
#2748 later runs prepare only once the plan is decided expert-parallel.
1c2910a to
3925d67
Compare
8bb2e35 to
cb62c2b
Compare
|
/ok to test 3925d67 |
|
meenchen-nv, on your two reviews: every item is answered in its thread. The duplicates the second review re-posted point to the same commits.
|
|
/claude review |
| checkpoint's own language model is a DeepSeek-V3 whose MoE holds per-expert Linears; here the name is | ||
| transformers' DeepSeek-V3 itself, since only the native class the adapter swaps in is built. | ||
| """ | ||
|
|
||
| from torch import nn | ||
| from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig, PreTrainedModel |
There was a problem hiding this comment.
[SUGGESTION] The stand-in's module-level DeepseekV3ForCausalLM is transformers' native class, so the new GPU test cannot detect a regression in the one thing the adapter exists to do.
What. patch_language_model_to_native replaces this name, and _native_for_the_prepared_config is then built with own_class is native_class — both isinstance branches return the same class. So test_kimi_vl_loads_expert_parallel_and_matches_the_whole_model would pass byte-for-byte if the swap (and the tie_weights shim, since it rides the same branch) became a no-op: language_model.model.layers[1].mlp.experts would already be the fused native experts, the per-expert tensors would still fuse via the native WeightConverter, and the logits would still match.
Why it matters. What the test does prove end-to-end is real and worth having — per-expert → fused 3-D fusion on each rank's own slice, vision tower/projector preserved, nothing unplaced, logit parity. But the swap itself is still covered only by the unit stand-in in tests/unit/torch/models/kimi_vl/test_ep.py, where the remote LM is a stub that is never loaded or forwarded. The no-class loud-failure case is likewise the only guard that the real Kimi-VL layout (an nn.ModuleList of per-expert Linears) never reaches the EP load. That leaves the reviewer's "end-to-end loading coverage" ask only partly closed, and in particular nothing exercises a real load of the per-expert module layout.
Suggested fix. Give the stand-in a module-level name the swap must actually replace, so building it is the regression signal. It never gets constructed on the passing path, so it needs no working forward:
class _OwnDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):
"""The checkpoint's own: routed experts as per-expert Linears, which EP cannot shard."""
def __init__(self, config):
super().__init__(config)
for layer in self.model.layers:
if hasattr(layer.mlp, "experts"):
layer.mlp.experts = nn.ModuleList(
nn.Linear(config.hidden_size, config.moe_intermediate_size, bias=False)
for _ in range(config.n_routed_experts)
)
# The name the adapter swaps for transformers' native class.
DeepseekV3ForCausalLM = _OwnDeepseekV3ForCausalLMWith that, a no-op swap builds the ModuleList layout and the expert-slice assertions fail loudly instead of passing. Keep the docstring note, but say the name is the checkpoint's own class rather than the native one.
Non-blocking — the current combination of unit and GPU coverage is defensible, but this is a few lines for a materially stronger test.
There was a problem hiding this comment.
Taken in 3b5d7e2. The stand-in's module-level DeepseekV3ForCausalLM is now the checkpoint's own class, which holds its routed experts as per-expert Linears. Only a swap that misses builds it, so a no-op swap would fail the GPU case's expert checks. Its CPU dry run still passes: the swap, the exact fusion and the logits.
|
/ok to test 3b5d7e2 |
3b5d7e2 to
d2321c1
Compare
cb62c2b to
b6e1f97
Compare
|
/ok to test d2321c1 |
d2321c1 to
cb90789
Compare
b6e1f97 to
abca417
Compare
|
/ok to test cb90789 |
Kimi-VL ships its own modeling, which builds the language model's MoE from an nn.ModuleList of per-expert Linears that transformers' expert- parallel machinery cannot shard. Its text backbone is DeepSeek-V3, so the adapter swaps in transformers' native deepseek_v3 language model (fused gate_up_proj/down_proj experts, and the plan to shard them) and keeps the bundled vision tower and projector from the checkpoint's code. The patch Kimi-VL and Kimi-K2.5 share lives in models/_kimi.py: the refusals prepare runs before it changes anything (transformers' own port, no text_config, no checkpoint path, untrusted code, a scheme on text_config alone), the import shims for symbols transformers 5 dropped, and the swap of the module-level DeepseekV3ForCausalLM, kept on a module reload through keep_patched, with a tie_weights wrapper for the kwargs newer transformers passes. deepseek_v3's expert-parallel plan becomes a shared BASE_MODEL_EP_PLAN, which the spec states as the plan prepare attaches, so a run at an EP degree of 1, which skips prepare, still finds the experts the checkpoint's own code builds and refuses to hold them whole on every rank. Moved out of the Kimi adapters change, whose history is archived. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Both Kimi checkpoints' code builds its routed experts at .mlp.experts, so the moe parameter varied nothing. The stand-in builds them there, and the test takes the model type only. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…e out The spec registers only its expert-parallel section, and is_moe's fallbacks do not recognize the native DeepseekV3MoE blocks prepare installs, so the HF export's MoE-block handling skips them until the sections are shared. The tie_weights shim drops missing_keys: a tied checkpoint would still load its embedding into both, with lm_head.weight reported missing (checked on transformers 5.18 and 5.5.4). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Deciding the plan applies the model spec to a copy of the config on every rank before the plan is known, and the Kimi adapters' prepare patches the checkpoint's modeling module for the process. A plan that then keeps the load off expert parallelism (replicated, or FSDP2) builds from the config as read, and the patched VLM class built transformers' native DeepseekV3ForCausalLM from the checkpoint's own text config, which it cannot take (on a stand-in, an AttributeError on pad_token_id). The swapped name now builds the native class only for the native DeepseekV3Config prepare puts on text_config, and the checkpoint's own class for its own config, so that load gets the language model it would with no patch. Both checkpoints' VLM classes use the name only to build their language model. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…e Kimi keeps the RuntimeError for a transformers below what the model's spec needs, which names the release. The remote-Kimi case expects whichever arm its transformers takes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…antized copy From meenchen-nv's review: the refusal of a quantization_config that only text_config declares advised declaring it on the top-level config too, which gets past this refusal only to meet the expert-parallel load's own refusal of a quantized checkpoint. It now names the way in, an unquantized copy. The module-scope comment also says why transformers is imported inside _kimi's functions: the optional [hf] extra, and the model registry imports this module. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From meenchen-nv's review: kimi_vl/ep.py's module docstring is now the adapter's contract, and its walkthrough moved to the PR description. The adapter's tests move to tests/unit/torch/models/kimi_vl/test_ep.py, mirroring kimi_vl/ep.py, as CONTRIBUTING asks; the directory is a package, so it can hold a test_ep.py beside the other adapters'. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From meenchen-nv's review: no test loaded checkpoint tensors through the Kimi-VL adapter. New 2-GPU case test_kimi_vl_loads_expert_parallel_and_ matches_the_whole_model writes a tiny Kimi-VL-shaped remote checkpoint (tests/_test_utils/torch/kimi_vl_tiny_modeling.py), its experts stored per expert, and loads it expert-parallel through the spec's prepare: each rank's native experts must equal the per-expert tensors fused, the vision tower and projector must load as stored, nothing may be left unplaced, and each rank's logits must match the whole model's on one GPU, with every rank's checks gathered under bounded collectives. Its non-distributed parts were checked on CPU: the swap, the exact fusion, the nested plan at the VLM root and the logits. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: the stand-in's module-level DeepseekV3ForCausalLM was transformers' native class, so the GPU case would pass even if the adapter's swap did nothing. It is now the checkpoint's own: routed experts as per-expert Linears, which only a swap that misses builds, and the case's expert checks then fail. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
abca417 to
dc200c3
Compare
cb90789 to
7849516
Compare
|
/ok to test 7849516 |
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: EP-sharded loading coverage is fixed, but real non-EP loading and forwarding after the process-wide patch remain untested.
Needs action:
- 💬 Author added the EP GPU case — extend
test_expert_parallel_load.pywith a real unprepared-config non-EP load/forward after preparation, checking checkpoint-owned weights and logits against an unpatched reference; see inline.
No action needed:
- ✔️ Resolved since the previous review: EP fusion/slice/logit coverage, the vacuous swap fixture, quantization advice, test placement, concise adapter contract, local-import rationale, and the description’s preparation scope.
| with torch.device("meta"): | ||
| assert type(remote.vlm(prepared).language_model) is DeepseekV3ForCausalLM | ||
| assert type(remote.vlm(as_read).language_model) is remote.remote_lm | ||
|
|
There was a problem hiding this comment.
Bot comment.
The added GPU case closes the EP-sharded loading part of the previous request, but this non-EP path still checks only class identity on meta. Its remote LM has no weights or forward, while the GPU reference loads copy.deepcopy(prepared) and therefore selects the native LM too. Moreover, the tiny checkpoint's as-read text config is already a native DeepseekV3Config, so simply loading it without prepared would not exercise own_class.
This matters in this PR independently: resolve_distribution_plan can prepare a copy and then return the original config for a non-EP plan. The compatibility shims and tie wrapper remain installed, and the checkpoint-owned LM—not the native LM used to justify the unreachable legacy-mask stub—then runs.
Please extend the existing loading case with a working checkpoint-owned backbone and a distinct remote text-config type. After preparing a copy, load the original config without EP, assert the checkpoint-owned class and saved values survive, and run logits against an unpatched reference. That would settle the remaining backward-compatibility coverage request without adding more mock-only tests.
There was a problem hiding this comment.
On the transformers releases this adapter supports, a later load off expert parallelism builds and runs the checkpoint's own language model, unchanged, so I have not added a real-load comparison:
- The swap chooses by config type (
_native_for_the_prepared_config): transformers' native class only for transformers' ownDeepseekV3Config, which only the copypreparemakes carries. Kimi-VL'sconfiguration_kimi_vl.py, like Kimi-K2.5'sconfiguration_deepseek.py, defines its ownDeepseekV3Config(PretrainedConfig)and buildstext_configfrom it, so the config as read builds the checkpoint's own class. This test pins that with a remote text-config type of its own: afterprepareon a copy, the as-read config builds exactly the checkpoint's language-model class, so its__init__, forward and state-dict names are the checkpoint's. - The legacy-mask stub never lands there: transformers 5.5.4 through 5.18.0 still define
_prepare_4d_causal_attention_mask, and the shims only fill names a release lacks, so the checkpoint's ownDeepseekV3Model.forwardcalls transformers' function. On those releases the shims add onlyis_torch_fx_available, returningFalse, which makes the remote code skiptorch.fx.wrap. - Nothing else
prepareinstalls reaches that load's weights or forward. Kimi-VL installs no decompressor and registers no weight conversion, and thetie_weightswrapper drops the keyword arguments and calls the checkpoint's owntie_weights, which the test above this one pins. - From [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748, the plan decision runs no
preparefor a plan off expert parallelism at all, only the build step, which installs no swap:test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code[fsdp2|whole]pins that the module keeps its ownDeepseekV3ForCausalLMand that the loader hooks are unchanged.
A real load and forward would run that same class through the same loader, comparing it with itself.
Note
Split out of #2748 at the reviewers' request (2026-10-30): the Kimi-VL adapter, the patch it shares with the Kimi-K2.5 adapter, and the shared deepseek_v3 plan. #2748 keeps the Kimi-K2.5 adapter.
Base:
shengliangx/hfep-base-11b-kimi, a branch (not a PR) that merges the PRs this one builds on, #2746 (the model specs) and #2747 with #3137 (on-the-fly decompression), with main, which holds the remote-code helpers (#3130) since they merged. The diff below is this PR's alone, and it merges after #2746 and #2747.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 onlyWhat does this PR do?
Type of change: New feature
Kimi-VL ships its own modeling. It builds the language model's MoE from an
nn.ModuleListof per-expert Linears, which transformers' expert-parallel machinery cannot shard. Its text backbone is DeepSeek-V3, so the adapter swaps in transformers' nativedeepseek_v3language model (fusedgate_up_proj/down_projexperts, and the plan to shard them). It keeps the bundled vision tower and projector from the checkpoint's own code.modelopt/torch/models/_kimi.py:text_config; a missing checkpoint path; untrusted remote code (get_class_from_dynamic_moduleruns the checkpoint's code without checkingtrust_remote_code); and a quantization scheme ontext_configalone, which the native text config drops.is_torch_fx_available, ...). They stub only code paths the native language model never runs.DeepseekV3ForCausalLMbuilds the native class from the nativeDeepseekV3Configprepareputs ontext_config, and the checkpoint's own class from its own config. Atie_weightswrapper covers therecompute_mappingnewer transformers passes. A remote module that defines no module-levelDeepseekV3ForCausalLMstops the load instead of passing for already patched.prepareon a copy of the config on every rank before the plan is known, and the patch lasts for the process. A plan that then keeps the load off expert parallelism (replicated, or FSDP2) builds from the config as read, so it gets the checkpoint's own language model, as with no patch. Before, the native class was handed the checkpoint's own text config, which it cannot take. The other effects ofprepareare harmless off expert parallelism: the import shims only add symbols the remote code imports, the staging copies its sibling modules, thetie_weightswrapper only drops keyword arguments for these untied checkpoints, and Kimi-K2.5's decompressor rewrites only the configs it marked, which are the prepared copies.keep_patched([OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130), on a module transformers executes again or loads again under another name. The reasons are in [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130's description.kimi.EPSpec.preparehook on thekimi_vlspec, so nothing outsidemodelopt/torch/modelslearns about Kimi.deepseek_v3's, and sharing them is left to a follow-up.prepareattaches (EPSpec.prepared_ep_plan, [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746): deepseek_v3's, now a sharedBASE_MODEL_EP_PLAN. So a run at an EP degree of 1, which skipsprepare, still finds the experts the checkpoint's own code builds, and refuses to hold them whole on every rank.The adapter's walkthrough (moved from
kimi_vl/ep.py's module docstring)Usage
Picked automatically by
apply_distributed_supportforkimi_vlconfigs undertorchrun; see #2754.Testing
tests/unit/torch/models/test_kimi_vl_ep.py, against a stand-in remote module, covers:DeepseekV3Configwith the EP plan;tie_weightswrapper;prepareran on a copy of its config, building the checkpoint's own language model (it failed before the fix);DeepseekV3ForCausalLM;prepareattaches naming the same experts as the plan the spec states.tests/unit/torch/utils/test_model_load_utils.py: a remote-code Kimi-VL, as its own code builds it, is refused when left whole on every rank at an EP degree of 1, withpreparenever run..mlp.experts.tests/unit/torch/modelsandtest_model_load_utils.pypass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.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.is_moemisses the nativeDeepseekV3MoEblocksprepareinstalls until the MoE and export sections are shared; the tie shim's says a tied checkpoint still loads, withlm_head.weightreported missing (checked on transformers 5.18 and 5.5.4). 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.tests/unit/torch/models,test_model_load_utils.py,test_ep_loader.py,test_unplaced_source_keys.py.preparemade (review item from [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746). The new test fails without the fix. 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.tests/unit/torch/models,test_model_load_utils.py,test_ep_loader.pyandtest_unplaced_source_keys.pypass on CPU, torch 2.14: 572 on transformers 5.18, and 447 passed, 56 skipped on 5.5.4.moe.config_levelshelper. [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 now raises the generic left-whole refusal as a ValueError and keeps the RuntimeError for a transformers below the spec's floor; the remote-Kimi refusal test expects whichever its transformers takes. 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.tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4.prepareonly for an expert-parallel plan. 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.test_kimi_vl_loads_expert_parallel_and_matches_the_whole_modelwrites a tiny Kimi-VL-shaped remote checkpoint (tests/_test_utils/torch/kimi_vl_tiny_modeling.py), with its experts stored per expert, and loads it expert-parallel through the spec'sprepare. Each rank's native experts must equal the per-expert tensors fused, the vision tower and projector must load as stored, nothing may be left unplaced, and each rank's logits must match the whole model's on one GPU. It passed on 2 GPUs and on 4 GPUs, with transformers 5.18.0, in a run of the whole GPU file at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 183bf95, which contains this PR. It loads through the spec'sprepareand the loader directly, not through the plan decision;text_configalone now names an unquantized copy as the way in;_kimi's module comment says why transformers is imported inside its functions, and the tests move totests/unit/torch/models/kimi_vl/test_ep.py;tests/unit/torch/modelsandtests/unit/torch/utilspass on transformers 5.18 and 5.5.4.DeepseekV3ForCausalLMis now the checkpoint's own, per-expert class, so a swap that misses fails the GPU case's expert checks; its CPU dry run passes. Pre-commit against main passes;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. At [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 4389363, which contains this head, the whole GPU file passed at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 4389363, 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. The Kimi-VL case passed on both, now with the tiny model's own per-expert language model, so a swap that missed would have failed it.test_model_load_utils.pynow holds the moved expert-parallel loader tests; its one import conflict kept both sides' imports, and the own diff is otherwise unchanged. 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. The Kimi-VL case passed on both; it loads 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.tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.[ep-half]skipped, as it needs 4; on 4 GPUs, all 6 passed. The Kimi-VL case passed on both; it loads through the spec and the loader directly, not through the plan decision.Before your PR is "Ready for review"
preparebefore the plan is decided, which shims the remote code's imports, patches its modeling module (which builds the native language model only for the configpreparemade, so a load the plan keeps off expert parallelism still builds the checkpoint's own), and can refuse the checkpoint, failing the decision. A run that would leave the MoE whole on every rank, with nothing asking for that, is refused, as other expert-parallel MoEs are. A single-process load is unchanged. [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 later runsprepareonly once the plan is decided expert-parallel.CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models, [1/2] of #2748's split: