Skip to content

target_loss_auto_stop: end training when the target is reached - #70

Open
scragnog wants to merge 26 commits into
koda-dernet:mainfrom
scragnog:pr/auto-stop
Open

target_loss_auto_stop: end training when the target is reached#70
scragnog wants to merge 26 commits into
koda-dernet:mainfrom
scragnog:pr/auto-stop

Conversation

@scragnog

Copy link
Copy Markdown

target_loss_auto_stop: end training when the target is reached

Part 4 of the training-core series — stacked on #69; until that merges this diff includes its commits.

With target-loss cruise control active, training previously held at the target LR-damped until epochs ran out — burning GPU-hours after the model was done. --target-loss-auto-stop ends the run (normal finalization: best/final saves intact) once the 5-epoch moving average crosses target_loss.

The feature took three fix rounds on our side, all squashed into this series so upstream gets the working version:

  • the flag was silently dropped at three points in the config pipeline (arg → factory → GUI)
  • the stop check lived only in the non-Fabric fallback loop; it now runs in the Fabric trainer too
  • wizard/preset/GUI-CLI passthrough completed

Off by default; zero behavior change unless the flag is set.

🤖 Generated with Claude Code

Rob Work PC and others added 26 commits July 29, 2026 11:34
Preprocessing now encodes BOTH the caption and the genre-tag text
variant into each .pt (backward-compatible optional keys); training
selects caption vs genre per sample per step according to genre_ratio
(auto-detected from preprocess_meta.json when unset) instead of the old
static per-sample choice at preprocess time. Old tensors keep working.

Co-authored-by: Gabriel <gadna166@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… GUI

conditioning_type flows from the dataset through collate into
ProgressWriter events; the training monitor shows live caption-vs-genre
usage. Also ports target-loss scale/EMA telemetry to the Fabric trainer
path (previously basic-loop only).

Co-authored-by: Gabriel <gadna166@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…hannel balancing, latent noise, t_bias

(from 5e400f6; scripts/ debug artifacts excluded)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rmup ramp

- MSE guarantees convergence to conditional expectation (required by flow matching)
- Huber converges to median, under-penalizes extreme flow vectors -> clipped highend
- Huber still available as option, just not default
- EMA warmup: effective_decay = min(target, (1+step)/(10+step)) so shadow catches up fast early
- ema_warmup_steps (default 2000): prevents EMA shadow from lagging early training
- At step 0 decay=0 (shadow=current weights), linearly ramps to target decay
- Wired through: defaults, configs, config_factory, CLI, wizard UI, frontend
- GUI: warmup field auto-hides when EMA decay=0
- No tracking until start_step, then snapshot current weights as shadow
- Mathematically clean: shadow is always a proper exponential average at full target decay
- Linear ramp was creating a Frankenstein blend of different averaging rates
- Renamed ema_warmup_steps -> ema_start_step across all 16 touchpoints
- Default 2000 steps: model learns basic structure before EMA locks in
… distribution

- training_step now stashes auxiliary metrics: raw/weighted loss, SNR weight stats,
  timestep mean, per-channel loss breakdown (max/min/ratio)
- drain_step_metrics() returns and clears the dict at log_every intervals
- Both Fabric and basic training loops log to TensorBoard: fidelity/*, ema/*
- Fixed timestep histogram to fire for flow_snr too (was min_snr only)
- Pipe fidelity dict + ema_active through progress writer for real-time GUI
- Frontend: unpack fidelity dict from progress messages into _tbScalars
- Added curated labels/colors for all new tags in mini-chart registry
- dynamic_channel_balance (opt-in toggle): blends 50/50 static VAE prior
  with running EMA of per-channel loss, so struggling channels get more gradient
- Wired through all config touchpoints: defaults, configs, factory, CLI, GUI
- GUI: checkbox auto-hides when channel_balance is off or legacy mode
- Fix: SNR weight charts were always 1.0 (logged post-normalization);
  now logs pre-normalization values + spread ratio
- New ema/effective_decay metric (0.0 when inactive, target decay when active)
- Added curated mini-chart labels for snr_weight_spread, ch_dynamic_blend_range,
  ema/effective_decay
…tric telemetry

- Add 3 new loss_fn options: pseudo_huber, x0_mse, x0_pseudo_huber
- x0_mse implements paper's intended formulation (implicit t² weighting)
- pseudo_huber: δ²(√(1+(e/δ)²) - 1) smooth Huber approximation
- Log fidelity/arrangement_loss (t>0.5) and fidelity/detail_loss (t≤0.5)
- Update config validation to accept all 5 loss_fn values
…prompts

- Expand --loss-fn choices to include pseudo_huber, x0_mse, x0_pseudo_huber
- Wizard shows delta prompt for all delta-using modes (huber, pseudo_huber, x0_pseudo_huber)
- Update review summary label to Huber/Pseudo-Huber delta
…art tags

- Add pseudo_huber, x0_mse, x0_pseudo_huber to loss function dropdown
- Update help panel with descriptions for all 5 modes
- Show huber delta input for all delta-using loss modes
- Add Arrangement Loss and Detail Loss to chart tag metadata
- build_train_namespace passes dynamic_channel_balance through (was
  silently reset to default for wizard and preset runs)
- presets recognise dynamic_channel_balance
- GUI CLI builder emits --loss-weighting unconditionally (backend
  default is flow_snr, so omitting 'none' trained with flow_snr),
  plus --loss-fn/--huber-delta/--t-bias/--latent-noise/
  --channel-balance/--dynamic-channel-balance/--vae-channel-prior/
  --legacy-loss/--ema-start-step, and snr-gamma for flow_snr too
- preset loader maps dynamic_channel_balance into the form

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e 'All The Levers submenu' on the GUI and Wizard
- Vendor Rose optimizer (stateless, zero VRAM) into sidestep_engine/vendor/rose.py
- Add SCAO (sparse curvature-aware) + Rose to build_optimizer() factory
- SCAO: requires pip install scao>=0.1.0, falls back to AdamW if missing
- Rose: vendored, always available, overrides compute_dtype to fp32
- Update _OPTIM_BPP: SCAO=14.0, Rose=0.5 (stateless)
- Add to CLI --optimizer-type choices
- Add to wizard menu with Rose LR guidance (0.01-0.1)
- Add to GUI frontend: both optimizer dropdowns, help panel, LR hints
- JS reactivity: Rose LR auto-switch (0.01) + hint display
- Add SCAO to dependency_check.py (Rose needs no dep check)
- Vendor Automagic from ostris/ai-toolkit (Apache 2.0) into vendor/automagic.py
- Merged optimizer_utils.py inline (Auto8bitTensor, copy_stochastic, etc.)
- Made QBytesTensor (optimum.quanto) optional -- not needed for bf16 LoRA
- Add to build_optimizer() factory, VRAM estimation (8.0 BPP), CLI, config
- GUI: both optimizer dropdowns, help panel, automagic-lr-hint div
- JS reactivity: auto-switch LR to 1e-6 + hint display
- Wizard: added to optimizer menu
- Verified: imports, factory, and 2-step optimization all pass
…ix Rose OOM

Root cause: previous presets changed 7+ variables from the proven config,
making optimizer comparison meaningless. All 4 runs converged to ~0.71
(vs 0.41 for good run) due to GA=4, Huber loss, SNR weighting, dim=256,
and 100-step patience — not optimizer issues.

Changes:
- All presets now use XL_SFT_Turbo base (GA=12, MSE, dim=512, 1000ep, patience=250)
- Only optimizer-specific settings differ: LR and scheduler
- Rose: gradient_checkpointing=true (was false, causing OOM on 32GB)
- SCAO: install_spec updated to git URL (not on PyPI)
- SCAO: installed via 'uv pip install git+https://github.com/whispering3/scao.git'
- Help text updated with correct SCAO install command
(from bedf53c; the four experimental Style_*/Optimised_* presets were
net-deleted later on the source branch and are not included)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(from 694b1cd, frontend logic only)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lr_scale_self_attn/cross_attn/mlp were absent from
build_train_namespace, preset field maps, and the GUI CLI builder, so
wizard and preset values silently fell back to 1.0.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…et_loss

