vram reduction for all models and slight it/s boost

[PlagueKind]

Use config dtype for training network when optimizer supports low precision

Previously, network.force_to() always used torch.float32, causing LoRA/LoKR adapter weights to be stored in float32 regardless of the configured model dtype. This resulted in unnecessarily high VRAM usage and slower iteration speeds during training.

Changes

BaseSDTrainProcess now selects the network dtype based on whether the configured optimizer supports low-precision training:

_fp32_optimizers = ["adam", "adamw", "sgd", "adagrad"]
_network_dtype = torch.float32 if self.train_config.optimizer.lower() in _fp32_optimizers else self.sd.torch_dtype
self.network.force_to(self.device_torch, dtype=_network_dtype)

Optimizers that use stochastic rounding or quantized states (automagic, automagic2, automagic3, adam8bit, adafactor, prodigy_8bit) will now run the network in the config dtype (typically bf16). Standard optimizers that require float32 precision for stable gradient accumulation (adam, adamw, sgd, adagrad) continue to use float32.

Impact

• ~50% VRAM reduction for the training network when using a low-precision-capable optimizer and a bf16 model. For a 256-module LoKR on Krea-2, this reduced network VRAM from ~9GB to ~4.5GB.
• Faster iteration speed due to reduced memory pressure and bf16 matmul throughput on modern CUDA hardware.
• No quality regression for supported optimizers — automagic3 uses stochastic rounding specifically to preserve update precision in low-precision dtypes.
• Safe fallback — float32-dependent optimizers are unaffected.

examples on 5060ti 16gb:

before change on krea2 I needed 70% offload not to oom during training. After i need 20%
before change on ideogram4 I was vram thrashing with no offloading and 1024, after I have about 2gb of headroom.