Skip to content

Make gradient checkpointing and offloading per-component#1476

Merged
dxqb merged 11 commits into
Nerogar:mergefrom
dxqb:split-offload
Jul 5, 2026
Merged

Make gradient checkpointing and offloading per-component#1476
dxqb merged 11 commits into
Nerogar:mergefrom
dxqb:split-offload

Conversation

@dxqb

@dxqb dxqb commented May 25, 2026

Copy link
Copy Markdown
Collaborator

@dxqb dxqb linked an issue May 25, 2026 that may be closed by this pull request
@dxqb

dxqb commented May 25, 2026

Copy link
Copy Markdown
Collaborator Author

Claude:

  • activation_offloading should default to False — it was always default-True but only took effect with layer offloading; now it works standalone, so a fresh fine-tune offloads activations out of the box.
  • "Layer Offload Fraction" is shown for CLIP encoders but is a no-op — the CLIP setup discards its conductor, so the value never drives anything; don't render the field for CLIP (or wire it up).

@dxqb dxqb added the preview merged in the preview branch label May 29, 2026
dxqb added a commit that referenced this pull request Jun 3, 2026
- BaseAnimaSetup: per-component checkpointing_or_offloading_enabled(),
  remove weight_list from create_autocast_context / disable_fp16_autocast_context
- AnimaFineTune/LoRASetup: latent_caching → image_caching / text_caching
- ModelType: add ANIMA to _MODEL_PARTS and supported_training_methods

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
dxqb added a commit to TheForgotten69/OneTrainer that referenced this pull request Jun 3, 2026
dxqb added a commit that referenced this pull request Jun 4, 2026
In the upstream TrainingTab.py (PR #1476), config is stored as
self.train_config on the view. In preview's Base*/controller pattern
it lives on controller.config. One call site in __setup_stable_diffusion_ui
was translated incorrectly during the bec207a merge.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
dxqb added a commit that referenced this pull request Jun 4, 2026
@dxqb dxqb mentioned this pull request Jun 6, 2026
3 tasks
@dxqb

This comment was marked as resolved.

dxqb added a commit that referenced this pull request Jun 14, 2026
…ent) into preview

# Conflicts:
#	modules/ui/ModelTab.py
#	modules/ui/TopBar.py
#	modules/ui/TrainUI.py
#	modules/ui/TrainingTab.py
dxqb added a commit that referenced this pull request Jun 14, 2026
Anima, Lens, and Ideogram setup files used the pre-#1476/#1462 4-arg
create_autocast_context/disable_fp16_autocast_context (with a weight-dtype
list), the old config.gradient_checkpointing.enabled() global check, and the
renamed config.latent_caching field. Update them to the current 3-arg
autocast helpers, per-part checkpointing via enable_checkpointing_for_*, and
config.image_caching/config.text_caching.
dxqb added a commit that referenced this pull request Jun 19, 2026
Add a _MODEL_PARTS table + ModelType.model_parts() as the single source
of truth for which components each model type has, keyed by TrainConfig
field names, and a ModelType.supported_training_methods() that enumerates
every type explicitly, raising on an unknown type rather than defaulting.

Collapse ModelTab's per-type __setup_*_ui methods into one __setup_ui that
derives the has_* widget flags from model_parts(), and collapse TopBar's
per-type training-method dispatch to build its dropdown from
supported_training_methods().

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
dxqb added 3 commits June 20, 2026 20:39
Rebased onto centralize-model-type: this is the offloading-only part of
split-offload, with the model-composition centralization (ModelType,
ModelTab, TopBar) excluded since it already landed separately.
Adds ModelType.model_parts()/has_multiple_text_encoders()/supported_training_methods()
as the single source of truth for which components a model has and which training
methods it supports. BaseModelTabView.build_content and TopBarController now derive
the Model tab layout and training-method dropdown generically from these, replacing
one hand-written per-model method / if-chain each.
Brings in the Ctk/PySide6 UI split (Base*View.py + *Controller.py) on
top of split-offload's per-component offloading UI. Resolved conflicts:

- ModelType.py: kept centralize-model-type's diffusion-model-first
  _MODEL_PARTS ordering (matches lora-output-formats).
- ModelTab.py/TopBar.py: deleted (already replaced by
  BaseModelTabView.py/TopBarController.py; split-offload never
  diverged on these files, so no porting needed here).
