[OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader

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

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.
 

modelopt/torch/utils/plugins/model_load_utils.py

@@ -71,6 +71,7 @@
     _num_experts_from_config,
     _routed_expert_keys,
     _routed_expert_modules,
+    config_levels,
     register_torch_a2a,
 )
 from modelopt.torch.utils.plugins.hf_checkpoint_utils import (
@@ -1624,14 +1625,22 @@ 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."""
+    # 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)
+
+
 def _ckpt_quant_method(cfg):
     """Return the checkpoint's native quantization method, or ``None``.
 
     Return the checkpoint's native quantization method (e.g. ``"fp8"`` for block-fp8 DeepSeek-V3),
     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 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:
@@ -1740,7 +1749,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" + (
@@ -1752,6 +1761,24 @@ 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.
+            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:
@@ -1843,8 +1870,43 @@ 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):
+            # 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. Imported here for the reason _decompressed_as_loaded gives.
+            from modelopt.torch.models import _onfly_decompress
+
+            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, 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 = (
+                []
+                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'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
     # the whole world (``dp``).
@@ -1918,7 +1980,16 @@ 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. A suffix is enough to tell them:
+        # check_compressed_keys refused any compressed tensor outside the experts.
+        if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]:
+            raise NotImplementedError(
+                f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter "
+                f"for, such as {carried[0]}: the export would copy them as stored, with no "
+                "quantization_config to read them. Decompress the checkpoint first."
+            )
 
         # 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.

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,129 @@ 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.<e>.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.spawn_materialize)
+    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.spawn_materialize = 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",
+        minversion="0.16",
+        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:
+        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