[diffusion][feature] Add offline SVDQuant W4A4 support (#6162)

Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/quantization/test_svdquant_tp_loading.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
+    layer = object.__new__(linear_type)
+    torch.nn.Module.__init__(layer)
+    layer.tp_rank = tp_rank
+    return layer
+
+
+def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
+    values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
+    return values.reshape(shape).to(dtype)
+
+
+def _method(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+    return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
+
+
+def test_column_parallel_loader_shards_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[32],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = loaded.narrow(parameter.output_dim, 32, 32)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_row_parallel_loader_shards_svdquant_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(RowParallelLinear, tp_rank=1)
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=64,
+        output_partition_sizes=[64],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_down": _values((128, 8), torch.bfloat16),
+        "smooth_factor": _values((128,), torch.bfloat16),
+    }
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        shard_size = parameter.shape[parameter.input_dim]
+        expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
+        torch.testing.assert_close(parameter, expected)
+
+
+def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
+    layer.total_num_heads = 4
+    layer.total_num_kv_heads = 2
+    layer.num_heads = 2
+    layer.num_kv_heads = 1
+    layer.num_kv_head_replicas = 1
+    layer.head_size = 8
+    layer.v_head_size = 8
+    method = _method(monkeypatch)
+    method.create_weights(
+        layer,
+        input_size_per_partition=128,
+        output_partition_sizes=[16, 8, 8],
+        input_size=128,
+        output_size=64,
+        params_dtype=torch.bfloat16,
+        weight_loader=layer.weight_loader,
+    )
+
+    loaded_tensors = {
+        "qweight": _values((64, 64), torch.int8),
+        "wscales": _values((8, 64), torch.float8_e4m3fn),
+        "proj_up": _values((64, 8), torch.bfloat16),
+        "wcscales": _values((64,), torch.bfloat16),
+    }
+    expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
+    for name, loaded in loaded_tensors.items():
+        parameter = getattr(layer, name)
+        parameter.weight_loader(parameter, loaded)
+        expected = torch.cat(
+            [
+                loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
+                for part in expected_global_outputs
+            ],
+            dim=parameter.output_dim,
+        )
+        torch.testing.assert_close(parameter, expected)