From b300f666568e6b83a1e8893a3e13fa71afa88aa5 Mon Sep 17 00:00:00 2001 From: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:45:55 -0400 Subject: [PATCH] [Bugfix] Fix HSDP compatibility with the new online FP8 linear method (#5677) Signed-off-by: roG0d --- docs/user_guide/diffusion/parallelism/hsdp.md | 36 +++++++++++++++++++ tests/diffusion/quantization/test_hsdp_fp8.py | 34 ++++++++++++++++-- vllm_omni/diffusion/quantization/hsdp_fp8.py | 15 +++++--- 3 files changed, 79 insertions(+), 6 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/hsdp.md b/docs/user_guide/diffusion/parallelism/hsdp.md index 96a357c86..8a8aa5ef9 100644 --- a/docs/user_guide/diffusion/parallelism/hsdp.md +++ b/docs/user_guide/diffusion/parallelism/hsdp.md @@ -64,6 +64,41 @@ omni = Omni( ) ``` +### Combined with Online FP8 Quantization + +HSDP can be combined with online FP8 quantization for diffusion models that +support both features. HSDP shards the quantized weights across GPUs to reduce +the per-GPU weight memory footprint. + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="", + quantization="fp8", + parallel_config=DiffusionParallelConfig( + use_hsdp=True, + hsdp_shard_size=2, + ), +) +``` + +For online serving: + +```bash +vllm serve \ + --omni \ + --quantization fp8 \ + --use-hsdp \ + --hsdp-shard-size 2 +``` + +The model must support both HSDP and online FP8 quantization. HSDP remains +incompatible with Tensor Parallelism. See the +[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and +hardware. + --- ## Example Script @@ -147,3 +182,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co 1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models 2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism 3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled +4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py index 482235ff2..ce939dbf3 100644 --- a/tests/diffusion/quantization/test_hsdp_fp8.py +++ b/tests/diffusion/quantization/test_hsdp_fp8.py @@ -36,15 +36,23 @@ class _ToyKernel: # --- 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): +def _make_fp8_toy_module( + out_features: int = 16, + in_features: int = 32, + *, + quant_method_cls=None, +): """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``. """ from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + if quant_method_cls is None: + quant_method_cls = Fp8LinearMethod + # Create the quant-method instance without calling __init__ (avoids # heavyweight upstream config dependencies). isinstance() still works. - qm = object.__new__(Fp8LinearMethod) + qm = object.__new__(quant_method_cls) qm.fp8_linear = _ToyKernel() # randn does not support float8 on CPU; zeros does and the test only @@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight(): assert module.weight.shape == (16, 32) +def test_rewrites_per_tensor_online_fp8_weight(): + """Regression test for vLLM's new online FP8 frontend (#45463).""" + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + module = _make_fp8_toy_module( + 16, + 32, + quant_method_cls=Fp8PerTensorOnlineLinearMethod, + ) + + 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) + w, *_ = module.quant_method.fp8_linear._get_layer_params(module) + assert w.shape == (32, 16) + assert w.stride() == (1, 32) + + def test_patched_get_layer_params_returns_transposed_view(): module = _make_fp8_toy_module(16, 32) prepare_fp8_layers_for_fsdp(module) diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py index 7c4666e29..6ffa8f2b4 100644 --- a/vllm_omni/diffusion/quantization/hsdp_fp8.py +++ b/vllm_omni/diffusion/quantization/hsdp_fp8.py @@ -33,9 +33,17 @@ import types from torch import nn from vllm.logger import init_logger from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod +from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, +) logger = init_logger(__name__) +_FP8_TRANSPOSED_WEIGHT_METHODS = ( + Fp8LinearMethod, + Fp8PerTensorOnlineLinearMethod, +) + def _build_transposed_get_layer_params(original_bound_method): """Return a ``_get_layer_params`` replacement that transposes ``w``. @@ -55,9 +63,8 @@ def _build_transposed_get_layer_params(original_bound_method): 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:`Fp8LinearMethod` - and whose weight is currently a non-contiguous transpose view, this - function: + For every layer whose quantization method stores FP8 weights as 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. @@ -76,7 +83,7 @@ def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: patched_kernel_ids: set[int] = set() for module in model.modules(): qm = getattr(module, "quant_method", None) - if not isinstance(qm, Fp8LinearMethod): + if not isinstance(qm, _FP8_TRANSPOSED_WEIGHT_METHODS): continue weight = getattr(module, "weight", None)