From 8933251a7e4fe986136673807ccd140a5c4d6cac Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 00:28:09 +0000 Subject: [PATCH] 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. --- studio/backend/core/inference/video.py | 73 ++++++++++++++++++++-- studio/backend/models/inference.py | 3 + studio/backend/routes/video.py | 1 + studio/backend/tests/test_video_backend.py | 24 +++++++ 4 files changed, 96 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 80feba8a90..b6ea024bf2 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2c07dcfe4d..0d82a34f1e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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") diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 6ad09af9de..1c6782409d 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -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"], diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index a8e70346e1..35eeec4bb4 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -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