- TrainUI.py/TrainingTab.py: deleted, with their per-component
  offloading UI (__create_offloading_widgets, VAE frame, "Async
  Offloading" switch) ported onto BaseTrainingTabView.py/
  BaseTrainUIView.py in place of the removed OffloadingWindow popup
  and its now-dead Base/Ctk/PySide6OffloadingWindowView.py +
  OffloadingWindowController.py split.
@dxqb dxqb changed the base branch from master to merge July 1, 2026 07:21
dxqb and others added 6 commits July 4, 2026 08:07
Add a ModelType.supported_training_methods() that enumerates every
model type explicitly, raising on an unknown type rather than
defaulting. Collapse ModelTab's per-type __setup_*_ui methods into one
build_content that derives the has_* widget flags from
ModelType.model_parts(), and collapse TopBar's per-type training-method
dispatch to build its dropdown from supported_training_methods().

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Resolves conflicts by keeping split-offload's per-component offload_fraction API
and reconciling it with centralize-model-type's ModelType/ModelTabController
centralization. Also ports the per-component checkpointing API (3-arg
enable_checkpointing_for_* calls, no global gradient_checkpointing.enabled()
guard) to Anima/Krea2, which were merged in as new files and had missed it, and
fixes the Krea2 16GB LoRA preset to the current per-part offload_fraction format.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
BaseAnimaSetup/BaseKrea2Setup were added wholesale by the centralize-model-type
merge and kept the old enabled()-guarded call pattern instead of the new
part-based one applied to Chroma/Qwen during conflict resolution.
# Conflicts:
#	modules/modelSetup/BaseChromaSetup.py
#	modules/modelSetup/BaseFlux2Setup.py
#	modules/modelSetup/BaseQwenSetup.py
#	modules/modelSetup/BaseZImageSetup.py
- default activation_offloading to False for fresh configs (it previously
  only took effect combined with layer offloading, so a fresh fine-tune
  now offloads activations out of the box, which enable_activation_offloading
  never did)
- hide the no-op "Layer Offload Fraction" control for CLIP text encoders,
  whose conductor is discarded in setup_optimizations
- narrow the caching_threads/layer-offloading guard in create_data_loader
  to the text-encoder parts that actually run in the caching dataloader's
  worker threads, via a new ModelType.text_encoder_parts() helper
- drop the dead model.vae.enable_gradient_checkpointing() call: the base
  SD setup only ever decodes the VAE under torch.no_grad(), so checkpointing
  it never had any effect; VAEs are small enough that FINE_TUNE_VAE doesn't
  need it either
…ding from train

- do not gate checkpointing or offloading on a part's train flag: grad can
  flow through a frozen part (embedding training runs grad through the frozen
  denoiser and text encoders), where checkpointing and offloading still save
  VRAM. checkpointing_enabled() -> gradient_checkpointing; offloading_enabled()
  -> offload_fraction > 0 or activation_offloading. Restores pre-branch
  behaviour for embedding training.
- reject offloading with checkpointing disabled (raise NotImplementedError)
  instead of silently forcing checkpointing on: the current offload modes need
  the use_reentrant recompute to service the backward pass.
- CLIP encoder checkpointing passes offload_enabled=False so a migrated
  offload_fraction can never build a self-activating conductor for a
  non-offloadable text encoder.
- hide the dead Layer Offload Fraction / Offload Activations controls for the
  SD/SDXL unet and Wuerstchen prior (no conductor), and gate the SD/SDXL native
  gradient checkpointing on checkpointing_enabled() so a dead offload toggle
  can no longer flip it on.
- derive ModelType.has_multiple_text_encoders() from model_parts() instead of a
  parallel is_*() chain.
- remove the dead LayerOffloadConductor.layer_offload_activated().

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@dxqb dxqb marked this pull request as ready for review July 4, 2026 12:29
These models always load their text encoder, but the checkpointing-refactor
merge on this branch didn't pick up the None-check removal from Nerogar#1579
because the surrounding lines had already diverged.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@dxqb dxqb merged commit 6fd9b7f into Nerogar:merge Jul 5, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

preview merged in the preview branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Activations offloading depends on layer offload fraction [Feat]: Separate offload settings for text encoder

1 participant