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 269134bb93f..943525b9e12 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -64,6 +64,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 74ab4a9486d..909b917236b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/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 ( @@ -1622,6 +1623,14 @@ 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``. @@ -1629,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, "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: @@ -1738,7 +1747,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 +1759,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: @@ -1841,8 +1868,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``). @@ -1916,7 +1978,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. diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 849305f4f6a..533147ac1b0 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,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..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 diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 55c1dc717a2..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", @@ -2162,7 +2165,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 +2219,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, @@ -2274,6 +2277,157 @@ 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, and a + config whose class keeps quantization_config, which the load cannot clear; one whose class only + forwards its text config's loads.""" + 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( + model_type=_NO_SPEC, + 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 + + # 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( + 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"): + 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(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) + + +@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, 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. 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 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.") + 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" # 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()} + 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( + model_type=_NO_SPEC, + num_experts=8, + num_hidden_layers=layers, + 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 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 + ) + + # ----------------------------------------------------------------------- _seed_parallel_states diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 46a5e7ef634..13fa3600e09 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -391,3 +391,83 @@ 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", + 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 + 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, "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", + 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)