ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Stub is_bf16_supported in the train-precision-modes test helper

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-07 06:57:41 +00:00
commit f98a06b72b
1 changed files with 5 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

6
studio/backend/tests/test_diffusion_dit_trainer.py
View file

@ -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)