diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 9c99de9406..344a5e66fa 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -69,6 +69,8 @@ def snapshot_backend_flags() -> Optional[dict]: matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None) if matmul is not None and hasattr(matmul, "allow_tf32"): state["matmul_tf32"] = bool(matmul.allow_tf32) + if matmul is not None and hasattr(matmul, "allow_fp16_accumulation"): + state["matmul_fp16_accum"] = bool(matmul.allow_fp16_accumulation) cudnn = getattr(torch.backends, "cudnn", None) if cudnn is not None: if hasattr(cudnn, "allow_tf32"): @@ -95,9 +97,9 @@ def restore_backend_flags(state: Optional[dict]) -> None: except Exception: # noqa: BLE001 — best-effort per-flag restore pass - _set( - getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32" - ) + matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None) + _set(matmul, "allow_tf32", "matmul_tf32") + _set(matmul, "allow_fp16_accumulation", "matmul_fp16_accum") cudnn = getattr(torch.backends, "cudnn", None) _set(cudnn, "allow_tf32", "cudnn_tf32") _set(cudnn, "benchmark", "cudnn_benchmark") @@ -179,6 +181,7 @@ def apply_speed_optims( "channels_last": False, "cudnn_benchmark": False, "tf32": False, + "fp16_accum": False, "fused_qkv": False, "compiled": False, "compiled_dequant": False, @@ -203,6 +206,21 @@ def apply_speed_optims( if on_cuda: applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) + # Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts, whose + # fp32-accumulate rate is halved; datacenter HBM parts gain nothing and keep the + # safer fp32 accumulate). Only affects fp16 matmuls -- bf16 loads were measured + # bit-identical with the flag on across every family (36/36 same-seed A/B cases), + # so on the quality-neutral tiers the flag engages only when the compute dtype is + # NOT fp16. On an fp16 pipeline (the pre-Ampere fallback dtype) the same harness + # measured real same-seed drift (mean 2-5% on SDXL / FLUX), so fp16 compute gets + # the 2x accumulate only under ``max``, the tier that already trades exactness for + # speed. Guarded by a per-family deny-list fed by the overflow validation harness + # and the UNSLOTH_DISABLE_FP16_ACCUM kill switch. + if on_cuda: + applied["fp16_accum"] = _enable_fp16_accumulation( + family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode + ) + # --- the compile lever, remapped per tier ---------------------------------------- # default = LIGHT compile: for a GGUF model, compile ONLY the dequant op chain # (~70-80% of eager GGUF time) -- cheap, VRAM-free, resolution-invariant; the @@ -330,6 +348,61 @@ def _enable_tf32(logger: Any) -> bool: return False +# Families the overflow validation harness (scripts/fp16_accum_validate.py) found to +# produce non-finite activations or NEW black frames under fp16 accumulation. Empty by +# measurement: across all six families the harness found no overflow anywhere -- bf16 +# loads are bit-identical with the flag on, and fp16 loads stay finite (their same-seed +# drift is why fp16 compute is additionally gated to the ``max`` tier below). +_FP16_ACCUM_DENY: frozenset[str] = frozenset() + + +def _enable_fp16_accumulation( + family: Any, + logger: Any, + *, + dtype: Any = None, + speed_mode: Optional[str] = None, +) -> bool: + """Turn on fp16-accumulated fp16 GEMMs for consumer GPUs, where they run ~2x the + fp32-accumulate rate (datacenter HBM parts are not throughput-nerfed, so they keep + the safer default). Gated on: the torch build exposing the flag (2.10+), a + consumer-class device, the family not being deny-listed by the overflow harness, + the UNSLOTH_DISABLE_FP16_ACCUM kill switch being unset, and -- when the pipeline + compute dtype IS fp16, the only case where the accumulator width changes results -- + the ``max`` tier (measured same-seed drift: mean 2-5%; bf16 loads are bit-identical + so they engage on any tier). The caller's snapshot/restore pair returns the + process-wide flag to its prior value on unload.""" + import os + + if os.environ.get("UNSLOTH_DISABLE_FP16_ACCUM", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ): + return False + name = str(getattr(family, "name", family or "")).lower() + if name in _FP16_ACCUM_DENY: + return False + if str(dtype).replace("torch.", "") == "float16" and speed_mode != SPEED_MAX: + return False + try: + import torch + + matmul = torch.backends.cuda.matmul + if not hasattr(matmul, "allow_fp16_accumulation"): + return False + from .diffusion_transformer_quant import _is_consumer_gpu + + if not _is_consumer_gpu(): + return False + matmul.allow_fp16_accumulation = True + return True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "fp16_accum", exc) + return False + + def _fuse_qkv(pipe: Any, logger: Any) -> bool: for owner in (pipe, getattr(pipe, "transformer", None)): fn = getattr(owner, "fuse_qkv_projections", None) diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index de73121e00..e6ce77116e 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -211,6 +211,7 @@ def test_speed_off_applies_nothing(monkeypatch): "fused_qkv": False, "compiled": False, "compiled_dequant": False, + "fp16_accum": False, } assert pipe.vae.mem_format is None and pipe.compiled is False # off must not touch any process-wide flag (bit-identical reference path). @@ -351,3 +352,154 @@ def test_apply_tolerates_missing_optims(monkeypatch): bare, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX ) assert applied["channels_last"] is False and applied["fused_qkv"] is False + + +# ── fp16 accumulation (consumer fp16-GEMM fast path) ────────────────────────── + + +def _stub_torch_fp16_accum( + monkeypatch, + *, + consumer = True, + with_flag = True, +): + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.channels_last = "channels_last" + matmul_attrs = {"allow_tf32": False} + if with_flag: + matmul_attrs["allow_fp16_accumulation"] = False + torch.backends = types.SimpleNamespace( + cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(**matmul_attrs)), + cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False), + ) + monkeypatch.setitem(sys.modules, "torch", torch) + import core.inference.diffusion_transformer_quant as tq + + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda device = None: consumer) + return torch + + +def test_snapshot_captures_fp16_accum_when_present(monkeypatch): + torch = _stub_torch_fp16_accum(monkeypatch) + torch.backends.cuda.matmul.allow_fp16_accumulation = True + snap = snapshot_backend_flags() + assert snap["matmul_fp16_accum"] is True + torch.backends.cuda.matmul.allow_fp16_accumulation = False + restore_backend_flags(snap) + assert torch.backends.cuda.matmul.allow_fp16_accumulation is True + + +def test_snapshot_skips_fp16_accum_on_older_torch(monkeypatch): + _stub_torch_fp16_accum(monkeypatch, with_flag = False) + snap = snapshot_backend_flags() + assert "matmul_fp16_accum" not in snap + restore_backend_flags(snap) # nothing to restore, no error + + +def test_fp16_accum_engages_on_consumer_cuda(monkeypatch): + torch = _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + applied = apply_speed_optims( + _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default" + ) + assert applied["fp16_accum"] is True + assert torch.backends.cuda.matmul.allow_fp16_accumulation is True + + +def test_fp16_accum_skipped_on_datacenter(monkeypatch): + torch = _stub_torch_fp16_accum(monkeypatch, consumer = False) + _stub_gguf_accel(monkeypatch) + applied = apply_speed_optims( + _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default" + ) + assert applied["fp16_accum"] is False + assert torch.backends.cuda.matmul.allow_fp16_accumulation is False + + +def test_fp16_accum_respects_kill_switch(monkeypatch): + _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + monkeypatch.setenv("UNSLOTH_DISABLE_FP16_ACCUM", "1") + applied = apply_speed_optims( + _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default" + ) + assert applied["fp16_accum"] is False + + +@pytest.mark.parametrize("value", ["TRUE", "Yes", "On", " true "]) +def test_fp16_accum_kill_switch_is_case_insensitive(monkeypatch, value): + # The documented safety escape hatch must honor the common boolean spellings, not only + # lowercase "1"/"true"/"yes": an operator setting UNSLOTH_DISABLE_FP16_ACCUM=TRUE to stop + # fp16-accumulation drift would otherwise be silently ignored. + _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + monkeypatch.setenv("UNSLOTH_DISABLE_FP16_ACCUM", value) + applied = apply_speed_optims( + _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default" + ) + assert applied["fp16_accum"] is False + + +def test_fp16_accum_respects_family_deny_list(monkeypatch): + _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + monkeypatch.setattr(ds_mod, "_FP16_ACCUM_DENY", frozenset({"fragile-family"})) + fam = types.SimpleNamespace(supports_torch_compile = True, name = "fragile-family") + applied = apply_speed_optims(_Pipe(), _target(), is_gguf = True, family = fam, speed_mode = "default") + assert applied["fp16_accum"] is False + + +def test_fp16_accum_skipped_when_flag_missing(monkeypatch): + _stub_torch_fp16_accum(monkeypatch, consumer = True, with_flag = False) + _stub_gguf_accel(monkeypatch) + applied = apply_speed_optims( + _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default" + ) + assert applied["fp16_accum"] is False + + +def test_fp16_accum_not_touched_off_cuda(monkeypatch): + torch = _stub_torch_fp16_accum(monkeypatch, consumer = True) + applied = apply_speed_optims( + _Pipe(), + _target(device = "mps"), + is_gguf = False, + family = _family(), + speed_mode = "eager", + ) + assert applied["fp16_accum"] is False + assert torch.backends.cuda.matmul.allow_fp16_accumulation is False + + +def test_fp16_accum_denied_on_fp16_dtype_below_max(monkeypatch): + # fp16 compute is where the accumulator width actually changes results (measured + # same-seed drift, mean 2-5%): the quality-neutral tiers must refuse it. + torch = _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + for mode in ("eager", "default"): + applied = apply_speed_optims( + _Pipe(), + _target(dtype = "float16"), + is_gguf = True, + family = _family(), + speed_mode = mode, + ) + assert applied["fp16_accum"] is False + assert torch.backends.cuda.matmul.allow_fp16_accumulation is False + + +def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch): + # max already trades exactness for speed (conv algos, max-autotune), so the 2x + # fp16 accumulate joins that tier for fp16 pipelines. + torch = _stub_torch_fp16_accum(monkeypatch, consumer = True) + _stub_gguf_accel(monkeypatch) + applied = apply_speed_optims( + _Pipe(with_compile = True, with_fuse = True), + _target(dtype = "float16"), + is_gguf = True, + family = _family(), + speed_mode = "MAX", + ) + assert applied["fp16_accum"] is True + assert torch.backends.cuda.matmul.allow_fp16_accumulation is True