Training, Optimizers, Device Handling & Z-Image Improvements

[daanforever]

Training, Optimizers, Device Handling & Z-Image Improvements

Overview

This PR bundles improvements across the training pipeline, optimizers, device handling, and Z-Image inference: timestep sampling and schedulers, Adafactor extensions, low-VRAM device handling, and a separate sampling model with quantization and UI fixes. It also adds fixes and refactors for Adafactor LR clamping, Z-Image loading/dtype, train checkpoint memory and optimizer state, PEFT LoRA loading, and text-embedding cache behaviour (shuffle_tokens, per-epoch shuffle, trigger_word).

Training & Schedulers

• blank_prompt_probability — Control how often Blank Prompt Preservation (BPP) runs to reduce step time while limiting degradation.
• Differential guidance — Separate metric and loss-graph category; LR graph with metric filtering.
• Timestep sampling — Gaussian/content-style distribution, configurable gaussian_std curriculum, timestep debug logging and mapping fixes in BaseSDTrainProcess.
• Warmup for cosine LR — Warmup support for cosine schedulers; fixes for SequentialLR.step() and T_0/T_max duplication with warmup; example config and guide (WARMUP_SCHEDULER_GUIDE.md).
• Fixed cycle training — Config and UI for fixed cycle; sigma parameter fix for weight calculation.
• Networks & config — rank_dropout/module_dropout, alpha for Peft LoRA, SNR compatibility for flow matching; validation for image/video scaling.
• UI — SimpleJob controls, docs, and types for new training options; loss graph and LR visualization.
• Train robustness — Free memory before checkpoint save to reduce OOM risk; validate optimizer state load and document save behaviour.

Optimizers (Adafactor)

• min_lr/max_lr and get_avg_learning_rate for finer control and logging.
• RMS tracking — Methods and weight-update RMS logging.
• Fixes — lr=0 with relative_step, uninitialized state, clamping to max_lr (including max_lr 1e-2→1e-4 and streamlined clamping logic); truncated normal sampling and spelling fix.

Device Handling

• safe_module_to_device in toolkit/util/device.py for consistent, low-VRAM-aware transfers.
• Conditional model transfer in ltx2, qwen_image, and wan22 based on low_vram (and related) flags instead of raw model.to(device).

Z-Image & Inference

• Separate sampling model — Own transformer with device handling (e.g. CPU when not generating) and optional quantization.
• Sampling model unload and LoRA on sampling transformer during generation.
• Quantization — Custom attribute for accuracy recovery adapter; assignment fix in quantize_model.
• Robustness — File cleanup with retries and logging; ZImageModel device transfer fixes.
• UI — Stop job API correctly marks job stopped and handles dead process; sample image viewer keyboard navigation fix.
• Loading & dtype — Load sampling transformer first and extract load helpers; normalize model paths (HF identifier warning); use dtype instead of deprecated torch_dtype for transformers; streamline embedding conversion and device management; optional debug_zimage_load to trace safetensors load/mmap.

Text Embedding Cache & Shuffle

• shuffle_tokens + cache_text_embeddings — Apply caption shuffle only when not caching; skip shuffling cached prompt embeds when shuffle_tokens is false in set_epoch_num and load_prompt_embedding; respect shuffle_tokens for cached embeddings so cache paths stay stable.
• Per-epoch shuffle — When cache_text_embeddings is true, shuffle token order along the sequence at the start of each new epoch; PromptEmbeds.shuffle_sequence(), dataset set_epoch_num / clear_cached_embeddings_memory, wired at epoch boundary in BaseSDTrainProcess; epoch_num synced at train start for correct resume.
• First segment/token fixed — In get_caption, when shuffle_tokens is on, keep the first comma-separated segment in place and shuffle the rest; in shuffle_sequence, keep index 0 fixed and permute only 1..seq_len-1.
• Memory & UX — Clear cached embeddings when not using cache; retain prompt embeddings when cached; logging when cached tokens are shuffled each epoch.
• SD trainer — Prefer trigger_word over blank when batch.prompt_embeds is None; require trigger_word when cache_text_embeddings is enabled and batch has no cached embeds (ValueError otherwise); reg batches keep using blank.

PEFT

• LoRA loading — Apply alpha key fix for all PEFT types when loading LoRA.

Commits

1. feat(training): timestep sampling, schedulers, networks, loss UI — blank_prompt_probability, differential guidance, LR graph; Gaussian/content-style timestep sampling, warmup for cosine LR; BaseSDTrainProcess timestep mapping, fixed cycle; rank_dropout/module_dropout, alpha for Peft LoRA, SNR flow matching; timestep debug logging, config and UI.

2. feat(optimizers): Adafactor min/max lr, RMS tracking, fixes — min_lr/max_lr, get_avg_learning_rate, RMS tracking; fixes for lr=0 with relative_step, uninitialized state, clamp to max_lr; truncated normal, weight update RMS logging.

3. refactor(device): safe_module_to_device and conditional model transfer — safe_module_to_device in toolkit/util/device.py; use in ltx2, qwen_image, wan22 for low_vram-aware handling.

4. feat(Z-Image): sampling model, quantization, stop job API — Separate sampling model, device handling, quantization; unload, LoRA on sampling transformer; accuracy recovery adapter, file cleanup with retries; stop job API and sample viewer fix.

5. fix(optimizers): adjust max_lr and improve learning rate clamping — Adafactor max_lr 1e-2→1e-4; streamlined clamping between min_lr and max_lr with relative_step.

6. refactor(z_image): streamline embedding conversion and device management — dtype instead of deprecated torch_dtype; normalize paths; load sampling transformer first, extract helpers; embedding conversion and device handling; debug_zimage_load.

7. fix(train): free memory before checkpoint save to reduce OOM risk — Free memory before checkpoint save; validate optimizer state load and document save behaviour.

8. fix(peft): apply alpha key fix for all peft types when loading LoRA — Alpha key fix for all PEFT types on LoRA load.

9. refactor(dataloader): remove redundant token shuffling logic — Redundant shuffling removed; shuffle_tokens skipped when caching; first segment/token fixed (get_caption + shuffle_sequence); set_epoch_num, clear_cached_embeddings_memory, shuffle at epoch start, logging; respect shuffle_tokens for cached embeds; trigger_word when cache_text_embeddings and no cache.

10. fix(adafactor): apply new min_lr/max_lr on restart after loading checkpoint — Ensures Adafactor applies updated min_lr/max_lr when resuming from checkpoint.

11. refactor(paths): rename normalize_model_path to normalize_path and strip whitespace — Centralised path normalisation and whitespace stripping in toolkit/paths.py.

12. feat(debug): add memory_debug util and use for CUDA load/unload logging — Adds toolkit/util/debug.py with memory_debug() context manager and debug config; wraps text encoder unload and Z-Image load stages for optional [DEBUG …] CUDA memory lines when logging.debug is enabled.

13. fix(sampling), feat(lora): sampling VRAM, single LoRA, path and debug wiring — Sampling: load sampling transformer on CPU, guarantee unload to CPU after generate_images, do not force unet to GPU. LoRA: single LoRA for training and sampling via shared parameters; free pretrained LoRA memory after load. Uses normalize_path in Z-Image and wires memory_debug in BaseSDTrainProcess LoRA load and Z-Image stages; train.example.yaml, base_model, network_mixins updates.