Fix silent lr_warmup_steps no-op under the default constant scheduler - #8464
Conversation
The four flow-matching families advertise lr_warmup_steps: 20 while the default lr_scheduler stays 'constant', which diffusers builds without reading num_warmup_steps. Advertise 'constant_with_warmup' with the warmup in FAMILY_TRAIN_DEFAULTS, and promote 'constant' with a positive warmup in normalized() as a backstop for direct callers. Fixes unslothai#8269.
for more information, see https://pre-commit.ci
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 67cc3882f1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
lr_scheduler and lr_warmup_steps are both checkpoint identity fields
(diffusion_checkpoint.py, _IDENTITY_LABELS), and the route builds the resume
identity from the normalised config (routes/training.py:2821). Rewriting the
pair inside normalized() therefore stranded every bundle written before the
rewrite: its manifest records "constant" with a positive warmup, replaying that
same config normalised to "constant_with_warmup", and mismatch_reason rejected
the resume as a changed learning-rate schedule. No config could express the
legacy pair, so those runs could not be continued at all.
The defaults table still advertises "constant_with_warmup" next to the warmup
steps for flux.1, qwen-image, flux.2-klein and flux.2-dev, which is what fixes
the silent no-op for a client building its request from
/training/diffusion/info. The warmup validation stays.
Replaces the promotion test with one that pins the resume identity round-trip
for ("constant", 20).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 282039d1d1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Both entries carried lr_warmup_steps: 20 with no lr_scheduler, so a request built from /training/diffusion/info still resolved to the "constant" default. The DiT and H3 trainers pass cfg.lr_scheduler and cfg.lr_warmup_steps straight to get_scheduler (diffusion_dit_trainer.py:2117, diffusion_h3_trainer.py:694), which drops num_warmup_steps for "constant", so neither family ramped. The invariant test now derives its family set from FAMILY_TRAIN_DEFAULTS rather than a hand-written list, so a family added with a warmup and no scheduler fails instead of being skipped.
|
@codex review |
1 similar comment
|
@codex review |
for more information, see https://pre-commit.ci
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
3 participants
Fixes #8269.
FAMILY_TRAIN_DEFAULTSrecommendslr_warmup_steps: 20for flux.1, qwen-image, flux.2-klein and flux.2-dev, butDiffusionLoraConfig.lr_schedulerdefaults to"constant", and diffusers'get_schedulerbuilds that schedule without ever readingnum_warmup_steps. So a config built from what/training/diffusion/infoadvertises never gets the ramp the defaults comment promises.Two changes in
diffusion_train_common.py:"lr_scheduler": "constant_with_warmup"next to their warmup steps, so the defaults table is coherent on its ownnormalized()promotes"constant"to"constant_with_warmup"whenlr_warmup_steps > 0, as a backstop for any direct caller that sets the pair by hand, and validates the warmup value while it is there. A zero warmup stays"constant", and an explicitly warmup-capable scheduler is left exactly as givenWhat the learning rate actually does (base LR 1e-4, warmup 20, diffusers 0.39.0):
Per family, whether the advertised defaults are honored in the schedule the trainer builds:
Tests:
studio/backend/tests/test_diffusion_warmup_defaults.py, five CPU-only cases covering the defaults invariant, the promotion, that explicit schedulers and zero-warmup configs are untouched, and warmup validation. Three of the five fail on main and pass with this change; the other two pin the unchanged behaviour.The fix was produced by Artemis Discovery (TurinTech) working from the issue report, scored against a harness measuring the fraction of families whose advertised defaults are honored (42.9 on main, 100 with this change). I reviewed every line and added the regression tests during that review.