817e14ad19
Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass.
111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unit tests for fp8 text-encoder casting (``diffusion_precision.py``).
|
|
|
|
Hermetic: torch + diffusers.hooks are stubbed via ``sys.modules`` so the gating and
|
|
the apply path run without a GPU or real diffusers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from core.inference.diffusion_precision import (
|
|
apply_fp8_text_encoder,
|
|
fp8_text_encoder_supported,
|
|
)
|
|
|
|
|
|
def _target(*, device = "cuda", dtype = "bfloat16"):
|
|
return types.SimpleNamespace(device = device, dtype = dtype)
|
|
|
|
|
|
def _stub_torch(monkeypatch, *, with_fp8 = True):
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16 = "bfloat16"
|
|
torch.float16 = "float16"
|
|
if with_fp8:
|
|
torch.float8_e4m3fn = "float8_e4m3fn"
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
return torch
|
|
|
|
|
|
def _stub_diffusers_hooks(monkeypatch, recorder):
|
|
hooks = types.ModuleType("diffusers.hooks")
|
|
casting = types.ModuleType("diffusers.hooks.layerwise_casting")
|
|
casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",)
|
|
|
|
def _apply(module, *, storage_dtype, compute_dtype, skip_modules_pattern):
|
|
recorder.append((module, storage_dtype, compute_dtype, skip_modules_pattern))
|
|
|
|
hooks.apply_layerwise_casting = _apply
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
|
|
|
|
# ── gating ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_fp8_supported_requires_cuda_bf16_and_fp8_dtype(monkeypatch):
|
|
_stub_torch(monkeypatch, with_fp8 = True)
|
|
assert fp8_text_encoder_supported(_target()) is True
|
|
assert fp8_text_encoder_supported(_target(device = "cpu")) is False
|
|
assert fp8_text_encoder_supported(_target(dtype = "float16")) is False
|
|
|
|
|
|
def test_fp8_unsupported_without_fp8_dtype(monkeypatch):
|
|
_stub_torch(monkeypatch, with_fp8 = False)
|
|
assert fp8_text_encoder_supported(_target()) is False
|
|
|
|
|
|
# ── apply ─────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_apply_disabled_returns_empty(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
assert apply_fp8_text_encoder(pipe, _target(), enable = False) == []
|
|
|
|
|
|
def test_apply_casts_all_present_text_encoders(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
recorder: list = []
|
|
_stub_diffusers_hooks(monkeypatch, recorder)
|
|
te1, te3 = object(), object()
|
|
# text_encoder + text_encoder_3 present, text_encoder_2 absent.
|
|
pipe = types.SimpleNamespace(text_encoder = te1, text_encoder_2 = None, text_encoder_3 = te3)
|
|
cast = apply_fp8_text_encoder(pipe, _target(), enable = True)
|
|
assert cast == ["text_encoder", "text_encoder_3"]
|
|
# Casts to fp8 storage with bf16 compute and skips norms.
|
|
assert {r[0] for r in recorder} == {te1, te3}
|
|
assert all(r[1] == "float8_e4m3fn" and r[2] == "bfloat16" for r in recorder)
|
|
|
|
|
|
def test_apply_unsupported_target_is_noop(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
recorder: list = []
|
|
_stub_diffusers_hooks(monkeypatch, recorder)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
assert apply_fp8_text_encoder(pipe, _target(device = "cpu"), enable = True) == []
|
|
assert recorder == []
|
|
|
|
|
|
def test_apply_tolerates_casting_failure(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
hooks = types.ModuleType("diffusers.hooks")
|
|
casting = types.ModuleType("diffusers.hooks.layerwise_casting")
|
|
casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",)
|
|
|
|
def _boom(module, **kwargs):
|
|
raise RuntimeError("fp8 not supported for this layer")
|
|
|
|
hooks.apply_layerwise_casting = _boom
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
# A casting failure leaves the encoder dense and reports nothing cast.
|
|
assert apply_fp8_text_encoder(pipe, _target(), enable = True) == []
|