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

Stub the torchao probe in the precision-mode capability tests

Browse source 
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).
This commit is contained in:
Daniel Han 2026-07-04 08:21:07 +00:00
commit 5df56c0804
1 changed files with 6 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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