From 5df56c080400c438f40fb4eae0d6486536fcbe98 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 08:21:07 +0000 Subject: [PATCH] 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):