Merge remote-tracking branch 'origin/image-generation' into fix/imggen-review-bugs
This commit is contained in:
oobabooga
commit
31ee7bae16
16 changed files with 297 additions and 56 deletions
|
|
@ -423,7 +423,9 @@ class DiffusionBackend:
|
|||
target = self._resolve_device_target(fam)
|
||||
if not dense_transformer_supported(target):
|
||||
return False
|
||||
scheme = select_transformer_quant_scheme(target, mode)
|
||||
scheme = select_transformer_quant_scheme(
|
||||
target, mode, family = getattr(fam, "name", None)
|
||||
)
|
||||
if scheme is None:
|
||||
return False
|
||||
source = resolve_prequant_source(
|
||||
|
|
@ -1329,7 +1331,7 @@ class DiffusionBackend:
|
|||
BEFORE the loader compiles the repeated block, so the order stays quantize ->
|
||||
compile -> placement."""
|
||||
# 1. Pre-quantized checkpoint, when one is configured for the resolved scheme.
|
||||
scheme = select_transformer_quant_scheme(target, mode)
|
||||
scheme = select_transformer_quant_scheme(target, mode, family = getattr(fam, "name", None))
|
||||
if scheme is None:
|
||||
# Bail BEFORE the (multi-GB) dense download: an explicit unsupported scheme
|
||||
# (e.g. fp8 on Ampere, nvfp4 off Blackwell) would otherwise materialise the
|
||||
|
|
@ -1368,7 +1370,14 @@ class DiffusionBackend:
|
|||
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
|
||||
)
|
||||
pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
|
||||
scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger)
|
||||
scheme = quantize_transformer(
|
||||
pipe,
|
||||
target,
|
||||
mode = mode,
|
||||
family = getattr(fam, "name", None),
|
||||
fast_accum = fast_accum,
|
||||
logger = logger,
|
||||
)
|
||||
if scheme is None:
|
||||
raise RuntimeError("transformer quant unsupported for this device/scheme")
|
||||
return pipe, scheme
|
||||
|
|
|
|||
|
|
@ -181,6 +181,11 @@ def install_compile_safe_patches() -> int:
|
|||
for cls, new_fn in _specs():
|
||||
if cls is None:
|
||||
continue
|
||||
# torch < 2.4 has no F.rms_norm: leave diffusers' original RMSNorm.forward in
|
||||
# place rather than installing a patch whose fast path would AttributeError.
|
||||
if cls is _RMSNorm and not hasattr(F, "rms_norm"):
|
||||
logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)")
|
||||
continue
|
||||
# Capture the live original BEFORE patching so the RMSNorm fast path can fall back
|
||||
# to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases.
|
||||
if cls is _RMSNorm:
|
||||
|
|
|
|||
|
|
@ -101,6 +101,30 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
|
|||
((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86
|
||||
)
|
||||
|
||||
# Families whose activation ranges break specific dense-quant schemes at the MODEL
|
||||
# level. The kernel smoke probe below cannot see this (it only proves the GEMM runs);
|
||||
# these were measured with the 28-pair prequant accuracy gate on a B200
|
||||
# (scripts/prequant_accuracy_gate.py) and reproduced with on-the-fly quantisation:
|
||||
# qwen-image + fp8 -> every frame black (mean luma 0.0000, SSIM 0.016 vs bf16). The
|
||||
# same per-row fp8 that matches bf16 on Z-Image / FLUX: Qwen's
|
||||
# activation outliers exceed even per-row fp8's dynamic range.
|
||||
# qwen-image + mxfp8 -> real semantic damage at 1024px (CLIP delta mean 0.0146, worst
|
||||
# cases 0.064 / 0.102 -- 2x the per-case bound).
|
||||
# qwen-image + nvfp4 -> LPIPS mean 0.51 vs bf16: unusable.
|
||||
# int8 dynamic (per-token) is excellent on Qwen (LPIPS mean 0.069 / SSIM 0.958), so the
|
||||
# auto ladder falls through to it. The deny also applies to an EXPLICIT request: a
|
||||
# scheme that renders black frames has no legitimate use, and returning None gives the
|
||||
# caller the same fallback contract as an unsupported scheme (GGUF build).
|
||||
_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
|
||||
"qwen-image": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}),
|
||||
"qwen-image-edit": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), # same DiT + activations
|
||||
}
|
||||
|
||||
|
||||
def _family_denied(family, scheme: str) -> bool:
|
||||
return scheme in _FAMILY_SCHEME_DENY.get(str(family or "").strip().lower(), ())
|
||||
|
||||
|
||||
# Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once.
|
||||
_SMOKE_CACHE: dict[tuple[str, str], bool] = {}
|
||||
|
||||
|
|
@ -201,18 +225,27 @@ def dense_transformer_supported(target: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]:
|
||||
def select_transformer_quant_scheme(
|
||||
target: Any,
|
||||
requested: Optional[str],
|
||||
family: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""The concrete scheme to apply, or None to fall back to GGUF.
|
||||
|
||||
``auto`` walks the per-arch ladder and returns the first scheme that passes a real
|
||||
quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are
|
||||
unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only
|
||||
if supported (else None -> GGUF), never silently swapped for a different one."""
|
||||
if supported (else None -> GGUF), never silently swapped for a different one.
|
||||
``family`` additionally applies the measured model-level deny list
|
||||
(``_FAMILY_SCHEME_DENY``): schemes that produce black frames or out-of-bar drift on
|
||||
that family are skipped by ``auto`` and refused when explicit."""
|
||||
requested = normalize_transformer_quant(requested)
|
||||
if requested is None or not dense_transformer_supported(target):
|
||||
return None
|
||||
device = str(getattr(target, "device", "cuda"))
|
||||
if requested != TQ_AUTO:
|
||||
if _family_denied(family, requested):
|
||||
return None
|
||||
return requested if _scheme_supported(requested, device) else None
|
||||
cap = _capability()
|
||||
if cap is None:
|
||||
|
|
@ -220,6 +253,8 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op
|
|||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
for scheme in _prefer_consumer_scheme(schemes, device):
|
||||
if _family_denied(family, scheme):
|
||||
continue
|
||||
if _scheme_supported(scheme, device):
|
||||
return scheme
|
||||
return None
|
||||
|
|
@ -385,6 +420,7 @@ def quantize_transformer(
|
|||
target: Any,
|
||||
*,
|
||||
mode: Optional[str],
|
||||
family: Optional[str] = None,
|
||||
min_features: int = DEFAULT_MIN_LINEAR_FEATURES,
|
||||
fast_accum: Optional[bool] = None,
|
||||
logger: Any = None,
|
||||
|
|
@ -396,7 +432,7 @@ def quantize_transformer(
|
|||
|
||||
``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None
|
||||
auto-detects (fast on consumer, precise on data-center), True/False force it."""
|
||||
scheme = select_transformer_quant_scheme(target, mode)
|
||||
scheme = select_transformer_quant_scheme(target, mode, family = family)
|
||||
if scheme is None:
|
||||
return None
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
|
|
|
|||
|
|
@ -486,6 +486,21 @@ def run_dit_lora_training(
|
|||
should_stop: Optional[StopCb] = None,
|
||||
) -> str:
|
||||
"""Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it."""
|
||||
cfg = config.normalized()
|
||||
spec = _SPECS.get(cfg.resolved_family)
|
||||
if spec is None:
|
||||
raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
|
||||
|
||||
# DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
|
||||
# explicitly asks for fp16 on a bf16-only family is refused rather than silently
|
||||
# upgraded, so the choice is never misrepresented. Validation runs before the heavy
|
||||
# imports so a host without diffusers still sees the real error.
|
||||
if cfg.mixed_precision == "fp16" and spec.force_bf16:
|
||||
raise ValueError(
|
||||
f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
|
||||
f"embedder internals. Set mixed precision to bf16."
|
||||
)
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
|
|
@ -493,11 +508,6 @@ def run_dit_lora_training(
|
|||
from peft import LoraConfig
|
||||
from peft.utils import get_peft_model_state_dict
|
||||
|
||||
cfg = config.normalized()
|
||||
spec = _SPECS.get(cfg.resolved_family)
|
||||
if spec is None:
|
||||
raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
|
||||
|
||||
rng = random.Random(cfg.seed)
|
||||
torch.manual_seed(cfg.seed)
|
||||
|
||||
|
|
@ -514,15 +524,6 @@ def run_dit_lora_training(
|
|||
save_on_stop = False
|
||||
return True
|
||||
|
||||
# DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
|
||||
# explicitly asks for fp16 on a bf16-only family is refused rather than silently
|
||||
# upgraded, so the choice is never misrepresented.
|
||||
if cfg.mixed_precision == "fp16" and spec.force_bf16:
|
||||
raise ValueError(
|
||||
f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
|
||||
f"embedder internals. Set mixed precision to bf16."
|
||||
)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
# The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
|
||||
# unsupported for real runs but keeps import/unit tests architecture-agnostic).
|
||||
|
|
@ -639,8 +640,11 @@ def run_dit_lora_training(
|
|||
(loss / cfg.gradient_accumulation_steps).backward()
|
||||
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
|
||||
|
||||
grad_norm: Optional[float] = None
|
||||
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
|
||||
torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
|
||||
# clip_grad_norm_ returns the PRE-clip total norm: the signal the grad-norm
|
||||
# chart wants (spikes stay visible even when clipping flattens the update).
|
||||
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
|
||||
optimizer.step()
|
||||
|
||||
running_loss += step_loss
|
||||
|
|
@ -661,6 +665,7 @@ def run_dit_lora_training(
|
|||
loss = round(step_loss, 5),
|
||||
avg_loss = round(running_loss / done, 5),
|
||||
learning_rate = cfg.learning_rate,
|
||||
grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
|
||||
samples_per_second = sps,
|
||||
peak_memory_gb = peak_gb or None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -340,8 +340,10 @@ def run_diffusion_lora_training(
|
|||
|
||||
# max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
|
||||
# passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
|
||||
grad_norm: Optional[float] = None
|
||||
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
|
||||
torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
|
||||
# clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal).
|
||||
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
|
||||
optimizer.step()
|
||||
lr_sched.step()
|
||||
|
||||
|
|
@ -366,6 +368,7 @@ def run_diffusion_lora_training(
|
|||
loss = round(step_loss, 5),
|
||||
avg_loss = round(running_loss / done, 5),
|
||||
learning_rate = lr_sched.get_last_lr()[0],
|
||||
grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
|
||||
samples_per_second = samples_per_second,
|
||||
peak_memory_gb = peak_gb or None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -384,9 +384,12 @@ def discover_image_caption_pairs(
|
|||
if sidecar.is_file():
|
||||
caption = sidecar.read_text(encoding = "utf-8").strip()
|
||||
break
|
||||
# 2. metadata row keyed by file name (basename or the name as written).
|
||||
# 2. metadata row keyed by file name (basename or the relative path; as_posix so a
|
||||
# Windows backslash path still matches the jsonl's forward-slash keys).
|
||||
if caption is None:
|
||||
caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root)))
|
||||
caption = meta_caption.get(img.name) or meta_caption.get(
|
||||
img.relative_to(root).as_posix()
|
||||
)
|
||||
# 3. dreambooth instance prompt.
|
||||
if caption is None and instance_prompt:
|
||||
caption = instance_prompt
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ def _idle_state() -> dict[str, Any]:
|
|||
"loss": None,
|
||||
"avg_loss": None,
|
||||
"learning_rate": None,
|
||||
"grad_norm": None,
|
||||
"num_images": None,
|
||||
"in_model_load": False,
|
||||
"output_dir": None,
|
||||
|
|
@ -82,16 +83,25 @@ def _idle_state() -> dict[str, Any]:
|
|||
"metric_steps": [],
|
||||
"metric_loss": [],
|
||||
"metric_lr": [],
|
||||
"metric_grad_norm": [],
|
||||
}
|
||||
|
||||
|
||||
def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None:
|
||||
"""Append one (step, loss, lr) point to the bounded history arrays on ``state``.
|
||||
def _append_metric(
|
||||
state: dict[str, Any],
|
||||
step: Any,
|
||||
loss: Any,
|
||||
lr: Any,
|
||||
grad_norm: Any = None,
|
||||
) -> None:
|
||||
"""Append one (step, loss, lr, grad_norm) point to the bounded history arrays on
|
||||
``state``.
|
||||
|
||||
Only records finite, positive-step points (mirrors the LLM trainer, which logs history
|
||||
only for step > 0 with a real loss). When the arrays hit ``_METRIC_CAP`` they are
|
||||
decimated in place (keep every other point) so appends stay bounded without losing the
|
||||
curve's shape. lr may be None (kept as None so the LR series can be sparse)."""
|
||||
curve's shape. lr / grad_norm may be None (kept as None so those series can be
|
||||
sparse while staying index-aligned with ``steps``)."""
|
||||
try:
|
||||
istep = int(step)
|
||||
except (TypeError, ValueError):
|
||||
|
|
@ -104,22 +114,34 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None
|
|||
return
|
||||
if floss != floss: # NaN guard
|
||||
return
|
||||
flr: Optional[float]
|
||||
try:
|
||||
flr = float(lr) if lr is not None else None
|
||||
except (TypeError, ValueError):
|
||||
flr = None
|
||||
|
||||
def _opt_float(v: Any) -> Optional[float]:
|
||||
try:
|
||||
return float(v) if v is not None else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
flr = _opt_float(lr)
|
||||
fgn = _opt_float(grad_norm)
|
||||
steps = state["metric_steps"]
|
||||
losses = state["metric_loss"]
|
||||
lrs = state["metric_lr"]
|
||||
gns = state["metric_grad_norm"]
|
||||
if len(steps) >= _METRIC_CAP:
|
||||
state["metric_steps"] = steps[::2]
|
||||
state["metric_loss"] = losses[::2]
|
||||
state["metric_lr"] = lrs[::2]
|
||||
steps, losses, lrs = state["metric_steps"], state["metric_loss"], state["metric_lr"]
|
||||
state["metric_grad_norm"] = gns[::2]
|
||||
steps, losses, lrs, gns = (
|
||||
state["metric_steps"],
|
||||
state["metric_loss"],
|
||||
state["metric_lr"],
|
||||
state["metric_grad_norm"],
|
||||
)
|
||||
steps.append(istep)
|
||||
losses.append(floss)
|
||||
lrs.append(flr)
|
||||
gns.append(fgn)
|
||||
|
||||
|
||||
class DiffusionTrainingService:
|
||||
|
|
@ -288,6 +310,7 @@ class DiffusionTrainingService:
|
|||
loss = ev.get("loss", s["loss"]),
|
||||
avg_loss = ev.get("avg_loss", s["avg_loss"]),
|
||||
learning_rate = ev.get("learning_rate", s["learning_rate"]),
|
||||
grad_norm = ev.get("grad_norm", s["grad_norm"]),
|
||||
message = "Training...",
|
||||
)
|
||||
# Fold optional perf fields (emitted by the trainers) so the UI can show
|
||||
|
|
@ -296,8 +319,14 @@ class DiffusionTrainingService:
|
|||
s["samples_per_second"] = ev.get("samples_per_second")
|
||||
if ev.get("peak_memory_gb") is not None:
|
||||
s["peak_memory_gb"] = ev.get("peak_memory_gb")
|
||||
# Retain a bounded (step, loss, lr) history for the live loss chart.
|
||||
_append_metric(s, ev.get("step"), ev.get("loss"), ev.get("learning_rate"))
|
||||
# Retain a bounded (step, loss, lr, grad_norm) history for the live charts.
|
||||
_append_metric(
|
||||
s,
|
||||
ev.get("step"),
|
||||
ev.get("loss"),
|
||||
ev.get("learning_rate"),
|
||||
ev.get("grad_norm"),
|
||||
)
|
||||
elif etype == "complete":
|
||||
# Reset in_model_load: a stop during model load emits complete without a
|
||||
# preceding model_load_completed, which would otherwise leave a stale
|
||||
|
|
|
|||
|
|
@ -734,12 +734,14 @@ class DiffusionTrainingStartResponse(BaseModel):
|
|||
|
||||
|
||||
class DiffusionMetricHistory(BaseModel):
|
||||
"""Paired step-indexed history arrays for the live training charts. ``lr`` entries may
|
||||
be null so a sparse learning-rate series still aligns with ``steps`` by index."""
|
||||
"""Paired step-indexed history arrays for the live training charts. ``lr`` and
|
||||
``grad_norm`` entries may be null so those sparse series still align with ``steps``
|
||||
by index."""
|
||||
|
||||
steps: List[int] = Field(default_factory = list)
|
||||
loss: List[float] = Field(default_factory = list)
|
||||
lr: List[Optional[float]] = Field(default_factory = list)
|
||||
grad_norm: List[Optional[float]] = Field(default_factory = list)
|
||||
|
||||
|
||||
class DiffusionTrainingStatusResponse(BaseModel):
|
||||
|
|
@ -754,6 +756,9 @@ class DiffusionTrainingStatusResponse(BaseModel):
|
|||
loss: Optional[float] = None
|
||||
avg_loss: Optional[float] = None
|
||||
learning_rate: Optional[float] = None
|
||||
# Pre-clip gradient norm from the trainer's progress events (None when clipping is
|
||||
# disabled), feeding the grad-norm chart.
|
||||
grad_norm: Optional[float] = None
|
||||
num_images: Optional[int] = None
|
||||
in_model_load: bool = False
|
||||
output_dir: Optional[str] = None
|
||||
|
|
|
|||
|
|
@ -1267,6 +1267,7 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s
|
|||
steps = snap.pop("metric_steps", []),
|
||||
loss = snap.pop("metric_loss", []),
|
||||
lr = snap.pop("metric_lr", []),
|
||||
grad_norm = snap.pop("metric_grad_norm", []),
|
||||
)
|
||||
return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history)
|
||||
|
||||
|
|
@ -1515,7 +1516,15 @@ def _image_record(
|
|||
caption = None
|
||||
break
|
||||
if caption is None:
|
||||
# Basename first, then the relative path as written in the jsonl (as_posix so a
|
||||
# Windows backslash path still matches forward-slash keys) -- the same lookup
|
||||
# order discover_image_caption_pairs uses.
|
||||
meta = meta_captions.get(image_path.name)
|
||||
if meta is None:
|
||||
try:
|
||||
meta = meta_captions.get(image_path.relative_to(folder).as_posix())
|
||||
except ValueError:
|
||||
meta = None
|
||||
if meta is not None:
|
||||
caption = meta
|
||||
source = "metadata"
|
||||
|
|
|
|||
|
|
@ -1764,7 +1764,9 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"):
|
|||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
# Resolve the scheme without the real GPU smoke probe, and configure no pre-quant
|
||||
# checkpoint so the dense materialise+quantise branch is the one exercised.
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: scheme)
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: scheme
|
||||
)
|
||||
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
|
||||
|
||||
def _quantize(pipe, target, *, mode, **kw):
|
||||
|
|
@ -1828,7 +1830,9 @@ def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeyp
|
|||
backend = DiffusionBackend()
|
||||
_force_cuda_target(backend, monkeypatch)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
|
||||
)
|
||||
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
|
||||
prequant_obj = object()
|
||||
loaded: dict = {"n": 0}
|
||||
|
|
@ -1958,7 +1962,9 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download(
|
|||
backend = DiffusionBackend()
|
||||
_force_cuda_target(backend, monkeypatch)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None)
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None
|
||||
)
|
||||
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -2005,7 +2011,9 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
|
|||
_force_cuda_target(backend, monkeypatch)
|
||||
fam = detect_family("unsloth/Z-Image-Turbo-GGUF")
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
|
||||
)
|
||||
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
|
||||
|
||||
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True
|
||||
|
|
@ -2016,10 +2024,14 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
|
|||
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
|
||||
# Unsupported scheme bails before the dense path (and so must the prefetch).
|
||||
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None)
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None
|
||||
)
|
||||
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
|
||||
# Device without dense support (e.g. non-CUDA) never widens.
|
||||
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
|
||||
monkeypatch.setattr(
|
||||
dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
|
||||
)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False)
|
||||
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
|
||||
|
||||
|
|
|
|||
|
|
@ -659,6 +659,17 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch):
|
|||
backend = _FakeBackend()
|
||||
backend.begin_load = _busy
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
# Pin the resolved device to cuda: the route only takes the arbiter for non-CPU
|
||||
# loads, so on a CPU-only host the ownership assert below would never hold.
|
||||
import types as _types
|
||||
|
||||
import core.inference.diffusion_device as devmod
|
||||
|
||||
monkeypatch.setattr(
|
||||
devmod,
|
||||
"resolve_diffusion_device_target",
|
||||
lambda: _types.SimpleNamespace(device = "cuda"),
|
||||
)
|
||||
resp = client.post(
|
||||
"/api/inference/images/load",
|
||||
json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"},
|
||||
|
|
|
|||
|
|
@ -433,7 +433,9 @@ def test_fp8_config_uses_per_row_granularity():
|
|||
|
||||
|
||||
def test_quantize_transformer_applies_and_marks(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8)
|
||||
monkeypatch.setattr(
|
||||
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
|
||||
)
|
||||
seen: dict = {}
|
||||
|
||||
def _mk(scheme, fast_accum = None):
|
||||
|
|
@ -458,13 +460,17 @@ def test_quantize_transformer_applies_and_marks(monkeypatch):
|
|||
|
||||
|
||||
def test_quantize_transformer_none_when_unsupported(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None)
|
||||
monkeypatch.setattr(
|
||||
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: None
|
||||
)
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
assert quantize_transformer(pipe, _target(), mode = "auto") is None
|
||||
|
||||
|
||||
def test_quantize_transformer_tolerates_failure(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8)
|
||||
monkeypatch.setattr(
|
||||
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_INT8
|
||||
)
|
||||
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg")
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
|
||||
|
|
@ -480,3 +486,60 @@ def test_quantize_transformer_tolerates_failure(monkeypatch):
|
|||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
# A quantise failure returns None (caller falls back to GGUF), never raises.
|
||||
assert quantize_transformer(pipe, _target(), mode = "int8") is None
|
||||
|
||||
|
||||
# ── family scheme deny (measured model-level breakage) ────────────────────────
|
||||
|
||||
|
||||
def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch):
|
||||
# B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the
|
||||
# Qwen DiT (per-row fp8 renders black frames on it; see _FAMILY_SCHEME_DENY) and
|
||||
# falls through the ladder to int8, which measures excellent on Qwen.
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image-edit") == TQ_INT8
|
||||
|
||||
|
||||
def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch):
|
||||
# An explicit fp8 request on qwen-image returns None (same contract as an
|
||||
# unsupported scheme: the caller builds the GGUF pipeline instead). int8 stays
|
||||
# honored on qwen, and fp8 stays honored on families outside the deny table.
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "fp8", family = "qwen-image") is None
|
||||
assert select_transformer_quant_scheme(_target(), "int8", family = "qwen-image") == TQ_INT8
|
||||
assert select_transformer_quant_scheme(_target(), "fp8", family = "z-image") == TQ_FP8
|
||||
|
||||
|
||||
def test_family_deny_no_family_keeps_ladder(monkeypatch):
|
||||
# Without a family (or an unknown one) the ladder is unchanged: fp8 first on B200.
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "sdxl") == TQ_FP8
|
||||
|
||||
|
||||
def test_quantize_transformer_threads_family(monkeypatch):
|
||||
# quantize_transformer passes the family down to the selector, so a denied
|
||||
# (family, scheme) pair never reaches torchao.
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
called = {}
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
|
||||
def _quantize(
|
||||
module,
|
||||
config,
|
||||
filter_fn = None,
|
||||
):
|
||||
called["scheme"] = True
|
||||
|
||||
tqz.quantize_ = _quantize
|
||||
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8-cfg"
|
||||
tqz.Float8DynamicActivationFloat8WeightConfig = lambda **kw: "fp8-cfg"
|
||||
tqz.PerRow = lambda: "per-row"
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
assert quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None
|
||||
assert called == {}
|
||||
|
|
|
|||
|
|
@ -291,6 +291,7 @@ export interface DiffusionMetricHistory {
|
|||
steps: number[];
|
||||
loss: number[];
|
||||
lr: Array<number | null>;
|
||||
grad_norm: Array<number | null>;
|
||||
}
|
||||
|
||||
// A snapshot of the current diffusion training job (GET /api/train/diffusion/status).
|
||||
|
|
|
|||
|
|
@ -481,7 +481,7 @@ function AdvancedSelect({
|
|||
return (
|
||||
<div className="flex flex-col gap-1">
|
||||
<div className="flex items-center justify-between gap-2">
|
||||
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
|
||||
<span className="flex shrink-0 items-center gap-1 whitespace-nowrap text-xs font-medium text-muted-foreground">
|
||||
{label}
|
||||
{hint && <InfoHint>{hint}</InfoHint>}
|
||||
</span>
|
||||
|
|
@ -1899,8 +1899,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
to GGUF (or nothing loaded) and otherwise show why it is unavailable. */}
|
||||
{!status?.loaded || status.model_kind === "gguf" ? (
|
||||
<AdvancedSelect
|
||||
label="GGUF compute"
|
||||
desc="Off runs the GGUF as-is. INT8/FP8/FP4 dequantise the transformer onto low-precision tensor cores for a faster step, at the cost of a larger download and more VRAM."
|
||||
label="Dtype"
|
||||
hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
|
||||
value={transformerQuant}
|
||||
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
|
||||
|
|
@ -1916,7 +1915,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
) : (
|
||||
<div className="flex items-center justify-between gap-2">
|
||||
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
|
||||
GGUF compute
|
||||
Dtype
|
||||
</span>
|
||||
<span className="text-xs text-muted-foreground/60">GGUF models only</span>
|
||||
</div>
|
||||
|
|
@ -2624,7 +2623,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
<p className="text-sm">
|
||||
{status?.loaded
|
||||
? "Enter a prompt and hit Generate."
|
||||
: "Select a model quant to load, then generate."}
|
||||
: "Select a diffusion model to load"}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ import type { TrainingSeriesPoint } from "@/features/training";
|
|||
// which are meaningless for diffusion LoRA training and showed as an empty card and an
|
||||
// "Evaluation not configured" placeholder. This is a diffusion-only two-card layout.
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-chart-card";
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card";
|
||||
// eslint-disable-next-line no-restricted-imports
|
||||
import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card";
|
||||
|
|
@ -43,14 +45,17 @@ function fullStepDomain(steps: number[]): [number, number] {
|
|||
return [min, max];
|
||||
}
|
||||
|
||||
// A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a
|
||||
// note under the loss card explaining why per-step loss looks noisy.
|
||||
// A diffusion-only metrics view: Training Loss and Learning Rate side by side, plus Grad
|
||||
// Norm (the pre-clip total gradient norm; spikes flag instability that raw loss noise
|
||||
// hides), with a note under the loss card explaining why per-step loss looks noisy.
|
||||
export function DiffusionCharts({
|
||||
lossHistory,
|
||||
lrHistory,
|
||||
gradNormHistory = [],
|
||||
}: {
|
||||
lossHistory: TrainingSeriesPoint[];
|
||||
lrHistory: TrainingSeriesPoint[];
|
||||
gradNormHistory?: TrainingSeriesPoint[];
|
||||
}): ReactElement | null {
|
||||
const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]);
|
||||
const smoothed = useMemo(
|
||||
|
|
@ -82,12 +87,24 @@ export function DiffusionCharts({
|
|||
[lrHistory],
|
||||
);
|
||||
|
||||
const gradNormData = useMemo(
|
||||
() =>
|
||||
compressSeries(
|
||||
gradNormHistory
|
||||
.filter((p) => Number.isFinite(p.value))
|
||||
.map((p) => ({ step: p.step, gradNorm: p.value, displayGradNorm: p.value })),
|
||||
MAX_RENDER_POINTS,
|
||||
),
|
||||
[gradNormHistory],
|
||||
);
|
||||
|
||||
const steps = useMemo(() => {
|
||||
const set = new Set<number>();
|
||||
for (const p of lossData) set.add(p.step);
|
||||
for (const p of lrData) set.add(p.step);
|
||||
for (const p of gradNormData) set.add(p.step);
|
||||
return Array.from(set).sort((a, b) => a - b);
|
||||
}, [lossData, lrData]);
|
||||
}, [lossData, lrData, gradNormData]);
|
||||
|
||||
const stepDomain = useMemo(() => fullStepDomain(steps), [steps]);
|
||||
const xAxisTicks = useMemo(
|
||||
|
|
@ -103,6 +120,10 @@ export function DiffusionCharts({
|
|||
() => buildYDomain(lrData.map((p) => p.displayLr)),
|
||||
[lrData],
|
||||
);
|
||||
const gradNormDomain = useMemo(
|
||||
() => buildYDomain(gradNormData.map((p) => p.displayGradNorm)),
|
||||
[gradNormData],
|
||||
);
|
||||
|
||||
const avgRaw =
|
||||
lossItems.length > 0
|
||||
|
|
@ -138,6 +159,15 @@ export function DiffusionCharts({
|
|||
xAxisTicks={xAxisTicks}
|
||||
scale="linear"
|
||||
/>
|
||||
{gradNormData.length > 0 && (
|
||||
<GradNormChartCard
|
||||
data={gradNormData}
|
||||
domain={gradNormDomain}
|
||||
visibleStepDomain={stepDomain}
|
||||
xAxisTicks={xAxisTicks}
|
||||
scale="linear"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -366,6 +366,13 @@ export function DiffusionTrainPanel({
|
|||
.map((step, i) => ({ step, value: h.lr[i] }))
|
||||
.filter((p): p is TrainingSeriesPoint => p.value != null);
|
||||
}, [status?.metric_history]);
|
||||
const gradNormHistory: TrainingSeriesPoint[] = useMemo(() => {
|
||||
const h = status?.metric_history;
|
||||
if (!h) return [];
|
||||
return h.steps
|
||||
.map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null }))
|
||||
.filter((p): p is TrainingSeriesPoint => p.value != null);
|
||||
}, [status?.metric_history]);
|
||||
|
||||
const onUpload = useCallback(async () => {
|
||||
const files = Array.from(fileInputRef.current?.files ?? []);
|
||||
|
|
@ -791,7 +798,17 @@ export function DiffusionTrainPanel({
|
|||
<>
|
||||
<div className="bg-card corner-squircle flex flex-col gap-3 rounded-3xl p-5 ring-1 ring-foreground/10">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-sm font-semibold capitalize">{status.status}</span>
|
||||
{/* A finished run should be unmistakable at a glance, so completed swaps
|
||||
the plain status word for a celebratory line in the success color. */}
|
||||
<span
|
||||
className={
|
||||
status.status === "completed"
|
||||
? "text-sm font-semibold text-emerald-600 dark:text-emerald-400"
|
||||
: "text-sm font-semibold capitalize"
|
||||
}
|
||||
>
|
||||
{status.status === "completed" ? "Training complete \u{1F389}" : status.status}
|
||||
</span>
|
||||
<span className="text-xs text-muted-foreground">
|
||||
{status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""}
|
||||
</span>
|
||||
|
|
@ -828,7 +845,11 @@ export function DiffusionTrainPanel({
|
|||
)}
|
||||
</div>
|
||||
|
||||
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
|
||||
<DiffusionCharts
|
||||
lossHistory={lossHistory}
|
||||
lrHistory={lrHistory}
|
||||
gradNormHistory={gradNormHistory}
|
||||
/>
|
||||
|
||||
{completed && (
|
||||
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue