iframe-proxy

vineethsaivs · 2026-08-12T16:40:31Z

Closes #8269.

The bug

FAMILY_TRAIN_DEFAULTS recommends lr_warmup_steps: 20 for the big flow-matching DiTs, and says why right above the table:

Warmup defaults: a short LR ramp keeps the first adapter updates from overshooting on the big flow-matching DiTs.

DiffusionLoraConfig.lr_scheduler defaults to "constant", and diffusers' get_scheduler returns before it ever reads num_warmup_steps for that schedule:

if name == SchedulerType.CONSTANT:
    return schedule_func(optimizer, last_epoch=last_epoch)

The trainers pass the value in good faith (diffusion_dit_trainer.py:2121, diffusion_h3_trainer.py:694), so nothing errors and nothing warns. The advertised ramp simply never happens.

One correction to the report

It names four families. There are six: ltx-2 and minimax-h3 carry the same lr_warmup_steps: 20 preset and are affected identically.

family          warmup   scheduler before      after
flux.1              20   constant              constant_with_warmup
qwen-image          20   constant              constant_with_warmup
flux.2-klein        20   constant              constant_with_warmup
flux.2-dev          20   constant              constant_with_warmup
ltx-2               20   constant              constant_with_warmup
minimax-h3          20   constant              constant_with_warmup
sdxl                 -   constant              constant      (no warmup preset, untouched)
z-image              -   constant              constant      (untouched)
krea-2               -   constant              constant      (untouched)

Fix

Carry the scheduler alongside the ramp. constant_with_warmup is already an accepted value in _LR_SCHEDULERS, and diffusers does read num_warmup_steps for it, so the preset now does what the comment says. Families that set no warmup keep constant and are not touched.

Why it survived

test_flow_families_carry_warmup_presets already existed and passes today: it asserts the warmup number is present and positive, but never that anything consumes it. Three tests added beside it, pinning the pairing rather than the value:

every family default with a positive lr_warmup_steps must name a scheduler that consumes warmup, and that scheduler must be one _LR_SCHEDULERS accepts
what /training/diffusion/info advertises must be constructible as-is and survive into a built DiffusionLoraConfig
families without a warmup preset must not have gained a scheduler override

                                                     before   after
test_a_warmup_preset_names_a_scheduler_that_uses_it   FAIL     pass
test_warmup_presets_survive_into_a_built_config       FAIL     pass
test_families_without_a_warmup_preset_keep_constant   pass     pass
test_flow_families_carry_warmup_presets (existing)    pass     pass

pytest studio/backend/tests/test_diffusion_train_extras.py gives 2 failed / 27 passed against main and 29 passed with this change.

Across every tests/*diffusion* file: 37 failed / 2077 passed on a clean tree and 37 failed / 2080 passed with this change, so the three added tests are the only difference. Those 37 are pre-existing and need a GPU (nf4 is not a CPU fallback); tests/test_mcp_server.py also fails to collect on a clean tree here for a missing fastmcp.

This changes training behaviour for those six families, in that the ramp now actually runs. That is the stated intent of the preset, so I read it as the fix rather than a new default, but if you would rather the presets be dropped than honoured, that is a one-line alternative and I am happy to switch.

FAMILY_TRAIN_DEFAULTS recommends lr_warmup_steps=20 for the big flow-matching DiTs, with a comment saying the ramp keeps the first adapter updates from overshooting. DiffusionLoraConfig.lr_scheduler defaults to "constant", and diffusers' get_scheduler returns before it ever reads num_warmup_steps for that schedule: if name == SchedulerType.CONSTANT: return schedule_func(optimizer, last_epoch=last_epoch) so the advertised ramp never happens. Carry "constant_with_warmup" alongside each warmup preset, which is already an accepted value in _LR_SCHEDULERS. Six families are affected, not the four the report names: ltx-2 and minimax-h3 set the same preset. sdxl, z-image and krea-2 set no warmup and are untouched. Closes unslothai#8269

vineethsaivs · 2026-08-15T21:04:36Z

oobabooga · 2026-08-17T01:39:18Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dd125cca69

ℹ️ 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".

chatgpt-codex-connector · 2026-08-17T01:41:24Z

+        "learning_rate": 1e-4,
+        "resolution": 512,
+        "lr_warmup_steps": 20,
+        "lr_scheduler": "constant_with_warmup",


Wire the scheduler presets into the Train UI

When any of these families is trained through the Studio UI, this new default is discarded: DiffusionTrainableFamily.defaults does not declare the scheduler or warmup fields, mergeFamilies() copies only rank/LR/resolution, and the panel initializes and sends lrScheduler="constant" with zero warmup. Consequently the primary training flow still performs no ramp for all six affected families despite /diffusion/info advertising one; the frontend needs to preserve and seed both new defaults.

Useful? React with 👍 / 👎.

Correct, and fixed in 47d586c. The panel narrowed the reported defaults to rank/lr/resolution in mergeFamilies(), initialised lrScheduler to "constant" with zero warmup, and then sent both fields explicitly on Start, so the pairing never reached a run started from the UI.

One detail differs from the report: DiffusionTrainableFamily.defaults is an untyped dict, so the values did reach the client. The blockers were the TypeScript defaults shape in features/images/api.ts, which enumerates its fields and had no scheduler or warmup, and mergeFamilies() itself.

The fix carries the pair through both merge arms, seeds it in the family re-seed effect, and resets it to constant/0 for a family with no warmup preset, so flux.1's 20 steps do not ride along into SDXL. Normalisation lives in diffusion-train-lr-schedule.ts: half a pair is dropped rather than seeded, since a warmup count under constant is exactly the inert state the backend pairing exists to prevent, and a scheduler the panel's Select cannot show (cosine_with_restarts, polynomial) is dropped rather than left rendering blank. The reported family owns the ramp outright, with no static-preset fallback, so a backend that drops a warmup preset drops the ramp here too. The LR schedule Select now marks the settings dirty like every other control, or the re-seed would replace a hand-picked schedule on the next family switch.

studio/frontend/tests/diffusion-train-warmup-preset.test.ts covers it: 8 tests, 3 of which fail against the previous panel. Full frontend suite 2015 passed, studio/backend/tests/test_diffusion_train_extras.py 29 passed.

oobabooga · 2026-08-17T04:06:33Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 47d586cefc

ℹ️ 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".

chatgpt-codex-connector · 2026-08-17T04:09:00Z

+      setLrScheduler(family.defaults.lrScheduler ?? "constant");
+      setLrWarmupSteps(family.defaults.lrWarmupSteps ?? 0);


Track LR schedule edits separately from other settings

After the user edits any numeric control—even Steps, rank, resolution, or seed—settingsDirty.current stays true, so this entire family-seeding block is skipped on subsequent family changes. As a result, switching from a warmed-up DiT to SDXL carries constant_with_warmup and 20 steps into SDXL, while switching in the opposite direction can retain constant and omit the new ramp; the LR schedule/warmup pair needs its own dirty state so unrelated edits do not suppress model-specific defaults.

Useful? React with 👍 / 👎.

Confirmed against a running Studio, and fixed in 2116951. settingsDirty is one flag over every numeric control, Steps and Seed included, so gating the ramp on it was wrong. Reproduced on the panel before the fix: open Train on SDXL, type a step count, switch to FLUX.1-dev, and the LR schedule stays Constant with the Warmup steps field hidden. That is the ramp this change exists to seed, silently not running.

The pair now has its own lrScheduleDirty ref, set by the LR schedule Select and by the Warmup steps input, and its re-seed sits outside the settingsDirty block. It gets a separate flag rather than joining the existing one because it is the only setting whose control disappears: under plain constant the Warmup steps field is not rendered, so a value carried past a family change is invisible rather than merely stale.

Same three cases read off the live panel after the fix:

SDXL -> edit Steps -> FLUX.1-dev Constant + warmup 20 (was: Constant, field hidden) FLUX.1-dev / Qwen-Image, untouched Constant + warmup 20 SDXL / Z-Image, untouched Constant hidden FLUX.1-dev -> pick Cosine -> SDXL Cosine decay 20 (user pick survives)

Two assertions added to diffusion-train-warmup-preset.test.ts, both failing against 47d586c. Frontend suite 2016 passed.

oobabooga · 2026-08-17T04:27:39Z

@codex review

chatgpt-codex-connector · 2026-08-17T04:30:38Z

Codex Review: Didn't find any major issues. Another round soon, please!

Reviewed commit: 2116951b65

ℹ️ 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".

oobabooga · 2026-08-17T04:34:21Z

Follow-up to my own change in 2116951, pushed as f457166.

Making "Warmup steps" visible by default put it in reach of numberField's shared dirty mark, which meant tuning the ramp froze the other family-seeded settings. Read off the live panel before the fix:

flux.1, untouched                LR 0.0001    warmup 20
flux.1, warmup tuned to 50       LR 0.0001    warmup 50
-> qwen-image                    LR 0.0001    warmup 50    <- should be 5e-5

qwen-image's own preset is 5e-5, so switching after touching the ramp trained it at double the recommended rate with nothing on the panel to say so. Before this PR the field was hidden under the default constant schedule, so it was not reachable.

numberField now takes an optional markDirty, and "Warmup steps" is the only caller that passes one: it charges the edit to lrScheduleDirty alone. Same sequence after the fix ends on LR 0.00005 / warmup 50, so the hand-typed ramp still survives the switch while the other three re-seed.

One added test, failing against 2116951. Frontend suite 2017 passed.

oobabooga · 2026-08-17T04:34:42Z

@codex review

chatgpt-codex-connector · 2026-08-17T04:37:48Z

Codex Review: Didn't find any major issues. Breezy!

Reviewed commit: f457166b0d

ℹ️ 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".

# Conflicts: # studio/backend/core/training/diffusion_train_common.py

oobabooga · 2026-08-17T06:44:35Z

Merged main in (773d3d1). #8464 landed the same backend fix about an hour ago, which put this branch in conflict, so the shape of this PR has changed.

The conflict was in FAMILY_TRAIN_DEFAULTS. Both sides made the same edit: lr_scheduler: "constant_with_warmup" on the six families carrying lr_warmup_steps: 20, identical keys and values, differing only in comment wording and key order. I resolved it to main's version, so diffusion_train_common.py is now byte-identical to main and the backend half of this PR is gone.

What is left is the part #8464 did not do. It fixed the advertisement; the Train panel still discards it. On main right now, /api/train/diffusion/info reports constant_with_warmup + 20 for those six families, mergeFamilies() narrows the reported defaults to rank/lr/resolution, and the panel sends lr_scheduler: "constant" with zero warmup on Start. So the ramp still does not run for anyone training through the UI, which is the primary flow. That is what the remaining diff wires up, and it is the same gap Codex raised here as P1.

Two of the three tests added here are now covered by test_diffusion_warmup_defaults.py from #8464, more generally than the versions here, so I dropped them. test_warmup_presets_survive_into_a_built_config is kept: #8464's helper filters train_defaults() down to the dataclass fields, so it does not catch a key added to FAMILY_TRAIN_DEFAULTS that DiffusionLoraConfig refuses, which would break any client posting the advertised defaults back verbatim.

Net diff against main is now 5 files: the frontend wiring, its tests, and that one backend test. Frontend suite 3489 passed, test_diffusion_train_extras.py + test_diffusion_warmup_defaults.py 32 passed.

oobabooga · 2026-08-17T06:44:51Z

@codex review

chatgpt-codex-connector · 2026-08-17T06:48:07Z

Codex Review: Didn't find any major issues. Bravo.

Reviewed commit: 773d3d1007

ℹ️ 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".

oobabooga · 2026-08-17T17:30:19Z

vineethsaivs requested a review from danielhanchen as a code owner August 12, 2026 16:40

chatgpt-codex-connector Bot reviewed Aug 17, 2026

View reviewed changes

Seed the diffusion warmup presets into the Train panel

47d586c

chatgpt-codex-connector Bot reviewed Aug 17, 2026

View reviewed changes

Give the LR schedule pair its own dirty flag

2116951

Keep a ramp edit from freezing the other family defaults

f457166

Merge remote-tracking branch 'origin/main' into review/pr-8593

773d3d1

# Conflicts: # studio/backend/core/training/diffusion_train_common.py

oobabooga merged commit 4b18f07 into unslothai:main Aug 17, 2026
34 of 36 checks passed

		setLrScheduler(family.defaults.lrScheduler ?? "constant");
		setLrWarmupSteps(family.defaults.lrWarmupSteps ?? 0);

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

Uh oh!

Conversation

vineethsaivs commented Aug 12, 2026

The bug

One correction to the report

Fix

Why it survived

Uh oh!

vineethsaivs commented Aug 15, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

oobabooga Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

oobabooga Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

chatgpt-codex-connector Bot commented Aug 17, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

chatgpt-codex-connector Bot commented Aug 17, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

chatgpt-codex-connector Bot commented Aug 17, 2026

Uh oh!

oobabooga commented Aug 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants