From c47c3ca2c7985ac23d991764cccd6be3db3995dd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:22:14 +0000 Subject: [PATCH 01/14] Load a checkpoint its model spec decompresses through the expert-parallel loader load_and_prepare_ep refused a compressed-tensors checkpoint at its format check. A config the on-the-fly decompressor marked now passes it, and on that path the load deletes quantization_config from its copy of the config at every level _ckpt_quant_method reads, so transformers builds no compressed-tensors quantizer to compete with the decompressor (from 5.13 it claims the experts' converters itself). Before anything loads, check_compressed_keys checks the names and dtypes the checkpoint stores: the experts alone compressed, every one of them, in floats bf16 holds. Both run inside the loader's collective checks: the format check and the deletion with the config checks, the stored-tensor check right after the checkpoint is made local, so a refusal on one rank fails every rank. Moved here from the on-the-fly decompression change, whose history is archived; the decompression itself is its own change below this one. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 ++++++++- .../unit/torch/utils/test_model_load_utils.py | 58 +++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 74ab4a9486d..730aea405a6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -66,6 +66,7 @@ from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.moe import ( _EXPERTS_STYLES, + _SUB_CONFIG_ATTRS, _choose_ep_degree, _ep_config, _num_experts_from_config, @@ -1622,6 +1623,13 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) +def _decompressed_as_loaded(cfg) -> bool: + """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint.""" + from modelopt.torch.models import _onfly_decompress + + return _onfly_decompress.is_marked(cfg) + + def _ckpt_quant_method(cfg): """Return the checkpoint's native quantization method, or ``None``. @@ -1738,7 +1746,7 @@ def load_and_prepare_ep( config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) quant_method = _ckpt_quant_method(hf_config) - if quant_method not in _EP_SOURCE_FORMATS: + if quant_method not in _EP_SOURCE_FORMATS and not _decompressed_as_loaded(hf_config): # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) loads = "unquantized checkpoints" + ( @@ -1750,6 +1758,16 @@ def load_and_prepare_ep( ) # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. hf_config = copy.deepcopy(hf_config) + if _decompressed_as_loaded(hf_config): + # The decompressor hands the experts over dequantized. Left on the config (Kimi-K2.5's + # config class copies its text config's to the top level), the checkpoint's + # quantization_config would build transformers' compressed-tensors quantizer, which from + # 5.13 claims the same expert converters first. Deleted rather than set to None: before + # 5.11 transformers takes any quantization_config attribute, None too, for a quantized + # checkpoint. Every level _ckpt_quant_method reads is cleared. + for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): + if cfg is not None and hasattr(cfg, "quantization_config"): + del cfg.quantization_config ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: @@ -1843,6 +1861,13 @@ def gpu() -> tuple[bool, bool, str | None]: # refused load downloads nothing. with _refused_together("the local checkpoint", alone=False): _, ckpt_path = ensure_local_checkpoint(ckpt_path) + if _decompressed_as_loaded(hf_config): + # Only the experts decompress as they load, and the config cannot say what else the + # checkpoint compressed, or in which dtypes: what it stores is checked before anything + # loads. + from modelopt.torch.models import _onfly_decompress + + _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 55c1dc717a2..5ef7bad4106 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2274,6 +2274,64 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): assert classes == [True] +def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( + ep_world_of_two, monkeypatch, tmp_path +): + """A model spec that installs the on-the-fly decompressor (Kimi-K2.5's) marks the config; that + checkpoint passes the format check and loads without its quantization_config, at every level, + which would build transformers' own compressed-tensors quantizer to compete with the + decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own. + So is one that compresses more than its experts, by the names it stores, before it loads.""" + from safetensors.torch import save_file + + from modelopt.torch.models import _onfly_decompress + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded = {} + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} + expert, attention = "model.layers.1.mlp.experts.0.up_proj", "model.layers.1.self_attn.o_proj" + for name, compressed in (("experts", [expert]), ("more", [expert, attention])): + stored = { + f"{m}.{n}": torch.ones(1, dtype=d) for m in compressed for n, d in stored_as.items() + } + (tmp_path / name).mkdir() + save_file(stored, str(tmp_path / name / "model.safetensors")) + quantization_config = {"quant_method": "compressed-tensors"} + text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) + llm = SimpleNamespace(quantization_config=quantization_config) + config = SimpleNamespace( + quantization_config=quantization_config, text_config=text, llm_config=llm + ) + plan = DistributionPlan(2, ep_size=2) + + with pytest.raises(NotImplementedError, match="cannot dequantize"): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) + _onfly_decompress._mark(config, tuple(stored_as)) + with pytest.raises(NotImplementedError, match="1 of them outside the experts"): + load_and_prepare_ep(str(tmp_path / "more"), plan, hf_config=config) + assert not loaded + with pytest.raises(_LoadedError): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) + + for level in (loaded["config"], loaded["config"].text_config, loaded["config"].llm_config): + assert not hasattr(level, "quantization_config") + assert config.quantization_config is text.quantization_config is quantization_config + assert llm.quantization_config is quantization_config + + # ----------------------------------------------------------------------- _seed_parallel_states From 47ee07e9798cacf2b3845584cffd34f22045925a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:52:05 +0000 Subject: [PATCH 02/14] Refuse a decompressed load whose model leaves compressed tensors unplaced The export copies the checkpoint tensors the model has no parameter for as stored, and the expert-parallel load of a decompressed checkpoint drops the quantization_config that would read a compressed one. An MTP layer's compressed experts, in no layer the model builds, would reach the export packed. The load now refuses once it knows the unplaced tensors, inside its collective checks. _ckpt_quant_method reads the same config levels as the deletion, from _SUB_CONFIG_ATTRS, and the deletion pops from the instance dict. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 18 +++++-- .../unit/torch/utils/test_model_load_utils.py | 53 +++++++++++++++++++ 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 730aea405a6..0725fd2c8f4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1637,7 +1637,7 @@ def _ckpt_quant_method(cfg): or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's ``text_config`` or ``llm_config``. """ - for c in (cfg, getattr(cfg, "text_config", None), getattr(cfg, "llm_config", None)): + for c in (cfg, *(getattr(cfg, attr, None) for attr in _SUB_CONFIG_ATTRS)): qc = getattr(c, "quantization_config", None) m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) if m: @@ -1766,8 +1766,8 @@ def load_and_prepare_ep( # 5.11 transformers takes any quantization_config attribute, None too, for a quantized # checkpoint. Every level _ckpt_quant_method reads is cleared. for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): - if cfg is not None and hasattr(cfg, "quantization_config"): - del cfg.quantization_config + if cfg is not None: + vars(cfg).pop("quantization_config", None) ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: @@ -1859,6 +1859,7 @@ def gpu() -> tuple[bool, bool, str | None]: dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The load and the export both read the local snapshot, made only once the checks pass, so a # refused load downloads nothing. + compressed_names: set[str] = set() with _refused_together("the local checkpoint", alone=False): _, ckpt_path = ensure_local_checkpoint(ckpt_path) if _decompressed_as_loaded(hf_config): @@ -1868,6 +1869,7 @@ def gpu() -> tuple[bool, bool, str | None]: from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). @@ -1941,7 +1943,15 @@ def gpu() -> tuple[bool, bool, str | None]: # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, # for the export to carry. - record_unplaced_source_keys(model, ckpt_path) + unplaced = record_unplaced_source_keys(model, ckpt_path) + # The export would copy them as stored, and this load dropped the quantization_config that + # reads a compressed one, such as an MTP layer's experts. + if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: + raise NotImplementedError( + f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " + f"for, such as {carried[0]}: the export would copy them as stored, with no " + "quantization_config to read them. Decompress the checkpoint first." + ) # The model carries the ep submesh, for calibration to re-install the all-to-all forwards # over (expert_parallel_mesh); transformers sets it differently by release. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5ef7bad4106..af3c19f90bf 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2332,6 +2332,59 @@ def from_pretrained(path, **kwargs): assert llm.quantization_config is quantization_config +@pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) +def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( + ep_world_of_two, monkeypatch, tmp_path, compressed +): + """The export copies the tensors the model has no parameter for as stored, and this load dropped + the quantization_config that reads a compressed one: an MTP layer's compressed experts are + refused, its uncompressed weights carried as before.""" + from safetensors.torch import save_file + + from modelopt.torch.models import _onfly_decompress + + class _PastError(Exception): + """Raised by the load's next check.""" + + def next_check(*args, **kwargs): + raise _PastError + + model = _model_with_decoder_at("model.") + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + model.config = copy.deepcopy(config) + return model + + mtp = "model.layers.2.mlp.experts.0.up_proj" # past the model's two layers + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} + names = list(stored_as) if compressed else ["weight"] + stored = {f"{m}.{n}": torch.ones(1, dtype=d) for m in (mtp,) for n, d in stored_as.items()} + save_file(stored, str(tmp_path / "model.safetensors")) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr( + model_load_utils, "unplaced_source_keys", lambda *args: [f"{mtp}.{n}" for n in names] + ) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + config = SimpleNamespace( + num_experts=8, + base_model_ep_plan=dict(_EP_PLAN), + quantization_config={"quant_method": "compressed-tensors"}, + ) + _onfly_decompress._mark(config, tuple(stored_as)) + + with pytest.raises(NotImplementedError if compressed else _PastError) as refused: + load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) + if compressed: + assert f"2 compressed tensors the model has no parameter for, such as {mtp}" in str( + refused.value + ) + + # ----------------------------------------------------------------------- _seed_parallel_states From 4ea763feef1d7f2971cfb22c81a637d0e8d17435 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:54:10 +0000 Subject: [PATCH 03/14] Let the loader tests' stand-in unplaced-key record return a list, as the real one does Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index af3c19f90bf..b1becc75ce2 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2162,7 +2162,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) class _Config(SimpleNamespace): @@ -2216,7 +2216,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = Qwen3MoeConfig( hidden_size=16, From 7e71d5d2b2b6a95007a83c6bc7f70deb06b1e0e6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:24:42 +0000 Subject: [PATCH 04/14] Pin that a real decompressed load places its experts' stored tensors The expert-parallel load refuses compressed tensors the model leaves unplaced, and relies on unplaced_source_keys placing the stored weight_packed / weight_scale / weight_shape of every expert the model builds. A CPU from_pretrained of a tiny compressed Qwen3-MoE, through the real decompressor, pins that: they resolve through the converters the load recorded (the stock mapping places none of them), and only the experts of a layer past num_hidden_layers stay unplaced. The two local modelopt.torch.models imports say why they are local, and the loader test's docstring says what its stand-in covers. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 3 +- .../unit/torch/utils/test_model_load_utils.py | 6 +- .../torch/utils/test_unplaced_source_keys.py | 77 +++++++++++++++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0725fd2c8f4..21a3a73d38c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1625,6 +1625,7 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st def _decompressed_as_loaded(cfg) -> bool: """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint.""" + # Importing modelopt.torch.models registers every model spec; only expert-parallel loads need them. from modelopt.torch.models import _onfly_decompress return _onfly_decompress.is_marked(cfg) @@ -1865,7 +1866,7 @@ def gpu() -> tuple[bool, bool, str | None]: if _decompressed_as_loaded(hf_config): # Only the experts decompress as they load, and the config cannot say what else the # checkpoint compressed, or in which dtypes: what it stores is checked before anything - # loads. + # loads. Imported here for the reason _decompressed_as_loaded gives. from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b1becc75ce2..9eb5f2d8289 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2337,8 +2337,10 @@ def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place ep_world_of_two, monkeypatch, tmp_path, compressed ): """The export copies the tensors the model has no parameter for as stored, and this load dropped - the quantization_config that reads a compressed one: an MTP layer's compressed experts are - refused, its uncompressed weights carried as before.""" + the quantization_config that reads a compressed one: an unplaced tensor under a stored name the + mark holds (an MTP layer's compressed experts) is refused, one under the plain ``weight`` name is + not. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in + test_unplaced_source_keys.py.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 46a5e7ef634..16a9368f23f 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -391,3 +391,80 @@ def test_mxfp4_expert_blocks_folded_in_at_load_are_placed(tmp_path): model = _load(tmp_path, quantization_config=Mxfp4Config(dequantize=True), dtype=torch.bfloat16) assert model._modelopt_unplaced_source_keys == [STRAY] + + +def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatch): + """The on-the-fly decompressor feeds each expert's stored ``weight_packed``/``weight_scale``/ + ``weight_shape`` into the fused weights. A built layer's are placed, through the converters + ``from_pretrained`` recorded (the stock mapping places none of them), so the expert-parallel load + passes them; only a layer the model does not build (an MTP layer) keeps its own, which that load + refuses to carry.""" + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) + import transformers.core_model_loading as cml + import transformers.modeling_utils as mu + from packaging.version import Version + from transformers import __version__ as transformers_version + + from modelopt.torch.models import _onfly_decompress + + if Version(transformers_version) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers >= 5.8") + # The decompressor's hooks are process-wide: undone after the test. + for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + monkeypatch.setattr(owner, name, getattr(owner, name)) + hidden, inter, experts = 64, 32, 2 + config = AutoConfig.for_model( + "qwen3_moe", + vocab_size=64, + hidden_size=hidden, + intermediate_size=64, + moe_intermediate_size=inter, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_experts=experts, + num_experts_per_tok=1, + tie_word_embeddings=False, + architectures=["Qwen3MoeForCausalLM"], + ) + group = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} + config.quantization_config = { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "config_groups": {"experts": {"targets": ["re:.*experts.*"], "weights": group}}, + } + compressor, scheme, stored_as, _ = _onfly_decompress._read_compression(config) + shapes = { + "gate_proj": (inter, hidden), + "up_proj": (inter, hidden), + "down_proj": (hidden, inter), + } + state, mtp = {}, [] + for layer in (0, 1): # layer 1 is past num_hidden_layers: an MTP layer, say + for e in range(experts): + for proj, (rows, cols) in shapes.items(): + weight = { + "weight": torch.ones(rows, cols), + "weight_scale": torch.ones(rows, cols // 32), + } + stem = f"model.layers.{layer}.mlp.experts.{e}.{proj}." + for name, tensor in compressor.compress(weight, scheme).items(): + state[stem + name] = tensor + if layer: + mtp.append(stem + name) + for name, tensor in AutoModelForCausalLM.from_config(config).state_dict().items(): + if ".experts." not in name: + state[name] = tensor.to(torch.bfloat16) + _save(tmp_path, config, state, config.quantization_config) + _onfly_decompress.install_onfly_decompress(config) + del config.quantization_config # as the expert-parallel load does once the hooks are in + + model = AutoModelForCausalLM.from_pretrained(str(tmp_path), config=config, dtype=torch.bfloat16) + if getattr(model, "_weight_conversions", None) is None: + pytest.skip("this transformers does not record the conversions from_pretrained ran") + + assert sorted(stored_as) == ["weight_packed", "weight_scale", "weight_shape"] + assert model_load_utils.unplaced_source_keys(model, str(tmp_path)) == sorted(mtp) From 56aec937003e73de97758fca4270da86ac91d5f5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:42:15 +0000 Subject: [PATCH 05/14] Load compressed experts expert-parallel on 2 GPUs, each rank decompressing its own A tiny Qwen3-MoE whose experts are compressed-tensors INT4 (Kimi-K2.5's format) or MXFP4 (Kimi-K3's), installed for on-the-fly decompression and loaded through load_and_prepare_ep: each rank must hold exactly the dense model's weights for its own experts, every stored tensor must be placed, and the sharded model's logits must match the dense model loaded whole. The MXFP4 case goes through compressed-tensors, which the decompressor reads, not transformers' own MXFP4 dequantization. Each rank's checks are gathered before any rank fails, under bounded collectives. Signed-off-by: Shengliang Xu --- .../torch/utils/test_expert_parallel_load.py | 131 +++++++++++++++++- 1 file changed, 130 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 849305f4f6a..1bdcbb4c957 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -17,20 +17,30 @@ ``load_and_prepare_ep`` shards an MoE's routed experts over the ranks, and the ``torch_a2a`` experts forward sends each token to the rank that holds its expert. Every rank runs its own batch through the -sharded model, and the result must match the whole model run on that batch alone. +sharded model, and the result must match the whole model run on that batch alone. A checkpoint whose +experts are compressed-tensors compressed loads the same way, each rank decompressing its own. """ +import json +from types import SimpleNamespace + import pytest import torch +import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version +from safetensors.torch import load_file, save_file transformers = pytest.importorskip("transformers") if Version(transformers.__version__) < Version("5.7"): pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7", allow_module_level=True) +import transformers.core_model_loading as cml +import transformers.modeling_utils as mu from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +from modelopt.torch.models import _onfly_decompress from modelopt.torch.utils.distributed import ( DistributionPlan, distribution_plan, @@ -101,6 +111,125 @@ def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree): torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2) +# The compressions the published Kimi checkpoints use, INT4 (Kimi-K2.5) and MXFP4 (Kimi-K3), through +# compressed-tensors (quant_method "compressed-tensors"), which the on-the-fly decompressor reads; not +# transformers' own MXFP4 (quant_method "mxfp4"). The experts draw values each format holds exactly, +# at a power-of-two scale bf16 holds too, so they decompress to the dense model's weights exactly. +_INT4 = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} +_E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6]) +_COMPRESSIONS = { + "int4": ("pack-quantized", _INT4, torch.arange(-8, 8) / 64, 1 / 64), + "mxfp4": ( + "mxfp4-pack-quantized", + {**_INT4, "type": "float", "scale_dtype": "torch.uint8"}, + torch.cat([_E2M1, -_E2M1[1:]]) / 16, + 1 / 16, + ), +} + + +def _save_compressed_tiny_qwen3_moe(dense_dir, packed_dir, compression): + """The tiny Qwen3-MoE twice: whole in bf16, and with its experts compressed by ``compression``.""" + fmt, weights, values, scale = _COMPRESSIONS[compression] + _save_tiny_qwen3_moe(dense_dir) + state = load_file(str(dense_dir / "model.safetensors")) + gen = torch.Generator().manual_seed(0) + for key in [k for k in state if ".experts." in k]: # ...mlp.experts..gate_proj.weight + draw = torch.randint(0, len(values), state[key].shape, generator=gen) + state[key] = values[draw].to(torch.bfloat16) + save_file(state, str(dense_dir / "model.safetensors"), metadata={"format": "pt"}) + + quantization_config = { + "quant_method": "compressed-tensors", + "format": fmt, + "config_groups": {"experts": {"targets": ["Linear"], "weights": weights, "format": fmt}}, + } + compressor, scheme, _, _ = _onfly_decompress._read_compression( + SimpleNamespace(quantization_config=quantization_config) + ) + packed = {} + for key, tensor in state.items(): + if ".experts." not in key: + packed[key] = tensor + continue + rows, cols = tensor.shape + scales = torch.full((rows, cols // 32), scale) + stored = compressor.compress({"weight": tensor.float(), "weight_scale": scales}, scheme) + for name, value in stored.items(): # floats as bf16, as the Kimi checkpoints store them + packed[key.removesuffix("weight") + name] = ( + value.bfloat16() if value.is_floating_point() else value + ) + packed_dir.mkdir() + save_file(packed, str(packed_dir / "model.safetensors"), metadata={"format": "pt"}) + config = json.loads((dense_dir / "config.json").read_text()) + config["quantization_config"] = quantization_config + (packed_dir / "config.json").write_text(json.dumps(config)) + + +def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_dir): + device = torch.device(f"cuda:{rank}") + problems = [] + # The decompressor's hooks are process-wide, and the workers outlive the job. + hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + try: + with bounded_collectives(seconds=120): + config = AutoConfig.from_pretrained(packed_dir) + _onfly_decompress.install_onfly_decompress(config) # as a model spec's prepare does + model = load_and_prepare_ep( + packed_dir, DistributionPlan(size, ep_size=size), hf_config=config + ).eval() + whole = AutoModelForCausalLM.from_pretrained(dense_dir, dtype=torch.bfloat16) + whole = whole.to(device).eval() + + # This rank's experts, decompressed as they loaded, are the dense model's, exactly. + per_rank = NUM_EXPERTS // size + owned = slice(rank * per_rank, (rank + 1) * per_rank) + for i, (layer, whole_layer) in enumerate(zip(model.model.layers, whole.model.layers)): + for name in ("gate_up_proj", "down_proj"): + local = _local(getattr(layer.mlp.experts, name)) + if local.shape[0] != per_rank: + problems.append(f"layer {i} {name}: {local.shape[0]} experts held") + elif not torch.equal(local, getattr(whole_layer.mlp.experts, name)[owned]): + problems.append(f"layer {i} {name}: rank {rank}'s experts differ") + # Every stored tensor of the experts it built went into them. + if model._modelopt_unplaced_source_keys: + problems.append(f"unplaced: {model._modelopt_unplaced_source_keys[:3]}") + set_distribution_plan(model, DistributionPlan(size, ep_size=size)) + + generator = torch.Generator().manual_seed(rank) + input_ids = torch.randint(0, 128, (2, 16), generator=generator).to(device) + with torch.no_grad(): + sharded = model(input_ids).logits.float() + reference = whole(input_ids).logits.float() + if not torch.allclose(sharded, reference, atol=2e-2, rtol=2e-2): + problems.append(f"logits differ by {(sharded - reference).abs().max().item():.3g}") + + gathered = [None] * size + dist.all_gather_object(gathered, problems) + finally: + mu.get_model_conversion_mapping, cml._materialize_copy = hooks + found = [problem for rank_problems in gathered for problem in rank_problems] + assert not found, "\n".join(found[:20]) + + +@pytest.mark.parametrize("compression", sorted(_COMPRESSIONS)) +def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own( + need_2_gpus, dist_workers, tmp_path, compression +): + """A checkpoint whose experts are compressed (INT4, MXFP4) loads expert-parallel through the + on-the-fly decompressor: each rank holds exactly the dense model's weights for its own experts, + every stored tensor is placed, and the sharded model matches the dense one loaded whole. Each + rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" + pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + if Version(transformers.__version__) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers 5.8") + if NUM_EXPERTS % dist_workers.world_size: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + dense_dir, packed_dir = tmp_path / "dense", tmp_path / "packed" + _save_compressed_tiny_qwen3_moe(dense_dir, packed_dir, compression) + dist_workers.run(_compressed_experts_load_on_their_own_ranks, str(dense_dir), str(packed_dir)) + + @pytest.mark.parametrize("ep_share", [1, 2], ids=["ep-world", "ep-half"]) def test_expert_parallel_model_matches_the_whole_model( need_2_gpus, dist_workers, tmp_path, ep_share From ce1671f49199512461d735297098b98dd401c710 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:54:03 +0000 Subject: [PATCH 06/14] Restore the decompressor's read hook where the tests undo its hooks The decompressor's keep-integers patch moved from _materialize_copy to spawn_materialize (#3137), so the 2-GPU job and the real-load unplaced test restore that hook. The GPU case's skip reason names the package it needs: compressed-tensors is not in CI's GPU image, where it skips. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 9 ++++++--- tests/unit/torch/utils/test_unplaced_source_keys.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 1bdcbb4c957..92ac059cb10 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -170,7 +170,7 @@ def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_di device = torch.device(f"cuda:{rank}") problems = [] # The decompressor's hooks are process-wide, and the workers outlive the job. - hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) try: with bounded_collectives(seconds=120): config = AutoConfig.from_pretrained(packed_dir) @@ -207,7 +207,7 @@ def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_di gathered = [None] * size dist.all_gather_object(gathered, problems) finally: - mu.get_model_conversion_mapping, cml._materialize_copy = hooks + mu.get_model_conversion_mapping, cml.spawn_materialize = hooks found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) @@ -220,7 +220,10 @@ def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own on-the-fly decompressor: each rank holds exactly the dense model's weights for its own experts, every stored tensor is placed, and the sharded model matches the dense one loaded whole. Each rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" - pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + pytest.importorskip( + "compressed_tensors", + reason="compressed-tensors is needed: the checkpoint's experts are compressed", + ) if Version(transformers.__version__) < Version("5.8"): pytest.skip("the on-the-fly decompression needs transformers 5.8") if NUM_EXPERTS % dist_workers.world_size: diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 16a9368f23f..4fd5da33406 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -412,7 +412,7 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc if Version(transformers_version) < Version("5.8"): pytest.skip("the on-the-fly decompression needs transformers >= 5.8") # The decompressor's hooks are process-wide: undone after the test. - for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) hidden, inter, experts = 64, 32, 2 config = AutoConfig.for_model( From ff20690977ba6bc3109dc4bc4c5de8a42b30f44c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:22:45 +0000 Subject: [PATCH 07/14] Leave no decompressing load flag behind the real-load unplaced test The real decompressing load sets the decompressor's load flag, which outlives the load; the test now runs it on a flag of its own. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_unplaced_source_keys.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 4fd5da33406..f3b89452821 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -414,6 +414,7 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc # The decompressor's hooks are process-wide: undone after the test. for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) + monkeypatch.setattr(_onfly_decompress, "_LOAD", {"decompressing": False}) hidden, inter, experts = 64, 32, 2 config = AutoConfig.for_model( "qwen3_moe", From bcf5a3783f4dc8133072a623b34dd43df9f932e4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:32:40 +0000 Subject: [PATCH 08/14] Refuse a config whose class keeps quantization_config, and say why the matchers hold From Claude's review: the load clears quantization_config from each config level's instance dict, which leaves a class attribute or property in place, and transformers' own quantizer would then compete with the decompressor. The load now refuses such a config inside its collective checks. Comments say why "weight" is left out of the compressed names (a float format's anchor, and a plain weight the export has always carried) and why a suffix tells a carried compressed tensor (check_compressed_keys refused any outside the experts). Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 12 +++++++++++- tests/unit/torch/utils/test_model_load_utils.py | 12 +++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 21a3a73d38c..851f42ee9c1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1769,6 +1769,13 @@ def load_and_prepare_ep( for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): if cfg is not None: vars(cfg).pop("quantization_config", None) + if hasattr(cfg, "quantization_config"): + # A class attribute or property, which the instance dict cannot clear. + raise NotImplementedError( + f"{type(cfg).__name__} keeps quantization_config off its instance, so " + "this load cannot clear it, and transformers' own quantizer would " + "compete with the decompressor. Decompress the checkpoint first." + ) ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: @@ -1870,6 +1877,8 @@ def gpu() -> tuple[bool, bool, str | None]: from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight + # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over @@ -1946,7 +1955,8 @@ def gpu() -> tuple[bool, bool, str | None]: # for the export to carry. unplaced = record_unplaced_source_keys(model, ckpt_path) # The export would copy them as stored, and this load dropped the quantization_config that - # reads a compressed one, such as an MTP layer's experts. + # reads a compressed one, such as an MTP layer's experts. A suffix is enough to tell them: + # check_compressed_keys refused any compressed tensor outside the experts. if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: raise NotImplementedError( f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 9eb5f2d8289..ed88abbc518 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2281,7 +2281,8 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( checkpoint passes the format check and loads without its quantization_config, at every level, which would build transformers' own compressed-tensors quantizer to compete with the decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own. - So is one that compresses more than its experts, by the names it stores, before it loads.""" + So is one that compresses more than its experts, by the names it stores, before it loads, and a + config whose class keeps quantization_config, which the load cannot clear.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress @@ -2331,6 +2332,15 @@ def from_pretrained(path, **kwargs): assert config.quantization_config is text.quantization_config is quantization_config assert llm.quantization_config is quantization_config + # A config class that keeps quantization_config as a class attribute cannot have it cleared. + typed_config = type("_Typed", (SimpleNamespace,), {"quantization_config": quantization_config}) + typed = typed_config(text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN)) + _onfly_decompress._mark(typed, tuple(stored_as)) + loaded.clear() + with pytest.raises(NotImplementedError, match="keeps quantization_config off its instance"): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=typed) + assert not loaded + @pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( From f5fcbe800dedb5b41fc816fbf0176efbcdc857b8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:46:25 +0000 Subject: [PATCH 09/14] Walk a config's levels through one helper, and check them once all are cleared From Claude's review: the top config and its nested language-model configs were walked three times, in _ckpt_quant_method, in the load's quantization_config deletion and in the decompressor. moe.config_levels, beside _SUB_CONFIG_ATTRS, now gives those levels to all three. (#3135 adds the same _SUB_CONFIG_ATTRS; the helper beside it conflicts there once, at the Kimi base's rebuild.) The deletion now clears every level before it checks any: a config class that forwards its text config's quantization_config through a property holds none once the text config's is gone, and is no longer refused; a class attribute, or a property of its own, still is. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 18 +++++-------- modelopt/torch/utils/moe.py | 6 +++++ .../torch/utils/plugins/model_load_utils.py | 25 ++++++++++--------- .../unit/torch/utils/test_model_load_utils.py | 15 ++++++++++- 4 files changed, 39 insertions(+), 25 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 1200032a42e..459b8ff6058 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -26,7 +26,7 @@ from safetensors import safe_open from modelopt.torch.utils import print_rank_0 -from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS, require_transformers +from modelopt.torch.utils.moe import config_levels, require_transformers def to_native_backbone_config( @@ -60,19 +60,13 @@ def to_native_backbone_config( return native -def _config_levels(hf_config) -> list: - """``hf_config`` and the nested language-model configs the expert-parallel load reads.""" - nested = (getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS) - return [c for c in (hf_config, *nested) if c is not None] - - def _find_quantization_config(hf_config): """Locate the compressed-tensors ``quantization_config`` (top-level, else on the nested language-model config) and return it as a plain dict, or ``None`` for a config without one.""" qc = next( ( found - for c in _config_levels(hf_config) + for c in config_levels(hf_config) if (found := getattr(c, "quantization_config", None)) is not None ), None, @@ -300,12 +294,12 @@ def _mark(hf_config, stored_as): def is_marked(hf_config) -> bool: """Whether :func:`install_onfly_decompress` marked ``hf_config``, or its nested language-model config, for its hooks.""" - return any(getattr(c, _MARK, None) is not None for c in _config_levels(hf_config)) + return any(getattr(c, _MARK, None) is not None for c in config_levels(hf_config)) def stored_names(hf_config) -> tuple[str, ...]: """The names the installed compressed-tensors stores each weight under for ``hf_config``'s scheme.""" - for c in _config_levels(hf_config): + for c in config_levels(hf_config): if (stored_as := getattr(c, _MARK, None)) is not None: return stored_as raise ValueError("The config is not marked for on-the-fly decompression: install it first.") @@ -323,7 +317,7 @@ def install_onfly_decompress(hf_config): raise TypeError("install_onfly_decompress: pass the model's config object, not a dict") # The loader sets each parameter from what the conversion hands it, uncast. Every level counts: # a backbone config derived from the text config (to_native_backbone_config) keeps its dtype. - for level in _config_levels(hf_config): + for level in config_levels(hf_config): dtype = getattr(level, "dtype", None) if dtype is not None and str(dtype).removeprefix("torch.") != "bfloat16": raise NotImplementedError( @@ -393,7 +387,7 @@ def patched(*args, **kwargs): ) # from_pretrained loads from a copy of the marked config; the mark must not reach the config # a later save writes, at any level. Cleared once the experts' converters are rewritten. - for c in _config_levels(config): + for c in config_levels(config): if hasattr(c, _MARK): delattr(c, _MARK) return out diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 0bc4d9d313c..ce2d92b79d3 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -60,6 +60,12 @@ def require_transformers(min_version: str, feature: str) -> None: _SUB_CONFIG_ATTRS = ("text_config", "llm_config") +def config_levels(hf_config) -> list: + """``hf_config`` and its nested language-model configs, those it has, which may carry its settings.""" + nested = (getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS) + return [c for c in (hf_config, *nested) if c is not None] + + def _ep_config(hf_config): """Return the config that carries the EP plan / expert count. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 851f42ee9c1..8681e70866e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -66,12 +66,12 @@ from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.moe import ( _EXPERTS_STYLES, - _SUB_CONFIG_ATTRS, _choose_ep_degree, _ep_config, _num_experts_from_config, _routed_expert_keys, _routed_expert_modules, + config_levels, register_torch_a2a, ) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -1638,7 +1638,7 @@ def _ckpt_quant_method(cfg): or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's ``text_config`` or ``llm_config``. """ - for c in (cfg, *(getattr(cfg, attr, None) for attr in _SUB_CONFIG_ATTRS)): + for c in config_levels(cfg): qc = getattr(c, "quantization_config", None) m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) if m: @@ -1766,16 +1766,17 @@ def load_and_prepare_ep( # 5.13 claims the same expert converters first. Deleted rather than set to None: before # 5.11 transformers takes any quantization_config attribute, None too, for a quantized # checkpoint. Every level _ckpt_quant_method reads is cleared. - for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): - if cfg is not None: - vars(cfg).pop("quantization_config", None) - if hasattr(cfg, "quantization_config"): - # A class attribute or property, which the instance dict cannot clear. - raise NotImplementedError( - f"{type(cfg).__name__} keeps quantization_config off its instance, so " - "this load cannot clear it, and transformers' own quantizer would " - "compete with the decompressor. Decompress the checkpoint first." - ) + levels = config_levels(hf_config) + for cfg in levels: + vars(cfg).pop("quantization_config", None) + # Checked once every level is cleared: a property forwarding to a nested level's is gone + # then, while a class attribute, or a property of its own, is not. + if kept := [cfg for cfg in levels if hasattr(cfg, "quantization_config")]: + raise NotImplementedError( + f"{type(kept[0]).__name__} keeps quantization_config off its instance, so this " + "load cannot clear it, and transformers' own quantizer would compete with the " + "decompressor. Decompress the checkpoint first." + ) ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ed88abbc518..f743c74aad2 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2282,7 +2282,8 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( which would build transformers' own compressed-tensors quantizer to compete with the decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own. So is one that compresses more than its experts, by the names it stores, before it loads, and a - config whose class keeps quantization_config, which the load cannot clear.""" + config whose class keeps quantization_config, which the load cannot clear; one whose class only + forwards its text config's loads.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress @@ -2341,6 +2342,18 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=typed) assert not loaded + # One whose class forwards the text config's, as Kimi-K2.5's copies it, loads once both are gone. + class _Forwarding(SimpleNamespace): + @property + def quantization_config(self): + return self.text_config.quantization_config + + text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) + forwarding = _Forwarding(text_config=text) + _onfly_decompress._mark(forwarding, tuple(stored_as)) + with pytest.raises(_LoadedError): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) + @pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( From 398d4655b649e9a5a984a741c2df090e96f9efba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:10:13 +0000 Subject: [PATCH 10/14] Build the unplaced-refusal loader test's model from the shared test model The EP loader tests now build their stand-in MoE from the shared model_with_decoder_at, and the file's own copy is gone. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index f743c74aad2..00417b0c702 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2374,7 +2374,7 @@ class _PastError(Exception): def next_check(*args, **kwargs): raise _PastError - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") class _Loader: @staticmethod From 0d65dc3aab82df12cd65245b5c77c4550c8f29d9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:13:02 +0000 Subject: [PATCH 11/14] Refuse an MTP layer's compressed experts before the load, not after it From Claude's review: the load refused the compressed tensors its model leaves unplaced only once from_pretrained had read the whole checkpoint. An MTP layer, the usual case, is stored past the decoder layers the EP config says the model builds, and the checkpoint's tensor names show that before anything loads. Its compressed experts are now refused there; the check after the load stays for any other shape. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 19 +++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 23 +++++++++++++++---- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8681e70866e..0f6cf400d82 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1881,6 +1881,25 @@ def gpu() -> tuple[bool, bool, str | None]: # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} + # An MTP layer is stored past the decoder layers the model builds, so its compressed + # experts would be refused only after the whole load, by the check that follows it. + # Refused here instead, before anything loads; that check stays for any other shape. + built = getattr(ep_cfg, "num_hidden_layers", None) + past = sorted( + name + for name in checkpoint_tensor_files(ckpt_path) + if name.rpartition(".")[2] in compressed_names + and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) + and built is not None + and int(layer[1]) >= built + ) + if past: + raise NotImplementedError( + f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " + f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " + "for them, and the export would copy them as stored, with no quantization_config " + "to read them. Decompress the checkpoint first." + ) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 00417b0c702..5d5c73453b9 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2355,14 +2355,19 @@ def quantization_config(self): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) -@pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) +@pytest.mark.parametrize( + ("compressed", "layers"), + [(False, 3), (True, 3), (True, 2)], + ids=["uncompressed", "compressed", "compressed-past-the-layers"], +) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( - ep_world_of_two, monkeypatch, tmp_path, compressed + ep_world_of_two, monkeypatch, tmp_path, compressed, layers ): """The export copies the tensors the model has no parameter for as stored, and this load dropped the quantization_config that reads a compressed one: an unplaced tensor under a stored name the mark holds (an MTP layer's compressed experts) is refused, one under the plain ``weight`` name is - not. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in + not. Compressed experts stored past the config's decoder layers are refused before anything + loads. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in test_unplaced_source_keys.py.""" from safetensors.torch import save_file @@ -2375,14 +2380,16 @@ def next_check(*args, **kwargs): raise _PastError model = model_with_decoder_at("model.") + loads = [] class _Loader: @staticmethod def from_pretrained(path, config, **kwargs): + loads.append(path) model.config = copy.deepcopy(config) return model - mtp = "model.layers.2.mlp.experts.0.up_proj" # past the model's two layers + mtp = "model.layers.2.mlp.experts.0.up_proj" # unplaced, as an MTP layer's are stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} names = list(stored_as) if compressed else ["weight"] stored = {f"{m}.{n}": torch.ones(1, dtype=d) for m in (mtp,) for n, d in stored_as.items()} @@ -2397,6 +2404,7 @@ def from_pretrained(path, config, **kwargs): monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( num_experts=8, + num_hidden_layers=layers, base_model_ep_plan=dict(_EP_PLAN), quantization_config={"quant_method": "compressed-tensors"}, ) @@ -2404,7 +2412,12 @@ def from_pretrained(path, config, **kwargs): with pytest.raises(NotImplementedError if compressed else _PastError) as refused: load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) - if compressed: + if layers == 2: + assert not loads + assert f"2 compressed tensors past the model's 2 decoder layers, such as {mtp}." in str( + refused.value + ) + elif compressed: assert f"2 compressed tensors the model has no parameter for, such as {mtp}" in str( refused.value ) From b0e0d6dc1a959dc284253e01ed52a39182bc86eb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:36:01 +0000 Subject: [PATCH 12/14] Refuse an MTP layer's experts from the names the compressed-key check read From Claude's review: the early refusal read the checkpoint's tensor names and every shard header a second time, right after check_compressed_keys had, and scanned them even for a config with no num_hidden_layers. It now takes the map check_compressed_keys returns, and skips the scan without a layer count. Its message says what it found, an MTP layer's typically, rather than assert the cause. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0f6cf400d82..909b917236b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1877,28 +1877,33 @@ def gpu() -> tuple[bool, bool, str | None]: # loads. Imported here for the reason _decompressed_as_loaded gives. from modelopt.torch.models import _onfly_decompress - _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + tensor_files = _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # An MTP layer is stored past the decoder layers the model builds, so its compressed # experts would be refused only after the whole load, by the check that follows it. - # Refused here instead, before anything loads; that check stays for any other shape. + # Refused here instead, before anything loads, from the names just read; that check + # stays for any other shape. transformers 5.18 builds no decoder layer past + # num_hidden_layers, so this refuses nothing that check would let through there. built = getattr(ep_cfg, "num_hidden_layers", None) - past = sorted( - name - for name in checkpoint_tensor_files(ckpt_path) - if name.rpartition(".")[2] in compressed_names - and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) - and built is not None - and int(layer[1]) >= built + past = ( + [] + if built is None + else sorted( + name + for name in tensor_files + if name.rpartition(".")[2] in compressed_names + and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) + and int(layer[1]) >= built + ) ) if past: raise NotImplementedError( f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " - f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " - "for them, and the export would copy them as stored, with no quantization_config " - "to read them. Decompress the checkpoint first." + f"decoder layers, such as {past[0]} (an MTP layer's, typically): the model builds " + "no parameter for them, and the export would copy them as stored, with no " + "quantization_config to read them. Decompress the checkpoint first." ) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over From c99f3dc7bdbf0a5d8697648e55ec6320cab5aa38 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:24:59 +0000 Subject: [PATCH 13/14] Give the decompressing loader tests' stand-in configs a model type no spec names The expert-parallel load is to read the config's model type for its spec checks, so each config standing in for a checkpoint's names one, as the moved loader tests' do. save_file is already imported at the top of the file the tests moved into. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_model_load_utils.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5d5c73453b9..684cc6a0f1d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1891,6 +1891,9 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( # ------------------------------------------------------------------------------- the plan +# The model type of the configs that stand in for a checkpoint's here: no model spec names it. +_NO_SPEC = "toy_moe" + _EP_PLAN = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", @@ -2284,8 +2287,6 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( So is one that compresses more than its experts, by the names it stores, before it loads, and a config whose class keeps quantization_config, which the load cannot clear; one whose class only forwards its text config's loads.""" - from safetensors.torch import save_file - from modelopt.torch.models import _onfly_decompress class _LoadedError(Exception): @@ -2315,7 +2316,10 @@ def from_pretrained(path, **kwargs): text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) llm = SimpleNamespace(quantization_config=quantization_config) config = SimpleNamespace( - quantization_config=quantization_config, text_config=text, llm_config=llm + model_type=_NO_SPEC, + quantization_config=quantization_config, + text_config=text, + llm_config=llm, ) plan = DistributionPlan(2, ep_size=2) @@ -2335,7 +2339,9 @@ def from_pretrained(path, **kwargs): # A config class that keeps quantization_config as a class attribute cannot have it cleared. typed_config = type("_Typed", (SimpleNamespace,), {"quantization_config": quantization_config}) - typed = typed_config(text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN)) + typed = typed_config( + model_type=_NO_SPEC, text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN) + ) _onfly_decompress._mark(typed, tuple(stored_as)) loaded.clear() with pytest.raises(NotImplementedError, match="keeps quantization_config off its instance"): @@ -2349,7 +2355,7 @@ def quantization_config(self): return self.text_config.quantization_config text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) - forwarding = _Forwarding(text_config=text) + forwarding = _Forwarding(model_type=_NO_SPEC, text_config=text) _onfly_decompress._mark(forwarding, tuple(stored_as)) with pytest.raises(_LoadedError): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) @@ -2369,8 +2375,6 @@ def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place not. Compressed experts stored past the config's decoder layers are refused before anything loads. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in test_unplaced_source_keys.py.""" - from safetensors.torch import save_file - from modelopt.torch.models import _onfly_decompress class _PastError(Exception): @@ -2403,6 +2407,7 @@ def from_pretrained(path, config, **kwargs): ) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( + model_type=_NO_SPEC, num_experts=8, num_hidden_layers=layers, base_model_ep_plan=dict(_EP_PLAN), From 92830fcba14d523cdd44545c2ae3714b67ce40e2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:49:59 +0000 Subject: [PATCH 14/14] Skip the compressed-expert load tests below compressed-tensors 0.16 Signed-off-by: Shengliang Xu #3137 now refuses an older compressed-tensors, whose compressors take no scheme: the tests that load through the real library skip there, as the decompressor's own do, rather than fail on the refusal. --- tests/gpu/torch/utils/test_expert_parallel_load.py | 1 + tests/unit/torch/utils/test_unplaced_source_keys.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 92ac059cb10..533147ac1b0 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -222,6 +222,7 @@ def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" pytest.importorskip( "compressed_tensors", + minversion="0.16", reason="compressed-tensors is needed: the checkpoint's experts are compressed", ) if Version(transformers.__version__) < Version("5.8"): diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index f3b89452821..13fa3600e09 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -400,7 +400,9 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc passes them; only a layer the model does not build (an MTP layer) keeps its own, which that load refuses to carry.""" pytest.importorskip( - "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + "compressed_tensors", + minversion="0.16", + reason="needs compressed-tensors (CI's unit session installs it)", ) import transformers.core_model_loading as cml import transformers.modeling_utils as mu