[Bugfix] Fix HSDP + FP8 online quantization compatibility (#4494)

Signed-off-by: Kwokho Tsui <kh_tsui@163.com>
Co-authored-by: Kwokho Tsui <kh_tsui@163.com>

tests/diffusion/quantization/test_hsdp_fp8.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HSDP/FSDP2 compatibility with online FP8 quantization."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+    _build_transposed_get_layer_params,
+    prepare_fp8_layers_for_fsdp,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# --- shared test doubles ---
+
+
+class _ToyKernel:
+    """Minimal kernel that mirrors ``FP8ScaledMMLinearKernel``."""
+
+    def __init__(self):
+        self.layer_param_names = ("weight", "weight_scale", "input_scale", "input_scale_ub")
+
+    def _get_layer_params(self, layer):
+        w, w_s, x_s, x_s_ub = self.layer_param_names
+        return (
+            getattr(layer, w),
+            getattr(layer, w_s),
+            getattr(layer, x_s, None),
+            getattr(layer, x_s_ub, None),
+        )
+
+
+# --- helper to build a toy module simulating online-FP8 post-load state ---
+
+
+def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
+    """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
+    ``quant_method`` is an ``Fp8OnlineLinearMethod`` with a ``_ToyKernel``.
+    """
+    from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+
+    # Create the quant-method instance without calling __init__ (avoids
+    # heavyweight upstream config dependencies).  isinstance() still works.
+    qm = object.__new__(Fp8OnlineLinearMethod)
+    qm.fp8_linear = _ToyKernel()
+
+    # randn does not support float8 on CPU; zeros does and the test only
+    # cares about shape/stride, not the actual weight values.
+    qweight = torch.zeros(out_features, in_features, dtype=torch.float8_e4m3fn)
+    weight = nn.Parameter(qweight.t())  # (in, out) non-contiguous column-major view
+
+    module = nn.Module()
+    module.quant_method = qm
+    module.weight = weight
+    module.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    # Set input_scale / input_scale_ub to verify _get_layer_params
+    # passes them through unchanged.
+    module.input_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    module.input_scale_ub = nn.Parameter(torch.ones(1, dtype=torch.float32))
+
+    return module
+
+
+# --- tests ---
+
+
+def test_transposed_get_layer_params():
+    import types
+
+    kernel = _ToyKernel()
+
+    # original bound method + patch + re-bind, matching prepare_fp8_layers_for_fsdp
+    original_bound = kernel._get_layer_params
+    patched_func = _build_transposed_get_layer_params(original_bound)
+    kernel._get_layer_params = types.MethodType(patched_func, kernel)
+
+    module = _make_fp8_toy_module(16, 32)
+    # Reproduce the storage rewrite from prepare_fp8_layers_for_fsdp:
+    # weight was qweight.t()  →  .t() recovers qweight (row-major contiguous).
+    module.weight = nn.Parameter(module.weight.data.t(), requires_grad=False)
+    assert module.weight.is_contiguous(), "qweight should be row-major contiguous"
+
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # patched method transposes row-major → column-major for Cutlass
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    # scales are passed through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_rewrites_non_contiguous_weight():
+    module = _make_fp8_toy_module(16, 32)
+    assert not module.weight.is_contiguous()
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 1
+    assert module.weight.is_contiguous()
+    assert module.weight.shape == (16, 32)
+
+
+def test_patched_get_layer_params_returns_transposed_view():
+    module = _make_fp8_toy_module(16, 32)
+    prepare_fp8_layers_for_fsdp(module)
+
+    kernel = module.quant_method.fp8_linear
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # Cutlass expects column-major B: (in, out), stride (1, in)
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    assert w.stride() == (1, 32)
+
+    # scales pass through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_contiguous_weights_skipped():
+    module = _make_fp8_toy_module(16, 32)
+    # make weight contiguous before calling
+    module.weight = nn.Parameter(module.weight.data.t().contiguous())
+    assert module.weight.is_contiguous()
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+
+
+def test_kernel_patched_only_once_per_instance():
+    module = _make_fp8_toy_module(16, 32)
+    kernel = module.quant_method.fp8_linear
+    original = kernel._get_layer_params
+
+    prepare_fp8_layers_for_fsdp(module)
+    first_patched = kernel._get_layer_params
+
+    # Second call with same kernel id must not stack patches
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+    assert kernel._get_layer_params is first_patched
+    assert kernel._get_layer_params is not original

vllm_omni/diffusion/distributed/hsdp.py

@@ -140,8 +140,13 @@ def apply_hsdp_to_model(
         fs_rank,
     )
 
+    # When the model contains FP8 parameters (online quantization), let FSDP
+    # keep the original storage dtype on all-gather instead of casting to
+    # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
+    # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
+    has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
     mp_policy = MixedPrecisionPolicy(
-        param_dtype=hsdp_config.param_dtype,
+        param_dtype=None if has_fp8_params else hsdp_config.param_dtype,
         reduce_dtype=hsdp_config.reduce_dtype,
         output_dtype=hsdp_config.output_dtype,
         cast_forward_inputs=False,

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -602,10 +602,17 @@ class DiffusersPipelineLoader:
             model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             with set_current_diffusion_config(self.od_config):
                 model = model_cls(od_config=self.od_config)
-            if not is_hsdp and target_device.type != "cpu":
+            # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak
+            # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only
+            # kernels inside load_weights via the layerwise loader, so when a quant
+            # config is set we initialize on the accelerator like the non-HSDP path;
+            # apply_hsdp_to_model shards GPU-resident params equally well.
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            if not hsdp_defer_to_cpu and target_device.type != "cpu":
                 model.to(target_device)
         else:
-            device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
+            hsdp_defer_to_cpu = is_hsdp and self.quant_config is None
+            device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device
             with device_ctx:
                 if load_format == "default":
                     model = initialize_model(self.od_config)
@@ -646,6 +653,22 @@ class DiffusersPipelineLoader:
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
+        # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
+        # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
+        # Rewrite affected layers in-place to row-major contiguous storage and
+        # shift the .t() to GEMM-call time. Layers using other quant methods or
+        # already-contiguous weights are left untouched.
+        if self.quant_config is not None:
+            from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+                prepare_fp8_layers_for_fsdp,
+            )
+
+            for _attr in ("transformer", "transformer_2"):
+                _trans = getattr(model, _attr, None)
+                if _trans is not None:
+                    prepare_fp8_layers_for_fsdp(_trans)
+
         # Collect all transformers to shard (some models have transformer_2 for MoE)
         transformers_to_shard = []
         transformer = getattr(model, "transformer", None)

vllm_omni/diffusion/quantization/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/quantization/hsdp_fp8.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HSDP/FSDP2 compatibility for online FP8 quantization.
+
+vllm's ``Fp8OnlineLinearMethod.process_weights_after_loading`` ends with
+``layer.weight = qweight.t()`` so that the Cutlass FP8 GEMM kernel sees its
+B operand as column-major ``[K, N]`` (the TN layout required by Hopper FP8
+``wgmma`` instructions). The resulting tensor is a non-contiguous transpose
+view of a ``(out_features, in_features)`` row-major storage.
+
+FSDP2 ``fully_shard`` rejects non-contiguous parameters because dim-0 sharding
+on a column-major view cannot be a contiguous memcpy. The two requirements are
+fundamentally at odds with the same physical buffer.
+
+This module reconciles them by separating the views: keep the parameter as
+the underlying ``(out, in)`` row-major contiguous storage (FSDP-friendly),
+and inject the equivalent ``.t()`` at the GEMM call site so the Cutlass
+kernel still receives a column-major B (zero-copy stride flip).
+
+The transpose is injected inside ``ScaledMMLinearKernel._get_layer_params``,
+which is the single place where ``apply_weights`` reads the weight tensor.
+The rest of ``apply_weights`` -- including its ``output_shape = w.shape[1]``
+computation and the eventual ``apply_scaled_mm(B=w, ...)`` call -- then
+operates on a tensor whose shape ``(K, N)`` and stride ``(1, K)`` match
+what the upstream FP8 kernel expects. No other upstream code is overridden,
+which keeps us robust against future changes inside ``apply_weights``.
+"""
+
+from __future__ import annotations
+
+import types
+
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+
+logger = init_logger(__name__)
+
+
+def _build_transposed_get_layer_params(original_bound_method):
+    """Return a ``_get_layer_params`` replacement that transposes ``w``.
+
+    ``original_bound_method`` is the kernel instance's existing
+    ``_get_layer_params`` (already bound to ``self``); we close over it and
+    apply ``.t()`` to the returned weight, leaving scales untouched.
+    """
+
+    def _get_layer_params(self, layer):
+        w, w_s, x_s, x_s_ub = original_bound_method(layer)
+        return w.t(), w_s, x_s, x_s_ub
+
+    return _get_layer_params
+
+
+def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
+    """Make online-FP8 linear layers in ``model`` FSDP2-compatible.
+
+    For every layer whose ``quant_method`` is an :class:`Fp8OnlineLinearMethod`
+    and whose weight is currently a non-contiguous transpose view, this
+    function:
+
+    1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major
+       contiguous storage so FSDP2 ``fully_shard`` accepts it.
+    2. Patches the per-layer GEMM kernel's bound ``_get_layer_params`` method
+       to return a ``.t()`` view of the weight, so ``apply_weights`` and the
+       downstream ``apply_scaled_mm`` continue to see a column-major
+       ``(in, out)`` B with zero copies.
+
+    Layers whose weight is already contiguous (e.g. Marlin FP8, offline-
+    quantized checkpoints) or that use a different quant method are skipped.
+
+    Returns:
+        Number of layers rewritten.
+    """
+    n_patched = 0
+    patched_kernel_ids: set[int] = set()
+    for module in model.modules():
+        qm = getattr(module, "quant_method", None)
+        if not isinstance(qm, Fp8OnlineLinearMethod):
+            continue
+
+        weight = getattr(module, "weight", None)
+        if weight is None or weight.is_contiguous():
+            continue
+
+        # ``weight`` here is ``qweight.t()`` (a non-contiguous (in, out) view
+        # of an (out, in) row-major qweight storage). ``weight.t()`` recovers
+        # that storage; ``.contiguous()`` is a zero-copy alias when the source
+        # is already row-major contiguous, but defensively materializes if
+        # upstream ever stacks another view.
+        contig = weight.data.t().contiguous()
+        new_param = nn.Parameter(contig, requires_grad=False)
+        module.weight = new_param
+
+        kernel = qm.fp8_linear
+        # Each linear layer is expected to have its own kernel instance, but
+        # guard against shared instances to avoid stacking multiple ``.t()``
+        # patches on the same object (which would compose into identity).
+        if id(kernel) not in patched_kernel_ids:
+            original_get = kernel._get_layer_params
+            kernel._get_layer_params = types.MethodType(_build_transposed_get_layer_params(original_get), kernel)
+            patched_kernel_ids.add(id(kernel))
+
+        n_patched += 1
+
+    if n_patched:
+        logger.info(
+            "Rewrote %d FP8 linear layer(s) into FSDP2-compatible storage; "
+            "transpose to column-major B is now applied at GEMM time.",
+            n_patched,
+        )
+    return n_patched