[shengliangxu]

Note

#2747 folded in (2026-11-03): at the user's request to fold the series backward, this PR now also carries #2747, the expert-parallel loader's side of the decompression, with its commits as they were. #2747 is closed and keeps its review history; its ticket is this one, OMNIML-6063.

Split out of #2747 at the reviewers' request (2026-10-30), and joined with it again: the on-the-fly decompression, modelopt/torch/models/_onfly_decompress.py, and the expert-parallel load of a checkpoint its model spec decompresses.

Base: main, which holds the expert-parallel loader (#3136) and its routing with the model spec (#2780) since they merged.

Series layout

• main ([OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 [1/2], [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131, [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132, [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134, [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135, [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140, [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067, [OMNIML-6071] Calibrate a distributed model on per-rank data #2740, [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 and [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 have merged)
  • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3] the loader's post-load refusals
  • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding → [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
  • [OMNIML-6063] [2/2] Decompress compressed-tensors MoE experts as they load, expert-parallel loads included #3137 [2/2] on-the-fly expert decompression, expert-parallel loads included (this PR; [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 folded in)
    • [OMNIML-6078] Add the Kimi-VL and Kimi-K2.5 expert-parallel adapters #3138 Kimi-VL and Kimi-K2.5 adapters ([OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 folded in)
      • [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 [2/2] Kimi-K3 adapter

What does this PR do?

Type of change: New feature

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(hf_config) rewrites the per-expert converters of the model built from the config it marks.
  • Each .weight source becomes the keys the checkpoint stores (weight_packed, weight_scale, ...). A DecompressExperts op 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 as None and are skipped.
  • It is format-agnostic: the compressed_tensors library does the math, driven by the checkpoint's own quantization_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.
  • It handles a checkpoint whose experts are its only compressed weights, in one config group that targets them. Several groups, or a group that is not the experts, are refused; so is a quantization_config with 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, under text_config.
  • Stored params that disagree on how many experts they hold are refused, as is a checkpoint key the decompression does not consume.
  • The rewritten converter reverses as the one it rewrote: the decompression is load-only, so a save (transformers' save_pretrained, ModelOpt's export) writes the experts back under the checkpoint's own .weight names (w1/w3/w2 for Kimi-K3).
  • The hooks are process-wide, but they rewrite only a marked config's converters, and only a load of such a model keeps integer tensors uncast: any other model loaded in the process keeps its stock converters and casts.
  • A second, differently compressed checkpoint in the same process is refused; the same checkpoint again is a no-op.
  • A config with nothing to decompress, or a raw dict, raises before anything is patched, and without needing 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.
  • It needs transformers 5.8: earlier releases' Concatenate looks up the sibling keys the decompression has already consumed.
  • It needs compressed-tensors 0.16, the first release whose compressors decompress one weight by its scheme (compression_param_names(scheme) and decompress(state_dict, scheme) as classmethods). An older one is refused before anything is read, naming the release and the upgrade.
• 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 (below). stored_names(hf_config) gives the names a marked config's checkpoint stores, and refuses an unmarked config by name.
• CI's unit nox sessions install compressed-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 for dequantize=True)? It does not fit an expert-parallel PTQ load:

• Its fused-experts path, a DecompressExperts op its update_weight_conversions prepends to the experts' converters, is recent: transformers 5.5 has none. The Kimi-K2.5 adapter runs from 5.8.
• That op decompresses every expert it is handed into one buffer sized for all of them, and has no case for the None the expert-parallel load hands it for each expert another rank owns. This decompression shards at load and decompresses only the rank's own experts.
• Before the load, the quantizer applies the checkpoint's compressed-tensors scheme to every module its config groups target (apply_quantization_config, compress_model). PTQ needs plain bf16 modules to place ModelOpt's quantizers on.

So the expert-parallel load deletes quantization_config on a decompressed load (below), and the experts are decompressed here.

The expert-parallel load of a decompressed checkpoint (folded in from #2747):

load_and_prepare_ep takes a checkpoint whose model spec decompresses its experts as they load (above), which its format check refused before.

• Format check. A config the decompressor marked passes the check; an unmarked compressed-tensors checkpoint is still refused.
• No competing quantizer. On the marked path the load deletes quantization_config from its copy of the config, at the top level and on text_config and llm_config: every level _ckpt_quant_method reads. So transformers builds no compressed-tensors quantizer to compete with the decompressor.
  • Kimi-K2.5's config class copies its text config's quantization_config to the top level, and from transformers 5.13 that quantizer claims the experts' converters itself.
  • It is deleted rather than set to None: before 5.11, transformers takes any quantization_config attribute, None too, for a quantized checkpoint.
  • The caller's config keeps its own.
• Stored tensors checked before anything loads (check_compressed_keys). A class target such as Linear does 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.
• No compressed tensor carried unread. The export copies the checkpoint tensors the model has no parameter for as stored, and this load dropped the config that would read a compressed one. So once the load knows those tensors, it refuses any that is compressed: an MTP layer's compressed experts, say, which pass the stored-tensor check but sit in no layer transformers builds. Uncompressed ones are carried as before. So are compressed tensors in a safetensors file the load never opens, such as a standalone mtp.safetensors, which the export copies as stored; a plain such file is carried.
• Collective checks. The format check and the deletion run inside the loader's collective checks ([OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136), and so does the stored-tensor check, right after the checkpoint is made local. The unplaced-tensor refusal runs inside the load's own collective block. So a refusal on one rank fails every rank.

Non-expert-parallel loads. A spec's prepare installs the decompressor and marks the config, and it runs only on expert-parallel routes:

• only an expert-parallel plan runs it, on every rank; the plan decision reads the plan the spec states instead (since [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748, now in [OMNIML-6078] Add the Kimi-VL and Kimi-K2.5 expert-parallel adapters #3138);
• an EP degree of 1 and a single process skip it;
• the FSDP2 load refuses a quantized checkpoint and a VLM before it loads.

So a marked config reaches from_pretrained only through this load. What outlives a prepare in the process is inert for other loads: the conversion-mapping hook rewrites only a marked config's converters.

Usage

Internal to the Kimi adapters (#3138, #2749), whose prepare installs the decompressor; the expert-parallel load does the rest.

Testing

• tests/unit/torch/models/test_onfly_decompress.py (49 cases) covers:

  • the rewritten converters, a converter source that is no weight refused at the rewrite, and a stored param without its weight refused at the load;
  • packed experts decompressing and fusing like an uncompressed checkpoint, through the loader's own steps;
  • only owned experts reaching the compressor, and stored params disagreeing on the expert count refused;
  • integer tensors keeping their dtype through transformers' own copy in a decompressing load, and a later load of an unmarked model casting as before (a real from_pretrained);
  • installing once, a second checkpoint compressed differently refused, and a config with nothing to decompress, a raw dict, or a load in a dtype other than bf16 at any level, refused;
  • the format taken from the one group, and a format with no compressor refused;
  • an unmarked config refused by name when its stored names are read;
  • the stored-tensor check over ten layouts, among them scales stored for other experts in equal number and a name the installed compressed-tensors does not give the scheme;
  • a VLM marked at either level;
  • real from_pretrained loads of a tiny compressed Mixtral, INT4 and MXFP4, with the real library: the experts load exactly, save back as the checkpoint's own w1/w3/w2 weights, 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.

• Skips below transformers 5.8, the release the decompression requires.

• Locally at the round-1 head:

  • pre-commit against main passes, mypy included;
  • on CPU, torch 2.14: 87 passed on transformers 5.18, and 45 passed, 3 skipped on 5.5.4. That is test_onfly_decompress.py plus the expert-parallel loader's test_ep_loader.py.

• 2026-10-30: round 2 (an unmarked config refused by name, the group's own format and a format with no compressor refused, a non-bf16 load refused, docstrings compressed), rebased onto [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.py on 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.

• 2026-10-30: rebased onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 b726fb3, restacked after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged, with no conflicts and the own diff 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.py and test_model_load_utils.py pass: 301 on transformers 5.18, and 237 passed, 21 skipped on 5.5.4.

• 2026-10-30: round 3 (cjluo-nv's review): the converter reverses as the one it rewrote, compressed weights must store the same names, integers stay uncast only in a decompressing load, bf16 at every config level, module-level imports. Each fix's test fails without it. 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.py and tests/unit/torch/export pass on transformers 5.18 and 5.5.4, apart from seven test_export_diffusers.py cases that fail the same way at the base in this environment.

• 2026-10-30: rebased onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 e251f6c (its module-scope imports and mixed-GPU refusal) with no conflicts and the own diff unchanged. Pre-commit against main passes; the same five test files pass on transformers 5.18 (308) and 5.5.4 (238 passed, 22 skipped).

• 2026-10-30: Claude's two suggestions on round 3 taken: a stored name the installed compressed-tensors does not give the scheme is refused, and so is a converter source that is no weight, where it is rewritten. Each test fails without its change. Pre-commit against main passes; the five test files pass on transformers 5.18 (309) and 5.5.4 (238 passed, 22 skipped).

• 2026-10-30: juhi's blocker fixed. From transformers 5.14 an expert-parallel load shards each expert's tensors through DtensorShardOperation.shard_tensor, which cast the packed integers to bf16 (a 2-GPU Kimi-K3 load on 5.18 failed in DecompressExperts). The keep-integers patch now sits on spawn_materialize, where every read starts; a new test shards an int32 tensor through transformers' own DtensorShardOperation on 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.py and test_model_load_utils.py pass 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.

• 2026-10-30: 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 [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 the spawn_materialize fix; 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.

• Review fixes carried from [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747:

  • _SUB_CONFIG_ATTRS is imported at module level;
  • the unit sessions install compressed-tensors on every torch, since the torch 2.8 and 2.9 legs it was kept out of are gone.

• 2026-10-30: rebased onto main 059588b through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 584c8b4, whose base now sits on main with [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 merged. No conflicts, and the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-10-30: restacked onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 56353ef (main 059588b). From Claude's review of [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747: check_compressed_keys returns the tensor name -> file map it read, for the loader to reuse (7a56a68); nothing else changed. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-10-31: restacked onto main d464d30 through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 ab06860. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-10-31: at 43f08b5, the decompressor ran in [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's compressed-expert GPU cases, [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.

• 2026-10-31: restacked onto main e352cac through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 930bb0f, which moved the expert-parallel loader tests into test_model_load_utils.py. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-10-31: at 9deea3b, over the rebuilt base, the decompressor ran in [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's compressed-expert GPU cases, [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.

• 2026-10-31, at c6b427a: on compressed-tensors 0.13.0 the decompressor's setup failed with TypeError: 'property' object is not callable. Before 0.16 a compressor states its parameter names as a property (to 0.14) or not at all (0.15), and decompresses through an instance built from the whole config. An older compressed-tensors is now refused up front, naming 0.16 and the upgrade (321ba9a, c6b427a), and the tests that need the real library skip below 0.16. test_a_compressed_tensors_before_per_scheme_compressors_is_refused_by_its_release fails without the check, on a stand-in compressor of 0.13's form. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. test_onfly_decompress.py against each compressed-tensors release, on CPU, torch 2.14, transformers 5.18:

  compressed-tensors	Passed	Skipped
  0.19.0, 0.16.0	54	0
  0.15.0.1, 0.13.0	45	9

• 2026-10-31, at 81bcef7, from Claude's review: the loader hook reads the decompressing flag first and looks at a source's dtype only in a decompressing load, so any other load passes through it untouched; test_integer_checkpoint_tensors_keep_their_dtype_on_load hands it a source with no dtype outside such a load, which raised before. The test's reverse name mapper is imported at module scope. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-01, at e72a810: restacked onto main cf38ec3 through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 4dd38a3. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; test_model_load_utils.py and test_onfly_decompress.py, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-01, at df2de9d: restacked onto main 51645db through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 35f1ac0; the tree is unchanged. Pre-commit against main passes, mypy included; test_model_load_utils.py and test_onfly_decompress.py, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-02, at fcd6c08: restacked onto main 2d4e8d0 through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 39f55f2. Main now holds [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135, whose _SUB_CONFIG_ATTRS in utils/moe.py is the one this PR added, word for word, so the PR no longer changes moe.py; the rest of the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-02, at f410618: restacked onto main 4f08a12, which now holds [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067, through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 c0041d0. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-02, at 59edbb4: restacked onto main 653ef61, which now holds [OMNIML-6071] Calibrate a distributed model on per-rank data #2740, through [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 d87f5bd, which now targets main. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-02, at c9dcaa9: rebased onto main bc5adc6, without [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136, and retargeted to main: its three files import nothing [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 adds. No conflicts; the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, and without transformers.

• 2026-11-02, at ac48ad2, from CodeRabbit's review: every in-function import in _onfly_decompress.py and its test now names its reason in a comment: optional dependencies, and core_model_loading being new in transformers 5. Comments only; pre-commit against main passes, and test_onfly_decompress.py passes on CPU, torch 2.14, on transformers 5.18, and skips on 5.5.4 and without transformers, as before.

• 2026-11-03, at a39cdb5: [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 folded in. This PR's branch was fast-forwarded to [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's head, so its commits are [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's as they were, on main a6cc167. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. On GPU, the whole GPU file passed at [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 7cfcd54, which contains this code, with transformers 5.18.0 and compressed-tensors 0.20.1: on 2 GPUs, 5 passed and [ep-half] skipped, as it needs 4; on 4 GPUs, all 6 passed, in the file's order and again in reverse.

• 2026-11-03, at 24bc5b6: the history was linearized onto main a6cc167, replaying this PR's and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's commits in order without the merge from main, so DCO no longer counts main's squash. The tree is identical to a39cdb5's, every commit is signed off, and pre-commit against main passes.

• 2026-11-03, at c3cfcc6, from CodeRabbit's review: the loader tests import _onfly_decompress, packaging's Version and transformers' version at module scope, since nothing makes them optional there; core_model_loading stays imported in its test, after the version check, with its reason beside it. Tests only; pre-commit against main passes, and test_model_load_utils.py and test_unplaced_source_keys.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-03, at 25a4cd8, from Claude's review: to_native_backbone_config moved to [OMNIML-6078] Add the Kimi-VL and Kimi-K2.5 expert-parallel adapters #3138, where the Kimi adapters call it, with its tests; this PR keeps the VLM case of the either-level load test. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.

• 2026-11-03, at eff418c, from cjluo-nv's review: a decompressing load refuses compressed tensors in the safetensors files it never opens, which the export copies as stored (ea3f31e; the test fails without it), and the real round trips give each expert scales of its own and store packed weights and scales across two shards in opposite expert order (eff418c), passing on transformers 5.8.0 and 5.18.0. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, and without transformers.

• 2026-11-03, at 2d8208d, from Claude's review: a decompressing load reads the checkpoint's shard headers once, as the unplaced-key check takes the map check_compressed_keys read. Pre-commit against main passes, mypy included; tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, and without transformers.

• 2026-11-03, at 86452a8, from kevalmorabia97's review: the tensor map reaches unplaced_source_keys by keyword and only when the load read one, so a two-argument replacement records its keys again (a new test pins both forms; it fails with the positional call), and the tests that write over a safetensors file they loaded clone its tensors first. Pre-commit against main passes, mypy included; tests/unit/torch/export, tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart from seven test_export_diffusers.py cases that fail the same way on main here, and without transformers.

• 2026-11-03, at 8feea6b: rebased onto main 7d44246, which now holds [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780. Three conflicts, in imports and context only: the loader's import list keeps the names [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 uses beside this PR's config_levels, the local-checkpoint step keeps [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780's hub options ahead of this PR's stored-tensor check, and the loader tests import _onfly_decompress beside get_spec. The tests' _NO_SPEC constant is no longer added here, as [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 added it word for word. The rest of the own diff is unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/export, tests/unit/torch/models and tests/unit/torch/utils, run in one process, pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4, apart from seven test_export_diffusers.py cases that fail the same way on main here, and tests/unit/torch passes without transformers.

• Earlier history and GPU results: see [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747.

Tests folded in from #2747:

• 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 to from_pretrained, whose config has no quantization_config at any level, llm_config's included, while the caller's config keeps its own.

  • The same checkpoint unmarked is refused.
  • One that also compresses an attention projection is refused before from_pretrained.
  • A compressed tensor the built model has no parameter for (an MTP layer's expert) is refused once the weights load; an uncompressed one gets past the check.

• 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:

  Checkpoints	Names	Names and dtypes
  Kimi-K2.5, K2.6, K2.7-Code, K3	0.5 s to 2.2 s (the most: K3's 497,220 names, read cold)	not timed
  Synthetic, K3's 247,500 compressed tensors in 10 shards (CPU)	0.5 s	1.0 s

• 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.

• tests/unit/torch/utils/test_unplaced_source_keys.py: a real decompressed from_pretrained on CPU (a tiny int4 pack-quantized Qwen3-MoE through the real decompressor): every built expert's stored tensors are placed through the recorded converters, and only a layer past num_hidden_layers stays unplaced.

• tests/gpu/torch/utils/test_expert_parallel_load.py::test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own[int4|mxfp4] (2 GPUs): a tiny Qwen3-MoE whose experts are compressed (INT4, MXFP4, through compressed-tensors) loads 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.

• [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's dated testing history, its 2- and 4-GPU runs included, stays on [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive; the decompressor is installed only by a model's prepare hook, and only a config a model spec marked takes the new load path.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A. compressed_tensors is imported only for compressed checkpoints, which already need it; only CI's unit sessions install it, for the tests.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A. The feature's entry lands with [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ✅ Approved on 2026-11-03 with [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 folded in (at eff418c). Before the fold, this PR's round 1 was approved on 2026-10-30.

Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models. #2747's split is now two PRs:

• [1/2] is [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130, merged;
• [2/2] is this PR, with [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 folded back in.

Jira: OMNIML-6063, which #2747 shared.

Summary by CodeRabbit

• New Features
  • Load supported compressed Mixture-of-Experts checkpoints through Transformers, with expert weights decompressed during loading and converted to bfloat16.
  • Expert-parallel loading distributes the relevant expert weights across workers while preserving integer tensor data types.
  • Checkpoint validation catches unsupported compression settings, inconsistent stored weights, and compressed tensors outside supported model layers.
  • Requires Transformers 5.8 or newer and compressed-tensors 0.16 or newer.

[shengliangxu]

/ok to test 3edfe3e

[shengliangxu]

/claude review

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 25a4cd8
• Latest commit, not reviewed: 8feea6b

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• Resume reviews
• Trigger review

Walkthrough

The change adds support for decompressing compressed-tensors MoE experts during Transformers loading. It validates configurations and checkpoint tensors, rewrites expert converters, preserves integer source dtypes during marked loads, and integrates these checks with expert-parallel loading.

Changes: Compressed Expert Decompression

Layer / File(s)	Summary
Compression configuration and checkpoint validation modelopt/torch/models/_onfly_decompress.py, tests/unit/torch/models/test_onfly_decompress.py	Config helpers locate supported compression settings. Checkpoint validation checks compressed tensor names, stored parameter sets, and dtypes.
Converter rewriting and expert decompression modelopt/torch/models/_onfly_decompress.py, tests/unit/torch/models/test_onfly_decompress.py	Installation marks compatible configs and rewrites expert converters. Converters decompress rank-owned experts before fusion. Marked loads preserve non-floating stored tensor dtypes. Tests cover converter behavior, config marking, and real checkpoint loading and saving.
Expert-parallel loader integration modelopt/torch/utils/moe.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_model_load_utils.py, tests/unit/torch/utils/test_unplaced_source_keys.py, tests/gpu/torch/utils/test_expert_parallel_load.py, noxfile.py	The loader handles marked checkpoints, removes quantization config from config levels, and rejects unsupported or unplaced compressed tensors. Tests cover loader validation, unplaced keys, and distributed loading. The unit Nox session installs compressed-tensors.

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant TransformersLoader
  participant ConverterMapping
  participant ExpertConverter
  participant Compressor
  TransformersLoader->>ConverterMapping: Select converters for marked config
  ConverterMapping->>ExpertConverter: Pass stored expert parameters
  ExpertConverter->>Compressor: Decompress rank-owned experts
  Compressor->>ExpertConverter: Return decompressed parameters
  ExpertConverter->>TransformersLoader: Emit fused bfloat16 weights

Loading

Suggested reviewers: cjluo-nv

Merge Risk: Moderate · previous review 25a4c

Some sharded compressed checkpoints can load incorrect expert weights silently. Resolve the ordering issue before merging.

Pre-merge checks: 6

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title names the main change: on-the-fly decompression of compressed-tensors MoE experts during loading, including expert-parallel loads. The changeset adds _onfly_decompress.py, rewrites per-exper…
Docstring Coverage		Docstring coverage is 85.71% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
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		No listed security anti-pattern was introduced. The changed modelopt Python files add no unsafe torch.load/numpy.load, hardcoded trust_remote_code=True, eval()/exec(), or # nosec usage. The .eval() ma…

Full details: Title check

Explanation

The title names the main change: on-the-fly decompression of compressed-tensors MoE experts during loading, including expert-parallel loads. The changeset adds _onfly_decompress.py, rewrites per-expert converters to decompress rank-owned experts to bf16, and extends expert-parallel loading in model_load_utils.py to support this path. The '[2/2]' tag and ticket ID are extra context that does not hide the main point.

Full details: Docstring Coverage

Explanation

Docstring coverage is 85.71% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 42 functions across 8 files.

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern was introduced. The changed modelopt Python files add no unsafe torch.load/numpy.load, hardcoded trust_remote_code=True, eval()/exec(), or # nosec usage. The .eval() matches are method calls. The only trust_remote_code=True match is pre-existing in the base revision. The PR does not change pyproject.toml or requirements.txt.

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

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

[claude]

[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"
        )

[shengliangxu]

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.

[claude]

[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.)

[shengliangxu]

I'm keeping the named list. I checked the alternatives on transformers 5.18 and 5.19:

• Copying everything doesn't work. WeightConverter uses __slots__ and has no __dict__, so vars(conv) raises TypeError. It isn't a dataclass there either, so dataclasses.replace is out.
• The mapping sets only two slots. Of the converter's slots, get_model_conversion_mapping sets only scope_prefix and base_model_prefix on what it returns. force_cpu comes from the constructor. The rewrite copies exactly these three.
• The other slots aren't the template's to carry:
  • The constructor derives compiled_sources and the _original_* patterns from the new patterns.
  • The load fills quantization_operation, collected_tensors, layer_targets and _was_used on its own copy of the rewritten converter.
• distributed_operation is never set. Both releases only initialise it to None. The experts are sharded at load because the first op is a MergeModulelist, and DecompressExperts is one.

If a later release starts setting another slot in the mapping, it will need adding here.

[claude]

[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.

[shengliangxu]

Fixed in bd40c48 with a refusal at install time:

• install_onfly_decompress raises NotImplementedError when the config's dtype is anything other than bfloat16 or auto. It says the experts decompress to bfloat16 and asks for a bfloat16 load.
• The expert-parallel load, the only caller, hands from_pretrained the config's dtype (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.

[claude]

[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."
        )

[shengliangxu]

Fixed in b58cfab, close to your suggestion:

• The format comes from the one config group's own format when it has one, and falls back to the top-level format.
• If compressed-tensors registers no compressor for that format, the install raises NotImplementedError and 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.

[claude]

[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.

[shengliangxu]

#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.

[claude]

[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 and DecompressExperts's docstring paragraph 2 say the same thing about subclassing MergeModulelist; one of them can go.
• The module docstring's last paragraph duplicates install_onfly_decompress's and to_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.

[shengliangxu]

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, _compressedify and to_native_backbone_config.
• The reason for subclassing MergeModulelist is now given only in _make_op_class's docstring. DecompressExperts has a one-line docstring.

I kept the comments on the subset-of-siblings pop and on deleting the mark before a save.

[claude]

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. _compressedify emits sources grouped by param name with the anchor group first, so input_dict's anchor keys arrive in gate-then-up order and out preserves it for the downstream MergeModulelist/Concatenate — gate/up fusion order is right. $-anchoring is what keeps .weight$ from claiming .weight_scale, and convert's marker handling stays consistent with the patterns _compressedify built. Sources the decompression can't consume are left in input_dict and raise, rather than reaching the merge in the wrong slot.
• Per-rank sharding. Subclassing MergeModulelist to get operations[0] recognized, then skipping None anchors, 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_native matches the real _materialize_copy(tensor, device, dtype, ...) positional order (pinned by test_a_copy_parameter_a_release_adds_passes_through), and is reached through spawn_materialize's call-time lookup. Scoping the skip to not t.is_floating_point() leaves fp8/e8m0 scales casting to bf16, which is value-exact.
• Mark lifecycle. from_pretrained deep-copies the config, so patched clearing _MARK strips it from the model's copy (keeping it out of a later save) while the caller's config stays marked for check_compressed_keys and for a re-prepare. The tests model this with explicit copy.deepcopy. The rewritten == 0 refusal 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 weight anchor, 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_keys set logic. .experts. correctly excludes shared_experts.; compressed-outside-experts, plain-expert-weight, and nothing-compressed all refuse. The dtype gate's F16/F32/F64 set 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_scale NVFP4 and fp32-weight_scale FP8 layouts, which matches the PR's stated Kimi INT4/MXFP4 scope and fails loudly with actionable text.
• Plugin laziness. compressed_tensors and transformers are both imported inside functions with the reason commented, _read_compression looks the config up before importing so a no-compression config needs neither, and moe.py still imports without transformers. No mode registration, modelopt_state schema, 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/unit budget. 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

[shengliangxu]

/ok to test 42fa4de

[shengliangxu]

Thanks for the review. Pushed 42fa4de, and the PR is rebased onto #3136's new head, which is on current main. Outcomes per suggestion:

1. Unmarked config: fixed (ad57437). It now raises a ValueError that names the problem.
2. Converter attribute whitelist: kept. The converter has __slots__, and only those three fields are set before the load (details in the thread).
3. bf16 contract: fixed (bd40c48). A config that loads in any other dtype is refused at install.
4. Group format: fixed (b58cfab). The group's own format wins, and a format with no compressor is refused.
5. _SUB_CONFIG_ATTRS readers: [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 switches them. This hunk matches [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135's.
6. compressed-tensors pin: left unpinned, like the other libraries the unit tests exercise.
7. Docstring volume: taken (42fa4de). The PR's added source lines drop from 505 to 447.

Each fix has a test that fails without it.

[shengliangxu]

/ok to test 8feea6b