diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py index a7307e167a1..e15a573b845 100644 --- a/tests/models/inkling/test_moe_weight_layout.py +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -187,6 +187,32 @@ def test_moe_loads_compressed_tensors_global_scale( assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"] +@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)]) +def test_moe_loads_channelwise_scale_for_tp( + projection: str, checkpoint_rows: int +) -> None: + param = torch.nn.Parameter(torch.empty(2, 4, 1)) + experts = SimpleNamespace( + **{f"{projection}_weight_scale": param}, + moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)), + ) + layer = SimpleNamespace( + experts=SimpleNamespace(routed_experts=experts), + _local_expert_slots=lambda: {0: 0, 2: 1}, + ) + checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1) + + loaded = moe.InklingMoE.load_expert_weight( + layer, f"experts.{projection}_weight_scale", checkpoint_scale + ) + + expected = checkpoint_scale[[0, 2]] + if projection == "w13": + expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2) + torch.testing.assert_close(param, expected.float()) + assert loaded == [f"experts.routed_experts.{projection}_weight_scale"] + + def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None: monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 32f7489a4c4..9ff3acaa14d 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -589,6 +589,9 @@ class InklingMoE(nn.Module): param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to( param.device ) + elif key == "w2_weight_scale" and weight.shape[-1] == 1: + # Per-output-channel scales are replicated across TP ranks. + param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype) elif key.startswith("w13"): # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the # fused param layout is [w1(gate); w3(up)]. The TP-local rows form