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

Merge diffusion-train-perf2: torchao probe stub in precision-mode tests

Browse source
This commit is contained in:
Daniel Han 2026-07-04 08:21:14 +00:00
commit 5fb79d39ad
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):