Make gradient checkpointing and offloading per-component#1476
Merged
Conversation
Collaborator
Author
|
Claude:
|
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
…-component) into preview
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
This comment was marked as resolved.
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.
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>
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.
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>
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fixes [Bug]: Activations offloading depends on layer offload fraction #1136 and [Feat]: Separate offload settings for text encoder #980
includes fix: Flux2 UI — sequence length and caption dropout labels no longer overlap #1460 because it was in the way
required for Microsoft Lens, because it has a text encoder that cannot be (easily) offloaded
some general testing
testing of SD VAE training because this PR touches it and it's rarely ever tested