(from d54ccb3; preset JSON retunes excluded — those presets are not in
this branch)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…peline

1. presets.py: PRESET_FIELDS whitelist didn't include the new field,
   so load_preset() stripped it. Added to both PRESET_FIELDS and PRESET_TYPES.
2. config_factory.py: TrainingConfigV2 constructor call didn't pass the
   field, so it always defaulted to False. Added import + kwarg.
3. Frontend: Missing checkbox element in HTML, missing chkMap entry in
   workspace-setup.js, and missing gatherFullConfig() entry in
   workspace-config.js. All three added.

Root cause: new config fields require updates at 6 locations. This closes
the full preset -> GUI -> config_factory -> TrainingConfigV2 -> trainer_loop
pipeline for target_loss_auto_stop.
The auto-stop logic existed only in run_basic_training_loop() inside
trainer_loop.py (the non-Fabric fallback path). But Lightning Fabric is
installed, so training ALWAYS runs through _train_fabric() in trainer.py,
which had its own complete epoch loop WITHOUT any auto-stop logic.

Every single training run from the GUI used the Fabric path. The
auto-stop code literally never executed — not because of pyc caching,
not because of config issues, but because it was in a function that
was never called.

Fixed by adding identical auto-stop logic to _train_fabric() after
the early stopping block.
…builder

Without these, auto-stop silently stayed off for wizard/preset runs and
GUI-generated CLI commands.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@coderabbitai

coderabbitai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

@scragnog, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 59 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 404710e1-0404-40df-9416-dac7e962d476

📥 Commits

Reviewing files that changed from the base of the PR and between ecd13bd and 10a68b9.

📒 Files selected for processing (32)
  • frontend/index.html
  • frontend/js/api-cli.js
  • frontend/js/defaults.json
  • frontend/js/training.js
  • frontend/js/workspace-behaviors.js
  • frontend/js/workspace-config.js
  • frontend/js/workspace-setup.js
  • presets/XL_SFT_Turbo.json
  • sidestep_engine/cli/args.py
  • sidestep_engine/core/adaptive_timestep.py
  • sidestep_engine/core/config_factory.py
  • sidestep_engine/core/configs.py
  • sidestep_engine/core/ema.py
  • sidestep_engine/core/lora_module.py
  • sidestep_engine/core/optim.py
  • sidestep_engine/core/timestep_sampling.py
  • sidestep_engine/core/trainer.py
  • sidestep_engine/core/trainer_loop.py
  • sidestep_engine/core/vram_estimation.py
  • sidestep_engine/data/preprocess.py
  • sidestep_engine/training_defaults.py
  • sidestep_engine/ui/config_panel.py
  • sidestep_engine/ui/dependency_check.py
  • sidestep_engine/ui/flows/common.py
  • sidestep_engine/ui/flows/review_summary.py
  • sidestep_engine/ui/flows/train_steps_advanced.py
  • sidestep_engine/ui/flows/train_steps_levers.py
  • sidestep_engine/ui/flows/train_steps_training.py
  • sidestep_engine/ui/presets.py
  • sidestep_engine/vendor/automagic.py
  • sidestep_engine/vendor/data_module.py
  • sidestep_engine/vendor/rose.py

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants