From 69b437fa21237bedb1393e8526c8ab31958c4965 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 07:31:30 +0000
Subject: [PATCH 1/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion_speed.py |  6 +++++-
 studio/backend/tests/test_diffusion_speed.py     | 11 +++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py
index dd2f2e664d..cb9101d03d 100644
--- a/studio/backend/core/inference/diffusion_speed.py
+++ b/studio/backend/core/inference/diffusion_speed.py
@@ -357,7 +357,11 @@ _FP16_ACCUM_DENY: frozenset[str] = frozenset()
 
 
 def _enable_fp16_accumulation(
-    family: Any, logger: Any, *, dtype: Any = None, speed_mode: Optional[str] = None
+    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
diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py
index 9dc50f4e28..5b645022e1 100644
--- a/studio/backend/tests/test_diffusion_speed.py
+++ b/studio/backend/tests/test_diffusion_speed.py
@@ -357,7 +357,12 @@ def test_apply_tolerates_missing_optims(monkeypatch):
 # ── fp16 accumulation (consumer fp16-GEMM fast path) ──────────────────────────
 
 
-def _stub_torch_fp16_accum(monkeypatch, *, consumer = True, with_flag = True):
+def _stub_torch_fp16_accum(
+    monkeypatch,
+    *,
+    consumer = True,
+    with_flag = True,
+):
     torch = types.ModuleType("torch")
     torch.bfloat16 = "bfloat16"
     torch.channels_last = "channels_last"
@@ -427,9 +432,7 @@ def test_fp16_accum_respects_family_deny_list(monkeypatch):
     _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"
-    )
+    applied = apply_speed_optims(_Pipe(), _target(), is_gguf = True, family = fam, speed_mode = "default")
     assert applied["fp16_accum"] is False
 
 

From 5df56c080400c438f40fb4eae0d6486536fcbe98 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 08:21:07 +0000
Subject: [PATCH 2/2] Stub the torchao probe in the precision-mode capability
 tests

train_precision_modes gates int8/fp8/mxfp8 on has_functional_torchao, and the
Backend CI runner does not install torchao, so the three capability-gating
tests collapsed to nf4/bf16/auto and failed. They exercise the CAPABILITY
gate, not torchao presence: stub the probe functional alongside the CUDA
capability patch. Validated with a torchao-blocked run (22 passed).
---
 studio/backend/tests/test_diffusion_dit_trainer.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index b58868c701..6d7736a106 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -230,10 +230,15 @@ def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
 
 def _patch_capability(monkeypatch, capability):
     # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor
-    # core capability (fp8 needs sm89+, mxfp8 needs sm100+).
+    # core capability (fp8 needs sm89+, mxfp8 needs sm100+). The torchao probe is stubbed
+    # functional so these tests exercise the CAPABILITY gate on hosts without torchao
+    # (the CPU-only CI runner does not install it).
     import torch
+
+    import core.training.diffusion_train_common as dtc
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
+    monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True)
 
 
 def test_train_precision_modes_blackwell_lists_mxfp8(monkeypatch):