P1: Anima loss=nan with Automagic or CAME optimizer

[wochenlong]

User report

A user reported that Anima training can quickly turn into loss=nan when using the Automagic or pytorch_optimizer.CAME optimizer.

Local report bundle (maintainer-only, not in repo):

• automagic.json / came.json (reporter TOML exports)
• screenshot showing finite loss then avr_loss=nan within the first few steps; log shows unet dtype: torch.bfloat16

Reported optimizer configs

CAME case:

• optimizer_type = "pytorch_optimizer.CAME"
• learning_rate = "2e-5"
• unet_lr = "2e-5"
• lr_scheduler = "constant_with_warmup"
• lr_warmup_steps = 100
• resolution = "768,768"
• save_precision = "bf16"

Automagic case:

• optimizer_type = "Automagic"
• learning_rate = "1e-6"
• unet_lr = "1e-6"
• lr_scheduler = "constant"
• resolution = "768,768"
• save_precision = "bf16"

Local reproduction attempt

Test environment (generic):

• Local Anima base model under sd-models/anima/
• Small private test dataset (not named in this issue)
• PyTorch: 2.6.0+cu124
• GPU: RTX 4090

Short 8-step smoke runs did not reproduce NaN locally:

• pytorch_optimizer.CAME: completed 8 steps, final avr_loss=0.113
• Automagic: completed 8 steps, final avr_loss=0.102

This suggests the issue is not an unconditional optimizer crash. It is likely triggered by a risky precision/optimizer combination plus user environment, data order, longer training, or PyTorch/runtime differences.

How to try reproducing the old failure mode

1. Use reporter configs above with Anima LoRA training.
2. Force the risky combo in the adapted *-sd-scripts.toml (or WebUI advanced): mixed_precision="bf16" and full_bf16=true (or full_fp16=true).
3. Run ≥ 50–200 steps (8-step smokes may stay finite).
4. Confirm adapted TOML: if full_bf16 is absent, backend mitigation is active (apply_anima_training_defaults strips it for CAME/Automagic).

On current main, step 2 is normally blocked unless you bypass the backend or use a build before the mitigation.

Suspected risk factor

The backend previously auto-promoted Anima mixed_precision=bf16 into full_bf16=true. That makes trainable LoRA weights/gradients run in bf16 too.

Upstream Anima docs recommend mixed_precision="bf16", but do not require full_bf16. They also mention that if loss becomes NaN, PyTorch should be 2.5 or newer.

For adaptive optimizers like Automagic and pytorch_optimizer.CAME, full half-precision trainable weights are a high-risk default.

Mitigation (implemented on main)

• Stop auto-enabling full_bf16 / full_fp16 for Anima just because mixed_precision is set.
• If Anima optimizer is Automagic or pytorch_optimizer.CAME, automatically remove full_bf16 / full_fp16 from submitted config and log a warning.
• Keep mixed_precision=bf16 valid, but keep trainable LoRA weights in fp32 for stability.
• Tests: tests/test_anima_training_defaults.py
• Docs: docs/anima-training.md (NaN troubleshooting)

Validation

• python -m unittest tests.test_anima_training_defaults tests.test_anima_train_wrapper tests.test_anima_backend_adapter
• Short CAME / Automagic smoke training (8 steps) with mitigation enabled

Follow-up

If users still report NaN after this mitigation, collect:

• PyTorch / CUDA / GPU model
• Full training TOML after backend adaptation (*-sd-scripts.toml)
• Step where loss first becomes NaN
• Whether full_bf16 / full_fp16 is still present in the adapted TOML

[wochenlong]

Update from user: switching the same Anima training case from bf16 to fp16 still produces avr_loss=nan; it only happens a few steps later.

This means the issue is not solved by changing bf16 to fp16. It also strengthens the suspicion that this is a high-risk optimizer/precision/model combination rather than a single bf16-only problem. fp16 is likely less stable for Anima because of the narrower numeric range.

