torchao 0.17 removed MXLinearConfig from prototype.mx_formats in favour of
MXFP8TrainingOpConfig.from_recipe shared with MoE training. _mxfp8_training_config
tries the 0.16 API first and falls back to the 0.17 one; both feed quantize_.
mxfp8 still degrades to bf16 with a warning when neither import resolves
- base_precision="mxfp8": torchao MX block-scaled float8 compute on the frozen
base linears (Blackwell sm100+, cuBLAS kernels). Applied after add_adapter like
fp8, never fatal, weights stay bf16 in memory. Measured 1.16x over compiled
bf16 on Z-Image at 1024px batch 4 (16k tokens/step); a wash at small token
counts, so it stays an explicit opt-in and auto never picks it.
- SDXL: regionally compile the U-Net's BasicTransformerBlocks through the same
never-fatal wrapper the DiT trainer uses. 1.35x steady state at 1024px batch 4
with same-seed loss parity (~1e-5 per step) and unchanged peak VRAM; ~30 s
one-time warmup. Steady-state samples/sec now excludes step 1, matching the
DiT trainer.
- /info: mxfp8 advertised only on sm100+; supports_compile now true for sdxl.
- NVFP4 training: not available in torchao 0.16 (no autograd path, no training
recipe), so NVFP4 stays an inference-only quant for now.
193 diffusion backend tests green; frontend build clean.
Inference:
- krea-2 DiffusionFamily (Krea2Pipeline / Krea2Transformer2DModel, base
krea/Krea-2-Turbo, bf16 only, no GGUF/sd.cpp mapping yet)
- Per-component pipeline loader (core/inference/diffusion_krea2.py): the krea
repo is exported with transformers 5.2, so the tokenizer config
(extra_special_tokens as a list, no slow-tokenizer vocab files) and the
text encoder rope settings (rope_parameters vs rope_scaling) need explicit
compat on the 4.x line; values are copied verbatim and equal the 4.x
Qwen3-VL defaults, so the math is unchanged. from_pretrained also
type-checks the tokenizer against the declared slow class, so the pipeline
is assembled through its constructor with the model_index init config
(is_distilled carries Turbo's fixed mu=1.15 schedule)
- Trust allowlist entry, curated picker entry + 8 step / cfg 0 defaults,
int8 exclusion token for the M=1 Krea2TimestepEmbedding projection
Training:
- krea-2 _FamilySpec in the DiT trainer: phased conditioning/transformer
load through the compat loader, shared Qwen-Image VAE latent path,
fixed-512 text embeds (static shapes, plain concat collate), inline 2x2
latent packing + shared position grid, the authors' recommended LoRA
target set and rank/alpha 32, lr 3e-4, 512px presets
- GPU smokes on B200: nf4 2.9 steps/s at 11.5 GB, bf16 3.4 steps/s at
30.1 GB, bf16 + regional compile 5.2 steps/s; adapter round-trip
generation verified
Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.
