[Model] Support llm-compressor Inkling NVFP4 weights (#49258)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 23, 2026 at 23:06 UTC
d65acd83d87b1ad8d13eaa9452510488de642543
3 files changed
+83
-10
tests/models/inkling/test_moe_weight_layout.py
+55
@@ -112,6 +112,30 @@ def test_inkling_mapper_maps_modelopt_exclusions() -> None:
112
assert not quant_config.is_layer_excluded("model.layers.3.mlp.experts")
113
114
115
+@pytest.mark.parametrize("projection", ["w13", "w2"])
116
+@pytest.mark.parametrize("nested", [False, True])
117
+@pytest.mark.parametrize(
118
+ "suffix",
119
+ [
120
+ "input_global_scale",
121
+ "weight_global_scale",
122
+ "weight_packed",
123
+ "weight_scale",
124
+ ],
125
+)
126
+def test_inkling_mapper_maps_compressed_tensors_expert_params(
127
+ projection: str, suffix: str, nested: bool
128
+) -> None:
129
+ projection_param = (
130
+ f"{projection}_weight.{suffix}" if nested else f"{projection}_{suffix}"
131
+ )
132
+ source = f"model.llm.layers.2.mlp.experts.{projection_param}"
133
+
134
+ mapped = _TmlForCausalLMBase.hf_to_vllm_mapper.apply_list([source])
135
+
136
+ assert mapped == [f"model.layers.2.mlp.experts.{projection}_{suffix}"]
137
+
138
+
139
@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)])
140
def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
141
experts = SimpleNamespace(
@@ -132,6 +156,37 @@ def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
156
assert loaded == [f"experts.routed_experts.{projection}_input_scale"]
157
158
159
+@pytest.mark.parametrize("projection", ["w13", "w2"])
160
+@pytest.mark.parametrize("scale_kind", ["input", "weight"])
161
+def test_moe_loads_compressed_tensors_global_scale(
162
+ projection: str, scale_kind: str
163
+) -> None:
164
+ param = torch.nn.Parameter(torch.empty(3, 2 if projection == "w13" else 1))
165
+ if projection == "w2":
166
+ param = torch.nn.Parameter(param.squeeze(1))
167
+ experts = SimpleNamespace(
168
+ **{f"{projection}_{scale_kind}_global_scale": param},
169
+ moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=0)),
170
+ )
171
+ layer = SimpleNamespace(
172
+ experts=SimpleNamespace(routed_experts=experts),
173
+ _local_expert_slots=lambda: {0: 0, 1: 1, 2: 2},
174
+ )
175
+ checkpoint_scale = torch.tensor([[1.0], [2.0], [3.0]])
176
+
177
+ loaded = moe.InklingMoE.load_expert_weight(
178
+ layer,
179
+ f"experts.{projection}_{scale_kind}_global_scale",
180
+ checkpoint_scale,
181
+ )
182
+
183
+ expected = (
184
+ checkpoint_scale.expand_as(param) if param.ndim == 2 else checkpoint_scale[:, 0]
185
+ )
186
+ torch.testing.assert_close(param, expected)
187
+ assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"]
188
+
189
+
190
def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None:
191
monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2)
192
monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1)
vllm/models/inkling/nvidia/model.py
+10
-1
@@ -378,11 +378,20 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
378
"language_model.lm_head.": "lm_head.",
379
},
380
orig_to_new_suffix={
381
- # NVFP4 scale
381
+ # ModelOpt NVFP4 scales
382
".w13_weight.scale": ".w13_weight_scale",
383
".w13_weight.scale2": ".w13_weight_scale_2",
384
".w2_weight.scale": ".w2_weight_scale",
385
".w2_weight.scale2": ".w2_weight_scale_2",
386
+ # Compressed tensors NVFP4 parameters
387
+ ".w13_weight.input_global_scale": ".w13_input_global_scale",
388
+ ".w13_weight.weight_global_scale": ".w13_weight_global_scale",
389
+ ".w13_weight.weight_packed": ".w13_weight_packed",
390
+ ".w13_weight.weight_scale": ".w13_weight_scale",
391
+ ".w2_weight.input_global_scale": ".w2_input_global_scale",
392
+ ".w2_weight.weight_global_scale": ".w2_weight_global_scale",
393
+ ".w2_weight.weight_packed": ".w2_weight_packed",
394
+ ".w2_weight.weight_scale": ".w2_weight_scale",
395
},
396
)
397
packed_modules_mapping = {
vllm/models/inkling/nvidia/moe.py
+18
-9
@@ -10,11 +10,10 @@ through vLLM's FusedMoE (which handles TP/EP); the sink experts run in
10
activates every sink) and always bf16 (the checkpoint excludes every
11
``shared_experts`` from quantization).
12
13
-NVFP4 routed experts reuse vLLM's ModelOpt NVFP4 fused-MoE method; excluded
14
-(bf16) layers fall back to the unquantized method. The checkpoint's fused
15
-stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` /
16
-``.input_amax`` aux tensors) are translated to the standard per-expert loads
17
-in :meth:`InklingMoE.load_expert_weight`.
13
+NVFP4 routed experts reuse vLLM's fused-MoE methods; excluded (bf16) layers
14
+fall back to the unquantized method. Checkpoint fused stacked tensors are
15
+translated to the standard per-expert loads in
16
+:meth:`InklingMoE.load_expert_weight`.
17
"""
18
19
from __future__ import annotations
@@ -575,11 +574,21 @@ class InklingMoE(nn.Module):
574
lids = [slots[g] for g in gids]
575
tp_rank = experts.moe_config.moe_parallel_config.tp_rank
576
578
- if key.endswith("_scale_2"):
577
+ if key.endswith(("_scale_2", "_global_scale")):
578
# Per-expert scalars, vectorized over the local experts. The
580
- # fused w13 param carries one slot per gate/up half.
581
- vals = weight[gids].float().to(param.device)
582
- param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals
579
+ # fused w13 params carry one slot per gate/up half. A single
580
+ # checkpoint value is shared by both halves.
581
+ vals = weight[gids].float().reshape(len(gids), -1)
582
+ target_width = math.prod(param.shape[1:])
583
+ if vals.shape[1] == 1:
584
+ vals = vals.expand(-1, target_width)
585
+ elif vals.shape[1] != target_width:
586
+ raise ValueError(
587
+ f"cannot load {tuple(weight.shape)} into {tuple(param.shape)}"
588
+ )
589
+ param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to(
590
+ param.device
591
+ )
592
elif key.startswith("w13"):
593
# Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the
594
# fused param layout is [w1(gate); w3(up)]. The TP-local rows form