Shift Qwen-Image training sigmas to the inference distribution
Qwen-Image's scheduler skips its static shift under use_dynamic_shifting, so the DiT trainer was drawing UNSHIFTED uniform-schedule sigmas for it (mean sigma 0.50) while inference always runs the exponential mu = log 3 shift plus the shift_terminal 0.02 stretch. Add a flow_shift config lever: "auto" (the new qwen-image default) rebuilds the training sigma table through the scheduler's own time_shift and stretch_shift_to_terminal so the draw matches the inference distribution exactly (mean sigma 0.72); a numeric value applies the standard linear shift s*u/(1+(s-1)*u); 1.0 keeps the historical identity behavior and stays the default for FLUX, Z-Image and Krea 2. The model timestep conditioning follows the shifted sigma, gathered in fp32 so bf16 rounding never skews it. Also wire two opt-in levers with off defaults: cfg_dropout (per-sample empty-prompt conditioning dropout, encoded alongside the captions before the text encoders are freed) and weighting_scheme="bell" (bsmntw-style mid-schedule Gaussian loss weighting normalized to mean 1). Verified with two 80-step rank-8 bf16 LoRA runs on Qwen/Qwen-Image (identity vs auto, same seed): both converge with finite decreasing loss and produce coherent same-seed previews. Unit tests cover the exact transform, the shifted sampling distribution, per-family defaults and config plumbing.
This commit is contained in:
Daniel Han
parent
5350a59cbf
commit
039049b05e
3 changed files with 374 additions and 14 deletions
|
|
@ -27,6 +27,7 @@ optimizer state and the frozen 4-bit base sit in VRAM during the loop.
|
|||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
|
|
@ -146,17 +147,62 @@ class _FamilySpec:
|
|||
|
||||
|
||||
# ── shared flow-matching helpers ──────────────────────────────────────────────
|
||||
def _gather_sigmas(scheduler, indices, device, dtype, n_dim):
|
||||
def _gather_sigmas(sigma_table, indices, device, dtype, n_dim):
|
||||
"""Gather per-sample sigmas for schedule ``indices`` and broadcast to ``n_dim``.
|
||||
Index-based (no per-item search): ``indices`` are the positions ``_sample_timesteps``
|
||||
drew from ``scheduler.timesteps``, and ``scheduler.sigmas`` is aligned with it, so this
|
||||
returns exactly what the diffusers ``get_sigmas`` timestep-matching helper would."""
|
||||
sigma = scheduler.sigmas[indices].to(device = device, dtype = dtype).flatten()
|
||||
drew from ``scheduler.timesteps``, and ``sigma_table`` (the scheduler's own sigmas, or
|
||||
the shifted copy from ``_training_sigma_table``) is aligned with it, so the identity
|
||||
table returns exactly what the diffusers ``get_sigmas`` helper would."""
|
||||
sigma = sigma_table[indices].to(device = device, dtype = dtype).flatten()
|
||||
while sigma.ndim < n_dim:
|
||||
sigma = sigma.unsqueeze(-1)
|
||||
return sigma
|
||||
|
||||
|
||||
def _training_sigma_table(scheduler, flow_shift):
|
||||
"""The sigma table training draws index into, per ``cfg.flow_shift``.
|
||||
|
||||
``1.0`` (every non-Qwen family's default) returns ``scheduler.sigmas`` unchanged: the
|
||||
historical behavior, correct for the families whose schedule already matches training
|
||||
convention. ``"auto"`` (the qwen-image default) reproduces the family's INFERENCE sigma
|
||||
distribution, which the scheduler never bakes into ``sigmas`` when
|
||||
``use_dynamic_shifting`` is true (the static ``shift`` at init is skipped, so Qwen-Image
|
||||
otherwise trains on unshifted uniform sigmas): apply the scheduler's own ``time_shift``
|
||||
at mu = ``max_shift`` (Qwen pins base_shift = max_shift = log 3, so the inference mu is
|
||||
constant at every resolution) followed by its ``stretch_shift_to_terminal`` -- using the
|
||||
scheduler's methods keeps the transform faithful across diffusers versions. A numeric
|
||||
value applies the standard linear shift s*u/(1+(s-1)*u) (musubi/kohya style
|
||||
discrete_flow_shift)."""
|
||||
sigmas = scheduler.sigmas
|
||||
if flow_shift == "auto":
|
||||
sc = scheduler.config
|
||||
if not getattr(sc, "use_dynamic_shifting", False):
|
||||
return sigmas # static-shift family: its init already baked the shift in
|
||||
mu = float(getattr(sc, "max_shift", None) or math.log(3.0))
|
||||
shifted = scheduler.time_shift(mu, 1.0, sigmas)
|
||||
if getattr(sc, "shift_terminal", None):
|
||||
# The stretch scales off the table's own final sigma, so it must see the full
|
||||
# descending schedule (as it does at inference set_timesteps), never a batch.
|
||||
shifted = scheduler.stretch_shift_to_terminal(shifted)
|
||||
return shifted
|
||||
s = float(flow_shift)
|
||||
if s == 1.0:
|
||||
return sigmas
|
||||
return s * sigmas / (1.0 + (s - 1.0) * sigmas)
|
||||
|
||||
|
||||
def _bell_loss_weights(num_train_timesteps):
|
||||
"""bsmntw-style bell loss-weight table over the training schedule: a Gaussian bell
|
||||
centered mid-schedule, floored at 0 and normalized to mean 1 (so the expected loss
|
||||
scale is unchanged). Indexed by round(sigma * num_train_timesteps)."""
|
||||
import torch
|
||||
steps = num_train_timesteps
|
||||
t = torch.arange(steps, dtype = torch.float32)
|
||||
w = torch.exp(-2.0 * ((t - steps / 2) / steps) ** 2)
|
||||
w = w - w.min()
|
||||
return w * (steps / w.sum())
|
||||
|
||||
|
||||
def _sample_timesteps(scheduler, batch_size, device):
|
||||
"""Logit-normal density timestep sampling (weighting_scheme='logit_normal'), returning
|
||||
(timesteps, indices) into the scheduler's schedule."""
|
||||
|
|
@ -1498,8 +1544,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
image_paths = [p for p, _ in pairs]
|
||||
captions = [c for _, c in pairs]
|
||||
uniq = sorted(set(captions))
|
||||
encoded = spec.encode_prompts(pipe, uniq, device)
|
||||
caption_embeds = {cap: emb for cap, emb in zip(uniq, encoded)}
|
||||
# CFG dropout swaps a sample's conditioning for the empty prompt, so its embedding must
|
||||
# be precomputed alongside the captions (the encoders are freed right after this).
|
||||
cfg_dropout = float(getattr(cfg, "cfg_dropout", 0.0) or 0.0)
|
||||
to_encode = uniq + ([""] if cfg_dropout > 0 and "" not in uniq else [])
|
||||
encoded = spec.encode_prompts(pipe, to_encode, device)
|
||||
caption_embeds = {cap: emb for cap, emb in zip(to_encode, encoded)}
|
||||
_free_text_encoders(pipe)
|
||||
gc.collect()
|
||||
if device == "cuda":
|
||||
|
|
@ -1596,6 +1646,19 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||||
cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
|
||||
)
|
||||
# Timestep-shift + loss-weighting setup (see _training_sigma_table). getattr defaults
|
||||
# keep an un-normalized config (tests, direct callers) on the historical behavior.
|
||||
flow_shift = getattr(cfg, "flow_shift", None)
|
||||
if flow_shift is None:
|
||||
flow_shift = "auto" if spec.family == "qwen-image" else 1.0
|
||||
sigma_table = _training_sigma_table(scheduler, flow_shift)
|
||||
shift_active = sigma_table is not scheduler.sigmas
|
||||
num_train_ts = scheduler.config.num_train_timesteps
|
||||
bell_weights = (
|
||||
_bell_loss_weights(num_train_ts).to(device)
|
||||
if str(getattr(cfg, "weighting_scheme", "none") or "none") == "bell"
|
||||
else None
|
||||
)
|
||||
# The LR schedule advances once per optimizer update, so warmup/decay are counted in
|
||||
# optimizer steps (matching the SDXL trainer; multiplying by the accumulation factor would
|
||||
# stretch warmup past the run and never reach the decay).
|
||||
|
|
@ -1652,11 +1715,24 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
|
||||
noise = torch.randn_like(latents)
|
||||
timesteps, t_indices = _sample_timesteps(scheduler, latents.shape[0], device)
|
||||
sigmas = _gather_sigmas(scheduler, t_indices, device, weight_dtype, latents.ndim)
|
||||
sigmas = _gather_sigmas(sigma_table, t_indices, device, weight_dtype, latents.ndim)
|
||||
if shift_active:
|
||||
# The model's timestep conditioning must follow the shifted sigma (the
|
||||
# scheduler convention is timestep = sigma * num_train_timesteps). Gather
|
||||
# in fp32 from the table so bf16 rounding never skews the conditioning.
|
||||
timesteps = (
|
||||
sigma_table[t_indices].to(device = device, dtype = torch.float32).flatten()
|
||||
* num_train_ts
|
||||
)
|
||||
noisy = (1.0 - sigmas) * latents + sigmas * noise
|
||||
|
||||
embeds = spec.collate(
|
||||
[caption_embeds[captions[i]] for i in idxs],
|
||||
[
|
||||
caption_embeds[""]
|
||||
if cfg_dropout > 0 and rng.random() < cfg_dropout
|
||||
else caption_embeds[captions[i]]
|
||||
for i in idxs
|
||||
],
|
||||
device,
|
||||
weight_dtype,
|
||||
pad_to = qwen_pad_to,
|
||||
|
|
@ -1666,7 +1742,17 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
|
|||
transformer, noisy, timesteps, sigmas, embeds, cfg, device, weight_dtype
|
||||
)
|
||||
target = noise - latents
|
||||
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
|
||||
if bell_weights is None:
|
||||
loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
|
||||
else:
|
||||
per = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
|
||||
w_idx = (
|
||||
(sigmas.flatten().float() * num_train_ts)
|
||||
.long()
|
||||
.clamp(0, num_train_ts - 1)
|
||||
)
|
||||
w = bell_weights[w_idx].view(-1, *([1] * (per.ndim - 1)))
|
||||
loss = (per * w).mean()
|
||||
(loss / cfg.gradient_accumulation_steps).backward()
|
||||
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
|
||||
|
||||
|
|
|
|||
|
|
@ -247,15 +247,25 @@ def get_trainer(family: str) -> Callable[..., str]:
|
|||
# Families absent here fall back to the DiffusionLoraConfig defaults.
|
||||
FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
|
||||
"sdxl": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 1024},
|
||||
"flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512},
|
||||
"qwen-image": {"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512},
|
||||
# Warmup defaults: a short LR ramp keeps the first adapter updates from overshooting on
|
||||
# the big flow-matching DiTs (whose logit-normal timestep draw concentrates loss mass
|
||||
# mid-schedule); the small warmups below are scaled for Studio's short-run step budgets.
|
||||
"flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20},
|
||||
"qwen-image": {
|
||||
"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512, "lr_warmup_steps": 20,
|
||||
},
|
||||
"z-image": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 768},
|
||||
# The Krea 2 authors' recommended starting point (their DreamBooth script defaults):
|
||||
# rank/alpha 32, lr 3e-4, 512px.
|
||||
"krea-2": {"lora_rank": 32, "learning_rate": 3e-4, "resolution": 512},
|
||||
# The upstream FLUX.2 DreamBooth references default to rank 16 / lr 1e-4.
|
||||
"flux.2-klein": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512},
|
||||
"flux.2-dev": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512},
|
||||
# The upstream FLUX.2 DreamBooth references default to rank 16 / lr 1e-4; FLUX.2's
|
||||
# uniform timestep draw benefits most from a warmup ramp.
|
||||
"flux.2-klein": {
|
||||
"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20,
|
||||
},
|
||||
"flux.2-dev": {
|
||||
"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -496,6 +506,19 @@ class DiffusionLoraConfig:
|
|||
# "fp8" (torchao float8 training on the frozen linears, Ada/Hopper/Blackwell + compile), or
|
||||
# "auto" (by free VRAM + GPU class). Non-nf4 modes need a dense base repo. SDXL ignores it.
|
||||
base_precision: str = "nf4"
|
||||
# Training-time timestep shift applied to the flow-matching sigma draw. None resolves
|
||||
# per family in normalized(): "auto" for qwen-image (reproduce the family's inference
|
||||
# sigma distribution: the scheduler's exponential time_shift at mu = max_shift, then the
|
||||
# shift_terminal stretch), 1.0 (identity, the historical behavior) for every other
|
||||
# family. A numeric value applies the standard linear shift s*u/(1+(s-1)*u); 1.0 is a
|
||||
# no-op. "auto" on a family without dynamic shifting falls back to identity.
|
||||
flow_shift: Optional[Any] = None # float | "auto" | None
|
||||
# Per-sample probability of replacing the caption conditioning with the empty prompt
|
||||
# (classifier-free-guidance dropout). 0.0 (default) disables it entirely.
|
||||
cfg_dropout: float = 0.0
|
||||
# Per-sample loss weighting over the drawn timestep: "none" (default, unweighted MSE)
|
||||
# or "bell" (bsmntw-style Gaussian bell centered mid-schedule, normalized to mean 1).
|
||||
weighting_scheme: str = "none"
|
||||
# How often to emit a progress event (in optimizer steps).
|
||||
log_every: int = 1
|
||||
# Optional explicit family override ("sdxl" / "flux.1" / ...); None = detect from base_model.
|
||||
|
|
@ -576,6 +599,34 @@ class DiffusionLoraConfig:
|
|||
f"{resolved_family}: its activations exceed fp8's range and corrupt the "
|
||||
f"trained result. Use 'nf4', 'int8', 'bf16', or 'auto'."
|
||||
)
|
||||
# flow_shift: None resolves to the family default ("auto" only for qwen-image, whose
|
||||
# scheduler skips its static shift under use_dynamic_shifting and would otherwise
|
||||
# train on unshifted uniform sigmas); an explicit value is validated and kept.
|
||||
flow_shift = self.flow_shift
|
||||
if flow_shift is None:
|
||||
flow_shift = "auto" if resolved_family == "qwen-image" else 1.0
|
||||
if isinstance(flow_shift, str):
|
||||
flow_shift = flow_shift.strip().lower()
|
||||
if flow_shift != "auto":
|
||||
try:
|
||||
flow_shift = float(flow_shift)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"flow_shift must be a positive number or 'auto', got {self.flow_shift!r}"
|
||||
) from exc
|
||||
if not isinstance(flow_shift, str):
|
||||
flow_shift = float(flow_shift)
|
||||
if flow_shift <= 0:
|
||||
raise ValueError("flow_shift must be > 0 (1.0 disables the shift), or 'auto'")
|
||||
try:
|
||||
cfg_dropout = float(self.cfg_dropout or 0.0)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"cfg_dropout must be a number, got {self.cfg_dropout!r}") from exc
|
||||
if not 0.0 <= cfg_dropout <= 1.0:
|
||||
raise ValueError("cfg_dropout must be between 0 and 1")
|
||||
weighting_scheme = str(self.weighting_scheme or "none").strip().lower()
|
||||
if weighting_scheme not in ("none", "bell"):
|
||||
raise ValueError("weighting_scheme must be one of none / bell")
|
||||
# A zero/negative gamma would zero out (or invert) the min-SNR weight and
|
||||
# silently train on a degenerate loss; None is the documented disable.
|
||||
if self.snr_gamma is not None and float(self.snr_gamma) <= 0:
|
||||
|
|
@ -604,6 +655,9 @@ class DiffusionLoraConfig:
|
|||
cache_variants = int(self.cache_variants),
|
||||
compile_transformer = compile_transformer,
|
||||
base_precision = base_precision,
|
||||
flow_shift = flow_shift,
|
||||
cfg_dropout = cfg_dropout,
|
||||
weighting_scheme = weighting_scheme,
|
||||
resolved_family = resolved_family,
|
||||
)
|
||||
|
||||
|
|
|
|||
220
studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py
Normal file
220
studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the DiT trainer's timestep-shift / CFG-dropout / loss-weighting levers.
|
||||
|
||||
CPU-only: cover the flow_shift config resolution (qwen-image defaults to "auto", every
|
||||
other family stays on the identity 1.0), the exact sigma transform for the auto and
|
||||
numeric modes, the shifted sampling distribution, and the bell weight table. The full
|
||||
training loop is exercised by the live GPU smokes, not here."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from core.training.diffusion_dit_trainer import (
|
||||
_bell_loss_weights,
|
||||
_gather_sigmas,
|
||||
_sample_timesteps,
|
||||
_training_sigma_table,
|
||||
)
|
||||
from core.training.diffusion_train_common import DiffusionLoraConfig
|
||||
|
||||
QWEN_SHIFT_TERMINAL = 0.02
|
||||
|
||||
|
||||
def _qwen_scheduler():
|
||||
# The Qwen/Qwen-Image scheduler config: shift=1.0 is SKIPPED at init because
|
||||
# use_dynamic_shifting is true, base_shift = max_shift = log 3 (constant inference mu),
|
||||
# exponential time shift, terminal stretch to 0.02.
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
return FlowMatchEulerDiscreteScheduler(
|
||||
num_train_timesteps = 1000,
|
||||
shift = 1.0,
|
||||
use_dynamic_shifting = True,
|
||||
base_shift = math.log(3.0),
|
||||
max_shift = math.log(3.0),
|
||||
shift_terminal = QWEN_SHIFT_TERMINAL,
|
||||
time_shift_type = "exponential",
|
||||
)
|
||||
|
||||
|
||||
def _flux_static_scheduler():
|
||||
# A static-shift scheduler (shift baked into sigmas at init, no dynamic shifting).
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
return FlowMatchEulerDiscreteScheduler(num_train_timesteps = 1000, shift = 3.0)
|
||||
|
||||
|
||||
# ── config resolution ─────────────────────────────────────────────────────────
|
||||
def test_flow_shift_defaults_per_family():
|
||||
qwen = DiffusionLoraConfig(
|
||||
base_model = "Qwen/Qwen-Image", data_dir = "d", output_dir = "o"
|
||||
).normalized()
|
||||
assert qwen.resolved_family == "qwen-image"
|
||||
assert qwen.flow_shift == "auto"
|
||||
flux = DiffusionLoraConfig(
|
||||
base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
|
||||
).normalized()
|
||||
assert flux.flow_shift == 1.0
|
||||
zimg = DiffusionLoraConfig(
|
||||
base_model = "Tongyi-MAI/Z-Image-Turbo", data_dir = "d", output_dir = "o"
|
||||
).normalized()
|
||||
assert zimg.flow_shift == 1.0
|
||||
|
||||
|
||||
def test_flow_shift_explicit_values_and_validation():
|
||||
cfg = DiffusionLoraConfig(
|
||||
base_model = "Qwen/Qwen-Image", data_dir = "d", output_dir = "o", flow_shift = 2.2
|
||||
).normalized()
|
||||
assert cfg.flow_shift == 2.2
|
||||
# String numerics from the Studio config path coerce; "auto" passes through.
|
||||
assert (
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "3.0"
|
||||
)
|
||||
.normalized()
|
||||
.flow_shift
|
||||
== 3.0
|
||||
)
|
||||
assert (
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "AUTO"
|
||||
)
|
||||
.normalized()
|
||||
.flow_shift
|
||||
== "auto"
|
||||
)
|
||||
with pytest.raises(ValueError, match = "flow_shift"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = 0.0
|
||||
).normalized()
|
||||
with pytest.raises(ValueError, match = "flow_shift"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "bogus"
|
||||
).normalized()
|
||||
|
||||
|
||||
def test_cfg_dropout_and_weighting_scheme_validation():
|
||||
cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o").normalized()
|
||||
assert cfg.cfg_dropout == 0.0
|
||||
assert cfg.weighting_scheme == "none"
|
||||
on = DiffusionLoraConfig(
|
||||
base_model = "b",
|
||||
data_dir = "d",
|
||||
output_dir = "o",
|
||||
cfg_dropout = 0.1,
|
||||
weighting_scheme = "bell",
|
||||
).normalized()
|
||||
assert on.cfg_dropout == 0.1
|
||||
assert on.weighting_scheme == "bell"
|
||||
with pytest.raises(ValueError, match = "cfg_dropout"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", cfg_dropout = 1.5
|
||||
).normalized()
|
||||
with pytest.raises(ValueError, match = "weighting_scheme"):
|
||||
DiffusionLoraConfig(
|
||||
base_model = "b", data_dir = "d", output_dir = "o", weighting_scheme = "sigma_sqrt"
|
||||
).normalized()
|
||||
|
||||
|
||||
def test_config_from_dict_plumbs_the_new_fields():
|
||||
from core.training.diffusion_train_common import _config_from_dict
|
||||
|
||||
cfg = _config_from_dict(
|
||||
{
|
||||
"base_model": "Qwen/Qwen-Image",
|
||||
"data_dir": "d",
|
||||
"output_dir": "o",
|
||||
"flow_shift": "auto",
|
||||
"cfg_dropout": 0.05,
|
||||
"weighting_scheme": "bell",
|
||||
}
|
||||
)
|
||||
assert cfg.flow_shift == "auto"
|
||||
assert cfg.cfg_dropout == 0.05
|
||||
assert cfg.weighting_scheme == "bell"
|
||||
|
||||
|
||||
# ── sigma table transforms ────────────────────────────────────────────────────
|
||||
def test_auto_table_matches_the_exact_qwen_transform():
|
||||
import torch
|
||||
|
||||
sched = _qwen_scheduler()
|
||||
table = _training_sigma_table(sched, "auto")
|
||||
base = sched.sigmas
|
||||
# Exponential shift at mu = log 3 with sigma exponent 1 is exp(mu)/(exp(mu) + 1/u - 1)
|
||||
# = 3u/(1 + 2u), then the terminal stretch maps the schedule's last sigma to 0.02.
|
||||
shifted = 3.0 * base / (1.0 + 2.0 * base)
|
||||
scale = (1.0 - shifted[-1]) / (1.0 - QWEN_SHIFT_TERMINAL)
|
||||
expected = 1.0 - (1.0 - shifted) / scale
|
||||
assert torch.allclose(table, expected, atol = 1e-6)
|
||||
# Fixed-point spot checks: sigma 1.0 stays 1.0, the terminal sigma lands on 0.02, and
|
||||
# the midpoint u = 0.5 rises to ~0.754 (3u/(1+2u) = 0.75 before the stretch).
|
||||
assert abs(float(table[0]) - 1.0) < 1e-6
|
||||
assert abs(float(table[-1]) - QWEN_SHIFT_TERMINAL) < 1e-6
|
||||
assert abs(float(table[499]) - 0.75427) < 1e-3
|
||||
# The table stays a valid descending schedule in (0, 1].
|
||||
assert bool((table[:-1] > table[1:]).all())
|
||||
|
||||
|
||||
def test_numeric_table_applies_the_linear_shift():
|
||||
import torch
|
||||
|
||||
sched = _qwen_scheduler()
|
||||
table = _training_sigma_table(sched, 2.2)
|
||||
base = sched.sigmas
|
||||
assert torch.allclose(table, 2.2 * base / (1.0 + 1.2 * base), atol = 1e-6)
|
||||
# u = 0.5 under shift s maps to s/(s+1).
|
||||
assert abs(float(table[499]) - 2.2 / 3.2) < 1e-3
|
||||
|
||||
|
||||
def test_identity_and_static_families_are_untouched():
|
||||
# flow_shift 1.0 must return the scheduler's own table object (no numeric drift for
|
||||
# FLUX / Z-Image / Krea 2), and "auto" on a static-shift scheduler is a no-op too:
|
||||
# its init already baked the shift into sigmas.
|
||||
sched = _qwen_scheduler()
|
||||
assert _training_sigma_table(sched, 1.0) is sched.sigmas
|
||||
static = _flux_static_scheduler()
|
||||
assert _training_sigma_table(static, "auto") is static.sigmas
|
||||
assert _training_sigma_table(static, 1.0) is static.sigmas
|
||||
|
||||
|
||||
def test_sampled_sigma_distribution_shifts_under_auto():
|
||||
import torch
|
||||
|
||||
torch.manual_seed(0)
|
||||
sched = _qwen_scheduler()
|
||||
auto_table = _training_sigma_table(sched, "auto")
|
||||
_, idx = _sample_timesteps(sched, 4096, "cpu")
|
||||
base = _gather_sigmas(sched.sigmas, idx, "cpu", torch.float32, 1)
|
||||
shifted = _gather_sigmas(auto_table, idx, "cpu", torch.float32, 1)
|
||||
# Unshifted logit-normal draws center at 0.5; the mu = log 3 shift + terminal stretch
|
||||
# pushes the mass toward high noise (mean ~0.72). Shift raises EVERY sample.
|
||||
assert abs(float(base.mean()) - 0.5) < 0.03
|
||||
assert float(shifted.mean()) > 0.68
|
||||
assert bool((shifted >= base - 1e-6).all())
|
||||
|
||||
|
||||
def test_gather_sigmas_broadcasts_to_ndim():
|
||||
import torch
|
||||
|
||||
sched = _qwen_scheduler()
|
||||
sig = _gather_sigmas(sched.sigmas, torch.tensor([0, 499, 999]), "cpu", torch.float32, 4)
|
||||
assert sig.shape == (3, 1, 1, 1)
|
||||
assert abs(float(sig[0].flatten()) - 1.0) < 1e-6
|
||||
|
||||
|
||||
# ── bell weighting ────────────────────────────────────────────────────────────
|
||||
def test_bell_weights_shape_peak_and_normalization():
|
||||
w = _bell_loss_weights(1000)
|
||||
assert w.shape == (1000,)
|
||||
assert float(w.min()) >= 0.0
|
||||
# Peak at mid-schedule, mean 1 so the expected loss scale is unchanged.
|
||||
assert int(w.argmax()) == 500
|
||||
assert abs(float(w.mean()) - 1.0) < 1e-5
|
||||
assert float(w[500]) > float(w[0])
|
||||
assert float(w[500]) > float(w[999])
|
||||
Loading…
Add table
Add a link
Reference in a new issue