Repository navigation
[OMNIML-6063] [18/21] Decompress compressed-tensors MoE experts on the fly as they load - #2747
shengliangxu wants to merge 13 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:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 65.96% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 47 functions across 4 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 PASS. The reviewed changes add no torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# noseccommen…Title check The title clearly and concisely describes the main change: on-the-fly decompression of compressed-tensors MoE experts during loading. Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Full details: Security Anti-Patterns
-
Explanation
PASS. The reviewed changes add no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# noseccomments. The PR does not modifypyproject.tomlorrequirements.txt; itscompressed-tensorsinstallation is limited tonoxfile.py.
-
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- 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.
95832dd to
e243f04
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-13-ep-export #2747 +/- ##
================================================================
Coverage ? 62.28%
================================================================
Files ? 658
Lines ? 73880
Branches ? 0
================================================================
Hits ? 46019
Misses ? 27861
Partials ? 0
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 e243f04 |
e243f04 to
d6afd32
Compare
|
/ok to test d6afd32 |
shengliangxu
changed the title
67355f4 to
b162532
Compare
|
/ok to test b162532 |
shengliangxu
changed the title
b162532 to
269d3d2
Compare
shengliangxu
changed the title
269d3d2 to
dddc2ee
Compare
|
/ok to test dddc2ee |
dddc2ee to
c668ec2
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if len(groups) != 1: | ||
| # Only the expert converters are rewritten, so another group's weights would load newly | ||
| # initialized, and which group holds the experts would be a guess. | ||
| raise NotImplementedError( | ||
| f"install_onfly_decompress: the checkpoint has {len(groups)} compressed-tensors config " | ||
| f"groups ({', '.join(groups) or 'none'}); on-the-fly decompression handles a single " | ||
| "group, compressing only the experts. Decompress the checkpoint first." | ||
| ) | ||
| [scheme] = groups.values() |
There was a problem hiding this comment.
[SUGGESTION] The multi-group refusal closes the main hazard from the previous round — thanks. Two residual edges in the same check:
-
The refusal message asserts the accepted shape is "a single group, compressing only the experts", but nothing verifies the one group actually targets the experts. A checkpoint whose single group targets, say,
re:.*self_attn.*is accepted, its scheme drivesparam_names, andpatched()rewrites the (uncompressed) expert converters to*.weight_packed$— matching no checkpoint key, so the experts load newly initialized. Not reachable from the Kimi adapters, but it is the one shape the check's own wording claims to exclude. Atargets-based guard (refuse when no target isLinear/re:matching an expert path) would make the invariant real rather than asserted. -
len(groups) != 1folds "zero groups" into the multi-groupNotImplementedError("has 0 … config groups (none); on-the-fly decompression handles a single group"). An emptyconfig_groupsis "nothing to decompress", which the module elsewhere treats as theValueErrorpath — routing it tolen(groups) == 0 → return Nonewould keep the two messages honest.
There was a problem hiding this comment.
Fixed in 2275f37: (1) the one group must target the experts, by the rule transformers' own compressed-tensors loader uses to find the experts' group: a class target such as Linear (the checkpoint's experts are Linear modules) or a re: target naming experts. A group targeting only attention is refused before anything is patched. The ignore list is not matched against module names, which install time does not have. (2) A config with no groups now raises ValueError (nothing to decompress) instead of the several-groups refusal; it is raised rather than returned as None so a process that already installed the hooks refuses it too, instead of returning without a mark. Both are new cases of the refusal test.
There was a problem hiding this comment.
Claude review (round 3) — [18/21] on-the-fly compressed-tensors MoE decompression
Scope: full review (the trigger comment carried no scoping instructions). All 7 changed files reviewed: modelopt/torch/models/_onfly_decompress.py (+338), _remote_code.py (+61), modelopt/torch/utils/plugins/model_load_utils.py (+8/-1), noxfile.py (+6/-1), and the three test files.
Findings: CRITICAL 0 · IMPORTANT 2 · SUGGESTION 1
Previous round resolved
Checked against the current code, not the replies:
_read_compressionnow refuses a checkpoint with several config groups (NotImplementedError) and unpacks the single group with[scheme] = groups.values(), so the "guess which group holds the experts" fallback and its test are gone. The new parametrized refusal test covers both the expert-group-among-others and no-expert-group shapes.DecompressExperts.convertnow raisesValueError: unconsumed checkpoint keys …instead ofout.update(input_dict), and_compressedify's non-.weightelse-branch is pinned bytest_a_source_the_decompression_does_not_consume_fails_the_load._mark's dict branch is gone — it is a plainsetattr.require_transformers("5.8", …)now runs before thetransformers.modeling_utils/core_model_loadingimports, so a 4.x install gets the actionable message.
This round
Both findings are about the mark's lifecycle — the single thing that decides whether a load gets the rewritten converters — and both fail silently (newly-initialized experts, or a competing quantizer) rather than loudly.
-
IMPORTANT — the
_decompressed_as_loadedwaiver permits the one state the decompressor cannot survive (inline). A config is only ever marked while it still carries a parseablequantization_config(_read_compressionreturnsNoneotherwise, and the already-installed early-return does not mark). So the waiver fires exactly whenquantization_configis still present — and nothing inload_and_prepare_epstrips it: thecopy.deepcopykeeps it, neither thefp8nor themxfp4branch overrides it, andfrom_pretrained_kwargs["config"] = hf_confighands it tofrom_pretrained, which builds aCompressedTensorsHfQuantizer. That contradicts_install_keep_int_native's own premise ("without anhf_quantizerthe loader casts every source…"). If instead the per-modelpreparealways strips it — asinstall_onfly_decompress's docstring says it does — thenquant_methodis alreadyNonehere and the waiver is dead code. The new test pins the un-stripped combination as passing. Stripping it on the marked path makes the two halves agree. -
IMPORTANT — the mark does not reach the model config for one of the two documented call orders (inline).
_markwrites top level only andpatched()reads top level only, butinstall_onfly_decompress's docstring explicitly accepts a config whosequantization_configlives undertext_config. Pass the VLM config and the mark lands on the VLM config, whileto_native_backbone_configderives the backbone config fromtext_configalone — the native config is unmarked,from_pretraineduses stock converters, the per-expert.weightkeys match nothing, and the experts load newly initialized. Re-installing on the derived config does not rescue it (_read_compression→None→returnwithout marking), andload_and_prepare_epcannot catch it either (the derived config has noquantization_config, so the format check passes). Passingvlm_cfg.text_configinstead works, and is the only order that makesis_marked'stext_configbranch reachable — so the two helpers in this file are coupled by an undocumented ordering. Marking both levels, or propagating the mark into_native_backbone_config, removes the trap.
Suggestion: the single-group check asserts "compressing only the experts" in its message but never checks the group's targets, and it folds zero groups into the multi-group refusal (inline).
Assessment
Risk: low-to-medium. The core mechanism continues to look right and is unusually well tested for monkeypatch-heavy code: driving transformers' own rename_source_key / add_tensor / convert objects in _load, the None-sharded expert case, the subset-of-stored-siblings case, the install-twice / different-checkpoint / malformed-config cases, and a real compressed_tensors INT4 round trip all earn their keep. Anchor-first new_sources ordering keeps gate-then-up correct through MergeModulelist → Concatenate, and the _materialize_copy wrapper is now thin enough that its process-wide permanence is a fair trade.
What is left is not the decompression math but its activation: the mark is a single ungated boolean that three places read with three different scopes (is_marked both levels, patched top level, _mark writes top level), it is consumed by delattr on first use, and every way of getting it wrong ends in a quiet load rather than an exception. Both IMPORTANT findings sit in that band and both are a few lines to close.
_remote_code.py now carries its test, and the staging wrap reads correctly (post-orig copy into the already-resolved stage dir, HF_MODULES_CACHE join for a relative resolve, idempotent, top-level *.py only as documented). The noxfile gate matches the real TORCH_VERSIONS keys, and the tf_pins install after it keeps each leg's transformers pin, so the real-library cases run only on the tf_latest × torch≥2.10 legs as the description says.
Verification note: transformers and compressed_tensors are not installed in this sandbox, so the WeightConverter field set, _materialize_copy's signature, compressor.decompress(dict, scheme)'s shape in the released library, and whether the loader really hands None over for non-local experts are all read from the diff, its comments and the tests. CI's tf_latest legs and the GPU Kimi-K3 run remain the real check on those. Finding 1 is verified against the checked-in model_load_utils.py; finding 2 rests on _onfly_decompress.py alone plus PretrainedConfig.to_dict/from_dict behavior, which the tests' _MARK not in model.config.to_dict() assertion corroborates.
🤖 Generated with Claude Code
|
/ok to test 95e1cfa |
|
/ok to test 0df6cca |
|
/claude review |
| # The experts' group as transformers' own compressed-tensors loader finds it: a class target | ||
| # (the checkpoint's experts are Linear modules) or a regex naming the experts. Another group | ||
| # leaves the experts stored as plain weights, which the rewritten converters would not find. | ||
| if not any(not t.startswith("re:") or "experts" in t for t in scheme.targets): | ||
| raise NotImplementedError( | ||
| f"install_onfly_decompress: the checkpoint's compressed-tensors config group targets " | ||
| f"{scheme.targets}, not the experts; on-the-fly decompression handles a checkpoint " | ||
| "whose experts are what it compresses. Decompress the checkpoint first." | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Algorithm] The guard checks targets but never qconfig.ignore — and for the class-target shape the Kimi checkpoints actually use, ignore is the only thing that limits the compression to the experts.
not t.startswith("re:") makes any class target pass unconditionally. The PR description says it plainly: Kimi's group is one class target "with an ignore list keeping the rest bf16". So the invariant this block claims to establish ("a checkpoint whose experts are what it compresses") is carried entirely by a field the code never reads, and two config shapes that violate it pass:
targets: ["Linear"], ignore: ["lm_head"]— the standard compressed-tensors W4A16 layout, where every Linear is compressed.patched()rewrites only converters with"experts"in their sources, andload_and_prepare_epnow deletesquantization_configat both levels on the marked path, so there is no quantizer to load the attention/denseweight_packedkeys either. Those layers load newly initialized, withfrom_pretrained's missing-keys warning rather than an error.targets: ["Linear"], ignore: ["re:.*experts.*"]— the experts are the one thing not compressed._compressedifystill rewrites their sources to*.weight_packed$, which match no checkpoint key, so the experts load newly initialized — exactly the failurepatched()'s own comment warns about for an unmarked model.
Both end in a silently garbage model. test_real_pack_quantized_experts_round_trip_exactly pins _int4(True, ["Linear"]) with no ignore at all as accepted, so the test suite enshrines shape 1's precondition rather than refusing it.
The config alone cannot resolve a class target without the module list, so the exact check needs one more input. Two options:
- Validate against the checkpoint instead of the config: the prepare hook already knows the checkpoint dir, so pass it through and refuse if
model.safetensors.index.jsoncarries aparam_names[0]key outside anexpertspath, or carries none inside one. That tests the real invariant (only expert weights are stored compressed) rather than a proxy for it. - If you want to stay config-only, require an expert-scoped regex target and have the Kimi specs opt into the class-target form explicitly (e.g.
install_onfly_decompress(cfg, experts_only_class_target=True)), so a checkpoint nobody has vetted cannot reach the rewrite by accident.
Either way, please drop "a class target (the checkpoint's experts are Linear modules)" from the comment or qualify it — as written it reads as though a class target implies expert-only, which is the inverse of what Linear means.
There was a problem hiding this comment.
Fixed in 577515b: the expert-parallel load now checks the names the checkpoint stores, before anything loads, which tests the invariant itself rather than a proxy for it. Every weight stored compressed (under a name the decompressor reads other than weight, such as weight_packed or weight_scale) must sit under an .experts. module, and every expert weight must be stored compressed. Both of your shapes are refused (with only lm_head ignored, attention is stored compressed too; with the experts ignored, they are stored as plain .weight), as are a checkpoint that compresses nothing and one that leaves a single expert plain. The check runs in load_and_prepare_ep right after the checkpoint is resolved to a local snapshot, over the tensor names the loader sees (checkpoint_tensor_files), and takes the stored names from the mark, which now holds the decompressor's stored param names instead of the format. I took this route over requiring a regex target because the real checkpoints show the config cannot answer it: Kimi-K2.5, K2.6, K2.7-Code and K3 each have one Linear group, and Kimi-K2.5's ignore list names neither its vision tower nor its projector, whose 335 tensors are stored plain anyway. Their stored names pass: K2.5, K2.6 and K2.7-Code store 69,120 compressed expert weights each and K3 247,296, none outside .experts. and no expert plain, and reading the names takes 0.5 s to 2.2 s against a load of minutes. The comment on the targets check no longer reads as if a class target meant experts only: it says a class target does not tell which modules it compressed, and points at the stored-name check. New tests: test_a_checkpoint_must_compress_its_experts_and_nothing_else (six layouts, each refusal failing without its clause), and through the load, a checkpoint that also compresses an attention projection is refused before from_pretrained.
| if qc is None and isinstance(hf_config, dict): | ||
| qc = hf_config.get("quantization_config") or (hf_config.get("text_config") or {}).get( | ||
| "quantization_config" | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] This raw-dict branch cannot lead anywhere good through the module's public entry point, and the way it fails leaves the process half-patched.
install_onfly_decompress ends with _mark(hf_config, fmt), which is a plain setattr since round 3 — so for a dict config it raises AttributeError: 'dict' object has no attribute '_modelopt_onfly_decompress'. By then _install_keep_int_native() has run and mu.get_model_conversion_mapping = patched has been assigned, so the caller gets an opaque AttributeError with both process-wide patches installed and no config marked — every subsequent load in the process then pays for the patches while the compressed checkpoint still loads with stock converters.
Either validate up front (if not hasattr(hf_config, "__dict__"): raise ValueError(...) before the first patch, or simply require a PretrainedConfig), or drop the dict branch here and in to_native_backbone_config so the module consistently takes config objects only — is_marked, _mark and patched all already assume that. The raw-dict parametrization in test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_config is testing a path no supported caller can complete.
There was a problem hiding this comment.
Fixed in 577515b: install_onfly_decompress now refuses a dict with a TypeError before it reads the config or patches anything (the load finds the hooks by a mark on the config object, which a dict cannot carry), and _find_quantization_config drops its raw-dict branch, which nothing reaches any more. The raw-dict lookup case is gone; the refused-before-anything-is-patched test gains a raw-dict config, which fails without the refusal. I kept to_native_backbone_config's dict branch: it takes a dict text_config under a config object, which completes (the mark then comes from the VLM config), and its own test covers it.
| for cfg in (hf_config, getattr(hf_config, "text_config", None)): | ||
| if cfg is not None and hasattr(cfg, "quantization_config"): | ||
| del cfg.quantization_config |
There was a problem hiding this comment.
[SUGGESTION] This loop covers hf_config and text_config, but _ckpt_quant_method 20 lines up walks (cfg, text_config, llm_config). A VLM that nests its backbone under llm_config therefore reports quant_method="compressed-tensors" from the shared helper while this deletion misses it entirely.
It isn't reachable through today's designed path — _onfly_decompress._find_quantization_config and is_marked don't look at llm_config either, so such a config can never get marked and the load is refused loudly by the check at line 2302. But the asymmetry is a trap for the next model spec: mark a config whose quantization_config only lives under llm_config and the compressed-tensors quantizer gets built anyway, which is precisely what this block exists to prevent.
Cheapest fix is to iterate the same three levels the file's own helper does:
for cfg in (hf_config, getattr(hf_config, "text_config", None), getattr(hf_config, "llm_config", None)):(and the matching level in is_marked / _find_quantization_config), so the two halves agree on where a backbone config can live.
There was a problem hiding this comment.
Fixed in 577515b: the deletion now walks the levels _ckpt_quant_method reads, llm_config included, so a quantization_config the format check could see at any level is gone from the config from_pretrained gets. The test gives the config an llm_config carrying one and checks all three levels are cleared while the caller's are not; it fails without the change. I left is_marked and _find_quantization_config at two levels: a config whose quantization_config sits only under llm_config cannot be marked, so its load is refused by the format check, as you note, and widening the mark's reach belongs with the spec that first needs it.
There was a problem hiding this comment.
Claude review (round 4) — [18/21] on-the-fly compressed-tensors MoE decompression
Scope: full review (the trigger comment carried no scoping instructions). All 7 changed files reviewed: modelopt/torch/models/_onfly_decompress.py (+361), _remote_code.py (+67), modelopt/torch/utils/plugins/model_load_utils.py (+17/-1), noxfile.py (+6/-1), and the three test files.
Findings: CRITICAL 0 · IMPORTANT 1 · SUGGESTION 2
Previous round resolved
Checked against the current code, not the replies:
- Round 3 #1 (the
_decompressed_as_loadedwaiver permitted a competing quantizer).load_and_prepare_epnow deletesquantization_configfrom its deepcopy, at the top level and ontext_config, on the marked path — with the reasons in the comment (Kimi-K2.5's config class copying it upward;delrather than= Nonefor pre-5.11). I re-checked that this cannot disturb the rest of the function:quant_methodis read once at line 2296, before the deletion, and its only later uses are thefp8andmxfp4branches, neither of which acompressed-tensorscheckpoint enters. The newtest_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passespins all three halves — refused unmarked, noquantization_configat either level in what reachesfrom_pretrained, caller's config untouched. - Round 3 #2 (the mark missed one of the two documented call orders).
to_native_backbone_confignow carries the mark to the config it derives, taking it from either level, andis_markedpluspatched'sdelattrboth read and clear both levels. Marking the VLM config and marking itstext_confignow reach the native backbone config equally, so the undocumented ordering coupling is gone. - Round 3 suggestion (the single-group check never looked at the group's targets). A targets check is now in place, and zero groups gets its own
ValueErrorinstead of being folded into the multi-group refusal. See below: the check is there, but one shape still slips through it.
This round
IMPORTANT — the targets check never consults qconfig.ignore, which is the only thing that makes the class-target shape expert-only (inline). not t.startswith("re:") passes any class target unconditionally, and the PR description states that Kimi's group is one class target "with an ignore list keeping the rest bf16". So the invariant the block's own error message asserts ("a checkpoint whose experts are what it compresses") rests entirely on a field the code never reads, and two shapes that violate it are accepted:
targets: ["Linear"], ignore: ["lm_head"]— the ordinary compressed-tensors W4A16 layout, where every Linear is compressed.patched()rewrites only converters with"experts"in their sources, and the new deletion means there is no quantizer left to load the attention/denseweight_packedkeys either, so those layers load newly initialized behind a missing-keys warning.targets: ["Linear"], ignore: ["re:.*experts.*"]— the experts are the one thing not compressed, their rewritten*.weight_packed$sources match nothing, and they load newly initialized: the exact failurepatched()'s comment warns about for an unmarked model.
test_real_pack_quantized_experts_round_trip_exactly uses _int4(True, ["Linear"]) with no ignore at all, so the suite pins shape 1's precondition as accepted rather than refusing it. A class target cannot be resolved from the config alone, so the exact check needs the checkpoint: the prepare hook knows the dir, and refusing when the weight index carries a param_names[0] key outside an experts path (or none inside one) tests the real invariant. The alternative is to require an expert-scoped regex and have the Kimi specs opt into the class-target form explicitly.
Suggestions: the raw-dict branch in _find_quantization_config leads to an AttributeError from _mark's setattr after both process-wide patches are installed, so a dict config leaves the process patched and unmarked (inline); and the quantization_config deletion walks two config levels where the file's own _ckpt_quant_method walks three, llm_config included (inline).
Assessment
Risk: low. Down from low-to-medium: both round-3 findings were about the mark's lifecycle, and both are closed in a way that makes the readers agree on scope rather than papering over the mismatch. The activation path now reads consistently end to end — install_onfly_decompress marks, to_native_backbone_config propagates, is_marked and patched read and clear both levels, load_and_prepare_ep waives the format check and strips the competing quantizer on exactly the marked path, and the caller's own config survives all of it. The decompression mechanism itself I have no remaining concerns about: anchor-first new_sources keeps gate-then-up ordering through MergeModulelist into Concatenate, the None-shard skip bounds peak memory at the per-rank shard, the subset-of-stored-siblings handling is explained and tested, and unconsumed sources raise instead of being passed down.
The one finding left is not in the mechanism but in its admission criteria, and its failure mode is the same silent one the earlier rounds kept landing on: a checkpoint the decompressor should refuse gets accepted and loads part of the model newly initialized. In practice only the Kimi adapters call this, which bounds the blast radius to a checkpoint someone deliberately pointed a Kimi spec at — that is why it is IMPORTANT rather than CRITICAL.
_remote_code.py still reads correctly (post-orig copy into the already-resolved stage dir, HF_MODULES_CACHE join for a relative resolve, shutil.copy onto the temp file so the mode comes from the source rather than mkstemp's 0600, atomic os.replace, idempotent, top-level *.py only as documented). The noxfile gate matches the real TORCH_VERSIONS keys, and the tf_pins install after it keeps each leg's transformers pin.
Verification note: neither transformers nor compressed_tensors is installed in this sandbox. The WeightConverter field set, _materialize_copy's signature, the classmethod shape of compressor.compression_param_names / compress / decompress in the released library, and whether the loader really hands None over for non-local experts are all read from the diff, its comments and the tests; CI's tf_latest × torch>=2.10 legs and the GPU Kimi-K3 run remain the real check on those. The finding above and both suggestions are verified against the checked-in source alone — the ignore field's existence on QuantizationConfig is the one library fact the finding leans on, so confirm that field name if you take the config-only route.
🤖 Generated with Claude Code
Expert-parallel PTQ loads a checkpoint with from_pretrained over the expert-parallel mesh, so each rank builds only its own experts. A checkpoint whose experts are compressed-tensors quantized -- Kimi-K2.5's INT4, Kimi-K3's MXFP4 -- cannot take that load as it is: transformers' fused-experts conversion expects each expert's plain .weight. _onfly_decompress rewrites those per-expert converters to read the checkpoint's stored keys (weight_packed, weight_scale, ...) and prepends a step that dequantizes each owned expert to bf16 before the merge, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk. It is format-agnostic: the compressed_tensors library does the math, driven by the checkpoint's own quantization_config. It needs transformers 5.8. It also holds the helper that rebuilds a VLM's bundled text config as a native transformers backbone config. _remote_code.stage_full_remote_closure makes transformers stage every module of a local checkpoint's remote code, not only those one import deep, which loading the Kimi checkpoints' own modeling needs. 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>
The decompress hooks are process-wide and keep the first checkpoint's compressor, scheme and stored params, so a second checkpoint compressed differently was decompressed with the first one's settings without a word. Installing again now compares the compression: the same one is a no-op, as is a config its installing prepare already stripped of quantization_config, and a different one raises. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The installed decompressor rewrote the per-expert converters of every model from_pretrained loaded afterwards. An uncompressed MoE in the same process had its .weight sources renamed to the first checkpoint's packed names, matched no stored key, and loaded its experts newly initialized without an error. Installing now marks the config it is given, a repeat install for the same checkpoint marks its config too, and only a model built from a marked config gets the rewritten converters. from_pretrained loads from a deep copy of the config, which keeps the mark; the rewrite removes it from that copy so a later save does not write it into config.json. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
They drive the real compressed-tensors registry and compressor, so they skip where the library is missing, as it is in a plain dev-test install. CI's unit session installs it; the skip reason now says so instead of leaving only an import error. With compressed-tensors 0.19.0 all four cases pass. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The real-library tests of on-the-fly expert decompression skip without compressed-tensors, and no CI environment installs it, so CI never ran them. The unit sessions now install it with the leg's torchvision pin, so the resolver keeps the leg's torch. compressed-tensors 0.19 needs torch>=2.10, so the torch 2.8 and 2.9 legs leave it out: there it would upgrade torch or, held back by the pin, fall back to an old 0.15 release. It stays out of the dev-test extra: the transformers-free partial_unit(torch) lane installs that extra, and compressed-tensors pulls in transformers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The two tests that restated _compressedify's source patterns now ride the packed-load test, parametrized over the stored-param sets (pack-quantized; weight and weight_scale): a swapped gate/up, an unanchored pattern and a decompress op that is not a MergeModulelist each fail it. The other-model test and the copy-from_pretrained test become one test over the three loads. The helper-level ValueError test goes: the install-level test asserts the same error and that nothing was patched. The non-weight-sources and native-backbone guard tests stay, since nothing else catches a regression in their branches. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… check A model spec that installs the on-the-fly decompressor marks its config. Kimi-K2.5's keeps its compressed-tensors quantization_config at the top level, where the decompressor reads it, so the expert-parallel load's format check would refuse a checkpoint it does decompress. A marked config now passes; an unmarked compressed-tensors checkpoint is still refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The decompressor read "no quantization_config" off a ValueError, but compressed_tensors raises pydantic's ValidationError, also a ValueError, for a quantization_config it cannot parse (say, one written by a newer compressed-tensors). Taken for none, the second install of a process returned without marking the config, and the checkpoint's experts would load newly initialized. Reading a config with none now returns None, so a parse error goes through. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ring patch transformers made moot The rewritten get_model_conversion_mapping forwards its arguments as it gets them: pinning two of them as keywords made a fourth positional collide with them. The rewritten expert converter keeps the scope transformers sets on a converter after building it (scope_prefix, base_model_prefix, force_cpu, whichever the release has), and the integer-keeping copy hook calls transformers' own copy instead of standing in for it. The patch that re-sorted experts numerically is gone: every transformers from 5.5 sorts the state dict by dot_natural_key before collecting, so experts already arrive in numeric order, and the patch only kept a process-wide id() map. stage_full_remote_closure gets a unit test, the install message prints on rank 0 only, and the module docstrings keep the contract, not the design notes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_read_compression took the config group whose targets mention "expert", else the first, as the experts' scheme. Only the expert converters are rewritten, so another group's weights loaded newly initialized, and a guessed group decompressed the experts with the wrong bits, group size or zero point. A checkpoint with more than one group is now refused before anything is patched. The Kimi checkpoints this path serves carry one group each (Kimi-K2.5, K2.6 and K2.7-Code INT4 pack-quantized, Kimi-K3 MXFP4), whose ignore list leaves everything but the routed experts in bf16. DecompressExperts handed a checkpoint key it did not consume on to the merge, which would fuse it into the wrong slot. It now raises, which transformers reports as a conversion error that fails the load. install_onfly_decompress checks the transformers 5.8 floor before it imports core_model_loading, so an older release gets that message rather than an ImportError, and _mark no longer writes a dict key no reader looks at. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Kimi-K2.5's config class copies its text config's quantization_config to the top level, where the text_config swap leaves it, so the expert-parallel load handed from_pretrained a marked config that still carried it. That builds transformers' compressed-tensors quantizer, whose own expert decompression (from 5.13) claims the per-expert converters before the decompressor's hook sees them. On the marked path the load now deletes quantization_config from its copy of the config, at the top level and on text_config; deleted rather than set to None, which transformers before 5.11 takes for a quantized checkpoint. The decompressor's mark has one scope now: a config counts as marked when it or its text_config carries the mark, the hook reads it that way and clears it at both levels, and to_native_backbone_config carries it to the backbone config it derives. The hook read the top level only, so a VLM whose text_config was marked loaded stock converters while the format check passed it, and a backbone loaded on its own from the derived config was never rewritten. The checkpoint's one config group must target the experts (a class target or a regex naming them, as transformers' own compressed-tensors loader finds the experts' group), and a quantization_config with no groups is refused as nothing to decompress rather than with the several-groups message. stage_full_remote_closure copies each module to a temp file in the stage dir and renames it into place, so a rank never imports a module another rank is still writing. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The on-the-fly decompressor rewrites only the experts' converters, and the expert-parallel load drops the quantization_config that would load any other compressed weight. The checkpoint's one config group can target a class such as Linear, which does not say which modules it compressed: the common compressed-tensors layout compresses every Linear but lm_head, and an ignore list can leave the experts plain instead. Either way part of the model loaded newly initialized behind a missing-keys warning. The load now checks the names the checkpoint stores before anything loads: every weight stored compressed must be an expert's, and every expert's weight compressed. The mark now holds the names a compressed weight is stored under, which the check reads off the config. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints pass it. install_onfly_decompress refuses a raw dict before it patches anything, since the load finds its hooks by a mark a dict cannot carry, and the load deletes the quantization_config of llm_config too, every level _ckpt_quant_method reads. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 577515b |
|
/claude review |
| return | ||
| orig = cml._materialize_copy | ||
|
|
||
| def _materialize_copy(tensor, device=None, dtype=None): | ||
| t = tensor[...] | ||
| if dtype is not None and dtype.is_floating_point and not t.is_floating_point(): | ||
| dtype = None # keep packed/shape/zero-point/g_idx integers native | ||
| return orig(t, device=device, dtype=dtype) |
There was a problem hiding this comment.
[IMPORTANT Algorithm] This keeps integer stored params native but lets float compression params be downcast to the load dtype before decompress sees them, which silently perturbs the dequantized weights for any format that stores a wider-than-bf16 scale.
nvfp4-pack-quantizedstoresweight_global_scaleas fp32;float-quantized/int-quantizedcheckpoints produced from an fp16 model storeweight_scaleas fp16. Both areis_floating_point(), sodtypeis left at the target bf16 and the cast happens — fp32→bf16 drops 16 mantissa bits, fp16→bf16 drops 3. The scale multiplies every element of the weight it governs, so the error lands on the whole dequantized tensor, with nothing in the output to show it happened.- The module's own docstring advertises format-agnosticism ("every format the library registers works"), and
_read_compressionaccepts whateverformatthe checkpoint declares, so this is reachable the moment someone points it at an NVFP4 compressed-tensors checkpoint. The two targeted formats happen to be safe (MXFP4'sweight_scaleis uint8 → preserved; Kimi-K2.5's INT4 scales are already bf16), which is exactly why it would go unnoticed.
Two ways out, either is fine:
- Preserve the source dtype for every stored compression param, not just the integer ones (the sibling params are tiny, so keeping them native costs almost nothing), or
- if only the bf16-scale formats are meant to be supported, say so in
_read_compressionand refuse a scheme whose stored float params are wider than the load dtype, so an unsupported format fails loudly instead of loading slightly wrong.
Separately: unlike the converter rewrite, this patch is not gated on is_marked — it stays installed process-wide and changes the cast behavior of every later from_pretrained in the process, including unmarked models. Worth gating or at least noting in the docstring.
| anchor_list = lists[param_names[0]] | ||
| present = [pn for pn in param_names if pn in lists] | ||
| # Non-local experts were sharded away at load (GroupedGemmParallel -> None); decompress | ||
| # only the local experts this rank owns. Keeps peak memory at the per-rank shard. | ||
| out[pat] = [ | ||
| compressor.decompress({pn: lists[pn][i] for pn in present}, scheme)[ | ||
| "weight" | ||
| ].to(torch.bfloat16) | ||
| for i in range(len(anchor_list)) | ||
| if anchor_list[i] is not None | ||
| ] |
There was a problem hiding this comment.
[SUGGESTION] Two small robustness gaps in the per-expert loop, both at a parallelism boundary where a mistake is silent:
- The sibling lists are indexed by the anchor's position (
lists[pn][i]) on the assumption that every stored param collected the same number of entries, in the same expert order. That holds when the checkpoint stores the sibling for every expert, but the comment just above explicitly allows a checkpoint to store only a subset of the scheme's params — if a subset is ever per-expert rather than per-checkpoint (one expert missing itsweight_scale), the lists shift by one and every later expert dequantizes against the wrong scale, with no error. A one-line guard pays for itself:anchor_list = lists[param_names[0]] present = [pn for pn in param_names if pn in lists] assert all(len(lists[pn]) == len(anchor_list) for pn in present), ( f"DecompressExperts: {stem} stored params disagree on expert count " f"{ {pn: len(lists[pn]) for pn in present} }" )
.to(torch.bfloat16)is hardcoded, so a model whoseconfig.dtypeis fp16 or fp32 gets bf16 expert params while the rest of the model keeps its own dtype. Since the loader's dtype cast happens on the sources (before this op), nothing downstream fixes it up. Either take the dtype from the op's construction site (the config is available inpatched) or refuse a non-bf16 load ininstall_onfly_decompress.
| def _find_quantization_config(hf_config): | ||
| """Locate the compressed-tensors ``quantization_config`` (top-level or under ``text_config``) and | ||
| return it as a plain dict, or ``None`` for a config without one.""" | ||
| qc = getattr(hf_config, "quantization_config", None) | ||
| if qc is None: | ||
| tc = getattr(hf_config, "text_config", None) | ||
| qc = getattr(tc, "quantization_config", None) if tc is not None else None | ||
| if qc is None: |
There was a problem hiding this comment.
[SUGGESTION] This module looks for the nested config under text_config only, while the rest of the EP loader treats text_config and llm_config as equally valid homes for it — modelopt/torch/utils/moe.py defines _SUB_CONFIG_ATTRS = ("text_config", "llm_config"), _ckpt_quant_method walks both, and the new code in load_and_prepare_ep deletes quantization_config from llm_config too ("every level _ckpt_quant_method reads").
So a VLM that carries its compressed-tensors config under llm_config reaches install_onfly_decompress and gets ValueError: config has no compressed-tensors quantization_config even though the loader's own helpers would have found it. The failure is loud, so this is not a bug today, but the asymmetry is easy to trip over when the next spec lands. is_marked and to_native_backbone_config have the same two-level view.
Reusing the existing constant would keep the three places in step:
from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS
...
for c in (hf_config, *(getattr(hf_config, a, None) for a in _SUB_CONFIG_ATTRS)):
| for py in glob.glob(os.path.join(src, "*.py")): | ||
| dst = os.path.join(stage_dir, os.path.basename(py)) | ||
| if not os.path.exists(dst): | ||
| # Every rank loading the checkpoint stages into this dir and imports from it: | ||
| # renamed into place (atomic), a file appears there whole, never half written. | ||
| fd, tmp = tempfile.mkstemp(dir=stage_dir) | ||
| os.close(fd) | ||
| shutil.copy(py, tmp) | ||
| os.replace(tmp, dst) | ||
| except OSError as e: # a failed copy only matters if a deeper import needs the file | ||
| warnings.warn(f"Could not stage the remote code of {src}: {e}") | ||
| return resolved |
There was a problem hiding this comment.
[SUGGESTION] The atomic rename is the right call; three follow-ups on the staging loop:
- The temp file leaks on failure. When
shutil.copyraises (theOSErrorbranch is exactly that case, and the new test exercises it), themkstempfile stays in the stage dir forever — an extension-less, partially-written copy of a modeling file sitting next to the modules transformers imports. Onefinally/exceptwithos.unlink(tmp)keeps the cache clean, and the existing test can assert the dir holds only.pyfiles. mkstempmode is 0600, andos.replacecarries that through, whereas transformers' ownshutil.copyleaves the source's 0644. On a cluster whereHF_HOME/HF_MODULES_CACHEis shared between users (or between a job user and an interactive user), the first rank to stage makes the module unreadable to everyone else and their import fails.os.chmod(tmp, 0o644 & ~current_umask)— or simply copying the source file's mode — avoids that.- A changed sibling is never refreshed.
if not os.path.exists(dst)means an edited checkpoint.pykeeps serving the previously staged copy, while transformers re-stages the entry module itself. During local iteration on a checkpoint's modeling code that shows up as edits that appear to have no effect. Comparing size/mtime (or content) before skipping would cost little here.
There was a problem hiding this comment.
Claude review — [18/21] on-the-fly decompression of compressed-tensors MoE experts
Scope: full review (trigger comment carried no extra scoping). 7 changed files, all reviewed: modelopt/torch/models/_onfly_decompress.py, modelopt/torch/models/_remote_code.py, modelopt/torch/utils/plugins/model_load_utils.py, noxfile.py, and the three test files.
Findings: CRITICAL: 0 · IMPORTANT: 2 · SUGGESTION: 3
Most impactful
-
[IMPORTANT Correctness]
patchedcan un-mark a config and rewrite nothing, silently (_onfly_decompress.py:340-355). Thedelattr(c, _MARK)runs before the rewrite loop, so the hook is single-shot per config object, and the loop's"experts" in pfilter can match zero converters without anything raising. Either wayfrom_pretrainedfalls back to the stock.weightconverters, logs unexpectedweight_packedkeys plus a missinggate_up_proj, and PTQ calibrates newly-initialized experts. The rest of this module goes out of its way to refuse exactly that outcome (check_compressed_keys, the unconsumed-keysValueError, theis_markedguard); a rewrite counter and moving thedelattrafter the loop closes the last hole. -
[IMPORTANT Algorithm]
_install_keep_int_nativeprotects integer stored params but not wide float ones (_onfly_decompress.py:255-262).weight_global_scale(fp32, nvfp4-pack-quantized) and an fp16weight_scaleareis_floating_point(), so they are cast to the bf16 load dtype beforecompressor.decompresssees them, and the lost mantissa bits land on every element of the dequantized weight. The two targeted formats are safe by accident (MXFP4's scale is uint8, Kimi-K2.5's is already bf16), but the module advertises itself as format-agnostic and_read_compressionaccepts any registeredformat. Either keep every stored param native (they are tiny) or refuse a scheme whose float params are wider than the load dtype. The same patch is also the one piece not gated onis_marked, so it changes casting for every laterfrom_pretrainedin the process.
Suggestions (non-blocking)
- Per-expert loop (
_onfly_decompress.py:190-200): no length check across the stored-param lists before index-aligning them, and a hardcoded.to(torch.bfloat16)regardless of the model dtype. _find_quantization_config/is_marked/to_native_backbone_configlook undertext_configonly, while_ckpt_quant_method,_ep_configand the newquantization_configdeletion all handlellm_configas well (_SUB_CONFIG_ATTRS)._remote_code.py:53-64: themkstemptemp file leaks when the copy fails, the staged module inherits 0600 (a sharedHF_MODULES_CACHEthen breaks other users' imports), andif not os.path.exists(dst)never refreshes an edited sibling.
Minor, not worth a thread: the noxfile gate is a denylist (torch_28, torch_29), so any future leg older than 2.10 added to TORCH_VERSIONS silently gets compressed-tensors; an allowlist of the 2.10+ ids would stay correct by construction.
What checked out
- Converter rewrite ordering:
new_sourcesgroups by stored-param name with the anchor first, andconvertiterateslist(input_dict.keys()), so gate-before-up survives into the anchor-keyed output for the downstreamConcatenate— including theformat: float-quantizedcase where the anchor is.weightitself (.weight_scaledoes notendswith(".weight"), so the anchor match stays unambiguous). No double-pop: siblings are drawn fromparam_names[1:], and non-anchor keys hitcontinuebefore anypop. check_compressed_keys:checkpoint_tensor_filesreturnstensor name -> file, so iterating it yields names as the signature expects. The three refusal clauses cover the layouts that matter, andstored != "weight"correctly keeps a plain expert weight out of the compressed set while still catching it asplain— are:.*shared_experts.*-targeted group slips past_read_compression's target check but is then caught here asoutside, since Kimi's shared experts carry no.experts.segment.load_and_prepare_ep: thequantization_configdeletion covers all three levels_ckpt_quant_methodreads, usesdelrather thanNonefor the pre-5.11 behavior, and operates on the deepcopy so the caller's config is untouched — the newtest_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passesasserts both halves.check_compressed_keysruns after_resolve_checkpoint_dirand beforefrom_pretrained, so a refusal leaves no partial load.- Patch target:
transformers.modeling_utils.get_model_conversion_mappingis the bindingfrom_pretrainedcalls, while modelopt's own export-side callers import fromtransformers.conversion_mappingand so keep the stock mapping. That separation is load-bearing and correct. - Tests are CPU-only, build tiny safetensors fixtures, and gate the real-library cases on
compressed_tensorsplus transformers 5.8 — appropriate fortests/unit.
Risk
Moderate. The feature is additive and reaches nothing unless a model spec installs it, so the blast radius is the Kimi EP-PTQ path; both IMPORTANT findings are silent-failure modes rather than things that would break a passing run today. Worth addressing before merge given that a silently mis-loaded expert produces a plausible-looking but wrong quantized checkpoint.
| return convs | ||
| # from_pretrained loads from a copy of the marked config; the mark must not reach the config | ||
| # a later save writes, at either level. | ||
| for c in (config, getattr(config, "text_config", None)): | ||
| if hasattr(c, _MARK): | ||
| delattr(c, _MARK) | ||
| out = [] | ||
| for c in convs: | ||
| # only the per-expert MoE converters (have `.weight` sources under `experts`) | ||
| if isinstance(c, WeightConverter) and any( | ||
| "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns | ||
| ): | ||
| out.append(_compressedify(c, compressor, scheme, param_names)) | ||
| else: | ||
| out.append(c) | ||
| return out |
There was a problem hiding this comment.
[IMPORTANT Correctness] The mark is consumed (deleted) on the first call, and the rewrite loop can match nothing — both failure modes are silent.
delattrruns before the loop, sopatchedis effectively single-shot per config object. If any transformers release callsget_model_conversion_mappingmore than once for onefrom_pretrained(or a VLM load reaches it once per sub-model), the second call sees an unmarked config, takes thereturn convsearly-exit, and the experts load from stock.weightconverters that match no stored key.- Even on the first call, if no converter matches the
"experts" in p and p.endswith((".weight", ".weight$"))filter — a model whose fused-MoE mapping names its sources differently, or a release that scopes the patterns (scope_prefix) soexpertsno longer appears insource_patterns—outis just a copy ofconvsand the function returns quietly.
In both cases the outcome is the one the module works hard to avoid elsewhere: from_pretrained logs unexpected weight_packed keys and a missing gate_up_proj, and PTQ then calibrates newly-initialized experts. Nothing raises.
Suggested fix: count the rewrites and fail loudly when there are none, and move the delattr after the rewrite so a config is only un-marked once its converters were actually rewritten:
out, rewritten = [], 0
for c in convs:
if isinstance(c, WeightConverter) and any(
"experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns
):
out.append(_compressedify(c, compressor, scheme, param_names))
rewritten += 1
else:
out.append(c)
if not rewritten:
raise RuntimeError(
"install_onfly_decompress: no per-expert WeightConverter to rewrite in this "
"model's conversion mapping; the compressed experts would load newly initialized."
)
# from_pretrained loads from a copy of the marked config; the mark must not reach the
# config a later save writes, at either level.
for c in (config, getattr(config, "text_config", None)):
if hasattr(c, _MARK):
delattr(c, _MARK)
return out
What does this PR do?
Type of change: New feature
Expert-parallel PTQ loads a checkpoint with
from_pretrainedover the expert-parallel mesh, so each rank builds only its own experts. A checkpoint whose experts are compressed-tensors quantized (Kimi-K2.5's INT4, Kimi-K3's MXFP4) cannot take that load as it is: transformers' fused-experts conversion expects each expert's plain.weight._onfly_decompressrewrites those per-expert converters to read the checkpoint's stored keys (weight_packed,weight_scale, ...). It prepends a step that dequantizes each owned expert to bf16 before the merge, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk.compressed_tensorslibrary does the math, driven by the checkpoint's ownquantization_config.Linear, or a regex naming the experts), and refuses one with several; aquantization_configwith no groups is refused as nothing to decompress. Kimi-K2.5, K2.6 and K2.7-Code (INT4) and Kimi-K3 (MXFP4) each carry one group, undertext_config, with an ignore list keeping the rest bf16.check_compressed_keys): every weight stored compressed must be an expert's, and every expert's weight stored compressed, since a class target such asLineardoes not say what was compressed. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints pass it.ValueErrorbefore anything is patched, and without needingcompressed_tensorsinstalled; a raw dict config raisesTypeErrorthere too, since the load finds the hooks by a mark only a config object carries. Aquantization_configthat fails to parse raises its own error, on a second install too, rather than pass for no config: the experts would otherwise load newly initialized.Concatenatelooks up the sibling keys the decompression has already consumed.text_configcarries it: the hook reads it that way and clears it at both levels, andto_native_backbone_configcarries it to the backbone config it derives.from_pretrainedthe checkpoint's compressed-tensorsquantization_config: on the marked path the expert-parallel load deletes it from its copy of the config, at the top level, ontext_configand onllm_config, every level_ckpt_quant_methodreads, so transformers builds no compressed-tensors quantizer to compete with the decompressor. Kimi-K2.5's config class copies its text config'squantization_configto the top level, and from transformers 5.13 that quantizer took over the expert converters itself._remote_code.stage_full_remote_closuremakes transformers stage every module of a local checkpoint's remote code, not only those one import deep, which loading the Kimi checkpoints' own modeling needs. It stages each module through a temp file renamed into place, so ranks staging the same checkpoint at once never import a half-written module.unitnox sessions installcompressed-tensorson torch 2.10 and later, the releases it supports, so the tests against the real library run there. It goes in no extra: it pulls in transformers, which CI's transformers-free lane must not get.Experts arrive in numeric order without help: every transformers from 5.5 sorts the state dict numerically (
dot_natural_key) before it collects tensors. The rewritten converter keeps the scope transformers sets on it (scope_prefix,base_model_prefix,force_cpu).Usage
Internal to the Kimi adapters ([19/21], [20/21]), called from their
preparehooks.Testing
New
tests/unit/torch/models/test_onfly_decompress.py(35 cases). It covers:compressed_tensorsblocked;compressed_tensorsinstalled, a checkpoint in several config groups, in one group that is not the experts, or in none refused, and an INT4 pack-quantized round trip. These 5 cases run in CI's unit sessions on torch 2.10 and later; elsewhere they skip, saying what they need.Mutation-checked: each test fails on at least one deliberate break of the code, and each of these fails its test: the stored-name check not called in the load,
llm_confignot cleared, each of the check's three refusal clauses removed, experts matched by substring (shared_expertsthen counts),weightcounted as compressed, and no raw-dict refusal.Skips below transformers 5.8, the release the decompressor requires, so CI's 5.5 job skips it and its 5.18 job runs it. Run locally on 5.12.
A second, differently compressed checkpoint in the same process raises rather than reusing the first one's decompression; the same checkpoint again is a no-op, and one whose
quantization_configfails to parse raises. Any other model loaded in the process keeps its stock converters: only the marked checkpoint's are rewritten.The source-pattern checks ride the packed-load test, the other-model and copy-load tests are one, and the helper-level
ValueErrortest is dropped. Each drop is checked by mutation.New
tests/unit/torch/models/test_remote_code.py: a three-module import chain staged in full, and installed once; a copy cut short leaves no module behind, and the next load stages it whole.tests/unit/torch/utils/test_ep_loader.py: the compressed-checkpoint case loads from a small safetensors checkpoint and runs the load tofrom_pretrained, whose config has noquantization_configat any level,llm_config's included, while the caller's config keeps its own; one that also compresses an attention projection is refused beforefrom_pretrained.The stored-name check, run read-only on the tensor names of the real Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints (their index and shard headers): all four pass. Reading the names takes, per rank:
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.The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.
On GPUs: the Kimi-K3 run ([20/21]) loads its MXFP4 experts through this path; Kimi-K3's prepare drops
quantization_configat both levels itself, so no compressed-tensors quantizer took its experts over.Before your PR is "Ready for review"
CONTRIBUTING.md: N/A.compressed_tensorsis imported only for compressed checkpoints, which already need it; only CI's unit sessions install it, for the tests.Additional Information
Part 18 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [17/21].
Summary by CodeRabbit