Close the A14B quant and offload gaps from review

Re-plan memory with the quant steady factor when the bf16 table forces
offload a quantised DiT would not need, mirroring the image dense-quant
path, and fall back to the bf16 plan when quant does not engage. Stream
the second expert under group offload (model and sequential already hook
every module). Fail the load cleanly when quant engages on only one
expert instead of running mixed precision with quant reported off.
Persist guidance_2 in the gallery recipe so A14B clips are reproducible.
This commit is contained in:
Daniel Han 2026-07-05 00:28:09 +00:00
commit 8933251a7e
4 changed files with 96 additions and 5 deletions

View file

@ -62,11 +62,12 @@ from .diffusion_speed import (
restore_backend_flags,
snapshot_backend_flags,
)
from .diffusion_auto_policy import build_resolved_record
from .diffusion_auto_policy import _QUANT_STEADY_FACTOR, build_resolved_record
from .diffusion_transformer_quant import (
dense_transformer_supported,
normalize_transformer_quant,
quantize_transformer,
select_transformer_quant_scheme,
)
from .video_families import (
VIDEO_CANCELLED_MSG,
@ -561,6 +562,43 @@ class VideoBackend:
companion_dense_mib = companion_mib,
requested_mode = normalize_memory_mode(memory_mode),
)
# Parity with the image dense-quant path: the bf16-table plan can force offload
# a quantised DiT would not need (offload also disables quant entirely). Re-plan
# with the scheme's steady factor and keep the resident placement when it fits;
# if quantisation later fails, the load falls back to this bf16 plan.
bf16_plan = plan
quant_replanned = False
if (
kind == "pipeline"
and plan.offload_policy != "none"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
and components is not None
):
scheme_preview = select_transformer_quant_scheme(
target, transformer_quant, family = fam.name
)
factor = _QUANT_STEADY_FACTOR.get(scheme_preview) if scheme_preview else None
if factor is not None:
quant_mib = int(
(components[0] * factor + components[1] + components[2]) * mib_per_gb
)
replanned = plan_diffusion_memory(
target = target,
device_memory = device_memory,
model_dense_mib = quant_mib,
runtime_headroom_mib = runtime_mib,
companion_dense_mib = None,
requested_mode = normalize_memory_mode(memory_mode),
)
if replanned.offload_policy == "none":
logger.info(
"video.transformer_quant: %s fits resident (%d MiB steady); "
"dropping the bf16 plan's '%s' offload",
scheme_preview, quant_mib, plan.offload_policy,
)
plan = replanned
quant_replanned = True
# ── build the pipeline.
pipeline_cls = getattr(diffusers, fam.pipeline_class)
@ -638,11 +676,23 @@ class VideoBackend:
)
if scheme is not None:
engaged.append(scheme)
# Report the scheme only if it engaged on every DiT; a partial quant (one expert
# dense, one quantised) would run the schedule at mismatched precision, so treat
# anything short of full coverage as not engaged.
if engaged and len(engaged) == len(views):
# Quant must engage on every DiT or none: the first expert is mutated in
# place, so a second-expert failure cannot fall back to dense (the schedule
# would run at mismatched precision with quant reported off). Fail the load
# cleanly instead; a full miss (nothing engaged) stays best-effort dense.
if engaged and len(engaged) < len(views):
del pipe
clear_gpu_cache()
raise RuntimeError(
f"transformer_quant={engaged[0]} engaged on only "
f"{len(engaged)}/{len(views)} experts; retry without quant."
)
if engaged:
transformer_quant_engaged = engaged[0]
# The quant-sized plan is only valid when quant actually engaged; a dense
# fallback must keep the conservative bf16 placement.
if quant_replanned and transformer_quant_engaged is None:
plan = bf16_plan
# ── optimisation layers, in the image backend's order: speed profile
# (compile must precede placement), attention (compile traces it),
@ -686,6 +736,19 @@ class VideoBackend:
attention_engaged = engaged
speed_optims = tuple(k for k, v in applied.items() if v)
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
if offload_policy == "group" and len(views) > 1:
# Group offload streams only ``pipe.transformer``; the second expert would
# otherwise sit resident (~57 GB bf16 on the A14B) and defeat the tier.
# model/sequential offload hook every top-level module, so only group needs
# this. Applied through the view so the helper streams transformer_2.
from .diffusion_memory import _apply_group_offload
for view in views[1:]:
if not _apply_group_offload(view, device, logger):
logger.warning(
"video.memory: group offload did not engage on the second "
"expert; it stays resident"
)
if not vae_tiling:
# Decode of a whole clip is the video memory peak; tiling is near-free
# in quality and keeps the decode bounded, so it is always on.

View file

@ -2412,6 +2412,9 @@ class GalleryVideo(BaseModel):
duration_s: float = Field(..., description = "Clip duration in seconds")
steps: int = Field(..., description = "Denoising steps")
guidance: float = Field(..., description = "Guidance scale")
guidance_2: Optional[float] = Field(
None, description = "Second-expert guidance scale (dual-expert families), if sent"
)
seed: int = Field(..., description = "Seed used")
has_audio: bool = Field(False, description = "Whether the MP4 carries an audio track")
model: Optional[str] = Field(None, description = "Model repo id that produced it")

View file

@ -194,6 +194,7 @@ async def generate_video(
"duration_s": result["duration_s"],
"steps": result["steps"],
"guidance": result["guidance"],
"guidance_2": request.guidance_2,
"seed": result["seed"],
"has_audio": result["has_audio"],
"model": result["repo_id"],

View file

@ -727,6 +727,30 @@ def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch):
assert status["transformer_quant"] == "int8"
def test_wan_a14b_partial_quant_fails_the_load(fake_runtime, monkeypatch):
# If the first expert quantises but the second does not, the pipe is left at
# mismatched precision with no way back (in-place mutation), so the load must
# fail cleanly rather than run mixed with quant reported off.
import core.inference.video as video_mod
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
outcomes = iter(["int8", None])
monkeypatch.setattr(
video_mod,
"quantize_transformer",
lambda view, target, *, mode, family, logger = None: next(outcomes),
)
backend = VideoBackend()
with pytest.raises(RuntimeError, match = "1/2 experts"):
backend.load_pipeline(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
model_kind = "pipeline",
transformer_quant = "int8",
)
assert backend.status()["loaded"] is False
def test_wan_ti2v_dense_quant_applies_to_single_dit(fake_runtime, monkeypatch):
# A single-DiT pipeline load quantises exactly one transformer.
import core.inference.video as video_mod