From 7e152ca64b6627eaf80eb7d94d271bd504f16b1a Mon Sep 17 00:00:00 2001 From: Kwokho Tsui <60513928+KwokhoTsui@users.noreply.github.com> Date: Fri, 19 Jun 2026 15:06:27 +0800 Subject: [PATCH] [Bugfix] Fix HSDP + FP8 online quantization compatibility (#4494) Signed-off-by: Kwokho Tsui Co-authored-by: Kwokho Tsui --- tests/diffusion/quantization/test_hsdp_fp8.py | 147 ++++++++++++++++++ vllm_omni/diffusion/distributed/hsdp.py | 7 +- .../model_loader/diffusers_loader.py | 27 +++- vllm_omni/diffusion/quantization/__init__.py | 2 + vllm_omni/diffusion/quantization/hsdp_fp8.py | 112 +++++++++++++ 5 files changed, 292 insertions(+), 3 deletions(-) create mode 100644 tests/diffusion/quantization/test_hsdp_fp8.py create mode 100644 vllm_omni/diffusion/quantization/__init__.py create mode 100644 vllm_omni/diffusion/quantization/hsdp_fp8.py diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py new file mode 100644 index 000000000..97279acf2 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 42675c7c0..3a056d6d4 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/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, diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index b2339843e..1aa40aae5 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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) diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/quantization/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py new file mode 100644 index 000000000..c2ddd68ae --- /dev/null +++ b/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