[Hardware][Ascend] Add Wan2.2 MXFP4 UOS, Smooth and W4A8 fallback (#7210)
Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>
This commit is contained in:
HAAZZZEEEE authored and GitHub committed
1 parent
4fab34c8f4
commit
becfefd440
20 files changed
+2435
-184
No files matched your search
@@ -6,22 +6,25 @@ W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
|
||||
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
|
||||
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
|
||||
|
||||
On Ascend, single-scale MXFP4 also supports Universal Optimal Scaling
|
||||
(UOS, `Qmax=7.25`), described in the single-scale scaling section below.
|
||||
|
||||
vLLM-Omni provides two quantization methods with different scale structures:
|
||||
|
||||
| Method | Scale structure | Mode | Use case |
|
||||
|--------|----------------|------|----------|
|
||||
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
|
||||
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
|
||||
| -------- | ---------------- | ------ | ---------- |
|
||||
| `mxfp4` | Single-scale (per-32 fine only) | Online + pre-quantized single-scale weights | Online baseline, or an existing single-scale checkpoint in the layout below |
|
||||
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | W4A4 with dual-level scales; offline mode loads calibrated Smooth tensors |
|
||||
|
||||
!!! tip "Recommended: `mxfp4_dualscale` offline"
|
||||
For production deployments, use the `mxfp4_dualscale` offline mode with a
|
||||
!!! tip "Existing W4A4-only DualScale deployments"
|
||||
For W4A4-only deployments, the existing `mxfp4_dualscale` offline mode uses a
|
||||
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
|
||||
calibrated `mul_scale` tensors from disk, providing measurably better accuracy
|
||||
than any online method. The one-time preprocessing cost amortises across all
|
||||
subsequent inference runs.
|
||||
calibrated `mul_scale` tensors from disk. Quality depends on the model,
|
||||
calibration recipe and BF16 fallback policy; evaluate the checkpoint against
|
||||
your quality target before deployment.
|
||||
|
||||
Use `mxfp4` online only for quick experimentation where preprocessing time
|
||||
is not acceptable and accuracy loss is tolerable.
|
||||
W4A8 layer/step fallback requires single-scale `mxfp4`, online or offline.
|
||||
DualScale supports W4A4 and BF16 routing only.
|
||||
|
||||
!!! warning "Online single-scale ≠ Offline dual-scale"
|
||||
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
|
||||
@@ -34,61 +37,244 @@ vLLM-Omni provides two quantization methods with different scale structures:
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
|
||||
| AMD ROCm (gfx950 / MI355X) | ✅ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ |
|
||||
The following matrix describes these vLLM-Omni implementations, rather than
|
||||
hardware support for the MXFP4 format in general.
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
| Device | Single-scale online OCP | Single-scale offline | UOS / W4A8 fallback | DualScale |
|
||||
| -------- | -------- | -------- | -------- | -------- |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ | ✅ |
|
||||
| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ | ❌ |
|
||||
| Other AMD architectures | ❌ | ❌ | ❌ | ❌ |
|
||||
| NVIDIA GPU | ❌ | ❌ | ❌ | ❌ |
|
||||
| Intel XPU | ❌ | ❌ | ❌ | ❌ |
|
||||
|
||||
Legend: `✅` supported by the implementation, `❌` unsupported.
|
||||
ROCm uses the existing AITER online path and requires `mxfp4_scale_alg=0`.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Wan2.2)
|
||||
|
||||
| Model | Online | Offline | Notes |
|
||||
|-------|--------|---------|-------|
|
||||
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
|
||||
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
|
||||
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
|
||||
| ------- | -------- | --------- | ------- |
|
||||
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Cascade experts resolve checkpoint metadata independently; explicit global runtime policy applies to both |
|
||||
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Per-expert configuration uses the same interface as T2V; validate I2V checkpoints separately |
|
||||
| Wan2.2-TI2V-5B | Requires validation | Requires validation | This guide does not establish checkpoint compatibility or generation quality for TI2V-5B |
|
||||
|
||||
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
|
||||
quantization quality, not model compatibility — both work on cascade (A14B) and
|
||||
single-transformer models alike, the same as `mxfp8` online:
|
||||
Both methods can be integrated with single-transformer and cascade models.
|
||||
This alone does not establish support or quality for every model variant:
|
||||
|
||||
- `mxfp4`: single-scale, lower overhead, simpler compute, online only
|
||||
- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
|
||||
- `mxfp4`: single-scale; online and pre-quantized weights, with optional W4A8 fallback on NPU
|
||||
- `mxfp4_dualscale`: dual-scale; online and offline, with BF16 fallback
|
||||
|
||||
**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
|
||||
by the merge script); there is no offline `mxfp4` single-scale format.
|
||||
The existing `merge_mxfp4_dualscale_checkpoint.py` produces **offline
|
||||
`mxfp4_dualscale`** checkpoints. Single-scale weights must satisfy the separate
|
||||
layout described below; changing a DualScale checkpoint's method name to
|
||||
`mxfp4` does not convert its weights. Prepare single-scale checkpoints with
|
||||
external tooling. Real exported checkpoint compatibility and model quality
|
||||
require validation of that checkpoint; the synthetic tests do not establish
|
||||
UOS export or video acceptance.
|
||||
|
||||
!!! note "Per-layer BF16 fallback in offline cascade models"
|
||||
The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
|
||||
An A14B offline DualScale checkpoint uses `quant_method: mxfp4_dualscale`. Most
|
||||
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
|
||||
retain their original BF16 weights and are listed in `ignored_layers` inside
|
||||
each transformer's `config.json`. The two transformers may have different
|
||||
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
|
||||
and rebuilds the config locally when they differ, so routing is always
|
||||
per-transformer-accurate.
|
||||
|
||||
!!! warning "TI2V-5B not supported"
|
||||
Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
|
||||
count makes it significantly more sensitive to 4-bit quantization noise,
|
||||
resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
|
||||
and resolves the storage format and BF16 layer policy independently.
|
||||
|
||||
## Configuration
|
||||
|
||||
### W4A8 fallback at selected layers and denoising steps (Ascend NPU)
|
||||
|
||||
For single-scale `mxfp4`, set `w4a8_fallback_steps` to a list of zero-based
|
||||
request denoising loop indices (not scheduler timestep values).
|
||||
Set `w4a8_fallback_layers` to exact runtime Linear paths, for example
|
||||
`["blocks.10.attn1.to_qkv", "blocks.12.ffn.net_2"]`. A selected layer uses W4A8
|
||||
at every step; at a selected step, all other quantized layers also use W4A8.
|
||||
The rules are combined with **OR**. W4A8 uses MXFP8 activations with MXFP4
|
||||
weights. Omitting both lists or passing `[]` keeps the existing W4A4 behavior. Existing
|
||||
`ignored_layers` take priority: **BF16 > (layer OR step) W4A8 > W4A4**.
|
||||
DualScale's existing leading BF16 block rule is unchanged; it does not enable
|
||||
W4A8.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="<Wan2.2-T2V-A14B-Diffusers>",
|
||||
quantization_config={
|
||||
"method": "mxfp4",
|
||||
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [0, 1, 38, 39],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
Both policies also work with offline single-scale `mxfp4`. Any nonempty
|
||||
W4A8 list with `mxfp4_dualscale`, online or offline, raises an explicit error
|
||||
at configuration construction. Empty lists preserve its existing W4A4 path.
|
||||
Each expert's checkpoint still determines its serialized
|
||||
format and BF16 layer routing. A serialized checkpoint config is the complete
|
||||
storage contract: `ignored_layers`, or its `modules_to_not_convert` alias,
|
||||
selects the BF16 layers, and omitting both means no explicit BF16 layers. A
|
||||
method-only checkpoint value such as `"mxfp4"` or
|
||||
`{"quant_method": "mxfp4"}` carries no storage policy and leaves a caller's
|
||||
online `ignored_layers` unchanged. An explicitly supplied quantization config
|
||||
retains its W4A8 layer and step lists when the loader rebuilds either expert's
|
||||
config from disk, including empty lists that disable checkpoint policies.
|
||||
Within an explicit runtime config, an omitted W4A8 list defaults to `[]`;
|
||||
individual omitted fields do not inherit saved policies. Without an explicit
|
||||
runtime config, auto-detection uses the saved checkpoint policies.
|
||||
|
||||
Layer names are relative to each expert (`blocks.10...`, without a
|
||||
`transformer.` or `transformer_2.` prefix). A global config applies the same
|
||||
list to both experts; per-component `transformer`/`transformer_2` configs can
|
||||
select different lists. Wan selects these expert names exactly: a `transformer`
|
||||
entry does not match `transformer_2`; an omitted expert uses `default` if set.
|
||||
A component mapped to `null` (or an unmatched component
|
||||
without a default) stays unquantized and requires a BF16 checkpoint; loading a
|
||||
quantized checkpoint for that disabled component raises an error. A top-level
|
||||
`quantization_config=None` retains checkpoint auto-detection.
|
||||
Names are exact, with no regex, wildcard, block expansion,
|
||||
or checkpoint-name remapping. Wan validates them against its runtime Linear
|
||||
modules and logs the selected paths/count and BF16 overrides per PP rank.
|
||||
Use `attn1.to_qkv` for fused self-attention; `attn1.to_q`, `to_k`, or `to_v`
|
||||
checkpoint names raise an error. Cross-attention `attn2.to_q/to_k/to_v` remains
|
||||
separate. Layer-only fallback does not require a denoise ForwardContext.
|
||||
|
||||
The text-to-video example accepts the JSON directly:
|
||||
|
||||
```bash
|
||||
ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model /path/to/Wan2.2-T2V-A14B-Diffusers \
|
||||
--quantization-config '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}' \
|
||||
--tensor-parallel-size 1 --ulysses-degree 2 --cfg-parallel-size 2 \
|
||||
--enforce-eager --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 40 --seed 42 --output wan22-w4a8.mp4
|
||||
```
|
||||
|
||||
Use either `--quantization` or `--quantization-config` in this example.
|
||||
Serving accepts the same JSON through `--diffusion-quantization-config`.
|
||||
|
||||
The list uses one request's loop index across both high- and low-noise experts,
|
||||
and both CFG branches see the same index. A new request starts at step zero.
|
||||
Indices outside the request's step range are never selected. A pipeline must
|
||||
publish `ForwardContext.denoise_step_idx`; requesting step fallback without
|
||||
that context raises an error instead of silently ignoring the list.
|
||||
|
||||
Single-scale W4A4 and W4A8 reuse the same prepared `weight` and `weight_scale`.
|
||||
Online loading quantizes the original BF16/FP16 weight once; offline loading
|
||||
packs the numeric FP4 values and reshapes the saved E8M0 scale without
|
||||
requantizing the weight. Only activation precision changes during forward.
|
||||
There is no DualScale decoding, second W4 representation or `w4a8_weight` /
|
||||
`w4a8_weight_scale` cache. No weight quantization runs inside the denoising loop.
|
||||
|
||||
The mixed FP8/FP4 GEMM uses BF16 output and a matching BF16 row bias for
|
||||
compatibility with the target A5 runtime. An FP16 caller receives that result
|
||||
converted back to FP16; its output shape and dtype are preserved.
|
||||
|
||||
The single-scale fallback validation target is eager Wan2.2 T2V A14B on A5.
|
||||
Quality scores, optimal fallback lists, real UOS exports and additional
|
||||
compile/offload combinations require separate evaluation. Existing W4A4-only
|
||||
DualScale still needs complete 512-channel input groups; use TP=1 with
|
||||
sequence/CFG parallelism for its initial multi-card run.
|
||||
|
||||
### `mxfp4` — Pre-Quantized Single-Scale Weights (Ascend NPU)
|
||||
|
||||
Set the following fields in each transformer's `config.json` when a
|
||||
single-scale checkpoint uses calibrated Smooth scaling:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": true,
|
||||
"require_smooth_scale": true,
|
||||
"ignored_layers": []
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Each quantized Linear must provide tensors under its model parameter name:
|
||||
|
||||
| Tensor | Checkpoint layout | Meaning |
|
||||
| -------- | ------------------- | --------- |
|
||||
| `weight` | `(N, K)` numeric FP4 values | Packed to FP4 after loading |
|
||||
| `weight_scale` | `(N, ceil(K / 32))`, `uint8` | Raw E8M0 exponent bytes, required |
|
||||
| `mul_scale` | `(K,)`, FP32 | Smooth pre-scale, required when declared above |
|
||||
|
||||
Both W4A4 and W4A8 quantize `x * mul_scale`. The Smooth tensor is converted
|
||||
to the layer's activation dtype after loading. Row-parallel layers load the
|
||||
corresponding input-channel shard. Pre-fused QKV weights must use one shared
|
||||
Smooth tensor for the fused projection.
|
||||
|
||||
For backward compatibility, `require_smooth_scale` defaults to `false`:
|
||||
an absent tensor uses identity scaling, while a present tensor is still loaded.
|
||||
Without a declaration, an absent Smooth tensor cannot be distinguished from a
|
||||
legacy unsmoothed checkpoint. Calibrated checkpoints should declare
|
||||
`require_smooth_scale: true` so missing or misnamed Smooth tensors fail loading;
|
||||
incorrect tensor shapes, non-floating tensors, NaN/Inf and nonpositive values
|
||||
also fail. Values that overflow or underflow to zero in the activation dtype
|
||||
are rejected during weight processing. The flag requires serialized single-scale
|
||||
weights and is rejected for online quantization. It does not make DualScale
|
||||
scales optional or change missing-parameter policies for other backends.
|
||||
Either the checkpoint or an explicit caller can require Smooth: a `true`
|
||||
declaration is retained when per-expert metadata rebuilds the configuration,
|
||||
including when the other source omits the flag or sets it to `false`.
|
||||
|
||||
Prepare the canonical numeric FP4 checkpoint before inference using external
|
||||
tooling. Runtime loading does not parse producer-specific export formats.
|
||||
Retain the producer recipe to identify the weight-generation algorithm and
|
||||
Smooth calibration; tensor storage alone does not establish UOS provenance.
|
||||
|
||||
### `mxfp4` — Single-Scale Block Scaling (OCP / UOS)
|
||||
|
||||
`mxfp4_scale_alg` selects the MXFP4 block-scaling method on NPU.
|
||||
UOS is a scaling rule, not a separate activation dtype or an activation-only
|
||||
algorithm. Both choices use FP4 E2M1 values with shared E8M0 scales.
|
||||
|
||||
| Mode | Weight preparation | A4 activation quantization | W4A8 fallback |
|
||||
| --- | --- | --- | --- |
|
||||
| Online | Selected scaling method, once at loading | Selected scaling method | Same W4; OCP MXFP8 A8 |
|
||||
| Offline | Load existing W4 and scales unchanged | Selected scaling method | Same W4; OCP MXFP8 A8 |
|
||||
|
||||
The available settings are:
|
||||
|
||||
| Setting | MXFP4 scaling method | NPU parameters |
|
||||
| --- | --- | --- |
|
||||
| `0` (default) | OCP MX | `scale_alg=0` |
|
||||
| `2` | UOS (`Qmax=7.25`) | `scale_alg=2, dst_type_max=7.25` |
|
||||
|
||||
Universal Optimal Scaling (UOS, `Qmax=7.25`) is derived in
|
||||
[MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://arxiv.org/abs/2607.24377),
|
||||
Section 4.1. The value `7.25` is the boundary used to select the shared
|
||||
power-of-two scale; the maximum finite E2M1 value remains `6`.
|
||||
This option applies UOS to Linear quantization. It does not implement the
|
||||
paper's full MXAttention pipeline or Pre-Normalization Quantization (PNQ).
|
||||
|
||||
Both use FP4 E2M1 with E8M0 scales, `axis=-1`, `block_size=32`,
|
||||
`round_mode="rint"`. W4A8 always quantizes the activation to
|
||||
`float8_e4m3fn` with `scale_alg=0` and the same grouping/rounding, without
|
||||
UOS parameters. A4 and A8 share the prepared W4 and scale. UOS is currently
|
||||
NPU-only; ROCm retains its existing AITER online algorithm.
|
||||
|
||||
An explicit UOS online configuration is
|
||||
`{"method": "mxfp4", "mxfp4_scale_alg": 2}`. For offline checkpoints the
|
||||
setting affects **only A4 activations**. Existing weights are never requantized:
|
||||
a UOS checkpoint can run UOS W4 + UOS A4 or UOS W4 + OCP A4, and both fall back to
|
||||
UOS W4 + OCP A8. Record the weight-generation algorithm separately from the
|
||||
runtime activation setting. Smooth remains active with either setting. The
|
||||
UOS configuration explicitly selects `2`; attention computation is unchanged.
|
||||
|
||||
### `mxfp4` — Single-Scale Online Mode
|
||||
|
||||
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
|
||||
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
|
||||
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
|
||||
calibration `mul_scale` is available. Applies equally to single-transformer
|
||||
and cascade (A14B) models — both transformers in a cascade receive the same
|
||||
quantization config automatically.
|
||||
Online mode quantizes BF16/FP16 weights once at loading and computes activation
|
||||
scales during forward. Each block of 32 K elements shares one E8M0 scale;
|
||||
no calibrated `mul_scale` is available. An explicit global quantization config
|
||||
applies to both cascade experts. Per-component configs can select different
|
||||
policies; checkpoint auto-detection resolves each expert independently.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
@@ -103,7 +289,7 @@ outputs = omni.generate(
|
||||
|
||||
```bash
|
||||
# Single-transformer or cascade model — same command
|
||||
python text_to_video.py --model <your-model> --quantization mxfp4
|
||||
python examples/offline_inference/text_to_video/text_to_video.py --model <your-model> --quantization mxfp4
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization mxfp4
|
||||
@@ -113,19 +299,19 @@ vllm serve <your-model> --omni --quantization mxfp4
|
||||
|
||||
Online DualScale mode computes both fine and coarse scales on the fly from BF16
|
||||
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
|
||||
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
|
||||
DualScale provides better quantization accuracy at higher compute cost.
|
||||
single-transformer and cascade (A14B) models with compatible Linear layers.
|
||||
Compare quality and performance using the intended model and workload.
|
||||
|
||||
The default configuration keeps the leading 5 transformer blocks in BF16
|
||||
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
|
||||
is sufficient to meet quality requirements and is the recommended setting.
|
||||
(`num_bf16_fallback_layers=5`). This is an implementation default; select and
|
||||
validate the BF16 policy against your own quality target.
|
||||
|
||||
```python
|
||||
omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
|
||||
```
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
|
||||
python examples/offline_inference/text_to_video/text_to_video.py --model <your-model> --quantization mxfp4_dualscale
|
||||
```
|
||||
|
||||
If accuracy debugging identifies additional precision-sensitive layers, they
|
||||
@@ -136,21 +322,22 @@ omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "mxfp4_dualscale",
|
||||
"ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
|
||||
"ignored_layers": ["blocks.10.attn1.to_qkv"], # explicit per-layer override
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
BF16 fallback routing in online mode applies two rules in priority order:
|
||||
|
||||
1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
|
||||
matches is kept in BF16 regardless of block index.
|
||||
1. **`ignored_layers`** (explicit per-layer override): a listed runtime Linear
|
||||
module is kept in BF16 regardless of block index. Use exact module paths,
|
||||
not arbitrary block prefixes or names ending in `.weight`.
|
||||
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
|
||||
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
|
||||
to `5` (recommended). Layers outside `blocks.N.*`
|
||||
(e.g. `condition_embedder`) are always quantized.
|
||||
to `5`. Layers outside `blocks.N.*` are quantized unless explicitly listed
|
||||
in `ignored_layers`.
|
||||
|
||||
### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
|
||||
### `mxfp4_dualscale` — DualScale Offline Mode
|
||||
|
||||
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
|
||||
preprocessing step converts the raw quantized output to the diffusers format
|
||||
@@ -169,8 +356,8 @@ inside `config.json`. At runtime, each layer's prefix is matched against
|
||||
Each quantized linear layer stores four tensors:
|
||||
|
||||
| Tensor | Shape | dtype | Description |
|
||||
|--------|-------|-------|-------------|
|
||||
| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
|
||||
| -------- | ------- | ------- | ------------- |
|
||||
| `weight` | `(N, K)` | float8_e4m3fn | Numeric FP4 values; packed into two values per byte after loading |
|
||||
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
|
||||
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
|
||||
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
|
||||
@@ -223,7 +410,7 @@ python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
|
||||
```
|
||||
|
||||
| Argument | Description |
|
||||
|----------|-------------|
|
||||
| ---------- | ------------- |
|
||||
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
|
||||
| `--original-model` | Root directory of the original BF16 diffusers model |
|
||||
| `--quant-path` | Root directory of the msModelSlim quantized output |
|
||||
@@ -251,14 +438,14 @@ The `quantization_config` injected into `config.json` for each transformer:
|
||||
```
|
||||
|
||||
`ignored_layers` lists every linear layer that retains its original BF16 weight,
|
||||
using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
|
||||
using vllm-omni runtime module paths (QKV-fused, FFN underscored, `to_out`
|
||||
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
|
||||
and may differ between `transformer` and `transformer_2` in a cascade model.
|
||||
|
||||
#### Step 3 — Serve
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
|
||||
python examples/offline_inference/text_to_video/text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
|
||||
|
||||
# Online serving
|
||||
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
|
||||
@@ -275,25 +462,32 @@ omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
|
||||
|
||||
## Parameters
|
||||
|
||||
### `mxfp4` (single-scale, online only)
|
||||
### `mxfp4` (single-scale, online + offline)
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `method` | str | — | `"mxfp4"` |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
|
||||
| `mxfp4_scale_alg` | int | `0` | MXFP4 block scaling: `0` OCP MX, `2` UOS (`Qmax=7.25`, NPU only). Online: W4 and A4; offline: A4 only. A8 remains OCP MXFP8 |
|
||||
| `is_checkpoint_mxfp4_serialized` | bool | `false` | Load pre-quantized single-scale weights and scales; declared in checkpoint `config.json` |
|
||||
| `require_smooth_scale` | bool | `false` | Require calibrated Smooth tensors for offline single-scale weights |
|
||||
| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16 |
|
||||
| `w4a8_fallback_layers` | list[str] | `[]` | Exact runtime Linear paths that always use A8; NPU only |
|
||||
| `w4a8_fallback_steps` | list[int] | `[]` | Zero-based denoise steps that use A8; NPU only |
|
||||
|
||||
### `mxfp4_dualscale` (dual-scale, online + offline)
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `method` | str | — | `"mxfp4_dualscale"` |
|
||||
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
|
||||
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
|
||||
| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
|
||||
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Validate the default of `5` against the intended model and quality target |
|
||||
| `w4a8_fallback_layers` | list[str] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
|
||||
| `w4a8_fallback_steps` | list[int] | `[]` | Must remain empty; DualScale does not support W4A8 fallback |
|
||||
|
||||
#### BF16 fallback priority (online mode)
|
||||
|
||||
```
|
||||
```text
|
||||
for each linear layer:
|
||||
if prefix in ignored_layers → BF16 (explicit override, highest priority)
|
||||
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
|
||||
@@ -307,8 +501,9 @@ unless they appear in `ignored_layers`.
|
||||
|
||||
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
|
||||
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
|
||||
available — all output partitions receive an identity pre-scale. No offline
|
||||
checkpoint format exists for this method.
|
||||
available — all output partitions receive an identity pre-scale. Offline
|
||||
single-scale checkpoints use the canonical numeric FP4 layout
|
||||
described above.
|
||||
|
||||
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
|
||||
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
|
||||
@@ -316,7 +511,7 @@ unless they appear in `ignored_layers`.
|
||||
explicit `ignored_layers` stay in BF16 for accuracy.
|
||||
|
||||
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
|
||||
**recommended for production** — loads four tensors per quantized layer: FP4
|
||||
loads four tensors per quantized layer: FP4
|
||||
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
|
||||
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
|
||||
layers have no quantization tensors and are routed via `ignored_layers`.
|
||||
@@ -325,21 +520,21 @@ unless they appear in `ignored_layers`.
|
||||
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
|
||||
float32 round-trip.
|
||||
|
||||
5. **Cascade model config propagation**: in a cascade model (transformer +
|
||||
transformer_2), vLLM-Omni reads each transformer's own `config.json` and
|
||||
rebuilds the quant config locally when `ignored_layers` differs between
|
||||
transformers, ensuring per-layer routing is accurate for each. The first
|
||||
transformer's config is propagated to `od_config` so the second transformer
|
||||
can reuse it as a starting point.
|
||||
5. **Cascade model configuration**: checkpoint auto-detection resolves each
|
||||
expert's own `config.json` independently; the first expert's detected config
|
||||
is not reused as the second expert's policy. An explicit global runtime
|
||||
config applies to both experts, while per-component configs select policies
|
||||
by the exact `transformer` / `transformer_2` names. Serialized checkpoint
|
||||
metadata still determines each expert's storage format and BF16 layers.
|
||||
|
||||
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
|
||||
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
|
||||
fused name (`attn1.to_qkv`), written automatically by the merge script.
|
||||
|
||||
7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
|
||||
DualScale offline method mitigates this with calibrated `mul_scale` smooth
|
||||
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
|
||||
off compression vs. accuracy for precision-sensitive layers.
|
||||
7. FP4 has fewer representable values than FP8, so evaluate quantization
|
||||
error and generation quality for the chosen recipe. Offline Smooth scaling
|
||||
and BF16 fallback can help with sensitive layers.
|
||||
`num_bf16_fallback_layers` applies only to online DualScale.
|
||||
|
||||
## Adapting MXFP4 for a New Model
|
||||
|
||||
@@ -350,27 +545,26 @@ a merge script for offline checkpoints.
|
||||
|
||||
### Step 1 — Discover runtime layer names
|
||||
|
||||
`ignored_layers` entries must match the **runtime parameter names** used inside
|
||||
vllm-omni, which may differ from the names stored in the diffusers checkpoint.
|
||||
The canonical source of truth is the model's own `named_parameters()`.
|
||||
`ignored_layers` entries use exact **runtime Linear module paths**, without
|
||||
`.weight`. Inspect the actual transformer instance in model initialization or
|
||||
a developer harness; `Omni` does not expose an `omni.pipeline.transformer` API.
|
||||
For a transformer instance named `transformer`:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
|
||||
# Load the model without quantization to inspect parameter names.
|
||||
omni = Omni(model="/path/to/your-model") # no --quantization flag
|
||||
for name, _ in omni.pipeline.transformer.named_parameters():
|
||||
if "weight" in name and "scale" not in name:
|
||||
for name, module in transformer.named_modules():
|
||||
if isinstance(module, LinearBase):
|
||||
print(name)
|
||||
```
|
||||
|
||||
Compare the printed names against the diffusers checkpoint keys
|
||||
Compare these module paths against the corresponding diffusers checkpoint keys
|
||||
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
|
||||
applies. Common patterns that differ in Wan2.2 (and may appear in other
|
||||
models):
|
||||
|
||||
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
|
||||
|---------------------------|------------------------|--------|
|
||||
| --------------------------- | ------------------------ | -------- |
|
||||
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
|
||||
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
|
||||
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
|
||||
@@ -382,8 +576,9 @@ fused runtime parameters.
|
||||
|
||||
!!! warning "Partial QKV fallback is not allowed"
|
||||
If your model fuses Q, K, V into a single layer, `ignored_layers` must
|
||||
include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
|
||||
`to_k`, `to_v` quantized) cannot be expressed at runtime because they share
|
||||
name the fused runtime module (for Wan, `attn1.to_qkv`). The source
|
||||
checkpoint must keep **all three or none** in BF16. A partial fallback
|
||||
(e.g. `to_q` in BF16 but `to_k`, `to_v` quantized) cannot be expressed at runtime because they share
|
||||
one `QKVParallelLinear`. The merge script enforces this and raises an error
|
||||
if only some of the trio appear as non-quantized.
|
||||
|
||||
@@ -392,12 +587,13 @@ fused runtime parameters.
|
||||
#### Online mode
|
||||
|
||||
Pass `ignored_layers` directly in the quantization config using the **runtime
|
||||
names** discovered in Step 1. No code changes to the model are required.
|
||||
names** discovered in Step 1. This assumes the model already passes its
|
||||
quantization config and module prefixes to compatible vLLM Linear layers.
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="/path/to/your-model",
|
||||
quantization={
|
||||
quantization_config={
|
||||
"method": "mxfp4_dualscale",
|
||||
"ignored_layers": [
|
||||
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
|
||||
@@ -409,9 +605,14 @@ omni = Omni(
|
||||
```
|
||||
|
||||
```bash
|
||||
# CLI does not support list-typed ignored_layers directly.
|
||||
# Use the Python API or set ignored_layers in config.json (offline).
|
||||
python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
|
||||
# The text-to-video example accepts list-valued fields through JSON.
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model /path/to/your-model \
|
||||
--quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
|
||||
|
||||
# Serving accepts the same configuration.
|
||||
vllm serve /path/to/your-model --omni \
|
||||
--diffusion-quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}'
|
||||
```
|
||||
|
||||
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
|
||||
@@ -436,9 +637,13 @@ the merge script is correct. The injected block:
|
||||
}
|
||||
```
|
||||
|
||||
To add a layer manually (e.g. to pin an additional layer to BF16 without
|
||||
re-running the merge script), edit `config.json` inside the transformer
|
||||
subfolder. Use runtime names, not diffusers checkpoint names.
|
||||
Adding an offline BF16 fallback layer requires both the original BF16 weight
|
||||
from the base model and a matching `ignored_layers` entry. Regenerate the
|
||||
checkpoint with the revised storage policy, restoring the original weight
|
||||
(and bias, if present) and removing that layer's quantization tensors.
|
||||
Editing `config.json` alone does not reconstruct BF16 weights from stored FP4
|
||||
values. For fused QKV, restore all three projections together. Use runtime
|
||||
module paths in `ignored_layers`.
|
||||
|
||||
### Step 3 — Write a merge script for offline mode
|
||||
|
||||
@@ -474,6 +679,7 @@ renaming rules and collect the results.
|
||||
!!! tip "Validate before serving"
|
||||
After producing the offline checkpoint, load it without a `--quantization`
|
||||
flag and verify that vLLM-Omni auto-detects the correct method. Check that
|
||||
the layer count reported in the startup log matches expectations: quantized
|
||||
layer count + `ignored_layers` count should equal total linear layer count.
|
||||
Any mismatch indicates a name-mapping bug in the merge script.
|
||||
each intended runtime Linear module resolves to the expected quantized or
|
||||
BF16 method and that every BF16 entry has original BF16 weights. Account
|
||||
for fused modules and, in online DualScale, the leading-block fallback rule;
|
||||
the raw length of `ignored_layers` alone is not a reliable layer count.
|
||||
@@ -384,13 +384,20 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
quantization_group = parser.add_mutually_exclusive_group()
|
||||
quantization_group.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
quantization_group.add_argument(
|
||||
"--quantization-config",
|
||||
type=json.loads,
|
||||
default=None,
|
||||
help='Quantization JSON, e.g. \'{"method":"mxfp4","w4a8_fallback_steps":[0,2]}\'.',
|
||||
)
|
||||
|
||||
# Distributed and parallel execution
|
||||
parser.add_argument(
|
||||
@@ -572,6 +579,8 @@ def main():
|
||||
omni_kwargs["flow_shift"] = args.flow_shift
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
if args.quantization_config is not None:
|
||||
omni_kwargs["quantization_config"] = args.quantization_config
|
||||
if args.cache_backend is not None:
|
||||
omni_kwargs["cache_backend"] = args.cache_backend
|
||||
omni_kwargs["cache_config"] = cache_config
|
||||
|
||||
@@ -1480,7 +1480,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path))
|
||||
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
|
||||
assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
|
||||
|
||||
@@ -1739,3 +1739,31 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
|
||||
builder(pipeline, deploy)
|
||||
else:
|
||||
builder(pipeline, user_deploy_config=deploy)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit", [False, True])
|
||||
def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
checkpoint = TransformerConfig.from_dict(
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"w4a8_fallback_steps": [37],
|
||||
}
|
||||
}
|
||||
)
|
||||
requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
||||
tf_model_config=checkpoint,
|
||||
quantization_config=requested,
|
||||
)
|
||||
assert cfg.quantization_config_is_auto_detected is not explicit
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint))
|
||||
cfg.enrich_config()
|
||||
restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
|
||||
assert restored.quantization_config_is_auto_detected is not explicit
|
||||
assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37])
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Regression tests for transformer quant-config auto-detection.
|
||||
|
||||
The loader path at pipeline_wan2_2.py carries two quantization contracts:
|
||||
@@ -27,6 +27,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
|
||||
|
||||
# od_config must remain unchanged — _create_transformer does not modify it.
|
||||
assert od_config.quantization_config is None
|
||||
|
||||
|
||||
def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
|
||||
import json
|
||||
|
||||
import torch
|
||||
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
|
||||
|
||||
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
|
||||
for component in ("transformer", "transformer_2"):
|
||||
directory = tmp_path / component
|
||||
directory.mkdir()
|
||||
config = dict(_MIN_CFG)
|
||||
expert_quant_config = (
|
||||
disk_quant_config[0 if component == "transformer" else 1]
|
||||
if isinstance(disk_quant_config, tuple)
|
||||
else disk_quant_config
|
||||
)
|
||||
if expert_quant_config is not None:
|
||||
config["quantization_config"] = expert_quant_config
|
||||
(directory / "config.json").write_text(json.dumps(config))
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
|
||||
|
||||
class FakeComponent(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
|
||||
|
||||
class FakeTransformer(FakeComponent):
|
||||
def __init__(self, quant_config=None, **kwargs):
|
||||
super().__init__()
|
||||
self.received_quant_config = quant_config
|
||||
|
||||
for module in (wan22_module, i2v_module):
|
||||
monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
|
||||
monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
|
||||
pipeline_cls = {
|
||||
"t2v": wan22_module.Wan22Pipeline,
|
||||
"i2v": i2v_module.Wan22I2VPipeline,
|
||||
"vace": vace_module.Wan22VACEPipeline,
|
||||
}[pipeline_kind]
|
||||
monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
|
||||
od_config = OmniDiffusionConfig(
|
||||
model=str(tmp_path),
|
||||
model_class_name=pipeline_cls.__name__,
|
||||
dtype=torch.bfloat16,
|
||||
quantization_config=active,
|
||||
boundary_ratio=0.875,
|
||||
flow_shift=12.0,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
od_config.enrich_config()
|
||||
pipeline = pipeline_cls(od_config=od_config)
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
def test_pipeline_constructor_routes_each_expert_quantization(
|
||||
pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
|
||||
):
|
||||
"""Run the actual pipeline constructors, expert factories and quant-method dispatch."""
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
active = build_quant_config(
|
||||
{
|
||||
"transformer": {
|
||||
"method": method,
|
||||
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [0, 1],
|
||||
},
|
||||
"transformer_2": {
|
||||
"method": method,
|
||||
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [],
|
||||
},
|
||||
}
|
||||
)
|
||||
disk_quant_config = None
|
||||
if offline:
|
||||
disk_quant_config = {
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.12.ffn.net_2"],
|
||||
"w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [9],
|
||||
}
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
|
||||
for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
|
||||
received = transformer.received_quant_config
|
||||
assert received.get_name() == method
|
||||
assert getattr(received, flag) == offline
|
||||
assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
|
||||
selected = f"blocks.{10 + index}.attn1.to_qkv"
|
||||
other = f"blocks.{11 - index}.attn1.to_qkv"
|
||||
assert received.w4a8_fallback_layers == [selected]
|
||||
assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
|
||||
assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
|
||||
if offline:
|
||||
assert isinstance(
|
||||
received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
|
||||
)
|
||||
assert pipeline.od_config.quantization_config is active
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit_none", [False, True])
|
||||
def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
FakeTransformer, captured = _make_fake_transformer()
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
|
||||
active = build_quant_config(components)
|
||||
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
|
||||
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
|
||||
assert captured[-1].get("quant_config") is None
|
||||
with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
|
||||
pipeline._create_transformer(
|
||||
{**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
|
||||
component="transformer_2",
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
FakeTransformer, captured = _make_fake_transformer()
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
|
||||
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
|
||||
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
|
||||
assert captured[-1]["quant_config"] is active.default_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize(
|
||||
"configured_expert, with_default",
|
||||
[
|
||||
("transformer", False),
|
||||
("transformer", True),
|
||||
("transformer_2", False),
|
||||
("transformer_2", True),
|
||||
],
|
||||
)
|
||||
def test_pipeline_partial_expert_config_does_not_match_sibling(
|
||||
pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
|
||||
):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
component_spec = {
|
||||
configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
}
|
||||
if with_default:
|
||||
component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
|
||||
active = build_quant_config(component_spec)
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
|
||||
sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
|
||||
selected = getattr(pipeline, configured_expert).received_quant_config
|
||||
other = getattr(pipeline, sibling).received_quant_config
|
||||
assert selected is active.component_configs[configured_expert]
|
||||
assert other is active.default_config
|
||||
if with_default:
|
||||
assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
|
||||
else:
|
||||
assert other is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize("different_ignored_layers", [False, True])
|
||||
def test_enriched_auto_quantization_preserves_each_expert_policy(
|
||||
pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
|
||||
):
|
||||
"""Startup auto-detection must not turn the first expert into a user override."""
|
||||
high = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ["blocks.0.ffn.net_2"],
|
||||
"w4a8_fallback_steps": [],
|
||||
"w4a8_fallback_layers": [],
|
||||
"mxfp4_scale_alg": 0,
|
||||
"require_smooth_scale": False,
|
||||
}
|
||||
low = {
|
||||
**high,
|
||||
"ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
|
||||
"w4a8_fallback_steps": [37],
|
||||
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
|
||||
"mxfp4_scale_alg": 2,
|
||||
"require_smooth_scale": True,
|
||||
}
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
|
||||
assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert.
|
||||
for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
|
||||
config = transformer.received_quant_config
|
||||
for policy in (
|
||||
"ignored_layers",
|
||||
"w4a8_fallback_steps",
|
||||
"w4a8_fallback_layers",
|
||||
"mxfp4_scale_alg",
|
||||
"require_smooth_scale",
|
||||
):
|
||||
assert getattr(config, policy) == expected[policy], policy
|
||||
@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cpu_platform(monkeypatch):
|
||||
def _cpu_pipeline_runtime(monkeypatch):
|
||||
# These lightweight pipelines have no accelerator or distributed groups.
|
||||
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
|
||||
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
|
||||
monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
|
||||
|
||||
|
||||
class _StubTransformer(nn.Module):
|
||||
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
|
||||
assert torch.equal(result, torch.full_like(latents, 10.0))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_first_request", [False, True])
|
||||
def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
get_forward_context,
|
||||
is_forward_context_available,
|
||||
override_forward_context,
|
||||
)
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
|
||||
|
||||
contexts: list[ForwardContext] = []
|
||||
events = []
|
||||
|
||||
class InjectedPredictionError(Exception):
|
||||
pass
|
||||
|
||||
class RecordingTransformer(_StubTransformer):
|
||||
def __init__(self, expert):
|
||||
super().__init__()
|
||||
self.expert = expert
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states, **kwargs):
|
||||
context = get_forward_context()
|
||||
assert context is contexts[-1]
|
||||
request_index = len(contexts) - 1
|
||||
step = context.denoise_step_idx
|
||||
positive = bool(encoder_hidden_states[0, 0, 0] > 0)
|
||||
events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
|
||||
if fail_first_request and request_index == 0 and step == 1 and not positive:
|
||||
raise InjectedPredictionError
|
||||
return (torch.ones_like(hidden_states) * (1 if positive else -1),)
|
||||
|
||||
def encode_prompt(**kwargs):
|
||||
context = get_forward_context()
|
||||
assert context.denoise_step_idx is None
|
||||
contexts.append(context)
|
||||
embeds = torch.ones(1, 8, 8)
|
||||
return embeds, -embeds
|
||||
|
||||
pipeline = _make_pipeline()
|
||||
pipeline.transformer = RecordingTransformer("high")
|
||||
pipeline.transformer_2 = RecordingTransformer("low")
|
||||
pipeline.scheduler = _StubScheduler([900, 500, 100])
|
||||
monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
|
||||
monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
|
||||
|
||||
# Use the real runner context manager, Wan forward/diffuse, and CFG
|
||||
# dispatch. Only model computation, scheduler math and unrelated KV I/O
|
||||
# are replaced; no test code publishes a denoise step.
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.pipeline = pipeline
|
||||
runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
|
||||
runner.vllm_config = None
|
||||
runner.cache_backend = None
|
||||
monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
|
||||
with override_forward_context(None):
|
||||
for request_index in range(2):
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="a fox walks",
|
||||
request_id=f"request-{request_index}",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
|
||||
),
|
||||
)
|
||||
|
||||
def execute():
|
||||
return runner._execute_request_list(
|
||||
[request],
|
||||
od_config=runner.od_config,
|
||||
allow_single_output=True,
|
||||
require_request_batch_support=False,
|
||||
record_name="test_wan_step_context",
|
||||
record_output_peak_memory=False,
|
||||
)
|
||||
|
||||
if fail_first_request and request_index == 0:
|
||||
with pytest.raises(InjectedPredictionError):
|
||||
execute()
|
||||
else:
|
||||
execute()
|
||||
assert not is_forward_context_available()
|
||||
|
||||
assert len(contexts) == 2 and contexts[0] is not contexts[1]
|
||||
expected = [
|
||||
(request, step, "high" if step == 0 else "low", positive, step in (0, 2))
|
||||
for request in range(2)
|
||||
for step in range(2 if fail_first_request and request == 0 else 3)
|
||||
for positive in (True, False)
|
||||
]
|
||||
assert events == expected
|
||||
|
||||
|
||||
class _StubDMDScheduler:
|
||||
def __init__(self) -> None:
|
||||
self.predict_clean_calls: list[tuple[float, float, float]] = []
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""
|
||||
|
||||
import sys
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config.load import LoadConfig
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.quantization.mxfp4_config import (
|
||||
DiffusionMXFP4Config,
|
||||
DiffusionMXFP4DualScaleMixedConfig,
|
||||
NPUMxfp4DualScaleLinearMethod,
|
||||
NPUMxfp4LinearMethod,
|
||||
NPUMxfp4OnlineLinearMethod,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _patch_tp_state(monkeypatch):
|
||||
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
|
||||
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
|
||||
class _CheckpointLayer(torch.nn.Module):
|
||||
def __init__(self, method, dtype=torch.bfloat16):
|
||||
super().__init__()
|
||||
method.create_weights(
|
||||
self,
|
||||
input_size_per_partition=512,
|
||||
output_partition_sizes=[2],
|
||||
input_size=512,
|
||||
output_size=2,
|
||||
params_dtype=dtype,
|
||||
weight_loader=default_weight_loader,
|
||||
)
|
||||
|
||||
def load_weights(self, weights):
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, value in weights:
|
||||
if name in params:
|
||||
params[name].weight_loader(params[name], value)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
|
||||
|
||||
def _load_checkpoint(layer, config, weights, monkeypatch):
|
||||
od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
|
||||
loader.load_weights(layer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
|
||||
# Exercise the shared implementation directly as well as the online MRO
|
||||
# below: checkpoint requirements must follow the configured storage mode.
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
|
||||
if serialized:
|
||||
with pytest.raises(ValueError, match="Required weights.*weight_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
|
||||
config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
|
||||
method = NPUMxfp4OnlineLinearMethod(config)
|
||||
layer = _CheckpointLayer(method)
|
||||
assert set(dict(layer.named_parameters())) == {"weight"}
|
||||
assert layer.weight.is_meta
|
||||
|
||||
weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
|
||||
packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
|
||||
generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
|
||||
quantized_inputs = []
|
||||
|
||||
def quantize(value, **kwargs):
|
||||
quantized_inputs.append(value.clone())
|
||||
expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
|
||||
if scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert kwargs == expected
|
||||
return packed_weight, generated_scale
|
||||
|
||||
# Keep real creation, lazy loading, processing, and strict loader checks;
|
||||
# substitute only the device quantization operator for this CPU test.
|
||||
npu = ModuleType("torch_npu")
|
||||
npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", npu)
|
||||
_load_checkpoint(layer, config, {"weight": weight}, monkeypatch)
|
||||
|
||||
assert len(quantized_inputs) == 1
|
||||
torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
|
||||
torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
|
||||
torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
|
||||
assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(quantized_inputs) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("require_smooth", [False, True])
|
||||
@pytest.mark.parametrize("has_smooth", [False, True])
|
||||
def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
scale = torch.linspace(0.5, 3.0, 512)
|
||||
if has_smooth:
|
||||
weights["mul_scale"] = scale
|
||||
else:
|
||||
weights.pop("mul_scale")
|
||||
|
||||
if require_smooth and not has_smooth:
|
||||
with pytest.raises(ValueError, match="mul_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_scale", ["shape", "name"])
|
||||
def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
if bad_scale == "shape":
|
||||
weights["mul_scale"] = torch.ones(511)
|
||||
with pytest.raises(AssertionError):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
weights["multiply_scale"] = weights.pop("mul_scale")
|
||||
with pytest.raises(ValueError, match="mul_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("required", [False, True])
|
||||
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
|
||||
def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
weights["mul_scale"][7] = bad_value
|
||||
with pytest.raises(ValueError, match="finite, strictly positive"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
|
||||
def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
weights["mul_scale"] = bad_scale
|
||||
with pytest.raises(ValueError, match="one-dimensional floating-point"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
|
||||
def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
layer = _CheckpointLayer(method, torch.float16)
|
||||
layer.mul_scale.data.fill_(bad_value)
|
||||
# Validation must fail before any device packing or GEMM.
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
|
||||
with pytest.raises(ValueError, match="finite, strictly positive"):
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "missing"),
|
||||
[
|
||||
("single", "weight"),
|
||||
("single", "weight_scale"),
|
||||
("dual", "weight"),
|
||||
("dual", "weight_scale"),
|
||||
("dual", "weight_dual_scale"),
|
||||
("dual", "mul_scale"),
|
||||
],
|
||||
)
|
||||
def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
|
||||
if mode == "single":
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
else:
|
||||
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
|
||||
method = NPUMxfp4DualScaleLinearMethod(config)
|
||||
layer = _CheckpointLayer(method)
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
|
||||
with pytest.raises(ValueError, match=missing):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
|
||||
from vllm.model_executor.layers.linear import RowParallelLinear
|
||||
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
|
||||
layer.tp_rank = rank
|
||||
checkpoint_scale = torch.arange(1024, dtype=torch.float32)
|
||||
RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
|
||||
torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("fallback", [False, True])
|
||||
@pytest.mark.parametrize("with_bias", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_single_scale_smooth_precedes_quantization_and_preserves_output(
|
||||
dtype, fallback, with_bias, scale_alg, monkeypatch
|
||||
):
|
||||
config = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
|
||||
)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
layer = _CheckpointLayer(method, dtype)
|
||||
weights = {
|
||||
"weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
|
||||
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
|
||||
"mul_scale": torch.linspace(0.5, 3.0, 512),
|
||||
}
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
|
||||
bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
|
||||
quantized_inputs = []
|
||||
|
||||
# Only the two device operators are substituted. Exercise real apply(),
|
||||
# step selection, Smooth multiplication, bias handling, and reshape logic.
|
||||
npu = ModuleType("torch_npu")
|
||||
|
||||
def quantize(value, **kwargs):
|
||||
quantized_inputs.append(value.clone())
|
||||
expected = dict(
|
||||
dst_type=torch.float8_e4m3fn if fallback else "fp4",
|
||||
axis=-1,
|
||||
block_size=32,
|
||||
round_mode="rint",
|
||||
scale_alg=0 if fallback else scale_alg,
|
||||
)
|
||||
if not fallback and scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert kwargs == expected
|
||||
return value, torch.ones((6, 16), dtype=torch.uint8)
|
||||
|
||||
def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
|
||||
if fallback:
|
||||
assert output_dtype == torch.bfloat16
|
||||
if bias is not None:
|
||||
assert bias.dtype == output_dtype
|
||||
assert bias.shape == (1, 2)
|
||||
result = x_q.float() @ weight.float()
|
||||
if bias is not None:
|
||||
result += bias.float()
|
||||
return result.to(output_dtype)
|
||||
|
||||
npu.__dict__.update(
|
||||
float8_e4m3fn=torch.float8_e4m3fn,
|
||||
float8_e8m0fnu=torch.uint8,
|
||||
float4_e2m1fn_x2="fp4",
|
||||
npu_dynamic_mx_quant=quantize,
|
||||
npu_quant_matmul=matmul,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", npu)
|
||||
with override_forward_context(ForwardContext()):
|
||||
set_forward_context_denoise_step_idx(0 if fallback else 1)
|
||||
result = method.apply(layer, x, bias)
|
||||
|
||||
expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
|
||||
assert len(quantized_inputs) == 1
|
||||
torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
|
||||
expected = expected_input.float() @ weights["weight"].float().T
|
||||
if bias is not None:
|
||||
expected += bias.float()
|
||||
if fallback:
|
||||
expected = expected.bfloat16()
|
||||
torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
|
||||
assert result.dtype == dtype
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
|
||||
|
||||
import pytest
|
||||
@@ -27,12 +27,42 @@ def test_mxfp4_config_get_name():
|
||||
assert DiffusionMXFP4Config.get_name() == "mxfp4"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3])
|
||||
def test_mxfp4_scale_alg_rejects_invalid_values(value):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="mxfp4_scale_alg"):
|
||||
build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("runtime_alg", [0, 2])
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg)
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"mxfp4_scale_alg": 2,
|
||||
"ignored_layers": ["proj_out"],
|
||||
},
|
||||
)
|
||||
assert resolved.is_checkpoint_mxfp4_serialized
|
||||
assert resolved.mxfp4_scale_alg == runtime_alg
|
||||
assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == []
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_defaults():
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
cfg = DiffusionMXFP4Config.from_config({})
|
||||
assert cfg.is_checkpoint_mxfp4_serialized is False
|
||||
assert cfg.ignored_layers == []
|
||||
assert cfg.w4a8_fallback_layers == []
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_serialized():
|
||||
@@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized():
|
||||
assert cfg.is_checkpoint_mxfp4_serialized is True
|
||||
|
||||
|
||||
def test_mxfp4_config_requires_offline_smooth():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True}
|
||||
assert DiffusionMXFP4Config.from_config(values).require_smooth_scale
|
||||
assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale
|
||||
with pytest.raises(ValueError, match="offline"):
|
||||
build_quant_config({"method": "mxfp4", "require_smooth_scale": True})
|
||||
with pytest.raises(ValueError, match="boolean"):
|
||||
DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1])
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"require_smooth_scale": True,
|
||||
},
|
||||
)
|
||||
assert resolved.require_smooth_scale
|
||||
assert resolved.is_checkpoint_mxfp4_serialized
|
||||
assert resolved.w4a8_fallback_steps == [1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disk_requirement", [None, False, True])
|
||||
def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
||||
disk = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
}
|
||||
if disk_requirement is not None:
|
||||
disk["require_smooth_scale"] = disk_requirement
|
||||
resolved = resolve_quant_config_from_disk(active, disk)
|
||||
assert resolved.require_smooth_scale
|
||||
assert resolved.ignored_layers == disk["ignored_layers"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}])
|
||||
def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(
|
||||
ignored_layers=["blocks.0.attn2.to_q"],
|
||||
w4a8_fallback_steps=[1],
|
||||
w4a8_fallback_layers=["blocks.1.attn1.to_qkv"],
|
||||
)
|
||||
|
||||
resolved = resolve_quant_config_from_disk(active, disk_qc)
|
||||
|
||||
assert resolved is active
|
||||
assert resolved.ignored_layers == ["blocks.0.attn2.to_q"]
|
||||
assert resolved.w4a8_fallback_steps == [1]
|
||||
assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"]
|
||||
|
||||
|
||||
def test_quantization_disk_string_rejects_mismatched_active_method():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = build_quant_config("mxfp4")
|
||||
|
||||
with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"):
|
||||
resolve_quant_config_from_disk(active, "mxfp4_dualscale")
|
||||
|
||||
|
||||
def test_quantization_disk_string_accepts_equivalent_method_alias():
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
class ActiveAliasConfig:
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "inc"
|
||||
|
||||
active = ActiveAliasConfig()
|
||||
|
||||
assert resolve_quant_config_from_disk(active, "auto-round") is active
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_ignored_layers():
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
@@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
|
||||
assert cfg.ignored_layers == ["proj_out"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"policy",
|
||||
[
|
||||
{"w4a8_fallback_steps": [0]},
|
||||
{"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
{"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
],
|
||||
)
|
||||
def test_dualscale_rejects_w4a8_configuration(serialized, policy):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
|
||||
|
||||
config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy}
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy)
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
DiffusionMXFP4DualScaleMixedConfig.from_config(config)
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
build_quant_config({"transformer": config})
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
resolve_quant_config_from_disk(None, config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
|
||||
@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)])
|
||||
def test_mxfp4_rejects_invalid_fallback_steps(method, steps):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="non-negative integer"):
|
||||
build_quant_config({"method": method, "w4a8_fallback_steps": steps})
|
||||
|
||||
|
||||
def test_mxfp4_fallback_steps_follow_request_context():
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
|
||||
|
||||
steps = [2, 0, 2]
|
||||
cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps})
|
||||
steps.append(1)
|
||||
assert cfg.w4a8_fallback_steps == [0, 2]
|
||||
with override_forward_context(None):
|
||||
assert not _is_w4a8_fallback_step([])
|
||||
with pytest.raises(RuntimeError, match="publishes its denoise step"):
|
||||
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
|
||||
for _request in range(2):
|
||||
with override_forward_context(ForwardContext()):
|
||||
for step, expected in [(0, True), (1, False), (2, True), (3, False)]:
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
# Both CFG passes and both experts read the same index.
|
||||
assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [
|
||||
expected,
|
||||
expected,
|
||||
]
|
||||
with pytest.raises(RuntimeError, match="publishes its denoise step"):
|
||||
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
|
||||
@pytest.mark.parametrize("steps", [[], [0, 2]])
|
||||
@pytest.mark.parametrize("already_serialized", [False, True])
|
||||
def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
layers = ["blocks.10.attn1.to_qkv"] if steps else []
|
||||
active = build_quant_config(
|
||||
{
|
||||
"method": method,
|
||||
flag: already_serialized,
|
||||
"w4a8_fallback_steps": steps,
|
||||
"w4a8_fallback_layers": layers,
|
||||
}
|
||||
)
|
||||
# Separate experts may have distinct high precision layers. Both must
|
||||
# retain the requested step policy during online/offline reconciliation.
|
||||
for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]:
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ignored,
|
||||
"w4a8_fallback_steps": [9],
|
||||
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
|
||||
},
|
||||
)
|
||||
assert getattr(resolved, flag)
|
||||
assert resolved.ignored_layers == ignored
|
||||
assert resolved.w4a8_fallback_steps == steps
|
||||
assert resolved.w4a8_fallback_layers == layers
|
||||
assert active.ignored_layers == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"runtime_policy",
|
||||
[{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}],
|
||||
)
|
||||
def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = build_quant_config({"method": "mxfp4", **runtime_policy})
|
||||
resolved_experts = []
|
||||
for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]):
|
||||
disk = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ignored,
|
||||
"w4a8_fallback_steps": [9],
|
||||
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
|
||||
}
|
||||
resolved = resolve_quant_config_from_disk(active, disk)
|
||||
assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", [])
|
||||
assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", [])
|
||||
assert resolved.ignored_layers == ignored
|
||||
detected = resolve_quant_config_from_disk(None, disk)
|
||||
assert detected.w4a8_fallback_steps == [9]
|
||||
assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"]
|
||||
resolved_experts.append(resolved)
|
||||
assert resolved_experts[0] is not resolved_experts[1]
|
||||
assert active.is_checkpoint_mxfp4_serialized is False
|
||||
assert active.ignored_layers == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
|
||||
)
|
||||
def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active_spec = {
|
||||
"method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
"w4a8_fallback_steps": [1] if method == "mxfp4" else [],
|
||||
"w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [],
|
||||
}
|
||||
if method == "mxfp4":
|
||||
active_spec["require_smooth_scale"] = True
|
||||
active = build_quant_config(active_spec)
|
||||
|
||||
resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True})
|
||||
|
||||
assert resolved is not active
|
||||
assert resolved.ignored_layers == []
|
||||
assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps
|
||||
assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers
|
||||
if method == "mxfp4":
|
||||
assert resolved.require_smooth_scale
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
|
||||
)
|
||||
@pytest.mark.parametrize("with_active", [False, True])
|
||||
def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = None
|
||||
disk_steps = [9] if method == "mxfp4" else []
|
||||
active_steps = [1] if method == "mxfp4" else []
|
||||
expected_steps = disk_steps
|
||||
if with_active:
|
||||
active = build_quant_config(
|
||||
{
|
||||
"method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
"w4a8_fallback_steps": active_steps,
|
||||
}
|
||||
)
|
||||
expected_steps = active_steps
|
||||
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"modules_to_not_convert": ["blocks.2.ffn.net_2"],
|
||||
"w4a8_fallback_steps": disk_steps,
|
||||
},
|
||||
)
|
||||
|
||||
assert resolved.ignored_layers == ["blocks.2.ffn.net_2"]
|
||||
assert resolved.w4a8_fallback_steps == expected_steps
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_quant_config integration
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
|
||||
assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
|
||||
|
||||
|
||||
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
|
||||
@pytest.mark.parametrize("steps", [[], [0, 2]])
|
||||
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps):
|
||||
"""A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
|
||||
cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
|
||||
cfg = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps
|
||||
)
|
||||
|
||||
assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
|
||||
|
||||
|
||||
def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch):
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
|
||||
cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2])
|
||||
|
||||
with pytest.raises(NotImplementedError, match="only supported on NPU"):
|
||||
cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
|
||||
|
||||
|
||||
def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
|
||||
"""MXFP4 on ROCm requires gfx950; any other arch must raise."""
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
@@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
|
||||
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
|
||||
assert layer.input_size_per_partition == _TP2_K // _TP2
|
||||
assert layer.output_size_per_partition == _TP2_N
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
|
||||
@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]])
|
||||
def test_mxfp4_rejects_invalid_fallback_layers(method, layers):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
||||
build_quant_config({"method": method, "w4a8_fallback_layers": layers})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4"])
|
||||
def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.component_config import resolve_component_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"]
|
||||
components = build_quant_config(
|
||||
{"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}}
|
||||
)
|
||||
names.append("blocks.11.attn1.to_qkv")
|
||||
assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]]
|
||||
assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == []
|
||||
detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]})
|
||||
assert detected.w4a8_fallback_layers == [names[0]]
|
||||
assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kind, serialized_flag",
|
||||
[("mxfp4", "is_checkpoint_mxfp4_serialized")],
|
||||
)
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_mxfp4_layer_and_step_union_routes_actual_methods(
|
||||
kind, serialized_flag, offline, scale_alg, monkeypatch, mocker
|
||||
):
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.quantization import build_quant_config, mxfp4_config
|
||||
|
||||
monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True)
|
||||
selected = "blocks.10.attn1.to_qkv"
|
||||
other = "blocks.11.attn2.to_q"
|
||||
preserved = "blocks.12.ffn.net_2"
|
||||
config = build_quant_config(
|
||||
{
|
||||
"method": kind,
|
||||
"mxfp4_scale_alg": scale_alg,
|
||||
serialized_flag: offline,
|
||||
"w4a8_fallback_layers": [selected, preserved],
|
||||
"ignored_layers": [preserved],
|
||||
"w4a8_fallback_steps": [1],
|
||||
}
|
||||
)
|
||||
# Real config dispatch creates separate methods for selected and other layers.
|
||||
linear = mocker.Mock(spec=LinearBase)
|
||||
assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod)
|
||||
x = torch.ones(2, 64, dtype=torch.bfloat16)
|
||||
layer = torch.nn.Module()
|
||||
for name in ("weight", "weight_scale", "mul_scale"):
|
||||
layer.register_buffer(name, torch.ones(64, dtype=x.dtype))
|
||||
calls = []
|
||||
|
||||
def record_a8(*args):
|
||||
calls.append("A8")
|
||||
return x
|
||||
|
||||
def record_a4(*args):
|
||||
calls.append("A4")
|
||||
return x
|
||||
|
||||
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8)
|
||||
for name in (selected, other):
|
||||
method = config.get_quant_method(linear, name)
|
||||
assert method.is_w4a8_fallback_layer == (name == selected)
|
||||
scales = (x, x)
|
||||
monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales)
|
||||
monkeypatch.setattr(method, "_quant_matmul", record_a4)
|
||||
for step in (0, 1, 2):
|
||||
with override_forward_context(ForwardContext(denoise_step_idx=step)):
|
||||
method.apply(layer, x)
|
||||
assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"]
|
||||
# A permanently selected layer works without a request context, even if steps are also configured.
|
||||
with override_forward_context(None):
|
||||
config.get_quant_method(linear, selected).apply(layer, x)
|
||||
assert calls[-1] == "A8"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("online", [False, True])
|
||||
@pytest.mark.parametrize("selected", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch):
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.quantization import mxfp4_config
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
|
||||
prepared = []
|
||||
packed = []
|
||||
|
||||
def quantize(x, **kwargs):
|
||||
prepared.append(kwargs)
|
||||
return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8)
|
||||
|
||||
def pack(x, dtype):
|
||||
packed.append(x)
|
||||
return x.to(dtype)
|
||||
|
||||
fake_npu = SimpleNamespace(
|
||||
float4_e2m1fn_x2=torch.int8,
|
||||
float8_e8m0fnu=torch.uint8,
|
||||
float8_e4m3fn=torch.float8_e4m3fn,
|
||||
npu_dtype_cast=pack,
|
||||
npu_dynamic_mx_quant=quantize,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
|
||||
config = mxfp4_config.DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=not online,
|
||||
w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [],
|
||||
w4a8_fallback_steps=[1],
|
||||
mxfp4_scale_alg=scale_alg,
|
||||
)
|
||||
cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod
|
||||
method = cls(config, prefix="blocks.0.attn1.to_qkv")
|
||||
layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16)
|
||||
layer.orig_dtype = torch.bfloat16
|
||||
if not online:
|
||||
layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False))
|
||||
layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False))
|
||||
method.process_weights_after_loading(layer)
|
||||
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(prepared) == int(online)
|
||||
expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
|
||||
if scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert prepared == ([expected] if online else [])
|
||||
assert len(packed) == int(not online)
|
||||
branches = []
|
||||
|
||||
def a4(x_q, x_scale, target, bias, dtype):
|
||||
assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers
|
||||
branches.append("A4")
|
||||
return torch.zeros(x_q.shape[0], 2, dtype=dtype)
|
||||
|
||||
def a8(x, weight, scale, bias, dtype):
|
||||
assert (weight.data_ptr(), scale.data_ptr()) == pointers
|
||||
branches.append("A8")
|
||||
return torch.zeros(x.shape[0], 2, dtype=dtype)
|
||||
|
||||
monkeypatch.setattr(method, "_quant_matmul", a4)
|
||||
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8)
|
||||
monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x))
|
||||
monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization"))
|
||||
monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking"))
|
||||
for step in (0, 1, 2):
|
||||
with override_forward_context(ForwardContext(denoise_step_idx=step)):
|
||||
method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16))
|
||||
assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"])
|
||||
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
|
||||
assert set(layer.state_dict()) == (
|
||||
{"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"}
|
||||
)
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
assert not hasattr(layer, "w4a8_weight_scale")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("online", [False, True])
|
||||
def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch):
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.quantization import mxfp4_config
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
|
||||
original_ones = torch.ones
|
||||
monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"}))
|
||||
calls = []
|
||||
|
||||
def quantize_dual(x, **kw):
|
||||
calls.append(kw.get("smooth_scale"))
|
||||
return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8)
|
||||
|
||||
fake_npu = SimpleNamespace(
|
||||
float4_e2m1fn_x2=torch.int8,
|
||||
npu_dtype_cast=lambda x, dtype: x.to(dtype),
|
||||
npu_format_cast=lambda x, *args, **kw: x,
|
||||
npu_dynamic_dual_level_mx_quant=quantize_dual,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
|
||||
config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig(
|
||||
is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[]
|
||||
)
|
||||
cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod
|
||||
method = cls(config)
|
||||
layer = torch.nn.Module()
|
||||
layer.input_size_per_partition = 512
|
||||
for name, value in {
|
||||
"weight": torch.ones(2, 512, dtype=torch.bfloat16),
|
||||
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
|
||||
"weight_dual_scale": torch.ones(2, 1, 1),
|
||||
"mul_scale": torch.ones(512),
|
||||
}.items():
|
||||
layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False))
|
||||
method.process_weights_after_loading(layer)
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(calls) == int(online)
|
||||
assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"}
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16))
|
||||
method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16))
|
||||
assert len(calls) == 1 + int(online)
|
||||
torch.testing.assert_close(calls[-1], layer.mul_scale)
|
||||
@@ -0,0 +1,237 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config.load import LoadConfig
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _patch_tp_state(monkeypatch):
|
||||
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
|
||||
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
|
||||
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
|
||||
prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer"
|
||||
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0)
|
||||
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
|
||||
class _WanCheckpoint(torch.nn.Module):
|
||||
load_weights = WanTransformer3DModel.load_weights
|
||||
_log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary
|
||||
|
||||
def __init__(self, config, *, num_blocks=1):
|
||||
super().__init__()
|
||||
blocks = []
|
||||
for block_idx in range(num_blocks):
|
||||
block = torch.nn.Module()
|
||||
block.attn1 = torch.nn.Module()
|
||||
block.attn1.to_qkv = QKVParallelLinear(
|
||||
hidden_size=512,
|
||||
head_size=32,
|
||||
total_num_heads=4,
|
||||
bias=False,
|
||||
params_dtype=torch.bfloat16,
|
||||
quant_config=config,
|
||||
prefix=f"blocks.{block_idx}.attn1.to_qkv",
|
||||
disable_tp=True,
|
||||
)
|
||||
blocks.append(block)
|
||||
self.blocks = torch.nn.ModuleList(blocks)
|
||||
|
||||
|
||||
def _checkpoint(mode, fused):
|
||||
if mode == "single":
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
||||
elif mode == "dual":
|
||||
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
|
||||
else:
|
||||
config = None
|
||||
model = _WanCheckpoint(config)
|
||||
weights = {}
|
||||
expected = {}
|
||||
for name, param in model.named_parameters():
|
||||
suffix = name.rsplit(".", 1)[1]
|
||||
output_dim = getattr(param, "output_dim", None)
|
||||
if output_dim is None:
|
||||
value = torch.full_like(param, 2.5)
|
||||
shards = [value] * 3
|
||||
else:
|
||||
parts = param.chunk(3, dim=output_dim)
|
||||
shards = [
|
||||
torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts)
|
||||
]
|
||||
value = torch.cat(shards, dim=output_dim)
|
||||
expected[name] = value
|
||||
if fused:
|
||||
weights[name] = value
|
||||
else:
|
||||
for shard_id, shard in zip(("q", "k", "v"), shards):
|
||||
weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard
|
||||
return model, config, weights, expected
|
||||
|
||||
|
||||
def _load(model, config, weights, monkeypatch):
|
||||
loader = DiffusersPipelineLoader(
|
||||
LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
|
||||
)
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
|
||||
loader.load_weights(model)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
|
||||
@pytest.mark.parametrize("fused", [False, True])
|
||||
def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch):
|
||||
model, config, weights, expected = _checkpoint(mode, fused)
|
||||
_load(model, config, weights, monkeypatch)
|
||||
for name, param in model.named_parameters():
|
||||
torch.testing.assert_close(param, expected[name], rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
|
||||
def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch):
|
||||
model, config, weights, _ = _checkpoint(mode, fused=False)
|
||||
weights.pop("blocks.0.attn1.to_k.weight")
|
||||
weights.pop("blocks.0.attn1.to_v.weight")
|
||||
with pytest.raises(ValueError, match="to_qkv.weight"):
|
||||
_load(model, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["single", "dual"])
|
||||
def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch):
|
||||
model, config, weights, _ = _checkpoint(mode, fused=False)
|
||||
weights.pop("blocks.0.attn1.to_k.weight_scale")
|
||||
weights.pop("blocks.0.attn1.to_v.weight_scale")
|
||||
with pytest.raises(ValueError, match="to_qkv.weight_scale"):
|
||||
_load(model, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["single", "dual"])
|
||||
@pytest.mark.parametrize("reverse", [False, True])
|
||||
def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch):
|
||||
model, config, weights, _ = _checkpoint(mode, fused=False)
|
||||
weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2
|
||||
if reverse:
|
||||
weights = dict(reversed(weights.items()))
|
||||
with pytest.raises(ValueError, match="must share the same Smooth tensor"):
|
||||
_load(model, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"])
|
||||
def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name):
|
||||
model, _, _, _ = _checkpoint("single", fused=True)
|
||||
loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))])
|
||||
assert loaded == set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fused", [False, True])
|
||||
def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch):
|
||||
model, config, weights, _ = _checkpoint("single", fused=fused)
|
||||
projection = "to_qkv" if fused else "to_q"
|
||||
weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1)
|
||||
with pytest.raises(ValueError, match="cannot load DualScale tensor"):
|
||||
_load(model, config, weights, monkeypatch)
|
||||
|
||||
|
||||
def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override():
|
||||
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"])
|
||||
model = _WanCheckpoint(config)
|
||||
# Uses the same validation called at the end of WanTransformer3DModel construction.
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
||||
for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"):
|
||||
config.w4a8_fallback_layers = [bad]
|
||||
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
||||
|
||||
|
||||
def test_wan_layer_policy_defers_only_other_pp_stage_paths():
|
||||
from vllm.model_executor.models.utils import PPMissingLayer
|
||||
|
||||
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"])
|
||||
model = _WanCheckpoint(config)
|
||||
model.blocks.append(PPMissingLayer())
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
||||
config.w4a8_fallback_layers = ["blocks.1"]
|
||||
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"])
|
||||
def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker):
|
||||
selected = "blocks.0.attn1.to_qkv"
|
||||
bf16 = "blocks.1.attn1.to_qkv"
|
||||
transformer_config = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
|
||||
w4a8_fallback_layers=[selected],
|
||||
)
|
||||
transformer_2_config = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
|
||||
ignored_layers=[bf16],
|
||||
w4a8_fallback_layers=[selected, bf16],
|
||||
)
|
||||
pipeline = torch.nn.Module()
|
||||
pipeline.transformer = _WanCheckpoint(transformer_config)
|
||||
pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2)
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config)
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config)
|
||||
|
||||
loader = DiffusersPipelineLoader(
|
||||
LoadConfig(),
|
||||
OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config),
|
||||
)
|
||||
events = []
|
||||
summaries = []
|
||||
|
||||
def process_weights(model, _target_device):
|
||||
events.append("process")
|
||||
for module in model.modules():
|
||||
method = getattr(module, "quant_method", None)
|
||||
if getattr(method, "is_w4a8_fallback_layer", False):
|
||||
module._already_called_process_weights_after_loading = True
|
||||
|
||||
def capture_summary(message, *args):
|
||||
rendered = message % args
|
||||
summaries.append(rendered)
|
||||
events.append(f"summary:{args[0]}")
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info",
|
||||
side_effect=capture_summary,
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is pipeline
|
||||
assert events == ["load", "process", "summary:transformer", "summary:transformer_2"]
|
||||
assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0]
|
||||
assert "BF16 overrides=0 []" in summaries[0]
|
||||
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0]
|
||||
assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1]
|
||||
assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1]
|
||||
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1]
|
||||
|
||||
|
||||
def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker):
|
||||
selected = "blocks.0.attn1.to_qkv"
|
||||
config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected])
|
||||
model = _WanCheckpoint(config)
|
||||
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
||||
logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info")
|
||||
|
||||
WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer")
|
||||
|
||||
message, *args = logger_info.call_args.args
|
||||
rendered = message % tuple(args)
|
||||
assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered
|
||||
@@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _local_model_paths(monkeypatch):
|
||||
# These tests exercise config transport, not model repository resolution.
|
||||
monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
|
||||
|
||||
|
||||
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
|
||||
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
|
||||
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import importlib.util
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -24,6 +26,51 @@ def _load_text_to_video():
|
||||
return module
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arguments", "expected"),
|
||||
[
|
||||
(["--quantization", "mxfp4"], {"quantization": "mxfp4"}),
|
||||
(
|
||||
["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'],
|
||||
{"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch):
|
||||
mod = _load_text_to_video()
|
||||
captured = {}
|
||||
|
||||
class StopBeforeModelLoadError(Exception):
|
||||
pass
|
||||
|
||||
def capture_omni(**kwargs):
|
||||
captured.update(kwargs)
|
||||
raise StopBeforeModelLoadError
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments])
|
||||
monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu"))
|
||||
monkeypatch.setattr(mod, "Omni", capture_omni)
|
||||
|
||||
with pytest.raises(StopBeforeModelLoadError):
|
||||
mod.main()
|
||||
|
||||
assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected
|
||||
|
||||
|
||||
def test_quantization_cli_options_are_mutually_exclusive(monkeypatch):
|
||||
mod = _load_text_to_video()
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'],
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
mod.parse_args()
|
||||
|
||||
assert exc.value.code == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("negative_prompt", "expected_negative_prompt"),
|
||||
[(None, None), ("", ""), ("blurry", "blurry")],
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
torch_npu = pytest.importorskip("torch_npu")
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2):
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.quantization.mxfp4_config import (
|
||||
DiffusionMXFP4Config,
|
||||
NPUMxfp4LinearMethod,
|
||||
NPUMxfp4OnlineLinearMethod,
|
||||
)
|
||||
|
||||
prefix = "blocks.10.attn1.to_qkv"
|
||||
steps = [1] if policy in ("step", "layer_and_step") else []
|
||||
layers = [prefix] if policy in ("layer", "layer_and_step") else []
|
||||
fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg}
|
||||
torch.manual_seed(42)
|
||||
layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype)
|
||||
if mode == "single":
|
||||
method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix)
|
||||
elif mode == "single_offline":
|
||||
method = NPUMxfp4LinearMethod(
|
||||
DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix
|
||||
)
|
||||
layer.orig_dtype = dtype
|
||||
values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0])
|
||||
layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight))
|
||||
layer.register_parameter(
|
||||
"weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False)
|
||||
)
|
||||
layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False))
|
||||
else:
|
||||
raise ValueError(f"Unsupported single-scale test mode: {mode}")
|
||||
method.process_weights_after_loading(layer)
|
||||
x = torch.randn(input_shape, device="npu", dtype=dtype)
|
||||
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
||||
method.process_weights_after_loading(layer)
|
||||
quantize = torch_npu.npu_dynamic_mx_quant
|
||||
activation_types = []
|
||||
|
||||
def record_quantize(x, **kwargs):
|
||||
dst = kwargs["dst_type"]
|
||||
expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg
|
||||
assert kwargs["scale_alg"] == expected_alg
|
||||
assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint"
|
||||
if expected_alg == 2:
|
||||
assert kwargs["dst_type_max"] == 7.25
|
||||
else:
|
||||
assert "dst_type_max" not in kwargs
|
||||
activation_types.append(dst)
|
||||
return quantize(x, **kwargs)
|
||||
|
||||
monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize)
|
||||
results = []
|
||||
for _request in range(2):
|
||||
with override_forward_context(ForwardContext()):
|
||||
for step in range(3):
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
output = method.apply(layer, x, layer.bias)
|
||||
torch.npu.synchronize()
|
||||
assert output.shape == (*input_shape[:-1], 128)
|
||||
assert output.dtype == dtype
|
||||
assert torch.isfinite(output).all()
|
||||
results.append(output.cpu())
|
||||
expected_types = [
|
||||
torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3)
|
||||
]
|
||||
assert activation_types == expected_types * 2
|
||||
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
assert not hasattr(layer, "w4a8_weight_scale")
|
||||
for first, second in zip(results[:3], results[3:]):
|
||||
torch.testing.assert_close(first, second, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_bias", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)])
|
||||
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
||||
def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch):
|
||||
_exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step")
|
||||
|
||||
|
||||
@pytest.mark.cards_1
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@torch.inference_mode()
|
||||
def test_single_scale_step_fallback_matches_independent_dense_reference(dtype):
|
||||
if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name():
|
||||
pytest.skip("MXFP4 kernels require an Ascend950 NPU")
|
||||
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.quantization.mxfp4_config import (
|
||||
DiffusionMXFP4Config,
|
||||
NPUMxfp4LinearMethod,
|
||||
)
|
||||
|
||||
# Synthetic serialized single-scale values and E8M0 exponent bytes.
|
||||
# Powers of two make both activation precisions exact: this independent
|
||||
# CPU dense oracle needs no production packing, dequantization or apply.
|
||||
# Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate.
|
||||
values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512)
|
||||
weight = torch.cat((values, values), dim=1)
|
||||
fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8)
|
||||
smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0)))
|
||||
bias = torch.tensor([0.25, -0.25]).repeat(64)
|
||||
x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64
|
||||
expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias)
|
||||
|
||||
layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype)
|
||||
layer.weight.copy_(weight)
|
||||
layer.bias.copy_(bias)
|
||||
layer.orig_dtype = dtype
|
||||
for name, value in {"weight_scale": fine, "mul_scale": smooth}.items():
|
||||
layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False))
|
||||
method = NPUMxfp4LinearMethod(
|
||||
DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2
|
||||
)
|
||||
)
|
||||
method.process_weights_after_loading(layer)
|
||||
prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
||||
assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"}
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
assert not hasattr(layer, "w4a8_weight_scale")
|
||||
x = x.to(device="npu", dtype=dtype)
|
||||
|
||||
with override_forward_context(ForwardContext()):
|
||||
for step in range(3):
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
actual = method.apply(layer, x, layer.bias).cpu()
|
||||
# A8's documented GEMM output is BF16 before restoring the caller
|
||||
# dtype; the reference does not reuse quantization or unpacking.
|
||||
reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype)
|
||||
torch.testing.assert_close(actual, reference, rtol=0, atol=0)
|
||||
assert actual.dtype == dtype
|
||||
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers
|
||||
|
||||
|
||||
@pytest.mark.cards_1
|
||||
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
||||
@pytest.mark.parametrize("policy", ["layer", "layer_and_step"])
|
||||
def test_mxfp4_layer_fallback(mode, policy, monkeypatch):
|
||||
"""Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix."""
|
||||
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy)
|
||||
|
||||
|
||||
@pytest.mark.cards_1
|
||||
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
||||
def test_mxfp4_ocp_compatibility(mode, monkeypatch):
|
||||
"""OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests."""
|
||||
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0)
|
||||
@@ -817,6 +817,8 @@ class _DiffusionConfigProjection:
|
||||
enable_stage_verification: bool = True
|
||||
prompt_file_path: str | None = None
|
||||
quantization_config: _QuantizationConfigType = None
|
||||
# Internal provenance, retained across config projection and worker transport.
|
||||
quantization_config_is_auto_detected: bool = False
|
||||
extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@field_validator("kv_transfer_config", mode="before")
|
||||
@@ -964,6 +966,8 @@ class _DiffusionConfigProjection:
|
||||
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
|
||||
)
|
||||
if should_use_checkpoint_config:
|
||||
if self.quantization_config is None:
|
||||
self.quantization_config_is_auto_detected = True
|
||||
self.quantization_config = quant_config
|
||||
if quant_method is not None:
|
||||
self.additional_config.setdefault("auto_detected_quant_method", quant_method)
|
||||
|
||||
@@ -1021,6 +1021,8 @@ class OmniDiffusionConfig:
|
||||
# str is resolved to {"method": <str>} internally.
|
||||
# Per-component: {"transformer": {"method": "fp8"}, "vae": None}
|
||||
quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
|
||||
# Internal provenance, retained across config projection and worker transport.
|
||||
quantization_config_is_auto_detected: bool = False
|
||||
# Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
|
||||
# does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
|
||||
# has already resolved to vLLM's ModelOpt FP8 linear method.
|
||||
@@ -1358,6 +1360,8 @@ class OmniDiffusionConfig:
|
||||
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
|
||||
)
|
||||
if should_use_checkpoint_config:
|
||||
if self.quantization_config is None:
|
||||
self.quantization_config_is_auto_detected = True
|
||||
self.quantization_config = tf_config.quant_config
|
||||
logger.info(
|
||||
"Auto-detected quantization '%s' from model config",
|
||||
|
||||
@@ -7,7 +7,7 @@ import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Generator, Iterable, Sequence
|
||||
from collections.abc import Callable, Generator, Iterable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
@@ -481,8 +481,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
|
||||
)
|
||||
|
||||
def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
|
||||
"""Return parameter names that should be covered by strict load checks."""
|
||||
all_parameter_names = {name for name, _ in model.named_parameters()}
|
||||
"""Return parameter names that should be covered by strict load checks.
|
||||
|
||||
A parameter with an initialized checkpoint default can explicitly set
|
||||
``is_checkpoint_optional``. It is still loaded when present on disk.
|
||||
"""
|
||||
all_parameter_names = {
|
||||
name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False)
|
||||
}
|
||||
sources = self._get_weight_sources(model)
|
||||
|
||||
# Keep strict behavior if no source metadata exists.
|
||||
@@ -853,9 +859,24 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
|
||||
del model
|
||||
return self.load_fresh_canonical_model()
|
||||
raise
|
||||
self._log_w4a8_fallback_load_summaries(model)
|
||||
self._attach_offload_startup_state(model)
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None:
|
||||
"""Ask discovered DiTs to report W4A8 state at the common load exit.
|
||||
|
||||
Each DiT derives readiness from its processed layers, so a deferred
|
||||
weight plan cannot be mistaken for the ordinary or HSDP load path.
|
||||
"""
|
||||
components = ModuleDiscovery.discover(model)
|
||||
for component_name, dit in zip(components.dit_names, components.dits):
|
||||
candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None)
|
||||
if callable(candidate):
|
||||
reporter = cast(Callable[[str], None], candidate)
|
||||
reporter(component_name)
|
||||
|
||||
@staticmethod
|
||||
def _request_offload_after_quant(model: nn.Module) -> int:
|
||||
"""Ask online-quant layers to return to host memory once quantized.
|
||||
@@ -1065,6 +1086,16 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
|
||||
# that have loaded weights tracking currently.
|
||||
if loaded_weights is not None:
|
||||
weights_not_loaded = weights_to_load - loaded_weights
|
||||
# Offline formats can require scales that older online loaders
|
||||
# were allowed to synthesize. Do not apply that legacy tolerance
|
||||
# to an explicitly required checkpoint tensor.
|
||||
required_missing = {
|
||||
name
|
||||
for name, param in model.named_parameters()
|
||||
if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False)
|
||||
}
|
||||
if required_missing:
|
||||
raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}")
|
||||
# NOTE: if the model is quantized, ignore not_loaded check for scale
|
||||
# weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
|
||||
# static activation `input_scale`, which the quant method may
|
||||
|
||||
@@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
|
||||
return {}
|
||||
|
||||
|
||||
def resolve_wan_transformer_quant_config(
|
||||
config: dict, quant_config: QuantizationConfig | None, component: str
|
||||
) -> QuantizationConfig | None:
|
||||
"""Resolve the expert before applying its checkpoint's storage contract."""
|
||||
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
# Wan experts are siblings: "transformer_2" must not inherit "transformer"
|
||||
# through the generic layer-prefix resolver. Select the whole expert name.
|
||||
component_quant_config = (
|
||||
quant_config.component_configs.get(component, quant_config.default_config)
|
||||
if isinstance(quant_config, ComponentQuantizationConfig)
|
||||
else quant_config
|
||||
)
|
||||
quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None
|
||||
resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config"))
|
||||
if quantization_disabled and resolved_quant_config is not None:
|
||||
raise ValueError(
|
||||
f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. "
|
||||
"Use a BF16 checkpoint for this component or enable its matching quantization method."
|
||||
)
|
||||
return resolved_quant_config
|
||||
|
||||
|
||||
def create_transformer_from_config(
|
||||
config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
|
||||
config: dict,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
*,
|
||||
component: str = "transformer",
|
||||
) -> WanTransformer3DModel:
|
||||
"""Create WanTransformer3DModel from config dict."""
|
||||
kwargs: dict = {}
|
||||
@@ -216,11 +244,7 @@ def create_transformer_from_config(
|
||||
if "pos_embed_seq_len" in config:
|
||||
kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
|
||||
|
||||
if "quantization_config" in config:
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
|
||||
|
||||
quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
|
||||
if quant_config is not None:
|
||||
kwargs["quant_config"] = quant_config
|
||||
if prefix:
|
||||
@@ -460,13 +484,13 @@ class Wan22Pipeline(
|
||||
# Initialize transformers with correct config (weights loaded via load_weights)
|
||||
if load_transformer:
|
||||
transformer_config = load_transformer_config(model, "transformer", local_files_only)
|
||||
self.transformer = self._create_transformer(transformer_config)
|
||||
self.transformer = self._create_transformer(transformer_config, component="transformer")
|
||||
else:
|
||||
self.transformer = None
|
||||
|
||||
if load_transformer_2:
|
||||
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
|
||||
self.transformer_2 = self._create_transformer(transformer_2_config)
|
||||
self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
|
||||
else:
|
||||
self.transformer_2 = None
|
||||
|
||||
@@ -504,10 +528,17 @@ class Wan22Pipeline(
|
||||
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
||||
)
|
||||
|
||||
def _create_transformer(self, config: dict) -> WanTransformer3DModel:
|
||||
def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
|
||||
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
return create_transformer_from_config(config, quant_config=quant_config)
|
||||
# Startup metadata describes the first expert, not a user policy for both.
|
||||
if getattr(self.od_config, "quantization_config_is_auto_detected", False):
|
||||
quant_config = None
|
||||
return create_transformer_from_config(
|
||||
config,
|
||||
quant_config=quant_config,
|
||||
component=component,
|
||||
)
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
|
||||
@@ -304,25 +304,10 @@ class Wan22I2VPipeline(
|
||||
# Transformers (weights loaded via load_weights)
|
||||
# Load config from model directory or HF Hub to get correct in_channels for I2V models
|
||||
transformer_config = load_transformer_config(model, "transformer", local_files_only)
|
||||
self.transformer = create_transformer_from_config(
|
||||
transformer_config,
|
||||
quant_config=od_config.quantization_config,
|
||||
)
|
||||
self.transformer = self._create_transformer(transformer_config, component="transformer")
|
||||
if self.has_transformer_2:
|
||||
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
|
||||
t2_quant = transformer_2_config.get("quantization_config")
|
||||
if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
|
||||
from vllm_omni.quantization.factory import build_quant_config
|
||||
|
||||
method = t2_quant["quant_method"]
|
||||
kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
|
||||
t2_quant = build_quant_config(method, **kwargs)
|
||||
else:
|
||||
t2_quant = None
|
||||
self.transformer_2 = create_transformer_from_config(
|
||||
transformer_2_config,
|
||||
quant_config=t2_quant,
|
||||
)
|
||||
self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
|
||||
else:
|
||||
self.transformer_2 = None
|
||||
|
||||
@@ -462,10 +447,13 @@ class Wan22I2VPipeline(
|
||||
image_embeds = self.image_encoder(pixel_values, output_hidden_states=True)
|
||||
return image_embeds.hidden_states[-2]
|
||||
|
||||
def _create_transformer(self, config: dict) -> WanTransformer3DModel:
|
||||
def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
|
||||
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
return create_transformer_from_config(config, quant_config=quant_config)
|
||||
# Startup metadata describes the first expert, not a user policy for both.
|
||||
if getattr(self.od_config, "quantization_config_is_auto_detected", False):
|
||||
quant_config = None
|
||||
return create_transformer_from_config(config, quant_config=quant_config, component=component)
|
||||
|
||||
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
|
||||
sampling_params_list = req.sampling_params_list
|
||||
|
||||
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
Wan22Pipeline,
|
||||
load_wan_weights_with_optional_gate,
|
||||
resolve_wan_transformer_quant_config,
|
||||
retrieve_latents,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
@@ -60,6 +61,8 @@ def create_vace_transformer_from_config(
|
||||
config: dict,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
*,
|
||||
component: str = "transformer",
|
||||
) -> WanVACETransformer3DModel:
|
||||
"""Create WanVACETransformer3DModel from config dict."""
|
||||
kwargs = {}
|
||||
@@ -97,6 +100,7 @@ def create_vace_transformer_from_config(
|
||||
kwargs["vace_layers"] = config["vace_layers"]
|
||||
if "vace_in_channels" in config:
|
||||
kwargs["vace_in_channels"] = config["vace_in_channels"]
|
||||
quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
|
||||
if quant_config is not None:
|
||||
kwargs["quant_config"] = quant_config
|
||||
if prefix:
|
||||
@@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
|
||||
|
||||
super().__init__(od_config=od_config, prefix=prefix)
|
||||
|
||||
def _create_transformer(self, config: dict) -> WanVACETransformer3DModel:
|
||||
def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel:
|
||||
"""Build VACE transformer. Respects od_config.quantization_config."""
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
return create_vace_transformer_from_config(config, quant_config=quant_config)
|
||||
# Startup metadata describes the first expert, not a user policy for both.
|
||||
if getattr(self.od_config, "quantization_config_is_auto_detected", False):
|
||||
quant_config = None
|
||||
return create_vace_transformer_from_config(config, quant_config=quant_config, component=component)
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
|
||||
@@ -18,7 +18,7 @@ from vllm.distributed import (
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.conv import Conv3dLayer
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.model_executor.models.utils import (
|
||||
@@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module):
|
||||
# ROPE helper
|
||||
self._cached_rope_emb = None
|
||||
self._cached_rope_resolution = None
|
||||
self._validate_w4a8_fallback_layers(quant_config)
|
||||
|
||||
def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None:
|
||||
"""Reject checkpoint aliases/typos before loading weights or running inference."""
|
||||
requested = tuple(getattr(quant_config, "w4a8_fallback_layers", []))
|
||||
self._w4a8_fallback_layers = requested
|
||||
self._local_w4a8_fallback_layers: tuple[str, ...] = ()
|
||||
if not requested:
|
||||
return
|
||||
modules = dict(self.named_modules())
|
||||
missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)]
|
||||
local_layers = [
|
||||
name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes)
|
||||
]
|
||||
self._local_w4a8_fallback_layers = tuple(local_layers)
|
||||
invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)]
|
||||
if invalid:
|
||||
raise ValueError(
|
||||
"w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; "
|
||||
f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, "
|
||||
"not checkpoint attn1.to_q/to_k/to_v names."
|
||||
)
|
||||
|
||||
def _log_w4a8_fallback_load_summary(self, component_name: str) -> None:
|
||||
"""Log post-load transforms/cache readiness, not numerical acceptance."""
|
||||
requested = self._w4a8_fallback_layers
|
||||
if not requested:
|
||||
return
|
||||
modules = dict(self.named_modules())
|
||||
local_layers = self._local_w4a8_fallback_layers
|
||||
selected_layers = [
|
||||
name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False)
|
||||
]
|
||||
bf16_overrides = [name for name in local_layers if name not in selected_layers]
|
||||
processed = [
|
||||
name
|
||||
for name in selected_layers
|
||||
if getattr(modules[name], "_already_called_process_weights_after_loading", False)
|
||||
]
|
||||
if not selected_layers:
|
||||
processing_state = "not-required"
|
||||
elif len(processed) == len(selected_layers):
|
||||
processing_state = "ready"
|
||||
else:
|
||||
processing_state = "not-ready"
|
||||
|
||||
def qualified(names: Iterable[str]) -> list[str]:
|
||||
return [f"{component_name}.{name}" for name in names]
|
||||
|
||||
logger.info(
|
||||
"Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; "
|
||||
"weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d",
|
||||
component_name,
|
||||
len(selected_layers),
|
||||
qualified(selected_layers),
|
||||
len(bf16_overrides),
|
||||
qualified(bf16_overrides),
|
||||
processing_state,
|
||||
len(processed),
|
||||
len(selected_layers),
|
||||
len(requested) - len(local_layers),
|
||||
)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
@@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module):
|
||||
}
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
single_scale_prefixes = {
|
||||
name
|
||||
for name, module in self.named_modules()
|
||||
if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod)
|
||||
}
|
||||
|
||||
def reject_dualscale_tensor(name: str) -> None:
|
||||
if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes:
|
||||
raise ValueError(
|
||||
f"Single-scale mxfp4 cannot load DualScale tensor {name}; "
|
||||
"changing quant_method does not convert a checkpoint."
|
||||
)
|
||||
|
||||
loaded_params: set[str] = set()
|
||||
loaded_qkv_shards: dict[str, set[str]] = {}
|
||||
qkv_smooth_scales: dict[str, torch.Tensor] = {}
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
name = weight_name_remapping.get(name, name)
|
||||
@@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module):
|
||||
# Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall
|
||||
# through to the else branch and are loaded directly.
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in original_name:
|
||||
if f"{weight_name}." not in original_name:
|
||||
continue
|
||||
lookup_name = original_name.replace(weight_name, param_name)
|
||||
lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1)
|
||||
reject_dualscale_tensor(lookup_name)
|
||||
# Skip weights that belong to PP stages other than this one
|
||||
if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
|
||||
break
|
||||
param = params_dict[lookup_name]
|
||||
if lookup_name.endswith(".mul_scale"):
|
||||
previous_scale = qkv_smooth_scales.get(lookup_name)
|
||||
if previous_scale is not None and not torch.equal(previous_scale, loaded_weight):
|
||||
raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}")
|
||||
if previous_scale is None:
|
||||
qkv_smooth_scales[lookup_name] = loaded_weight.clone()
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(original_name)
|
||||
if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False):
|
||||
# Shared input-channel Smooth is a complete tensor even
|
||||
# when supplied under a single Q/K/V source name.
|
||||
loaded_params.add(lookup_name)
|
||||
else:
|
||||
shards = loaded_qkv_shards.setdefault(lookup_name, set())
|
||||
shards.add(shard_id)
|
||||
if shards == {"q", "k", "v"}:
|
||||
loaded_params.add(lookup_name)
|
||||
break
|
||||
else:
|
||||
# diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
|
||||
@@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module):
|
||||
if is_pp_missing_parameter(lookup_name, self):
|
||||
continue
|
||||
|
||||
reject_dualscale_tensor(lookup_name)
|
||||
if lookup_name not in params_dict:
|
||||
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
|
||||
continue
|
||||
@@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module):
|
||||
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
loaded_params.add(original_name)
|
||||
loaded_params.add(lookup_name)
|
||||
loaded_params.update((original_name, lookup_name))
|
||||
|
||||
return loaded_params
|
||||
@@ -446,6 +446,33 @@ def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> b
|
||||
return False
|
||||
|
||||
|
||||
_MXFP4_SERIALIZED_FLAGS = {
|
||||
"mxfp4": "is_checkpoint_mxfp4_serialized",
|
||||
"mxfp4_dualscale": "is_checkpoint_serialized",
|
||||
}
|
||||
|
||||
|
||||
def _normalize_serialized_mxfp4_layer_policy(qc_method: str, qc_kwargs: dict[str, Any]) -> None:
|
||||
"""Canonicalize the BF16 layer policy owned by serialized MXFP4 metadata.
|
||||
|
||||
A serialized checkpoint config is a complete storage contract: an omitted
|
||||
layer list means no explicit BF16 layers. Method-only online metadata does
|
||||
not carry that contract and must leave caller-provided layer policy intact.
|
||||
"""
|
||||
method = _normalize_quant_method_alias(qc_method)
|
||||
if method is None:
|
||||
return
|
||||
serialized_flag = _MXFP4_SERIALIZED_FLAGS.get(method)
|
||||
if serialized_flag is None or not qc_kwargs.get(serialized_flag):
|
||||
return
|
||||
|
||||
ignored_layers = qc_kwargs.get("ignored_layers")
|
||||
if not ignored_layers:
|
||||
ignored_layers = qc_kwargs.get("modules_to_not_convert")
|
||||
qc_kwargs["ignored_layers"] = ignored_layers or []
|
||||
qc_kwargs.pop("modules_to_not_convert", None)
|
||||
|
||||
|
||||
def resolve_quant_config_from_disk(
|
||||
quant_config: QuantizationConfig | None,
|
||||
disk_qc: dict[str, Any] | str | None,
|
||||
@@ -460,22 +487,24 @@ def resolve_quant_config_from_disk(
|
||||
- quant_config is None: auto-detect from disk_qc (full build).
|
||||
- Methods mismatch: raise ValueError — prevents silent weight corruption.
|
||||
- Disk marks serialized but quant_config is online: rebuild from disk.
|
||||
- ignored_layers differ: rebuild from disk (per-transformer BF16 routing).
|
||||
- Serialized MXFP4 metadata owns the complete BF16 layer policy; omitted
|
||||
ignored_layers means the checkpoint default (no explicit BF16 layers).
|
||||
- Method-only metadata carries no storage policy and preserves caller
|
||||
ignored_layers and runtime fallback policies.
|
||||
"""
|
||||
if disk_qc is None:
|
||||
return quant_config
|
||||
|
||||
if isinstance(disk_qc, str):
|
||||
if quant_config is None:
|
||||
logger.info("Auto-detected quantization from config.json: method=%s", disk_qc)
|
||||
return build_quant_config(disk_qc)
|
||||
return quant_config
|
||||
qc_method = disk_qc
|
||||
qc_kwargs: dict[str, Any] = {}
|
||||
else:
|
||||
if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
|
||||
return quant_config
|
||||
qc_method = disk_qc["quant_method"]
|
||||
qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
|
||||
|
||||
if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc:
|
||||
return quant_config
|
||||
|
||||
qc_method: str = disk_qc["quant_method"]
|
||||
qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"}
|
||||
_normalize_serialized_mxfp4_layer_policy(qc_method, qc_kwargs)
|
||||
|
||||
if quant_config is None:
|
||||
logger.info(
|
||||
@@ -494,6 +523,26 @@ def resolve_quant_config_from_disk(
|
||||
"Pass a matching --quantization flag or omit it for auto-detection."
|
||||
)
|
||||
|
||||
if isinstance(disk_qc, str):
|
||||
return quant_config
|
||||
|
||||
# Serialized checkpoint metadata owns storage and BF16 layer routing. The
|
||||
# active config owns runtime step/layer policies, including explicit empty
|
||||
# lists. Preserve them when each cascade transformer rebuilds from disk.
|
||||
for policy in ("w4a8_fallback_steps", "w4a8_fallback_layers"):
|
||||
if hasattr(quant_config, policy):
|
||||
qc_kwargs[policy] = list(getattr(quant_config, policy))
|
||||
if hasattr(quant_config, "mxfp4_scale_alg"):
|
||||
# Runtime activation policy is independent of how offline W4 was made.
|
||||
qc_kwargs["mxfp4_scale_alg"] = quant_config.mxfp4_scale_alg
|
||||
if hasattr(quant_config, "require_smooth_scale"):
|
||||
disk_requires_smooth = qc_kwargs.get("require_smooth_scale", False)
|
||||
if not isinstance(disk_requires_smooth, bool):
|
||||
raise ValueError("require_smooth_scale must be a boolean")
|
||||
# Either a calibrated checkpoint or an explicit caller can require
|
||||
# Smooth. Per-expert rebuilds must never weaken that requirement.
|
||||
qc_kwargs["require_smooth_scale"] = quant_config.require_smooth_scale or disk_requires_smooth
|
||||
|
||||
if _disk_marks_serialized(qc_kwargs, quant_config):
|
||||
logger.info(
|
||||
"config.json marks checkpoint as serialized; switching to offline %s mode.",
|
||||
@@ -501,6 +550,14 @@ def resolve_quant_config_from_disk(
|
||||
)
|
||||
return build_quant_config(qc_method, **qc_kwargs)
|
||||
|
||||
if (
|
||||
"require_smooth_scale" in qc_kwargs
|
||||
and hasattr(quant_config, "require_smooth_scale")
|
||||
and qc_kwargs["require_smooth_scale"] != quant_config.require_smooth_scale
|
||||
):
|
||||
logger.info("config.json Smooth requirement differs from active config; rebuilding quant_config.")
|
||||
return build_quant_config(qc_method, **qc_kwargs)
|
||||
|
||||
# AutoRound MXFP checkpoints use data_type="mx_fp" instead of
|
||||
# is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
|
||||
# selected according to the checkpoint's bit width.
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers.
|
||||
|
||||
Architecture mirrors mxfp8_config.py:
|
||||
@@ -41,7 +41,7 @@ Dual-scale (W4A4_MXFP4_DUALSCALE):
|
||||
and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast.
|
||||
|
||||
Checkpoint tensor shapes for dual-scale:
|
||||
weight : (N, K) float8_e4m3fn – FP4 packed
|
||||
weight : (N, K) float8_e4m3fn – numeric FP4 values (packed after loading)
|
||||
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits)
|
||||
weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert)
|
||||
mul_scale : (K,) float32 – per-input-channel activation pre-scale
|
||||
@@ -51,9 +51,9 @@ Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/q
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
from vllm.logger import init_logger
|
||||
@@ -83,9 +83,54 @@ from vllm_omni.quantization.mxfp8_config import (
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _validate_w4a8_fallback_steps(steps: list[int] | None) -> list[int]:
|
||||
if steps is None:
|
||||
return []
|
||||
if not isinstance(steps, list) or any(type(step) is not int or step < 0 for step in steps):
|
||||
raise ValueError("w4a8_fallback_steps must be a list of non-negative integer denoise step indices.")
|
||||
return sorted(set(steps))
|
||||
|
||||
|
||||
def _validate_w4a8_fallback_layers(layers: list[str] | None) -> list[str]:
|
||||
"""Names are exact runtime Linear paths, not checkpoint aliases or patterns."""
|
||||
if layers is None:
|
||||
return []
|
||||
if not isinstance(layers, list) or any(
|
||||
not isinstance(name, str) or re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*(?:\.[A-Za-z_0-9]+)*", name) is None
|
||||
for name in layers
|
||||
):
|
||||
raise ValueError("w4a8_fallback_layers must be a list of non-empty exact runtime Linear paths.")
|
||||
return sorted(set(layers))
|
||||
|
||||
|
||||
def _is_w4a8_fallback_step(steps: list[int]) -> bool:
|
||||
"""Read the request's zero-based denoise index; never keep state on a Linear."""
|
||||
if not steps:
|
||||
return False
|
||||
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
get_forward_context,
|
||||
is_forward_context_available,
|
||||
)
|
||||
|
||||
step = get_forward_context().denoise_step_idx if is_forward_context_available() else None
|
||||
if step is None:
|
||||
raise RuntimeError("W4A8 step fallback requires a pipeline that publishes its denoise step in ForwardContext.")
|
||||
return step in steps
|
||||
|
||||
|
||||
def _validate_smooth_scale(scale: torch.Tensor) -> None:
|
||||
"""Reject malformed calibration before loading or activation dtype conversion."""
|
||||
if scale.ndim != 1 or not scale.is_floating_point():
|
||||
raise ValueError("MXFP4 mul_scale must be a one-dimensional floating-point Smooth tensor.")
|
||||
if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()):
|
||||
raise ValueError("MXFP4 mul_scale must contain finite, strictly positive Smooth values.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -105,10 +150,24 @@ class DiffusionMXFP4Config(QuantizationConfig):
|
||||
self,
|
||||
is_checkpoint_mxfp4_serialized: bool = False,
|
||||
ignored_layers: list[str] | None = None,
|
||||
w4a8_fallback_steps: list[int] | None = None,
|
||||
require_smooth_scale: bool = False,
|
||||
w4a8_fallback_layers: list[str] | None = None,
|
||||
mxfp4_scale_alg: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
|
||||
self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
|
||||
if not isinstance(require_smooth_scale, bool):
|
||||
raise ValueError("require_smooth_scale must be a boolean")
|
||||
if require_smooth_scale and not self.is_checkpoint_mxfp4_serialized:
|
||||
raise ValueError("require_smooth_scale requires an offline MXFP4 checkpoint")
|
||||
self.require_smooth_scale = require_smooth_scale
|
||||
if type(mxfp4_scale_alg) is not int or mxfp4_scale_alg not in (0, 2):
|
||||
raise ValueError("mxfp4_scale_alg must be 0 (OCP MX) or 2 (UOS with dst_type_max=7.25).")
|
||||
self.mxfp4_scale_alg = mxfp4_scale_alg
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
@@ -139,6 +198,10 @@ class DiffusionMXFP4Config(QuantizationConfig):
|
||||
return cls(
|
||||
is_checkpoint_mxfp4_serialized=is_serialized,
|
||||
ignored_layers=ignored_layers,
|
||||
w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
|
||||
w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
|
||||
require_smooth_scale=config.get("require_smooth_scale", False),
|
||||
mxfp4_scale_alg=config.get("mxfp4_scale_alg", 0),
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
@@ -155,9 +218,15 @@ class DiffusionMXFP4Config(QuantizationConfig):
|
||||
return UnquantizedLinearMethod()
|
||||
if current_omni_platform.is_npu():
|
||||
if self.is_checkpoint_mxfp4_serialized:
|
||||
return NPUMxfp4LinearMethod(self)
|
||||
return NPUMxfp4OnlineLinearMethod(self)
|
||||
return NPUMxfp4LinearMethod(self, prefix=prefix)
|
||||
return NPUMxfp4OnlineLinearMethod(self, prefix=prefix)
|
||||
if current_omni_platform.is_rocm():
|
||||
if self.mxfp4_scale_alg != 0:
|
||||
raise NotImplementedError("MXFP4 UOS quantization is currently only supported on NPU (Ascend).")
|
||||
if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
|
||||
raise NotImplementedError(
|
||||
"MXFP4 W4A8 step/layer fallback is currently only supported on NPU (Ascend)."
|
||||
)
|
||||
gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName
|
||||
if "gfx950" not in gcn_arch:
|
||||
raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}")
|
||||
@@ -176,6 +245,56 @@ class DiffusionMXFP4Config(QuantizationConfig):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _npu_quantize_mxfp4(x: torch.Tensor, scale_alg: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize weight or activation using the same explicit FP4 algorithm."""
|
||||
import torch_npu
|
||||
|
||||
# Universal Optimal Scaling (Qmax=7.25), MXAttention Section 4.1:
|
||||
# https://arxiv.org/abs/2607.24377
|
||||
scale_kwargs = {"dst_type_max": 7.25} if scale_alg == 2 else {}
|
||||
return torch_npu.npu_dynamic_mx_quant(
|
||||
x,
|
||||
dst_type=torch_npu.float4_e2m1fn_x2,
|
||||
axis=-1,
|
||||
block_size=32,
|
||||
round_mode="rint",
|
||||
scale_alg=scale_alg,
|
||||
**scale_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _npu_w4a8_matmul(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""MXFP8 activations with prepared single-scale MXFP4 weights (N, K)."""
|
||||
import torch_npu
|
||||
|
||||
x_q, x_scale = torch_npu.npu_dynamic_mx_quant(
|
||||
x, dst_type=torch_npu.float8_e4m3fn, axis=-1, block_size=32, round_mode="rint", scale_alg=0
|
||||
)
|
||||
# Use the BF16 output/bias pair supported by the mixed FP8/FP4 kernel
|
||||
# in the target A5 runtime. Restore the caller's dtype afterwards.
|
||||
if bias is not None:
|
||||
bias = bias.to(torch.bfloat16).reshape(1, -1)
|
||||
output = torch_npu.npu_quant_matmul(
|
||||
x_q,
|
||||
weight.transpose(0, 1),
|
||||
weight_scale.transpose(0, 1),
|
||||
scale_dtype=torch_npu.float8_e8m0fnu,
|
||||
x2_dtype=torch_npu.float4_e2m1fn_x2,
|
||||
pertoken_scale=x_scale,
|
||||
pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
|
||||
bias=bias,
|
||||
output_dtype=torch.bfloat16,
|
||||
group_sizes=[0, 0, 32],
|
||||
)
|
||||
return output.to(output_dtype)
|
||||
|
||||
|
||||
class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||||
"""NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints.
|
||||
|
||||
@@ -187,7 +306,8 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||||
NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: DiffusionMXFP4Config) -> None:
|
||||
def __init__(self, quant_config: DiffusionMXFP4Config, *, prefix: str = "") -> None:
|
||||
self.is_w4a8_fallback_layer = prefix in quant_config.w4a8_fallback_layers
|
||||
self.quant_config = quant_config
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
|
||||
@@ -210,11 +330,14 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
# BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights.
|
||||
layer.register_parameter(
|
||||
"weight",
|
||||
ModelWeightParameter(
|
||||
data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype),
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
@@ -234,12 +357,42 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||||
),
|
||||
)
|
||||
|
||||
# Older single-scale checkpoints have no Smooth tensor. Keep their
|
||||
# identity transform while allowing calibrated checkpoints to load it.
|
||||
def smooth_weight_loader(
|
||||
param: ModelWeightParameter, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
_validate_smooth_scale(loaded_weight)
|
||||
if not callable(weight_loader):
|
||||
raise ValueError("Loading MXFP4 Smooth requires a callable weight_loader")
|
||||
weight_loader(param, loaded_weight, *args, **kwargs)
|
||||
|
||||
layer.register_parameter(
|
||||
"mul_scale",
|
||||
ModelWeightParameter(
|
||||
data=torch.ones(input_size_per_partition, dtype=torch.float32),
|
||||
input_dim=0,
|
||||
output_dim=None,
|
||||
weight_loader=smooth_weight_loader,
|
||||
),
|
||||
)
|
||||
setattr(layer.mul_scale, "ignore_warning", True)
|
||||
setattr(layer.mul_scale, "is_checkpoint_optional", not self.quant_config.require_smooth_scale)
|
||||
if self.quant_config.is_checkpoint_mxfp4_serialized:
|
||||
setattr(layer.weight_scale, "is_checkpoint_required", True)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
import torch_npu
|
||||
|
||||
smooth_param = cast(torch.Tensor, layer.mul_scale)
|
||||
_validate_smooth_scale(smooth_param)
|
||||
smooth = smooth_param.data.to(dtype=cast(torch.dtype, layer.orig_dtype)).contiguous()
|
||||
# Finite FP32 values can overflow or underflow in the activation dtype.
|
||||
_validate_smooth_scale(smooth)
|
||||
|
||||
# NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose.
|
||||
w = layer.weight
|
||||
if w.dtype != torch_npu.float4_e2m1fn_x2:
|
||||
@@ -258,15 +411,31 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||||
|
||||
replace_parameter(layer, "weight", w)
|
||||
replace_parameter(layer, "weight_scale", s)
|
||||
# Keep loader metadata while converting the calibration tensor once.
|
||||
layer.mul_scale.data = smooth.to(device=cast(torch.Tensor, w).device)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
|
||||
# --- NPU MXFP4 ops — shared with online path via inheritance ---
|
||||
|
||||
def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
import torch_npu
|
||||
def _apply_inner(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
ori_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
# Online quantization has no calibrated Smooth tensor. For offline
|
||||
# weights the same pre-scale must precede both A4 and A8 quantization.
|
||||
if (mul_scale := getattr(layer, "mul_scale", None)) is not None:
|
||||
x = x * mul_scale.to(x.dtype)
|
||||
if self.is_w4a8_fallback_layer or _is_w4a8_fallback_step(self.quant_config.w4a8_fallback_steps):
|
||||
return _npu_w4a8_matmul(
|
||||
x, cast(torch.Tensor, layer.weight), cast(torch.Tensor, layer.weight_scale), bias, ori_dtype
|
||||
)
|
||||
return super()._apply_inner(layer, x, bias, ori_dtype)
|
||||
|
||||
# No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2.
|
||||
return torch_npu.npu_dynamic_mx_quant(x)
|
||||
def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return _npu_quantize_mxfp4(x, self.quant_config.mxfp4_scale_alg)
|
||||
|
||||
def _quant_matmul(
|
||||
self,
|
||||
@@ -318,8 +487,6 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
import torch_npu
|
||||
|
||||
if layer.weight.device == torch.device("meta"):
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty_like(layer.weight, device=layer._load_device),
|
||||
@@ -331,8 +498,10 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
|
||||
layer.register_parameter("weight", weight)
|
||||
initialize_single_dummy_weight(layer.weight)
|
||||
|
||||
# NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2.
|
||||
weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight)
|
||||
# Prepare once; offline checkpoints keep their existing weight algorithm.
|
||||
weight_fp4, weight_scale_raw = _npu_quantize_mxfp4(
|
||||
cast(torch.Tensor, layer.weight), self.quant_config.mxfp4_scale_alg
|
||||
)
|
||||
|
||||
# Weight stays (N, K) — no pre-transpose for FP4 packed format.
|
||||
# Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul.
|
||||
@@ -497,7 +666,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
"""NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints.
|
||||
|
||||
Checkpoint tensors and their canonical post-load shapes:
|
||||
weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte)
|
||||
weight : (N, K) float8_e4m3fn – numeric FP4 values; packed into NZ after loading
|
||||
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2)
|
||||
weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N)
|
||||
mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration)
|
||||
@@ -509,7 +678,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
Reference: MindIE-SD W4A4MXFP4DualQuantLinear.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Any) -> None:
|
||||
def __init__(self, quant_config: DiffusionMXFP4DualScaleMixedConfig) -> None:
|
||||
self.quant_config = quant_config
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
|
||||
@@ -532,7 +701,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
# FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16
|
||||
# Checkpoint FP4 values are stored numerically in FP8; BF16 preserves them exactly.
|
||||
layer.register_parameter(
|
||||
"weight",
|
||||
ModelWeightParameter(
|
||||
@@ -587,6 +756,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
),
|
||||
)
|
||||
setattr(layer.mul_scale, "ignore_warning", True)
|
||||
setattr(layer.weight_scale, "is_checkpoint_required", True)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
@@ -594,7 +764,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
|
||||
import torch_npu
|
||||
|
||||
# float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29).
|
||||
# Numeric FP4 values → packed float4_e2m1fn_x2 → NZ hardware format (format ID 29).
|
||||
# NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices.
|
||||
w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2)
|
||||
w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8)
|
||||
@@ -631,7 +801,10 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||||
"""
|
||||
if ori_dtype not in (torch.bfloat16, torch.float16):
|
||||
x = x.to(torch.bfloat16)
|
||||
x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale)
|
||||
# The dual-level activation kernel requires x and Smooth to have the
|
||||
# same dtype. Prepared calibration scales are BF16, including online.
|
||||
smooth_scale = cast(torch.Tensor, layer.mul_scale).to(x.dtype)
|
||||
x_q, l0_scale, l1_scale = self._quantize_activation(x, smooth_scale)
|
||||
return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype)
|
||||
|
||||
def _quantize_activation( # type: ignore[override]
|
||||
@@ -785,11 +958,20 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
|
||||
is_checkpoint_serialized: bool = False,
|
||||
ignored_layers: list[str] | None = None,
|
||||
num_bf16_fallback_layers: int = 5,
|
||||
w4a8_fallback_steps: list[int] | None = None,
|
||||
w4a8_fallback_layers: list[str] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.is_checkpoint_serialized = is_checkpoint_serialized
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.num_bf16_fallback_layers = num_bf16_fallback_layers
|
||||
self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps)
|
||||
self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers)
|
||||
if self.w4a8_fallback_steps or self.w4a8_fallback_layers:
|
||||
raise ValueError(
|
||||
"mxfp4_dualscale does not support W4A8 fallback. Use mxfp4 with single-scale weights; "
|
||||
"changing a DualScale checkpoint's method name does not convert its storage format."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
@@ -822,6 +1004,8 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
|
||||
is_checkpoint_serialized=is_serialized,
|
||||
ignored_layers=ignored_layers,
|
||||
num_bf16_fallback_layers=num_bf16_fallback_layers,
|
||||
w4a8_fallback_steps=config.get("w4a8_fallback_steps"),
|
||||
w4a8_fallback_layers=config.get("w4a8_fallback_layers"),
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
|
||||
Reference in new issue
Block a user