New evidence screenshot was provided in chat: finite loss until around step 7, then avr_loss=nan from around step 8 onward.

[wochenlong]

Latest local validation update:

Using the same parameters as the user report, local training does not produce NaN. The run reached at least step 23/800 with finite avr_loss around 0.095-0.106, and the training monitor/TensorBoard loss curves also remained finite.

This rules out "the reported parameter set always causes NaN" as the primary explanation.

Current direction:

• Parameters alone are probably not the root cause.
• The remaining likely causes are environment/runtime differences, such as PyTorch/CUDA version, GPU/driver behavior, cached latents or text encoder outputs, dataset file differences, or stale runtime state.
• Keep the current mitigation for high-risk precision settings (Automagic / CAME should not use full fp16/bf16 trainable weights), but continue collecting environment details from affected users.

Useful next evidence to request from affected users:

• torch.__version__
• CUDA version and GPU model
• Whether cached *_anima.npz / *_anima_te.npz files were generated locally or copied from another environment
• The adapted *-sd-scripts.toml
• The first step where avr_loss becomes NaN

[wochenlong]

Additional user environment details:

• GPU: NVIDIA RTX 3060 Laptop GPU
• Python: 3.10.11
• PyTorch: reported as 2.7.0 + CUDA 12.8
• AdamW8bit does not produce NaN in the same scenario, but the loss trends upward instead of decreasing.

Interpretation update:

• This further reduces the likelihood that the GUI parameters or Anima training wrapper are the direct root cause.
• The issue appears optimizer/runtime sensitive: Automagic and pytorch_optimizer.CAME can NaN, while AdamW8bit avoids NaN but still does not learn well.
• RTX 3060 Laptop has a much narrower memory/thermal/power envelope than the local RTX 4090 validation machine, so precision, CUDA kernel behavior, cached tensors, and optimizer state stability may differ materially.
• Since PyTorch is reported as newer than the upstream minimum (>=2.5), the remaining likely causes are CUDA 12.8 runtime differences, laptop GPU behavior, dataset/cache-specific tensors, or an optimizer combination that is simply not robust for this Anima setup.

Next evidence to request:

• Exact torch.__version__ string printed by Python, not just installer/package label.
• NVIDIA driver version.
• Whether NaN still happens after deleting all *_anima.npz and *_anima_te.npz caches and rebuilding them locally.
• Adapted *-sd-scripts.toml after backend rewriting.
• A short run with mixed_precision=bf16, no full_bf16, and AdamW8bit at lower LR, to distinguish optimizer instability from general training divergence.

[wochenlong]

Additional update:

The user also reports NaN when using the "Shentong" optimizer, which likely maps to Prodigy / a Prodigy-style adaptive optimizer in the UI.

This changes the working hypothesis:

• The problem is no longer isolated to Automagic or pytorch_optimizer.CAME.
• It may affect most adaptive/non-AdamW optimizers in this user's environment.
• AdamW8bit avoids NaN but the loss rises, which suggests the run is numerically unstable or not learning correctly even under the safer optimizer.

Current stronger hypothesis:

• The root cause is likely environment/data/cache/model specific, with optimizers only changing how quickly the instability becomes visible.
• Possible triggers include corrupted/stale Anima latent/text-encoder caches, model file mismatch, problematic samples/batches, CUDA 12.8 + RTX 3060 Laptop behavior, or a training setup that is unstable for this dataset at the current resolution/repeat/LR.

Suggested next triage steps:

1. Delete all cached *_anima.npz and *_anima_te.npz files in the training dataset, then regenerate caches locally on the affected machine.
2. Verify hashes/sizes of Anima base model, Qwen3, and VAE against a known-good machine.
3. Run a minimal 20-50 step test with AdamW8bit, mixed_precision=bf16, no full bf16/fp16, lower LR, and no preview generation.
4. If NaN or rising loss persists, identify the first batch/image that triggers the divergence by reducing dataset size or fixing dataloader order.
5. Collect the exact adapted *-sd-scripts.toml and the first NaN step for each optimizer.