[Feature][Diffusion] Allow online INT8 quantization with DLO AllGather (#6573)

Signed-off-by: brandneway <gyuan4892@gmail.com>
2026-08-29 09:50:30 +00:00
1 parent d9980c75e8
commit 5a5bc55c25
@@ -273,12 +273,15 @@ def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
    assert loader.take_host_weight_plan() is None


-def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
+def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
    return OmniDiffusionConfig(
        model="",
        dtype=torch.float32,
        quantization_config="fp8",
-        parallel_config=DiffusionParallelConfig(data_parallel_size=2),
+        parallel_config=DiffusionParallelConfig(
+            data_parallel_size=dp_size,
+            sequence_parallel_size=1,
+        ),
        enable_distributed_layerwise_offload=True,
        dlo_use_allgather=True,
    )
@@ -824,6 +827,93 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
    assert loader.take_host_weight_plan() is None


+def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
+    """A DLO group of one runs no weight collective, so an otherwise
+    unvalidated online method must not be rejected by the AllGather gate."""
+
+    class UnsupportedOnlineMethod:
+        uses_meta_device = True
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config(dp_size=1)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnsupportedOnlineMethod()
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == []
+    assert calls == [("load", True), "process"]
+
+
 def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

@@ -847,10 +937,54 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
        ),
    )

-    with pytest.raises(ValueError, match="per-tensor FP8 linears"):
+    with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
        loader.load_model(load_device="cpu")


+def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
+    mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
+    NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
@@ -245,6 +245,160 @@ class TestDistributedLayerwiseOffloadHook:
        )
        torch.testing.assert_close(rank0_block.weight_scale, expected_scale)

+    def test_allgather_reconstructs_online_int8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-INT8 tensors byte-for-byte.
+
+        One block carries both finalized online-INT8 layouts: the NPU method
+        leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
+        per-channel scale, while the CUDA method leaves a transposed-view
+        weight (stride (1, K)) with a scalar fp32 scale.  The int8 and fp32
+        dtype groups must survive the flat-shard AllGather independently.
+        """
+
+        class OnlineInt8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical = torch.arange(1.0, 13.0).reshape(3, 4)
+                # NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
+                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
+                self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
+                # CUDA online int8: Parameter(qweight.t().data) -> transposed
+                # view (stride (1, K)) over the (N, K) storage, scalar scale
+                self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
+                self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+
+        current_block = OnlineInt8Block()
+        rank0_block = OnlineInt8Block()
+        rank1_block = OnlineInt8Block()
+        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+        def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+            return DistributedLayerwiseOffloadHook(
+                next_block=block,
+                device=torch.device("cpu"),
+                dp_group=object(),
+                dp_size=2,
+                rank=rank,
+                copy_stream=DummyStream(),
+                comm_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        rank0_hook = make_hook(rank0_block, 0)
+        rank1_hook = make_hook(rank1_block, 1)
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineInt8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        for name, tensor in rank0_block.named_parameters():
+            assert tensor.stride() == expected_strides[name]
+            torch.testing.assert_close(
+                tensor.float(),
+                expected[name].float(),
+                rtol=0,
+                atol=0,
+            )
+
+    def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.
+
+        One block carries both finalized online-MXFP8 layouts: the NPU method
+        leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
+        (K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
+        (XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
+        view over a contiguous (K/32, N) buffer.  The fp8 and e8m0 dtype
+        groups must survive the flat-shard AllGather independently.
+        """
+
+        class OnlineMxfp8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical = torch.arange(1.0, 13.0).reshape(3, 4)
+                # Build scales from raw e8m0 bit patterns (1 byte/elem).
+                npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
+                kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
+                # NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
+                # weight, contiguous (K_groups/2, N, 2) e8m0 scale
+                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
+                self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
+                # vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
+                # scale kept as a transposed view (stride (1, K/32)) over the
+                # contiguous (K/32, N) storage
+                self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
+                self.kernel_weight_scale = nn.Parameter(
+                    kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
+                )
+
+        current_block = OnlineMxfp8Block()
+        rank0_block = OnlineMxfp8Block()
+        rank1_block = OnlineMxfp8Block()
+        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+        def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+            return DistributedLayerwiseOffloadHook(
+                next_block=block,
+                device=torch.device("cpu"),
+                dp_group=object(),
+                dp_size=2,
+                rank=rank,
+                copy_stream=DummyStream(),
+                comm_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        rank0_hook = make_hook(rank0_block, 0)
+        rank1_hook = make_hook(rank1_block, 1)
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineMxfp8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        for name, tensor in rank0_block.named_parameters():
+            assert tensor.stride() == expected_strides[name]
+            # Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
+            # version-sensitive, while contiguous().view(uint8) is an exact
+            # byte-level comparison for both contiguous and transposed-view
+            # parameters.
+            torch.testing.assert_close(
+                tensor.detach().contiguous().view(torch.uint8),
+                expected[name].contiguous().view(torch.uint8),
+                rtol=0,
+                atol=0,
+            )
+
    def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
        """DTensor wrapper should be preserved through prefetch/offload cycle."""
        current_block = TinyBlock(_make_values(1.0))