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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-26 08:01:43 +00:00
commit 8c3554e38c
3 changed files with 53 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

4
studio/backend/core/inference/diffusion.py
View file

@ -623,9 +623,7 @@ class DiffusionBackend:
pipe_kwargs["token"] = hf_token
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
pipe.to(device)
scheme = quantize_transformer(
pipe, target, mode = mode, fast_accum = fast_accum, logger = logger
)
scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger)
if scheme is None:
raise RuntimeError("transformer quant unsupported for this device/scheme")
return pipe, scheme

38
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -63,13 +63,35 @@ _SMOKE_CACHE: dict[tuple[str, str], bool] = {}
# torch.cuda.get_device_name(), so the workstation "A4000" is not mistaken for the
# data-center "A40". Anything not here -- GeForce, workstation RTX, or an unknown name --
# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus.
_DATACENTER_GPU_TOKENS = frozenset({
"B200", "B100", "GB200", "GB300", "GB10", # Blackwell data center
"H200", "H100", "H800", "H20", # Hopper data center
"A100", "A800", "A30", "A40", "A16", "A10", "A2", # Ampere data center
"L40", "L40S", "L4", "L20", "L2", # Ada data center
"V100", "P100", "P40", "T4", # legacy data center
})
_DATACENTER_GPU_TOKENS = frozenset(
{
"B200",
"B100",
"GB200",
"GB300",
"GB10", # Blackwell data center
"H200",
"H100",
"H800",
"H20", # Hopper data center
"A100",
"A800",
"A30",
"A40",
"A16",
"A10",
"A2", # Ampere data center
"L40",
"L40S",
"L4",
"L20",
"L2", # Ada data center
"V100",
"P100",
"P40",
"T4", # legacy data center
}
)
def _is_consumer_gpu(device: Any = None) -> bool:
@ -85,7 +107,6 @@ def _is_consumer_gpu(device: Any = None) -> bool:
import re
import torch
name = torch.cuda.get_device_name(device).upper()
except Exception: # noqa: BLE001 — no torch / no device -> assume consumer
return True
@ -224,7 +245,6 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
# quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations).
try:
from torchao.float8 import Float8MMConfig
return Float8DynamicActivationFloat8WeightConfig(
mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum))
)

40
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -212,31 +212,37 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
def _stub_device_name(monkeypatch, name):
torch = types.ModuleType("torch")
torch.cuda = types.SimpleNamespace(get_device_name=lambda device=None: name)
torch.cuda = types.SimpleNamespace(get_device_name = lambda device = None: name)
monkeypatch.setitem(sys.modules, "torch", torch)
@pytest.mark.parametrize("name", [
"NVIDIA GeForce RTX 5090",
"NVIDIA GeForce RTX 4090",
"NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40
"NVIDIA RTX 6000 Ada Generation",
"NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC)
])
@pytest.mark.parametrize(
"name",
[
"NVIDIA GeForce RTX 5090",
"NVIDIA GeForce RTX 4090",
"NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40
"NVIDIA RTX 6000 Ada Generation",
"NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC)
],
)
def test_is_consumer_gpu_true(monkeypatch, name):
_stub_device_name(monkeypatch, name)
assert tq._is_consumer_gpu() is True
@pytest.mark.parametrize("name", [
"NVIDIA B200",
"NVIDIA H100 80GB HBM3",
"NVIDIA A100-SXM4-80GB",
"NVIDIA A40", # data-center Ampere (distinct token from RTX A4000)
"NVIDIA L40S",
"NVIDIA L4",
"Tesla V100-SXM2-16GB",
])
@pytest.mark.parametrize(
"name",
[
"NVIDIA B200",
"NVIDIA H100 80GB HBM3",
"NVIDIA A100-SXM4-80GB",
"NVIDIA A40", # data-center Ampere (distinct token from RTX A4000)
"NVIDIA L40S",
"NVIDIA L4",
"Tesla V100-SXM2-16GB",
],
)
def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name):
_stub_device_name(monkeypatch, name)
assert tq._is_consumer_gpu() is False