Repository navigation
[OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load - #3137
shengliangxu wants to merge 20 commits into
Conversation
|
/ok to test 3edfe3e |
|
/claude review |
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review: 3
Please check the settings in the CodeRabbit UI or the
 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| stored_as = next( | ||
| getattr(c, _MARK) for c in _config_levels(hf_config) if getattr(c, _MARK, None) is not None | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] Bare next() with no default raises StopIteration when the config is not marked.
check_compressed_keys is documented to run "the expert-parallel load" (#2747), i.e. after install_onfly_decompress. But patched deletes _MARK from every level of the config it rewrote, so calling this after the model has been built — or simply calling it on a config that was never passed to install_onfly_decompress — raises an unexplained StopIteration instead of a message naming the real mistake. A StopIteration escaping inside a generator/comprehension frame higher up the stack would be swallowed as an early stop rather than reported.
stored_as = next(
(getattr(c, _MARK) for c in _config_levels(hf_config) if getattr(c, _MARK, None)), None
)
if stored_as is None:
raise ValueError(
"check_compressed_keys: config is not marked for on-the-fly decompression; "
"call install_onfly_decompress first"
)There was a problem hiding this comment.
Fixed in ad57437. check_compressed_keys now reads the mark through stored_names(hf_config). When no config level carries the mark, it raises ValueError("The config is not marked for on-the-fly decompression: install it first.") instead of StopIteration. test_an_unmarked_config_has_no_stored_names_to_check fails without the change.
| source_patterns=new_sources, | ||
| target_patterns=tgt, | ||
| operations=[DecompressExperts(), *conv.operations], | ||
| ) | ||
| # What else transformers sets on a converter when it maps a model -- the sub-model scope it | ||
| # matches keys under, whether it loads on the CPU -- whichever of them this release has. | ||
| for attr in ("scope_prefix", "base_model_prefix", "force_cpu"): | ||
| if hasattr(conv, attr): | ||
| setattr(new, attr, getattr(conv, attr)) |
There was a problem hiding this comment.
[SUGGESTION] The attribute whitelist silently drops anything else transformers set on the converter.
The rewrite builds a fresh WeightConverter from source_patterns / target_patterns / operations and then copies back exactly three named attributes. Anything else the mapping (or a release in the declared 5.8–5.18 support window) sets on a converter is lost without a diagnostic. distributed_operation is the one that matters most here: the whole design depends on the loader handing non-local experts over as None (GroupedGemmParallel.shard_tensor), and if a release ever carries that on the converter template rather than resolving it per entry, dropping it means every rank decompresses every expert — a silent correctness-and-memory regression, not an error.
test_a_rewritten_converter_keeps_the_scope_transformers_set_on_it pins the three names, so it cannot catch a fourth being added. Copying what is actually there inverts the default to "carry forward":
new = WeightConverter(
source_patterns=new_sources,
target_patterns=tgt,
operations=[DecompressExperts(), *conv.operations],
)
# Carry over whatever else this release set on the converter (sub-model scope, CPU loading,
# distributed op), rather than a name whitelist a new field would fall out of.
for attr, value in vars(conv).items():
if attr not in ("source_patterns", "target_patterns", "operations"):
setattr(new, attr, value)(If WeightConverter is a dataclass in all supported releases, dataclasses.replace(conv, source_patterns=..., operations=...) is even smaller.)
There was a problem hiding this comment.
I'm keeping the named list. I checked the alternatives on transformers 5.18 and 5.19:
- Copying everything doesn't work.
WeightConverteruses__slots__and has no__dict__, sovars(conv)raisesTypeError. It isn't a dataclass there either, sodataclasses.replaceis out. - The mapping sets only two slots. Of the converter's slots,
get_model_conversion_mappingsets onlyscope_prefixandbase_model_prefixon what it returns.force_cpucomes from the constructor. The rewrite copies exactly these three. - The other slots aren't the template's to carry:
- The constructor derives
compiled_sourcesand the_original_*patterns from the new patterns. - The load fills
quantization_operation,collected_tensors,layer_targetsand_was_usedon its own copy of the rewritten converter.
- The constructor derives
distributed_operationis never set. Both releases only initialise it toNone. The experts are sharded at load because the first op is aMergeModulelist, andDecompressExpertsis one.
If a later release starts setting another slot in the mapping, it will need adding here.
| # only the local experts this rank owns. Keeps peak memory at the per-rank shard. | ||
| out[pat] = [ | ||
| compressor.decompress({pn: lists[pn][i] for pn in present}, scheme)[ | ||
| "weight" | ||
| ].to(torch.bfloat16) | ||
| for i in range(len(anchor_list)) | ||
| if anchor_list[i] is not None |
There was a problem hiding this comment.
[SUGGESTION] The hardcoded bf16 output is an unenforced contract with the load's dtype.
Two places assume the model is being loaded in bf16: this .to(torch.bfloat16), and check_compressed_keys's refusal of stored floats "which bf16 cannot hold exactly ... the dtype the experts decompress to". Nothing checks it. _install_keep_int_native only protects integer tensors, so with from_pretrained(..., dtype=torch.float16) the stored weight_scale is cast to the param dtype before decompress sees it, and the decompressed experts come back bf16 while every other param is fp16 — a dtype mismatch that surfaces far from here (at the first forward, or when ModelOpt places quantizers), with nothing pointing back at the decompression.
Cheapest fix is to fail at install time, where the config is in hand:
dtype = getattr(hf_config, "dtype", None) or getattr(hf_config, "torch_dtype", None)
if dtype is not None and dtype not in (torch.bfloat16, "bfloat16"):
raise NotImplementedError(
f"install_onfly_decompress: the experts decompress to bf16; a {dtype} load would "
"mix dtypes. Load the checkpoint in bfloat16."
)Alternatively, derive the target dtype in _make_op_class and drop the bf16 assumption from check_compressed_keys too — but a loud refusal is the smaller diff.
There was a problem hiding this comment.
Fixed in bd40c48 with a refusal at install time:
install_onfly_decompressraisesNotImplementedErrorwhen the config'sdtypeis anything other than bfloat16 orauto. It says the experts decompress to bfloat16 and asks for a bfloat16 load.- The expert-parallel load, the only caller, hands
from_pretrainedthe config'sdtype(bfloat16 when unset), so this is the dtype that load runs in.
The not-bf16 case of test_a_config_with_nothing_to_decompress_leaves_the_loader_alone is a float16 config. It fails without the check.
| fmt = qconfig.format | ||
| fmt = fmt.value if hasattr(fmt, "value") else str(fmt) | ||
| compressor = BaseCompressor.get_value_from_registry(fmt) | ||
| param_names = tuple(compressor.compression_param_names(scheme)) |
There was a problem hiding this comment.
[SUGGESTION] A scheme-level / mixed-precision format degrades into an opaque registry error.
QuantizationConfig.format is not always a single registered compressor name in current compressed_tensors: once any scheme carries its own format, the top-level value becomes CompressionFormat.mixed_precision (and in some releases the field is a list). str(fmt) then yields "mixed-precision" or "['pack-quantized']", and BaseCompressor.get_value_from_registry raises a registry lookup error — the user sees an internal compressed_tensors traceback rather than any of this module's deliberate "Decompress the checkpoint first" refusals, which is the whole point of the single-group / expert-targets gating just above.
Since exactly one group is already enforced, preferring the scheme's own format covers it:
fmt = getattr(scheme, "format", None) or qconfig.format
fmt = fmt.value if hasattr(fmt, "value") else str(fmt)
if fmt not in BaseCompressor.registered_names():
raise NotImplementedError(
f"install_onfly_decompress: compressed-tensors format {fmt!r} has no registered "
"compressor; on-the-fly decompression cannot read it. Decompress the checkpoint first."
)There was a problem hiding this comment.
Fixed in b58cfab, close to your suggestion:
- The format comes from the one config group's own
formatwhen it has one, and falls back to the top-levelformat. - If compressed-tensors registers no compressor for that format, the install raises
NotImplementedErrorand asks for the checkpoint to be decompressed first. That covers "mixed-precision" and a list-valued field too.
test_the_format_is_the_one_group_s_own_and_must_have_a_compressor covers both halves:
the-group-names-it: the top level says "mixed-precision" and the group says "pack-quantized".none-registered: the format has no compressor.
Each case fails without its half of the fix.
| # The nested configs a VLM's expert-parallel plan and expert count are read from. A spec attaching | ||
| # its plan anywhere else would be read as having none, and the MoE planned as a dense model. | ||
| _SUB_CONFIG_ATTRS = ("text_config", "llm_config") |
There was a problem hiding this comment.
[SUGGESTION] The new constant is not the source of truth it claims to be.
Its docstring says these are "The nested configs a VLM's expert-parallel plan and expert count are read from", but _ep_config (line 71) and _num_experts_from_config (line 85) — the two functions that actually do those reads, immediately below — still hardcode ("text_config", "llm_config") inline. So the file now carries three copies of the same tuple, and the only consumer of the named one is _onfly_decompress. That is exactly the drift CONTRIBUTING's "keep a single source of truth" is about: adding a third nested attr later means finding all three sites.
I understand the PR note that this hunk is byte-identical to the model-specs PR so the two merge cleanly — but the cleanup is a one-line change in each function (for _attr in _SUB_CONFIG_ATTRS: / *(getattr(hf_config, a, None) for a in _SUB_CONFIG_ATTRS)) and would be worth doing in whichever of the two lands second, if not here.
There was a problem hiding this comment.
#3135 does that cleanup. It adds this same constant with the same comment, and switches _ep_config and _num_experts_from_config to it. Its spec check reads the constant too. This hunk stays byte-for-byte identical to #3135's so the two merge cleanly in either order. Once #3135 is in, every reader uses the one tuple.
| # with the torchvision pin so the resolver keeps the leg's torch. Not in dev-test: the | ||
| # transformers-free partial_unit(torch) lane installs that extra, and compressed-tensors pulls in | ||
| # transformers. | ||
| session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", "compressed-tensors") |
There was a problem hiding this comment.
[SUGGESTION] compressed-tensors is unpinned, so CI tracks whatever release is newest on the day.
The module is coupled to a specific slice of that library's API — BaseCompressor.get_value_from_registry(fmt), compressor.compression_param_names(scheme), and compressor.decompress(state_dict, scheme)["weight"] as a classmethod. All three are internal-ish surfaces that have moved before (decompress signatures and the QuantizationScheme.format field both changed within the 0.x line). Unpinned, a new release breaks this leg on an unrelated PR, and the failure lands on whoever happens to push that morning.
A floor plus a compatible-release ceiling keeps it reproducible and makes the supported range explicit, the way TRANSFORMERS_VERSIONS already does for transformers:
| session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", "compressed-tensors") | |
| session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", "compressed-tensors~=0.13.0") |
(Substitute the version you actually validated against — the PR says the real-library tests pass locally, so that release is the one to pin.)
There was a problem hiding this comment.
I'm leaving it unpinned, as the repo does for the libraries its unit tests exercise:
dev-testpins the test tooling (pytest, coverage, xdist) but nottimmortorch-geometric.- The partial lane installs
megatron-coreunpinned.
Users install compressed-tensors next to their checkpoints, usually at its newest release. If that release changes one of the three calls used here, CI should show it, not a pin. The real-library tests pass with compressed-tensors 0.19.0.
|
|
||
| Each rank decompresses only its own experts into the transformers 5.x fused-experts layout, so | ||
| expert-parallel PTQ runs on the original compressed checkpoint, with no decompressed copy on disk. | ||
| The math is the ``compressed_tensors`` library's, driven by the checkpoint's ``quantization_config``: | ||
| ``compressor.compression_param_names(scheme)`` names the stored keys of one weight, and | ||
| ``compressor.decompress(...)["weight"]`` dequantizes them, so every format the library registers | ||
| works whose stored floating-point tensors bf16 holds exactly (:func:`check_compressed_keys`). It | ||
| needs ``compressed_tensors`` importable and transformers 5.8 or newer (earlier releases' | ||
| ``Concatenate`` looks up the sibling keys the decompression consumed). | ||
|
|
||
| :func:`install_onfly_decompress` rewrites the per-expert converters ``get_model_conversion_mapping`` | ||
| returns, for the model built from the config it marks: each ``.weight`` source becomes the stored | ||
| keys, and a ``DecompressExperts`` op ahead of the merge dequantizes each expert. | ||
| :func:`to_native_backbone_config` builds a VLM's transformers-native backbone config. | ||
| """ |
There was a problem hiding this comment.
[SUGGESTION] Docstring and comment volume runs well past the coding standard.
CONTRIBUTING's coding standards ask for "a one-line summary with optionally a short description plus only the Args:/Returns: entries whose meaning isn't obvious", comments "to one or two lines", and state that "root-cause analysis belongs in the PR description, not the source". This file is ~475 lines of which roughly 150 are prose, and a lot of it is the PR description restated: the transformers-5.8 Concatenate rationale appears in the module docstring, _make_op_class's docstring and DecompressExperts's docstring; the MergeModulelist-subclass-so-the-loader-shards explanation appears three times; install_onfly_decompress's docstring is a 12-line essay on process-wide hook semantics.
Worth a compression pass, keeping the genuinely non-obvious why comments (the marker/$ handling, the subset-of-siblings pop, the mark-deletion-before-save) and moving the design narrative to the PR body where it already lives. Two concrete redundancies to collapse:
_make_op_class's docstring paragraph 2 andDecompressExperts's docstring paragraph 2 say the same thing about subclassingMergeModulelist; one of them can go.- The module docstring's last paragraph duplicates
install_onfly_decompress's andto_native_backbone_config's own docstrings.
Non-blocking, but it compounds — every future reader pays for it, and the skills guidance in this repo makes the same point about text that has to be re-read on each use.
There was a problem hiding this comment.
Taken in 42fa4de, with no code change. The PR's added source lines drop from 505 to 447:
- The module docstring is now a short summary plus the transformers 5.8 note. The paragraph that repeated the function docstrings is gone.
- These docstrings are down to a few lines each:
install_onfly_decompress,check_compressed_keys,_install_keep_int_native,_read_compression,_compressedifyandto_native_backbone_config. - The reason for subclassing
MergeModulelistis now given only in_make_op_class's docstring.DecompressExpertshas a one-line docstring.
I kept the comments on the subset-of-siblings pop and on deleting the mark before a save.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (trigger comment was a bare /claude review). 4 changed files, all reviewed: modelopt/torch/models/_onfly_decompress.py (475 added), modelopt/torch/utils/moe.py, noxfile.py, and tests/unit/torch/models/test_onfly_decompress.py. Diffed two-dot against origin/shengliangx/hfep-6a2-ep-loader-core.
Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 6
| # | Where | Suggestion |
|---|---|---|
| 1 | _onfly_decompress.py:436 |
bare next() → StopIteration when the config isn't marked (the mark is deleted by patched, so the ordering contract is easy to get wrong) |
| 2 | _onfly_decompress.py:262 |
_compressedify's three-name attribute whitelist silently drops any other field a release sets on the converter; distributed_operation is the one the per-expert None sharding rests on |
| 3 | _onfly_decompress.py:222 |
hardcoded .to(torch.bfloat16) and check_compressed_keys's bf16-exactness rule are an unenforced contract with the load's dtype |
| 4 | _onfly_decompress.py:154 |
a scheme-level / mixed-precision format degrades to an opaque compressed_tensors registry error instead of this module's deliberate refusal |
| 5 | utils/moe.py:57 |
_SUB_CONFIG_ATTRS isn't used by _ep_config / _num_experts_from_config right below it — three copies of the same tuple |
| 6 | noxfile.py:73 |
compressed-tensors unpinned, so every unit leg tracks the newest release of a library this module uses at an internal-ish API surface |
What I traced and found correct
- Decompress-then-fuse dataflow.
_compressedifyemits sources grouped by param name with the anchor group first, soinput_dict's anchor keys arrive in gate-then-up order andoutpreserves it for the downstreamMergeModulelist/Concatenate— gate/up fusion order is right.$-anchoring is what keeps.weight$from claiming.weight_scale, andconvert'smarkerhandling stays consistent with the patterns_compressedifybuilt. Sources the decompression can't consume are left ininput_dictand raise, rather than reaching the merge in the wrong slot. - Per-rank sharding. Subclassing
MergeModulelistto getoperations[0]recognized, then skippingNoneanchors, correctly bounds peak memory at the rank's own shard. The positional expert-count check closes the real foot-gun (one short sibling list would pair every later expert's weight with the next one's scale). - Integer preservation.
_install_keep_int_nativematches the real_materialize_copy(tensor, device, dtype, ...)positional order (pinned bytest_a_copy_parameter_a_release_adds_passes_through), and is reached throughspawn_materialize's call-time lookup. Scoping the skip tonot t.is_floating_point()leaves fp8/e8m0 scales casting to bf16, which is value-exact. - Mark lifecycle.
from_pretraineddeep-copies the config, sopatchedclearing_MARKstrips it from the model's copy (keeping it out of a later save) while the caller's config stays marked forcheck_compressed_keysand for a re-prepare. The tests model this with explicitcopy.deepcopy. Therewritten == 0refusal correctly covers the silent newly-initialized-experts case. - Install idempotence. Re-install for the same checkpoint marks and returns without double-wrapping (which, with a
weightanchor, would have decompressed twice); a differently-compressed second checkpoint, a config with nothing to decompress, and a raw dict all raise before anything is patched. check_compressed_keysset logic..experts.correctly excludesshared_experts.; compressed-outside-experts, plain-expert-weight, and nothing-compressed all refuse. The dtype gate'sF16/F32/F64set is right for the bf16 target (fp8/e8m0/bf16 all round-trip exactly). Worth noting as a known limitation rather than a defect: this refuses the common fp32-weight_global_scaleNVFP4 and fp32-weight_scaleFP8 layouts, which matches the PR's stated Kimi INT4/MXFP4 scope and fails loudly with actionable text.- Plugin laziness.
compressed_tensorsandtransformersare both imported inside functions with the reason commented,_read_compressionlooks the config up before importing so a no-compression config needs neither, andmoe.pystill imports without transformers. No mode registration,modelopt_stateschema, config schema, or public__init__export is touched, so no restore/migration surface here. - Tests. 39 cases that drive transformers' own loader objects rather than asserting on mocks, with the real-library cases skipping with a reason where the dep is absent. CPU-only and fast — fits the
tests/unitbudget. No prompt-injection or review-steering text anywhere in the diff.
Risk: low. Additive, behind a private module installed only from a model's prepare hook, with the only process-wide side effect (_materialize_copy) scoped to a cast that just one kind of checkpoint asks for. The one structural thing to keep an eye on across the declared 5.8–5.18 transformers range is #2 — the converter-attribute whitelist is the place where a release bump could regress sharding without an error.
🤖 Generated with Claude Code
4b6470c to
4c59a6b
Compare
3edfe3e to
42fa4de
Compare
|
/ok to test 42fa4de |
|
Thanks for the review. Pushed 42fa4de, and the PR is rebased onto #3136's new head, which is on current main. Outcomes per suggestion:
Each fix has a test that fails without it. |
|
/ok to test 1a90be5 |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: earlier blockers are resolved; the substantial new test suite still needs human sign-off.
Needs action:
- Review and sign off on the 1,001-line test addition in
tests/unit/torch/models/test_onfly_decompress.py. - Move the test’s local
build_reverse_name_mapperimport to module scope, or document a concrete reason to keep it lazy (non-blocking;CONTRIBUTING.md).
No action needed:
- ✔️ Resolved: expert-ID alignment, nested dtype validation, reverse naming, sequential-load isolation, and sharded integer preservation. The source diff is now within budget; versioned EP2 results are recorded.
1a90be5 to
7a56a68
Compare
|
/ok to test 7a56a68 |
75bc2d6 to
ab06860
Compare
7a56a68 to
43f08b5
Compare
|
/ok to test 43f08b5 |
7747c5d to
930bb0f
Compare
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. install_onfly_decompress rewrites the per-expert converters of the model built from the config it marks: each .weight source becomes the stored keys (weight_packed, weight_scale, ...), and a DecompressExperts op ahead of the merge dequantizes each owned expert to bf16 with the compressed_tensors library, driven by the checkpoint's own quantization_config. Non-local experts arrive as None and are skipped. A checkpoint must compress its experts alone, in one config group, in floats bf16 holds exactly; check_compressed_keys checks the stored names and dtypes. The hooks are process-wide but rewrite only a marked config's converters. to_native_backbone_config rebuilds a VLM's bundled text config as a native transformers backbone config, carrying the mark. utils/moe.py gains _SUB_CONFIG_ATTRS, the nested configs the expert- parallel load reads, word for word as the model specs' change adds it. CI's unit sessions install compressed-tensors, for the tests against the real library. Moved out of the on-the-fly decompression change, whose history is archived. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
modelopt.torch.utils.moe imports nothing from the model registry at module level, so the decompressor can import the constant at the top of the file: importing the registry, moe, the decompressor or a Kimi adapter on its own shows no cycle. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The torch 2.8 and 2.9 legs it was kept out of are gone from the unit sessions, so the condition and the comment about older legs are dead. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
check_compressed_keys read the names the mark holds through a bare next(), so a config the decompressor never marked, or one whose mark the hook has already cleared, raised an unexplained StopIteration. stored_names() reads them, and refuses an unmarked config with a ValueError that says what is missing. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…h no compressor compressed-tensors writes "mixed-precision" at a config's top level once a group names its own format, and the registry lookup of that name failed inside the library. The one group's own format is read first, and a format the library registers no compressor for is refused here, with this module's "decompress the checkpoint first". Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The experts decompress to bf16, and the loader sets each parameter from what the conversion hands it, without a cast: a checkpoint loading in fp16 or fp32 would end up with bf16 experts beside the rest of the model, failing far from here. install_onfly_decompress refuses such a config before anything is patched. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
from_pretrained keeps the converters it ran as model._weight_conversions, and a save reverses them: transformers' save_pretrained, and ModelOpt's export, which names the expanded experts back after the checkpoint. The rewritten converter reversed into its stored names (weight_packed, weight_scale, ...), so the export's expert-leaf arity check refused it and the export fell back to in-memory names: gate_proj/up_proj/down_proj for a checkpoint whose experts are w1/w3/w2, Kimi-K3's among them. The rewritten converter is now a WeightConverter subclass that reverses as the converter it rewrote: the decompression is load-only. Real from_pretrained loads of a tiny compressed Mixtral, INT4 and MXFP4, replace the converter-only round trip: the experts load exactly, save back as the checkpoint's own w1/w3/w2 weights, and export under those names. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The load hands the decompression each expert's stored tensors by the order it collects them, not by expert id, so equal counts do not make a weight and a scale the same expert's: packed weights for experts 0 and 1 beside scales for experts 1 and 2 would pair expert 0's weight with expert 1's scale. check_compressed_keys, which still sees the names, now refuses a checkpoint unless every compressed weight stores the same names; a name no weight stores (Kimi-K3's weight_global_scale) is not read. A scale stored without its packed weight is refused the same way. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The copy patch kept every integer source native for every later load in the process, so an unmarked model whose checkpoint stores an integer for a float parameter would load it uncast. The patched conversion mapping, which from_pretrained asks for before it reads a tensor, now records whether the load under way is one whose converters it rewrote, and only that load keeps integers native. A real load of an unmarked model after installing casts as before. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The install checked only the top-level config's dtype. A VLM config with none there and float16 on its text config passed, and the backbone config to_native_backbone_config derives from that text config keeps float16, with the experts decompressed to bf16 beside it. Every level must now be bf16 or unset; "auto" is refused too, since it reads the dtype from the checkpoint. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
print_rank_0, require_transformers and safetensors' safe_open add no optional dependency at import. compressed-tensors, the version-dependent transformers modules and the checkpoint plugin stay lazy. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… read The stored names come from the installed compressed-tensors, not the checkpoint: a name a newer release added under a compressed weight would match no rewritten source, load as an unexpected key, and leave the experts decompressed without it. check_compressed_keys now refuses it. A per-expert converter source that is no weight was handed to the decompress op, which refused it from a loader thread as an unconsumed key; _compressedify now refuses it where the converter is rewritten. stored_names says what it returns. Comments tightened to keep the change within its size. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From transformers 5.14 an expert-parallel load reads each expert's tensors through DtensorShardOperation.shard_tensor, which casts them to the parameter's dtype as it shards them onto the mesh, and never reaches the _materialize_copy the keep-integers patch sat on. A compressed checkpoint's packed weights (INT4's int32, MXFP4's uint8) then reached the decompression as bf16 and failed there, as a 2-GPU Kimi-K3 load on transformers 5.18 did. The patch now sits on spawn_materialize, where the loader starts both the copy and the shard, in every release from 5.8: a decompressing load drops the float dtype for an integer source there, judging a safetensors slice by its stored dtype without reading it. A test shards an int32 tensor onto a one-rank CPU mesh through transformers' own DtensorShardOperation; it fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…loor first From Claude's review: _compressedify refuses a source that is no expert weight, so a dense weight fused with the experts' cannot be rewritten to stored names it does not have and load newly initialized. The stored-name check refuses only a weight's names the installed compressed-tensors does not give the scheme, so an expert's uncompressed bias, loaded by its own converter, passes. The install checks the transformers floor before it needs compressed-tensors, so the first error names what to fix first. Comments tightened. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A real decompressing load leaves the load flag set until the next load, and the fixture copied whatever an earlier test left, so the integer test failed when it ran after one in the same session. The fixture now starts each test from a clean flag and restores the module's own after. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From Claude's review: the unknown-name check covers weight* names only, which every compressor uses today, and says so. The rewrite's refusal of a source that is no expert weight explains the bias case: a bias belongs in a converter of its own, and grouped with the experts' weights it has no stored key either. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
torch.distributed.tensor is a core import with no optional dependency, so the shard-path test imports it with the module's other imports. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
No behaviour change: shorter module, converter and hook docstrings, and the source-dtype check folded into the hook it serves. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… reuse From Claude's review of #2747: the expert-parallel loader read the checkpoint's tensor names and shard headers again, for a check of its own, right after this one read them. The map is now returned instead of thrown away. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
43f08b5 to
9deea3b
Compare
|
/ok to test 9deea3b |
Note
Split out of #2747 at the reviewers' request (2026-10-30): the on-the-fly decompression itself,
modelopt/torch/models/_onfly_decompress.py, with its tests. #2747 keeps the expert-parallel loader's side.Base: #3136 (#2741 [2/3], the expert-parallel loader), for
require_transformers. The diff below is this PR's alone, and it merges after #3136.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
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.install_onfly_decompress(hf_config)rewrites the per-expert converters of the model built from the config it marks..weightsource becomes the keys the checkpoint stores (weight_packed,weight_scale, ...). ADecompressExpertsop ahead of the merge dequantizes each expert this rank owns to bf16, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk. Experts other ranks own arrive asNoneand are skipped.compressed_tensorslibrary does the math, driven by the checkpoint's ownquantization_config. The format is the config group's own when it names one; the top level then says "mixed-precision". A format compressed-tensors registers no compressor for is refused.quantization_configwith no groups, since there is nothing to decompress. Kimi-K2.5, K2.6 and K2.7-Code (INT4) and Kimi-K3 (MXFP4) each carry one group, undertext_config.save_pretrained, ModelOpt's export) writes the experts back under the checkpoint's own.weightnames (w1/w3/w2for Kimi-K3).compressed_tensors. So does a config that loads in a dtype other than bf16, the dtype the experts decompress to, at any of its levels.Concatenatelooks up the sibling keys the decompression has already consumed.check_compressed_keys(hf_config, ckpt_dir)reads the checkpoint's stored tensor names and dtypes from its shard headers. Every compressed weight must be an expert's, every expert's weight must be compressed, every compressed weight must store the same names (the load pairs them by position), all of them names the installed compressed-tensors gives the scheme, and no compressed tensor may be a float bf16 cannot hold exactly. The expert-parallel load runs it ([OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747).stored_names(hf_config)gives the names a marked config's checkpoint stores, and refuses an unmarked config by name.to_native_backbone_configrebuilds a VLM's bundled text config as a native transformers backbone config, carrying the mark.utils/moe.pygains_SUB_CONFIG_ATTRS, the nested configs the expert-parallel load reads. It is word for word as the model specs' change adds it, so the two merge without a conflict.unitnox sessions installcompressed-tensors, 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.Why not transformers' own compressed-tensors quantizer (
run_compressed=False, which 5.18 deprecates fordequantize=True)? It does not fit an expert-parallel PTQ load:DecompressExpertsop itsupdate_weight_conversionsprepends to the experts' converters, is recent: transformers 5.5 has none. The Kimi-K2.5 adapter runs from 5.8.Nonethe expert-parallel load hands it for each expert another rank owns. This decompression shards at load and decompresses only the rank's own experts.apply_quantization_config,compress_model). PTQ needs plain bf16 modules to place ModelOpt's quantizers on.So the expert-parallel load deletes
quantization_configon a decompressed load (#2747), and the experts are decompressed here.Usage
Internal to the Kimi adapters (#2748, #2749), called from their
preparehooks.Testing
tests/unit/torch/models/test_onfly_decompress.py(50 cases) covers:from_pretrained);from_pretrainedloads of a tiny compressed Mixtral, INT4 and MXFP4, with the real library: the experts load exactly, save back as the checkpoint's ownw1/w3/w2weights, and export under those names. These and the other cases with the real library run in CI's unit sessions; elsewhere they skip, saying what they need.test_onfly_decompress.pyplus the expert-parallel loader'stest_ep_loader.py.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: 146 passed on transformers 5.18, and 94 passed, 9 skipped on 5.5.4; pre-commit against main passes, mypy included. Each fix's test fails without it.test_onfly_decompress.py,test_ep_loader.py,test_unplaced_source_keys.py,test_hf_checkpoint_utils.pyandtest_model_load_utils.pypass: 301 on transformers 5.18, and 237 passed, 21 skipped on 5.5.4.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.DtensorShardOperation.shard_tensor, which cast the packed integers to bf16 (a 2-GPU Kimi-K3 load on 5.18 failed inDecompressExperts). The keep-integers patch now sits onspawn_materialize, where every read starts; a new test shards an int32 tensor through transformers' ownDtensorShardOperationon a one-rank CPU mesh and fails without the fix. Claude's suggestions taken: experts' weights only rewritten, an expert's bias passes the stored-name check, the transformers floor checked first, wording. Rebased onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 1491b66. 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 (313) and 5.5.4 (239 passed, 22 skipped), with compressed-tensors 0.19.0. The 2-GPU run of the compressed expert-parallel load on 5.18 ([OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's GPU case) is pending on our cluster.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 [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 e2e7486 (on this PR's 0d7866e), loading each case at EP2 through thespawn_materializefix; on one GPU both skip, as they need two. Since then this PR gained only the module-scope test import (b56c3e7) and shorter docstrings (48b922c), with no behaviour change; source is now 495 added lines, within the ~500 budget._SUB_CONFIG_ATTRSis imported at module level;compressed-tensorson every torch, since the torch 2.8 and 2.9 legs it was kept out of are gone.tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.check_compressed_keysreturns thetensor name -> filemap it read, for the loader to reuse (7a56a68); nothing else changed. Pre-commit against main passes, mypy included;tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.tests/unit/torch/modelsandtests/unit/torch/utilspass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.[int4]and[mxfp4], which 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.test_model_load_utils.py. No conflicts; the own diff is 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.[int4]and[mxfp4], which 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"
preparehook.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 of expert-parallel PTQ for Hugging Face MoE models, [2/3] of #2747's split: