diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index f937add782..dd2f2e664d 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -208,12 +208,18 @@ def apply_speed_optims( # 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 -- the bf16 DiT paths are - # untouched -- so it engages wherever fp16 compute appears (SDXL pipelines, fp16 - # LoRA adapters). Guarded by a per-family deny-list fed by the overflow validation - # harness and the UNSLOTH_DISABLE_FP16_ACCUM kill switch. + # 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) + applied["fp16_accum"] = _enable_fp16_accumulation( + family, logger, dtype = getattr(target, "dtype", None), speed_mode = speed_mode + ) # --- the compile lever, remapped per tier ---------------------------------------- # default = LIGHT compile: for a GGUF model, compile ONLY the dequant op chain @@ -343,19 +349,25 @@ def _enable_tf32(logger: Any) -> bool: # Families the overflow validation harness (scripts/fp16_accum_validate.py) found to -# produce non-finite activations or visible drift under fp16 accumulation. Empty until -# a family actually fails: the gate below already restricts the flag to consumer GPUs -# and fp16 GEMMs, and the DiT families run bf16 compute (unaffected by this flag). +# 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) -> bool: +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, - and the UNSLOTH_DISABLE_FP16_ACCUM kill switch being unset. The caller's - snapshot/restore pair returns the process-wide flag to its prior value on unload.""" + 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() in ("1", "true", "yes"): @@ -363,6 +375,8 @@ def _enable_fp16_accumulation(family: Any, logger: Any) -> bool: 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 diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index dadaf73b93..9dc50f4e28 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -453,3 +453,36 @@ def test_fp16_accum_not_touched_off_cuda(monkeypatch): ) 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