Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config

Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.

quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.
This commit is contained in:
Daniel Han 2026-06-28 06:07:11 +00:00
commit daad898d05
2 changed files with 11 additions and 7 deletions

View file

@ -27,7 +27,7 @@ REPO = "unsloth/Z-Image-Turbo-GGUF"
GGUF = "z-image-turbo-Q4_K_M.gguf"
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/probe_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images"
def _psnr(a, b):
@ -39,17 +39,19 @@ _LPIPS = {"fn": None}
def _lpips(ref_arr, arr):
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable."""
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable.
Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a
resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM."""
try:
import torch
import lpips
if _LPIPS["fn"] is None:
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
def t(x):
t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
return t.cuda()
return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
@ -111,7 +113,7 @@ def _quant_config(name):
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
except (TypeError, AttributeError):
return MXDynamicActivationMXWeightConfig()
raise ValueError(name)

View file

@ -265,7 +265,9 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
except (TypeError, AttributeError):
# TypeError: older torchao without the explicit dtype knobs.
# AttributeError: a torch build without torch.float8_e4m3fn.
return MXDynamicActivationMXWeightConfig()
raise ValueError(f"unknown transformer quant scheme '{scheme}'")