Repository navigation
[OMNIML-6063] [18/21] Decompress compressed-tensors MoE experts on the fly as they load - #2747
shengliangxu wants to merge 14 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.22% 62.33% +0.10%
=================================================================
Files 656 658 +2
Lines 73733 73955 +222
=================================================================
+ Hits 45883 46097 +214
- Misses 27850 27858 +8
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/ok to test 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}
| 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.
Pushing back on (2), and postponing (1) and (3). (2): shutil.copy is copyfile plus copymode, so it gives the temp file the source's permission bits before the rename, and the staged module gets the checkpoint file's mode, not mkstemp's 0600: with a 0644 source under umask 022, the temp file goes from 0600 to 0644 at the copy and is staged 0644 (Python 3.12). (1) and (3) are real but small: a failed copy leaves an extension-less temp file, which neither transformers nor Python imports, beside the warning that reports the failure; and a sibling edited after it was staged keeps its staged copy until the stage dir is cleared, while transformers refreshes the modules it stages itself. They are postponed because [19/21] (#2748) changes this module in the same round, and an edit here would conflict with it.
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.
|
/ok to test 1948e08 |
|
/claude review |
There was a problem hiding this comment.
Claude review (round 6) — [18/21] on-the-fly compressed-tensors MoE decompression
Scope: full review (the trigger comment carried no scoping instructions). 7 changed files, all reviewed: modelopt/torch/models/_onfly_decompress.py (+459), _remote_code.py (+67), modelopt/torch/utils/plugins/model_load_utils.py (+28/-1), noxfile.py (+6/-1), and the three test files.
Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
Previous round resolved. Both round-5 IMPORTANT findings are closed, checked against the current code rather than the replies:
patchedcould un-mark a config and rewrite nothing, silently. There is now arewrittencounter with aRuntimeErrorwhen no per-expert converter matched, and thedelattrloop moved after the rewrite — so the mark survives the refusal and a marked model whose mapping has nothing to rewrite fails loudly instead of loading newly-initialized experts. The clearing walks_config_levels, so both nested levels are covered.- Wide float stored params were cast to bf16 before
decompresssaw them. Closed by refusal rather than protection:check_compressed_keysnow reads the dtype of every tensor the compressed experts store, from the shard headers, and refusesF16/F32/F64. I checked the set is the right one —BF16,F8_E4M3/F8_E5M2,F8_E8M0and the integer dtypes are all held exactly by bf16, so accepting them is correct, and an nvfp4 fp32weight_global_scaleis caught. The(F32) … bf16 cannot hold exactlycase pins it.
Round-5 suggestions: the per-expert length check is in (the counts ValueError), _config_levels now drives _find_quantization_config / is_marked / the delattr off _SUB_CONFIG_ATTRS so llm_config is handled everywhere the loader handles it, and _remote_code.py's staged module now takes its mode from the source via shutil.copy.
Round 4's open finding — the targets check never consulting qconfig.ignore — is also answered, and well: not t.startswith("re:") deliberately admits the class-target shape, and check_compressed_keys then tests the real invariant against the stored names. I traced both shapes that worried me there and both are now refused: targets: ["Linear"], ignore: ["lm_head"] trips outside (the attention stems carry no .experts.), and ignore: ["re:.*experts.*"] trips plain plus not compressed. Checking the checkpoint instead of the config was the right call.
This round. Three non-blocking items, all inline:
patched's converter filter uses the unanchored"experts" in pwhilecheck_compressed_keysuses".experts.", so the two disagree onshared_experts— in the silent direction, and past both therewrittencounter and the stored-key check. Unreachable for the targeted models (transformers' DeepSeek-V3 mapping has no shared-expert converter), one character class to close.- The hardcoded
.to(torch.bfloat16)is now the sole justification for the new dtype refusal, butload_and_prepare_ephonours the config's owndtype(model_load_utils.py:2383into:2434) rather than forcing bf16 — so a non-bf16 config makes the cast and the refusal disagree in both directions. _remote_code.py'strysits outside the staging loop: one bad sibling abandons the rest of the closure, and themkstemptemp file leaks when the copy fails.
What checked out.
- Converter rewrite.
new_sourcesgroups by stored-param name with the anchor first, so gate-then-up ordering survivesconvert'slist(input_dict.keys())into the anchor-keyed output and on throughMergeModulelistintoConcatenate. Thefloat-quantizedcase where the anchor is.weightstays unambiguous (.weight_scaledoes notendswith(".weight")). Siblings are drawn fromparam_names[1:]and non-anchor patternscontinuebefore anypop, so no double-pop; the subset-of-stored-siblings handling (Kimi-K3 storing noweight_global_scale) is explained and tested; unconsumed sources raise rather than reaching the merge. - Expert count pairing. The new
len(lists[pn]) != len(anchor_list)check guards exactly the hazard positional pairing creates, and theNone-shard skip after it keeps peak memory at the per-rank shard. check_compressed_keys.checkpoint_tensor_filesreturnsname -> file, so iterating it yields names as the signature expects;stored != "weight"correctly keeps a plain expert weight out ofcompressedwhile still catching it asplain; thefloat-quantizedlayout (where"weight"is inparam_names) classifies correctly in both loops; the three refusal clauses cover the layouts that matter.- Install/mark lifecycle.
get_value_from_registryreturns the registered class, so thecompression != installedtuple comparison is stable for the same checkpoint and the install-twice case is a genuine no-op rather than a spuriousRuntimeError.require_transformers("5.8", ...)runs before thecore_model_loadingimport._read_compressionreturningNonebefore importingcompressed_tensorskeeps the no-config path dependency-free. - Loader integration.
quant_methodis read once at :2336, before the deletion, and its only later uses are thefp8andmxfp4branches a compressed-tensors checkpoint never enters. The deletion covers all three levels_ckpt_quant_methodreads, usesdelrather than= Nonefor the pre-5.11 behaviour, and runs on the deepcopy — the newtest_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passesasserts the refused-unmarked, no-quantization_config-at-any-level and caller's-config-untouched halves together. - Patch target separation.
_onfly_decompresspatchestransformers.modeling_utils.get_model_conversion_mapping(the bindingfrom_pretrainedcalls) whilemodel_load_utils.py:39imports the stock function fromtransformers.conversion_mappingfor the export-side mapping — which is what the export wants, since the model it describes is decompressed. Load-bearing and correct. - Laziness.
_onfly_decompressimports onlyosandtorchat module level; transformers,compressed_tensorsandsafetensorsare all function-local, andmodelopt/torch/models/__init__.pydoes not pull it in. - noxfile. The
("torch_28", "torch_29")gate matches the realTORCH_VERSIONSkeys and thetf_pinsinstall lands after it, so each leg keeps its transformers pin and the real-library cases run only on thetf_latestlegs with torch 2.10+. (It is a denylist, so a future pre-2.10 leg would silently getcompressed-tensors; an allowlist of the 2.10+ ids would stay correct by construction. Previously noted, still not worth a thread.) - Tests. CPU-only, tiny safetensors fixtures, real-library cases gated on
compressed_tensorsplus transformers 5.8. Thespawn_materializeassertion intest_integer_checkpoint_tensors_keep_their_dtype_on_loadis a nice touch — it pins that the wrapper is reached through transformers' own call path, not just directly.
Risk: low. Additive, and nothing reaches it unless a model spec installs it, so the blast radius is the Kimi EP-PTQ path. Six rounds in, the mechanism and its activation both read consistently end to end, and the shift this round — from protecting a wide-float stored param to refusing it against the checkpoint's own shard headers — is the right trade for a path whose failure mode is a plausible-looking but wrong quantized checkpoint. The three remaining items are latent edges rather than things that would affect a passing run today.
Verification note: neither transformers nor compressed_tensors is installed in this sandbox. The WeightConverter field set, _materialize_copy's signature, the classmethod shape of compression_param_names / decompress, and whether the loader really hands None over for non-local experts are read from the diff, its comments and the tests; CI's tf_latest legs with torch 2.10+ and the GPU Kimi-K3 run remain the real check on those. Finding 2 is verified against the checked-in model_load_utils.py; findings 1 and 3 stand on the new files alone.
🤖 Generated with Claude Code
| if isinstance(c, WeightConverter) and any( | ||
| "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns | ||
| ): |
There was a problem hiding this comment.
[SUGGESTION] This filter uses the unanchored substring "experts", while check_compressed_keys classifies modules with the dot-anchored ".experts." (lines 430, 432). The two disagree on exactly one family of names — *_experts, i.e. shared_experts — and in the direction that fails silently:
- a converter whose sources are
...mlp.shared_experts.gate_proj.weightmatches"experts" in p, so_compressedifyrewrites it to...shared_experts.gate_proj.weight_packed$; - Kimi's
ignorelist leaves the shared experts bf16, so the checkpoint stores...shared_experts.gate_proj.weightand nothing matches the rewritten source — the target param is never filled and loads newly initialized behind a missing-keys warning; check_compressed_keyscannot catch it:".experts."does not matchshared_experts., so that stem lands in neithercompressednorexpert_weights, andoutside/plainboth stay empty;- the
rewrittencounter does not catch it either — the routed-expert converter also matches, so the count is non-zero.
This is not reachable for the checkpoints this PR targets: in transformers' DeepSeek-V3 mapping the shared experts are plain nn.Linears loaded by name with no WeightConverter, so the only *experts* converters are the routed ones. That is why this is a SUGGESTION and not blocking. But the filter is the admission criterion for the rewrite, and the fix is one character class:
| if isinstance(c, WeightConverter) and any( | |
| "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns | |
| ): | |
| if isinstance(c, WeightConverter) and any( | |
| ".experts." in p and p.endswith((".weight", ".weight$")) | |
| for p in c.source_patterns | |
| ): |
That also makes the rewrite scope and the stored-key check agree on one definition of "an expert", which is worth having given that check_compressed_keys is the only thing standing behind the class-target (targets: ["Linear"] + ignore) shape.
There was a problem hiding this comment.
Postponed: agreed the two definitions should match, but nothing the unanchored match would misfile reaches it today. The hooks are installed only by the Kimi-K2.5 and Kimi-K3 specs, whose expert mappings (deepseek_v3's, and the kimi_linear one _register_kimi_k3_conversions registers) convert routed experts alone; anchoring on .experts. belongs with the first spec whose mapping also fuses shared_experts weights, as transformers' minimax_m3_vl does.
| 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] The hardcoded .to(torch.bfloat16) is now load-bearing in a second place, which makes it worth a guard rather than just a comment.
check_compressed_keys's new dtype refusal (lines 450-456) is justified entirely by this cast — "bf16, the dtype the experts decompress to". But load_and_prepare_ep does not force bf16: model_load_utils.py:2383 computes dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 and passes it straight to from_pretrained (model_load_utils.py:2434). So for a checkpoint whose config declares dtype: float16 the two halves disagree in both directions:
- the refusal fires on an
F16weight_scalethat the load would in fact have kept exactly, so a checkpoint that would decompress correctly is refused; - the dense half loads fp16 while these experts come out bf16, so the fused expert param is either double-rounded on assignment or errors on the dtype mismatch — and nothing reports it.
Low reachability (the four targeted Kimi checkpoints are all bf16), hence SUGGESTION. Two cheap ways to close it: thread the load dtype into _make_op_class / check_compressed_keys so one value drives both the cast and the "cannot hold exactly" set, or — if bf16 really is the only supported output — assert it where the premise is made, e.g. refuse in install_onfly_decompress when the config declares a non-bf16 dtype. Either makes the invariant checkable instead of implicit across two files.
There was a problem hiding this comment.
Postponed: valid, but the hooks are installed only by the Kimi-K2.5 and Kimi-K3 specs, whose checkpoints load in bf16, as you note, so today the cast and the check_compressed_keys refusal agree. Making the premise one checked value, by refusing a non-bf16 config dtype in install_onfly_decompress or by threading the load dtype into DecompressExperts and the refusal, is a follow-up.
| 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}") |
There was a problem hiding this comment.
[SUGGESTION] The shutil.copy-onto-mkstemp rewrite addressed the 0600 mode from the last round, but two edges remain in the error path, and both come from the try sitting outside the loop rather than inside it:
- One bad sibling abandons the rest. The first
OSError— an unreadable*.py, a full disk, a permission problem on one file — exits the whole loop, so the siblings after it are never staged. Since the point of this helper is to stage the full closure, a partial stage reproduces theFileNotFoundErrorfromget_relative_import_filesthat the docstring describes, just one import deeper. - The temp file leaks. If
shutil.copyraises,os.replacenever runs andtmpstays in the stage dir. It is not a*.py, so it does not breakget_relative_import_files, but with every rank staging into a sharedHF_MODULES_CACHEthe droppings accumulate silently.
Moving the try inside the loop with a finally that unlinks fixes both — stage what you can, warn per file, leave nothing behind:
for py in glob.glob(os.path.join(src, "*.py")):
dst = os.path.join(stage_dir, os.path.basename(py))
if os.path.exists(dst):
continue
# 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.
tmp = None
try:
fd, tmp = tempfile.mkstemp(dir=stage_dir)
os.close(fd)
shutil.copy(py, tmp)
os.replace(tmp, dst)
tmp = None
except OSError as e: # a failed copy only matters if a deeper import needs the file
warnings.warn(f"Could not stage {py} of the remote code of {src}: {e}")
finally:
if tmp is not None:
with contextlib.suppress(OSError):
os.unlink(tmp)There was a problem hiding this comment.
Postponed: both edges are real, but only after a copy has already failed and warned, and an import that needs a file left unstaged still fails loudly with FileNotFoundError. Moving the try into the loop, with a finally that unlinks the temp file, is a follow-up rather than a change to this approved slice.
|
/ok to test 69f3234 |
|
/ok to test e8a3ef9 |
|
/ok to test 5801d1b |
|
/ok to test f3f0d07 |
|
/ok to test a60fddb |
|
/ok to test 2ab4b71 |
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>
Three ways the on-the-fly decompression could still load a compressed checkpoint wrong without a word are now refused: - The conversion-mapping hook rewrote whatever per-expert converters it found, possibly none, and cleared the config's mark before rewriting. A marked model whose mapping has no per-expert converter to rewrite, whose experts would load newly initialized, is now refused, and the mark is cleared only once the converters are rewritten. - The load may cast a source tensor to the model's dtype before the decompression reads it; only integers are kept as stored. The expert-parallel load's check of what the checkpoint stores now also reads the dtypes of the compressed experts' tensors and refuses one bf16 cannot hold exactly (fp16, fp32, fp64), such as NVFP4's fp32 global scale. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints store int32, uint8 and bf16 there. - DecompressExperts pairs each expert's stored tensors by position, and now refuses stored params that disagree on how many experts they hold. The decompressor also finds its compressed-tensors config and mark on llm_config, as the rest of the expert-parallel load does, and the integer-keeping loader hook says why it stays process-wide. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 008b1f9 |
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, for any format whose stored floats bf16 holds exactly; a checkpoint whose compressed experts store fp16, fp32 or fp64 (such as NVFP4's fp32 global scale) is refused, since the load may cast it to bf16 before decompressing.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. It also reads, from the shard headers, the dtype of every tensor the compressed experts store, and refuses a float bf16 cannot hold exactly. Kimi-K2.5, K2.6 and K2.7-Code store int32 packed weights and shapes with bf16 scales, and Kimi-K3 uint8 packed weights and scales (one shard header each), so all four pass that too.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_configandllm_config, the levels the rest of the expert-parallel load reads; the hook refuses a marked model whose conversion mapping has no per-expert converter to rewrite (its experts would load newly initialized), and clears the mark at every level only once it has rewritten them.to_native_backbone_configcarries the mark 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(39 cases). It covers:compressed_tensorsblocked;llm_config, read as ontext_config;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, no raw-dict refusal, no refusal of a mapping with nothing to rewrite, the mark cleared before the rewrite, no expert-count check,text_configas the only nested level, and no dtype check.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, which stores int32 packed weights and bf16 scales as the real ones do, 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 stored names, and with them the compressed experts' dtypes, 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