Repository navigation
[OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader - #2747
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 |
|---|---|---|
| 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), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec commen… |
| Title check | |
The title accurately identifies the main change: loading checkpoints with on-the-fly decompression through the expert-parallel loader. Although the wording is slightly awkward, it is specific and rela… |
| 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.
Full details: Title check
-
Explanation
The title accurately identifies the main change: loading checkpoints with on-the-fly decompression through the expert-parallel loader. Although the wording is slightly awkward, it is specific and related to the changeset.
Finishing Touches
-
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
2142693 to
2e134ec
Compare
95832dd to
e243f04
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-13-ep-export #2747 +/- ##
================================================================
Coverage ? 62.46%
================================================================
Files ? 658
Lines ? 74106
Branches ? 0
================================================================
Hits ? 46294
Misses ? 27812
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 |
2e134ec to
fdd217a
Compare
e243f04 to
d6afd32
Compare
|
/ok to test d6afd32 |
shengliangxu
changed the title
67355f4 to
b162532
Compare
0a3e056 to
67f1c8f
Compare
|
/ok to test b162532 |
shengliangxu
changed the title
67f1c8f to
2f18815
Compare
b162532 to
269d3d2
Compare
shengliangxu
changed the title
269d3d2 to
dddc2ee
Compare
|
/ok to test dddc2ee |
dddc2ee to
c668ec2
Compare
68f620d to
5c5cde0
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): | ||
| if cfg is not None: | ||
| vars(cfg).pop("quantization_config", None) | ||
| if hasattr(cfg, "quantization_config"): | ||
| # A class attribute or property, which the instance dict cannot clear. | ||
| raise NotImplementedError( | ||
| f"{type(cfg).__name__} keeps quantization_config off its instance, so " | ||
| "this load cannot clear it, and transformers' own quantizer would " | ||
| "compete with the decompressor. Decompress the checkpoint first." | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] The new refusal is order-dependent: the top level is popped and checked before text_config/llm_config have been popped. Any config class that exposes quantization_config at the top level as a forwarding class attribute — a @property returning self.text_config.quantization_config, which is exactly the "copies its text config's to the top level" shape the comment above describes, just implemented on the class rather than in __init__ — hits the hasattr while the sub-config still holds the dict, and the load refuses a checkpoint it could have handled. Clear all levels first, then verify: after text_config's dict is gone, such a property raises inside hasattr, which is False, and from_pretrained sees nothing to build a quantizer from.
levels = [
cfg
for cfg in (hf_config, *(getattr(hf_config, a, None) for a in _SUB_CONFIG_ATTRS))
if cfg is not None
]
for cfg in levels:
vars(cfg).pop("quantization_config", None)
for cfg in levels: # checked only once every level is cleared: a forwarding
if hasattr(cfg, "quantization_config"): # property reads a level cleared later
raise NotImplementedError(
f"{type(cfg).__name__} keeps quantization_config off its instance, so "
"this load cannot clear it, and transformers' own quantizer would "
"compete with the decompressor. Decompress the checkpoint first."
)Two lines, and it makes the check exhaustive rather than order-sensitive — which is the property the block exists for. Non-blocking: the real Kimi configs you ran do the copy in __init__, so nothing is broken today, and the _Typed test case uses SimpleNamespace, which has no forwarding, so it would not catch a regression here either way. (Also note levels is what _onfly_decompress._config_levels already computes — see the duplication point raised last round.)
There was a problem hiding this comment.
Taken in dc0a618. The load now clears quantization_config from every level first and checks afterwards. A class that forwards its text config's through a property then holds none, while a class attribute or a property of its own is still refused.
The test now has a config class with such a forwarding property: it loads. It fails if the check runs level by level.
| # The export would copy them as stored, and this load dropped the quantization_config that | ||
| # reads a compressed one, such as an MTP layer's experts. A suffix is enough to tell them: | ||
| # check_compressed_keys refused any compressed tensor outside the experts. | ||
| if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: | ||
| raise NotImplementedError( | ||
| f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " | ||
| f"for, such as {carried[0]}: the export would copy them as stored, with no " | ||
| "quantization_config to read them. Decompress the checkpoint first." | ||
| ) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] This refusal is sound only as long as model._weight_conversions is populated. _conversion_plans (line 188) is (_conversion_plan(model), _conversion_plan(model, loaded) if loaded else None), and the first of those goes through the patched get_model_conversion_mapping — whose patched has already delattr'd _MARK from model.config during from_pretrained, so is_marked is false and it returns the stock converters, which carry .experts.*.weight sources and resolve none of weight_packed/weight_scale/weight_shape. The rewritten converters are therefore visible only through _weight_conversions.
On a release that does not record it, every built expert's stored tensors land in unplaced, and this check refuses a perfectly good checkpoint with a message that says the model has no parameter for them — and len(carried) is then in the hundreds of thousands for Kimi-K3. The possibility is not hypothetical in this repo's own judgement: test_unplaced_source_keys.py:80's pre-existing _load helper, and your new test_compressed_experts_decompressed_at_load_are_placed, both pytest.skip on exactly _weight_conversions is None, so the one test that pins this path silently stops covering it on such a release while production hard-refuses.
Cheap fix: skip the refusal when the record is missing, so the behaviour degrades to the pre-PR carry rather than to a false block —
unplaced = record_unplaced_source_keys(model, ckpt_path)
# Only the record from_pretrained kept of the converters it ran resolves a compressed
# expert's stored names: the stock mapping has .weight sources. Without it every expert's
# would look unplaced, so there is nothing to tell an MTP layer's from them.
if getattr(model, "_weight_conversions", None) is None:
compressed_names = set()or, if you would rather keep the block loud, raise a message that names the real cause (this transformers does not record the conversions from_pretrained ran) instead of blaming the checkpoint. Non-blocking, and moot if every transformers ≥ 5.8 records it — which I could not check here (neither transformers nor compressed_tensors is importable in this sandbox).
There was a problem hiding this comment.
I checked this against every release the decompression supports, and it does not happen there.
convert_and_load_state_dict_in_model ends with model._weight_conversions = [c for c in weight_mapping if c.was_used()], unconditionally. That holds in transformers 5.8.0 (core_model_loading.py line 1436), 5.13.0 (line 1700), 5.14.0 (line 1735) and 5.18.0 (line 1798). The decompression requires 5.8, which install_onfly_decompress checks first. The rewritten converters match the stored keys, so they are always among the ones recorded.
test_compressed_experts_decompressed_at_load_are_placed pins that on a real load: it fails when unplaced_source_keys reads the stock mapping alone. Its skip on a missing record is the file's general guard for older releases, which the decompression does not run on.
A guard for a missing record would be dead code on every supported release, so I'd rather not add it.
There was a problem hiding this comment.
Claude review (post-split round 2) — [3/3] decompressed checkpoints through the expert-parallel loader
Scope: full review — the trigger comment carried no scoping instructions. All 4 changed files opened (+385/-6): modelopt/torch/utils/plugins/model_load_utils.py (+49/-3), tests/gpu/torch/utils/test_expert_parallel_load.py (+133/-1), tests/unit/torch/utils/test_ep_loader.py (+125/-2), tests/unit/torch/utils/test_unplaced_source_keys.py (+78). Diff taken two-dot against origin/shengliangx/hfep-11a1-decompress-experts. I also read the base-branch _onfly_decompress.py (_config_levels, _mark/is_marked/stored_names, install_onfly_decompress, patched, check_compressed_keys) and unplaced_source_keys / _conversion_plans / record_unplaced_source_keys / _refused_together for the contracts the diff consumes.
Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2
| № | Location | Finding |
|---|---|---|
| 1 | model_load_utils.py:1812 |
The new class-attribute refusal is order-dependent: the top level is checked before the sub-configs are popped, so a forwarding quantization_config property would be refused instead of handled. Clear every level, then check every level. |
| 2 | model_load_utils.py:2024 |
The unplaced-compressed refusal resolves expert stored names only through model._weight_conversions; on a release that does not record it, every built expert looks unplaced and a good checkpoint is hard-refused with a message that blames the checkpoint. The one test covering this path skips on exactly that condition. |
Neither blocks merge. Approving.
Since the last round
One commit, e2e7486b, closing findings 1 and 2 of post-split round 1. I checked the result rather than the message:
- Finding 1 (silent
vars().popmiss) is closed, and closed the right way. Thehasattrre-check turns the one silent failure the block exists to prevent — a class attribute orattribute_mapindirection surviving the instance-dict pop, after whichfrom_pretrainedbuilds the competing compressed-tensors quantizer anyway — into aNotImplementedErrornaming the class. It raises on the deepcopy only, so the caller's config is untouched on refusal, and it sits inside_refused_together("the expert-parallel load's checks", alone=False), which re-raises on every rank (model_load_utils.py:498-513), so the refusal is collective. The_Typedleg oftest_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passespins it and assertsnot loaded; it fails without the check, since the old pop would have let the load through. My finding 1 is the residual ordering nit on the same lines, not a reopening. - Finding 2 (unexplained
- {"weight"}and suffix-only matching) is closed as comments, which is what I asked for. Both are accurate against the base:check_compressed_keys'outsideclause refuses any compressed stem without.experts., which is what makeskey.rpartition(".")[2]safe, and- {"weight"}keeps a scheme whosecompression_param_namesstarts with"weight"(naive-quantized) from flagging every plain unplaced weight. The exclusion cannot open a hole: such a scheme still storesweight_scale, so an MTP expert is caught by that name, and a scheme storing onlyweightis refused by thenot compressedclause earlier. - Finding 3 (the config-level walk spelled out three times) is not addressed — still non-blocking, still true, now four times counting the new loop. Folded into finding 1, whose fix materializes the list naturally, rather than re-raised as its own thread.
The rest of the diff is byte-identical to round 1, so what I re-traced is what this commit could disturb plus the seams I had not pinned down.
What I traced and found sound
- The mark survives the deletion, and the load-bearing ordering is intact.
_MARKis a separate attribute, andpatchedcaptures(compressor, scheme, param_names, fmt)at install time rather than re-reading the config at mapping time, so droppingquantization_configcannot disarm the decompressor.patcheddelattrs_MARKfrom every level duringfrom_pretrained, socompressed_namesmust stay captured at line 1907; re-deriving it at 2027 would raisestored_names'ValueError. This commit moves neither point. - The deletion cannot disturb the format bookkeeping.
quant_methodis computed at 1791, before the deletion, so the== "mxfp4"branch below is unaffected and noMxfp4Config(dequantize=True)is injected to compete. A marked config whose ownpreparealready droppedquantization_config(Kimi-K3's) takes thequant_method is Nonebranch and still gets the deletion,check_compressed_keysandcompressed_names— the checks hinge on the mark, not on the config being quantized, which is the right hinge. _ckpt_quant_method's_SUB_CONFIG_ATTRSrefactor is behaviour-identical (("text_config", "llm_config")), so the cleared set at 1812 is exactly the read set at 1660 and stays so if the tuple grows.- Collective uniformity of all three refusals. Format check + deletion in
_refused_together("the expert-parallel load's checks");check_compressed_keys+compressed_namesin"the local checkpoint", right afterensure_local_checkpoint(a rank-local, deterministic header read, no collectives); the unplaced refusal in"the expert-parallel load". The last matters most, since that block is followed by_seed_parallel_states' collectives and a rank-dependent raise there would strand the ranks that carried on —unplacedcomes from the same headers and parameter names on every rank (EP sharding renames nothing) andrecord_unplaced_source_keysreturns a sorted list, solen(carried)andcarried[0]are rank-identical. - The refusal cannot fire falsely on a non-expert key. A
*.weight_scale-shaped tensor outside the experts was already refused asoutside, and aweight*sibling under a compressed stem that the scheme does not name by theunknownpass — so the suffix test can only match an expert-shaped stem. - Nothing a currently-working load can reach changes.
compressed_namesis empty on every unmarked path, socarriedis empty; the only other delta is oneand not _decompressed_as_loaded(...)that can only widen what the load accepts.record_unplaced_source_keysalready returnedlist[str]in the base — this PR only starts consuming it, hence the two monkeypatches moving fromNoneto[]. Nomodelopt_state, mode registration, config schema or public API touched, and no new hard import:_decompressed_as_loadedimports inside the function and is short-circuited behindquant_method not in _EP_SOURCE_FORMATSon the format-check path. - Tests. The GPU test is the strongest addition: dense and compressed built from the same tensors, with values INT4 and MXFP4 hold exactly at a power-of-two scale bf16 also holds, then
torch.equalagainst the dense model's fusedgate_up_proj/down_projslice for each rank's own experts — an exact assertion rather than a tolerance hiding a bug — plus_modelopt_unplaced_source_keysempty, plus end-to-end logits. It gathers per-rank problems withall_gather_objectunderbounded_collectives(seconds=120)before any rank asserts, so a regression fails instead of hanging CI, and restores(get_model_conversion_mapping, spawn_materialize)in afinally— needed both becausedist_workersoutlive the job and because that restore is what lets the second parametrization install again without trippinginstall_onfly_decompress' "compressed differently"RuntimeError._LOAD["decompressing"]is left unrestored but ends atFalse(the post-load_conversion_planre-enterspatchedwith the mark cleared) andspawn_materializeis restored anyway, so it is inert.test_compressed_experts_decompressed_at_load_are_placedis the one case running a real decompressedfrom_pretrained, and it pins that only the past-num_hidden_layerslayer stays unplaced.
No prompt-injection-shaped content in the diff.
Assessment
Risk: low. 49 added source lines, all behind a mark only an expert-parallel prepare sets, each of the three refusals backed by a test that fails without it, and every refusal placed where it fails all ranks together — two of them before anything downloads or loads. This commit strictly converts a silent miss into a loud refusal, the direction this file's standard asks for. Both suggestions are robustness-on-other-releases rather than defects in what you measured; the property to preserve on any future edit of this stretch remains that compressed_names is captured before from_pretrained.
Verification note: neither transformers nor compressed_tensors is importable in this sandbox and I did not execute the suite here. Whether every supported PretrainedConfig keeps quantization_config in its instance dict (suggestion 1), whether every transformers >= 5.8 records model._weight_conversions (suggestion 2), and whether from_pretrained honours an explicitly-passed config without re-reading config.json are read from the diff, the base source and the tests — CI's transformers legs and the GPU Kimi-K3 run remain the real check. Everything stated about this round's commit is verified against the checked-in source alone.
🤖 Generated with Claude Code
|
/ok to test dc0a618 |
|
/claude review |
| unplaced = record_unplaced_source_keys(model, ckpt_path) | ||
| # The export would copy them as stored, and this load dropped the quantization_config that | ||
| # reads a compressed one, such as an MTP layer's experts. A suffix is enough to tell them: | ||
| # check_compressed_keys refused any compressed tensor outside the experts. | ||
| if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: | ||
| raise NotImplementedError( | ||
| f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " | ||
| f"for, such as {carried[0]}: the export would copy them as stored, with no " | ||
| "quantization_config to read them. Decompress the checkpoint first." | ||
| ) | ||
|
|
There was a problem hiding this comment.
[SUGGESTION] This refusal is correct, but it lands after from_pretrained has placed the whole checkpoint — for the motivating case (an MTP layer's compressed experts on a Kimi-scale checkpoint) the user pays a full multi-hundred-GB load, across every rank, before learning the load cannot proceed.
The information for the common shape is already in hand before anything loads: check_compressed_keys walks every compressed stem from the shard headers, and the EP config carries num_hidden_layers. A compressed stem whose .layers.<i>. index is at or past that count is exactly the MTP case, and refusing it there would cost nothing on top of the walk that already happens (0.5–2.2 s per the PR's measurements).
Not blocking, and I see the argument for keeping check_compressed_keys architecture-agnostic — unplaced_source_keys is the authoritative answer and only the built model can give it, so this check has to stay as the backstop either way. But a cheap pre-load refusal for the layer-index shape would turn the expected failure mode from "refused after the load" into "refused before the download", and this one would remain for anything the index heuristic misses.
(Checked and fine: record_unplaced_source_keys returns list[str], the suffix match is unambiguous because check_compressed_keys has already refused any compressed tensor outside .experts., every rank computes the same unplaced under EP — non-local experts still resolve to a parameter name that exists on the rank — and the raise sits inside _refused_together("the expert-parallel load"), so one rank refusing fails all.)
There was a problem hiding this comment.
Taken in 9bc105a. Before anything loads, the loader now refuses compressed tensors stored at a layer index at or past the EP config's num_hidden_layers, which is where an MTP layer sits; the error names the first such tensor and the layer count. It reads only the checkpoint's tensor names, the same walk check_compressed_keys makes, once more. The refusal after the load stays for any other shape, and a config without num_hidden_layers skips the early check.
It lives in the expert-parallel loader, which has the EP config, so check_compressed_keys stays architecture-agnostic. On transformers 5.18 no model builds a decoder layer past num_hidden_layers: DeepSeek-V3, GLM-4-MoE and Step-3.7 ignore their MTP layers on load.
The loader test's new [compressed-past-the-layers] case checks the refusal and that nothing loads; it fails without the change. The other two cases now put the stored experts within the layer count, so the early check passes them, and fail if it refuses every layer.
|
/ok to test 9bc105a |
|
/claude review |
| built = getattr(ep_cfg, "num_hidden_layers", None) | ||
| past = sorted( | ||
| name | ||
| for name in checkpoint_tensor_files(ckpt_path) | ||
| if name.rpartition(".")[2] in compressed_names | ||
| and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) | ||
| and built is not None | ||
| and int(layer[1]) >= built | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Performance] This scans every shard header a second time. check_compressed_keys (line 1963) already called checkpoint_tensor_files(ckpt_dir) and built the full tensor name -> file dict; calling it again here re-reads the index plus every shard header. On the checkpoints the description measures that is 0.5–2.2 s of cold header reads duplicated, on every rank (both calls sit inside the same _refused_together block), and it rebuilds a ~497k-entry dict for K3.
Two smaller points in the same comprehension:
built is not Noneis the last condition, so whenep_cfghas nonum_hidden_layersthe regex still runs once per tensor name and the whole scan is wasted. It is loop-invariant — hoisting it out skips the scan entirely in that case.check_compressed_keysalready has thecompressedstem set and receiveshf_config, so the layer-bound question is the same question it is answering ("which compressed tensors would load wrong"). Folding this refusal in there, or having it return the names/stems it collected for reuse here, removes the second scan and keeps the two compressed-key refusals in one place.
Concretely, the cheapest version that keeps the check where it is:
tensor_names = _onfly_decompress.check_compressed_keys(hf_config, ckpt_path)
compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"}
built = getattr(ep_cfg, "num_hidden_layers", None)
past = (
sorted(
name
for name in tensor_names
if name.rpartition(".")[2] in compressed_names
and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name))
and int(layer[1]) >= built
)
if built is not None
else []
)(with check_compressed_keys returning its tensor_files — it is the only caller-visible change, and it is in #3137.)
There was a problem hiding this comment.
| if past: | ||
| raise NotImplementedError( | ||
| f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " | ||
| f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " | ||
| "for them, and the export would copy them as stored, with no quantization_config " | ||
| "to read them. Decompress the checkpoint first." | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] This early refusal is strictly broader than the post-load one it front-runs: the carried check at line 2096 refuses compressed tensors the model has no parameter for, while this one refuses any compressed expert tensor whose layer index is >= num_hidden_layers, placed or not. Today the two coincide — no transformers release builds decoder layers past num_hidden_layers, so a num_nextn_predict_layers MTP block is never placed — but if one ever does (DeepSeek-V3/Kimi-style MTP built as a real module), this refuses a load the carried check would have let through, and the only way out is to re-materialize the checkpoint.
Not worth restructuring for; what would help is making the message say what was actually observed rather than asserting a cause, since (an MTP layer) is an inference from the index:
| if past: | |
| raise NotImplementedError( | |
| f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " | |
| f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " | |
| "for them, and the export would copy them as stored, with no quantization_config " | |
| "to read them. Decompress the checkpoint first." | |
| ) | |
| if past: | |
| raise NotImplementedError( | |
| f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " | |
| f"decoder layers, such as {past[0]} (an MTP layer's, typically): the model builds " | |
| "no parameter for them, and the export would copy them as stored, with no " | |
| "quantization_config to read them. Decompress the checkpoint first." | |
| ) |
There was a problem hiding this comment.
Taken in 23eb10d. The message now reports what the check saw: the tensors lie past the model's decoder layers, "an MTP layer's, typically", and the model builds no parameter for them. The comment above it states the assumption: on transformers 5.18, which builds no decoder layer past num_hidden_layers, this refuses nothing that the check after the load would let through.
|
LGTM: the loader change is correct as far as I traced it and the new early refusal is tested; a human still has to sign off on two edits to existing tests. Needs action:
No action needed:
|
|
/ok to test 23eb10d |
|
/claude review |
|
LGTM: the one new commit closes my last round's suggestion — the early MTP refusal now reuses the tensor-name map Needs action:
No action needed:
|
|
/ok to test 7d78ca1 |
…llel loader load_and_prepare_ep refused a compressed-tensors checkpoint at its format check. A config the on-the-fly decompressor marked now passes it, and on that path the load deletes quantization_config from its copy of the config at every level _ckpt_quant_method reads, so transformers builds no compressed-tensors quantizer to compete with the decompressor (from 5.13 it claims the experts' converters itself). Before anything loads, check_compressed_keys checks the names and dtypes the checkpoint stores: the experts alone compressed, every one of them, in floats bf16 holds. Both run inside the loader's collective checks: the format check and the deletion with the config checks, the stored-tensor check right after the checkpoint is made local, so a refusal on one rank fails every rank. Moved here from the on-the-fly decompression change, whose history is archived; the decompression itself is its own change below this one. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…aced The export copies the checkpoint tensors the model has no parameter for as stored, and the expert-parallel load of a decompressed checkpoint drops the quantization_config that would read a compressed one. An MTP layer's compressed experts, in no layer the model builds, would reach the export packed. The load now refuses once it knows the unplaced tensors, inside its collective checks. _ckpt_quant_method reads the same config levels as the deletion, from _SUB_CONFIG_ATTRS, and the deletion pops from the instance dict. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…the real one does Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The expert-parallel load refuses compressed tensors the model leaves unplaced, and relies on unplaced_source_keys placing the stored weight_packed / weight_scale / weight_shape of every expert the model builds. A CPU from_pretrained of a tiny compressed Qwen3-MoE, through the real decompressor, pins that: they resolve through the converters the load recorded (the stock mapping places none of them), and only the experts of a layer past num_hidden_layers stay unplaced. The two local modelopt.torch.models imports say why they are local, and the loader test's docstring says what its stand-in covers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ssing its own A tiny Qwen3-MoE whose experts are compressed-tensors INT4 (Kimi-K2.5's format) or MXFP4 (Kimi-K3's), installed for on-the-fly decompression and loaded through load_and_prepare_ep: each rank must hold exactly the dense model's weights for its own experts, every stored tensor must be placed, and the sharded model's logits must match the dense model loaded whole. The MXFP4 case goes through compressed-tensors, which the decompressor reads, not transformers' own MXFP4 dequantization. Each rank's checks are gathered before any rank fails, under bounded collectives. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The decompressor's keep-integers patch moved from _materialize_copy to spawn_materialize (#3137), so the 2-GPU job and the real-load unplaced test restore that hook. The GPU case's skip reason names the package it needs: compressed-tensors is not in CI's GPU image, where it skips. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The real decompressing load sets the decompressor's load flag, which outlives the load; the test now runs it on a flag of its own. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…e matchers hold From Claude's review: the load clears quantization_config from each config level's instance dict, which leaves a class attribute or property in place, and transformers' own quantizer would then compete with the decompressor. The load now refuses such a config inside its collective checks. Comments say why "weight" is left out of the compressed names (a float format's anchor, and a plain weight the export has always carried) and why a suffix tells a carried compressed tensor (check_compressed_keys refused any outside the experts). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…e cleared From Claude's review: the top config and its nested language-model configs were walked three times, in _ckpt_quant_method, in the load's quantization_config deletion and in the decompressor. moe.config_levels, beside _SUB_CONFIG_ATTRS, now gives those levels to all three. (#3135 adds the same _SUB_CONFIG_ATTRS; the helper beside it conflicts there once, at the Kimi base's rebuild.) The deletion now clears every level before it checks any: a config class that forwards its text config's quantization_config through a property holds none once the text config's is gone, and is no longer refused; a class attribute, or a property of its own, still is. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…odel The EP loader tests now build their stand-in MoE from the shared model_with_decoder_at, and the file's own copy is gone. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: the load refused the compressed tensors its model leaves unplaced only once from_pretrained had read the whole checkpoint. An MTP layer, the usual case, is stored past the decoder layers the EP config says the model builds, and the checkpoint's tensor names show that before anything loads. Its compressed experts are now refused there; the check after the load stays for any other shape. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… read From Claude's review: the early refusal read the checkpoint's tensor names and every shard header a second time, right after check_compressed_keys had, and scanned them even for a config with no num_hidden_layers. It now takes the map check_compressed_keys returns, and skips the scan without a layer count. Its message says what it found, an MTP layer's typically, rather than assert the cause. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… spec names The expert-parallel load is to read the config's model type for its spec checks, so each config standing in for a checkpoint's names one, as the moved loader tests' do. save_file is already imported at the top of the file the tests moved into. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 6b87b69 |
Note
Split at the reviewers' request (2026-10-30): this PR is now [3/3], the expert-parallel loader's side of on-the-fly decompression.
stage_full_remote_closure, with [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748'skeep_patched.Base: #3137. The diff below is this PR's alone, and it merges after #3137.
Series layout
main([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130, [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132, [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 have merged)fsdp2_wrapmoved intofsdp2_utilsWhat does this PR do?
Type of change: New feature
load_and_prepare_eptakes a checkpoint whose model spec decompresses its experts as they load (#3137), which its format check refused before.quantization_configfrom its copy of the config, at the top level and ontext_configandllm_config: every level_ckpt_quant_methodreads. So transformers builds no compressed-tensors quantizer to compete with the decompressor.quantization_configto the top level, and from transformers 5.13 that quantizer claims the experts' converters itself.None: before 5.11, transformers takes anyquantization_configattribute,Nonetoo, for a quantized checkpoint.check_compressed_keys). A class target such asLineardoes not say what was compressed, and the load drops the config that would load any other compressed weight. So every weight stored compressed must be an expert's, and every expert's weight stored compressed. No compressed tensor may be a float that bf16 cannot hold exactly.Non-expert-parallel loads. A spec's
prepareinstalls the decompressor and marks the config, and it runs only on expert-parallel routes:So a marked config reaches
from_pretrainedonly through this load. What outlives apreparein the process is inert for other loads: the conversion-mapping hook rewrites only a marked config's converters.Usage
Internal to the Kimi adapters (#2748, #2749), whose
prepareinstalls the decompressor; the expert-parallel load does the rest.Testing
tests/unit/torch/utils/test_model_load_utils.py: the compressed-checkpoint case loads from a small safetensors checkpoint, which stores int32 packed weights and bf16 scales as the real ones do. It runs the load tofrom_pretrained, whose config has noquantization_configat any level,llm_config's included, while the caller's config keeps its own.from_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:
On GPUs, the Kimi-K3 run ([OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749) loads its MXFP4 experts through this path.
Locally at the round-1 head:
tests/unit/torch/utilsandtests/unit/torch/modelspass on CPU, torch 2.14: 927 on transformers 5.18, and 861 passed, 27 skipped on 5.5.4.2026-10-30: round 2 (unplaced compressed tensors refused,
_ckpt_quant_methodreads the_SUB_CONFIG_ATTRSlevels, the deletion pops from the instance dict), rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 42fa4de, on [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 4c59a6b (main 99c3834);tests/unit/torch/models/test_onfly_decompress.py+tests/unit/torch/utils/test_ep_loader.py+test_hf_checkpoint_utils.pyon CPU, torch 2.14: 149 passed on transformers 5.18, and 97 passed, 9 skipped on 5.5.4; pre-commit against main passes, mypy included. The refusal's test fails without it.2026-10-30: round 3.
test_unplaced_source_keys.pygains a real decompressedfrom_pretrainedon CPU (a tiny int4 pack-quantized Qwen3-MoE through the real decompressor): every built expert's stored tensors are placed through the recorded converters, so the new refusal passes them, and only a layer pastnum_hidden_layersstays unplaced. It fails with the stock mapping alone and without the decompressor. The two local imports state their reason.test_onfly_decompress.py,test_ep_loader.py,test_hf_checkpoint_utils.pyandtest_unplaced_source_keys.pyon CPU, torch 2.14: 163 passed on transformers 5.18, and 110 passed, 10 skipped on 5.5.4; pre-commit against main passes, mypy included.2026-10-30: rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 7a57fe4, which moved onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 b726fb3 after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged. No conflicts; the own diff is unchanged. Pre-commit against main 62a5c48 passes, mypy included. On CPU, torch 2.14,
test_onfly_decompress.py,test_ep_loader.py,test_unplaced_source_keys.py,test_hf_checkpoint_utils.pyandtest_model_load_utils.pypass: 305 on transformers 5.18, and 240 passed, 22 skipped on 5.5.4.2026-10-30: rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 ee9dc75 (cjluo-nv's round there) with the own diff unchanged. Pre-commit against main passes, mypy included. On CPU, torch 2.14,
test_onfly_decompress.py,test_ep_loader.py,test_unplaced_source_keys.py,test_hf_checkpoint_utils.py,test_model_load_utils.pyandtests/unit/torch/exportpass on transformers 5.18 and 5.5.4, apart from seventest_export_diffusers.pycases that fail the same way at the base in this environment.2026-10-30: rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 a005dab, on [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 e251f6c. One textual conflict, in
model_load_utils.py's moe import list: both_EXPERTS_STYLES([OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136) and_SUB_CONFIG_ATTRS(this PR) kept; the own diff is unchanged. Pre-commit against main passes;test_onfly_decompress.py,test_ep_loader.py,test_unplaced_source_keys.py,test_hf_checkpoint_utils.pyandtest_model_load_utils.pypass on transformers 5.18 (312) and 5.5.4 (241 passed, 23 skipped).2026-10-30: rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 c0022f6 with the own diff unchanged; pre-commit against main passes; the five test files pass on transformers 5.18 (313) and 5.5.4 (241 passed, 23 skipped).
2026-10-30: a 2-GPU case,
tests/gpu/torch/utils/test_expert_parallel_load.py::test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4|mxfp4]: a tiny Qwen3-MoE whose experts are compressed (INT4, MXFP4 through compressed-tensors), loaded at EP2 through the decompressor; each rank's experts must equal the dense model's exactly, every stored tensor must be placed, and the logits must match the dense model loaded whole. compressed-tensors is not in CI's GPU image, so it skips there (the skip names the package); it runs on our cluster, where the run on transformers 5.18 with compressed-tensors 0.19.0 is pending. Also from Claude's review: a config whose class keepsquantization_configis refused, since the load cannot clear it. Rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 0d7866e ([OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 1491b66). Pre-commit against main passes; the five unit test files pass on transformers 5.18 (317) and 5.5.4 (242 passed, 23 skipped).2026-10-30: the 2-GPU case ran on two GPUs, with transformers 5.18.0 and compressed-tensors 0.19.0,
tests/gpu/torch/utils/test_expert_parallel_load.py::test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4]and[mxfp4]passed at this PR's e2e7486, loading each case at EP2 through thespawn_materializefix; on one GPU both skip, as they need two. Since then: rebased onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137's no-behaviour-change commits, the three config walks sharemoe.config_levels(dc0a618; the same levels as before), and the class-heldquantization_configcheck runs once every level is cleared, which changes nothing for the configs the case uses. [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 adds the same_SUB_CONFIG_ATTRStomoe.py, soconfig_levelsbeside it conflicts once where the two meet (resolved at the Kimi base's rebuild). Pre-commit against main passes; the five unit test files pass on transformers 5.18 (317) and 5.5.4 (242 passed, 23 skipped).Re-cut on 2026-10-29 onto [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741, and split on 2026-10-30; the history before is archived. Review fixes then:
_SUB_CONFIG_ATTRSis imported at module level;2026-10-30: rebased onto main 059588b through [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 1a90be5 and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 584c8b4. [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's loader tests now build their stand-in MoE from the shared
model_with_decoder_at, and the unplaced-refusal test here does too. From Claude's review (4428815680): an MTP layer's compressed experts, stored at or past the EP config'snum_hidden_layers, are refused before anything loads, from the checkpoint's tensor names. The refusal after the load stays for any other shape. The loader test's new[compressed-past-the-layers]case fails without the change, and its other cases fail if every layer is refused. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. The compressed-expert GPU cases passed on 2 and 4 GPUs at dc0a618, in [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749's run of the whole GPU file (transformers 5.18.0, compressed-tensors 0.19.0); this round adds only the early check, which their checkpoints pass.2026-10-30: restacked onto [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 7a56a68 and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 56353ef. From Claude's review (4429033004, 4429033879): the early MTP refusal reads the map
check_compressed_keysreturns instead of every shard header again, skips the scan withoutnum_hidden_layers, and its message says what it found rather than assert the cause. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. At this head, with transformers 5.18.0 and compressed-tensors 0.19.0,test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4]and[mxfp4]passed on 2 GPUs and on 4 GPUs, in a run of the whole GPU file.2026-10-31: restacked onto main d464d30 through [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 43f08b5 and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 ab06860. [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 shortened the comment before the local-checkpoint guard; this PR's
compressed_namesline now follows the shortened comment, and the own diff is otherwise unchanged. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-31: at 7d78ca1,
test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4]and[mxfp4]passed on 2 GPUs and on 4 GPUs, with transformers 5.18.0 and compressed-tensors 0.19.0, in a run of the whole GPU file at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 63c8267, which contains it.2026-10-31: restacked onto main e352cac through [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 9deea3b and [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 930bb0f. [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 moved the expert-parallel loader tests from
test_ep_loader.pyintotest_model_load_utils.py, and this PR's loader-test changes moved with them, unchanged. Its stand-in configs now name a model type no spec registers (6b87b69), as [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746's load reads the config's model type first. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.2026-10-31: at 6b87b69, over the rebuilt base,
test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4]and[mxfp4]passed on 2 GPUs and on 4 GPUs, with transformers 5.18.0 and compressed-tensors 0.19.0, in a run of the whole GPU file at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 a60c01c, which contains it.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models, [3/3] of #2747's split. It merges after #3137.