From f98a06b72b3acb63eb127e0ec31329d766057183 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:57:41 +0000 Subject: [PATCH] Stub is_bf16_supported in the train-precision-modes test helper The three train_precision_modes capability tests patched is_available and get_device_capability but not is_bf16_supported, so on the CPU-only CI runner (where the real probe is False) the dense modes collapsed to nf4 and the fp8 / mxfp8 assertions failed; they only passed on a bf16 GPU dev box. An Ada or Blackwell GPU is by definition bf16-capable, so the helper must stub it True to exercise the capability gate the tests target. --- studio/backend/tests/test_diffusion_dit_trainer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 0e8e8a0370..c02342b021 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -265,12 +265,16 @@ 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+). 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). + # (the CPU-only CI runner does not install it). is_bf16_supported must be stubbed True + # too: the dense modes gate on it, and an Ada/Blackwell GPU is by definition bf16-capable, + # so without this the modes collapse to nf4 on a CPU runner where the real probe is False + # (the test otherwise only passes on a bf16 GPU host). import torch import core.training.diffusion_train_common as dtc monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True)