Merge dfe1da4d36 into fa53f71ec0

* fix: improve numerical stability by conditionally using float32 in block computations

* doc: update README for improvement stability for fp16 training on Anima in version 0.10.3

README-ja.md

@@ -50,6 +50,9 @@ Stable Diffusion等の画像生成モデルの学習、モデルによる画像
 
 ### 更新履歴
 
+- **Version 0.10.3 (2026-04-02):**
+    - Animaでfp16で学習する際の安定性をさらに改善しました。[PR #2302](https://github.com/kohya-ss/sd-scripts/pull/2302) 問題をご報告いただいた方々に深く感謝します。
+
 - **Version 0.10.2 (2026-03-30):**
     - SD/SDXLのLECO学習に対応しました。[PR #2285](https://github.com/kohya-ss/sd-scripts/pull/2285) および [PR #2294](https://github.com/kohya-ss/sd-scripts/pull/2294) umisetokikaze氏に深く感謝します。
         - 詳細は[ドキュメント](./docs/train_leco.md)をご覧ください。

README.md

@@ -47,6 +47,9 @@ If you find this project helpful, please consider supporting its development via
 
 ### Change History
 
+- **Version 0.10.3 (2026-04-02):**
+    - Stability when training with fp16 on Anima has been further improved. See [PR #2302](https://github.com/kohya-ss/sd-scripts/pull/2302) for details. We deeply appreciate those who reported the issue.
+
 - **Version 0.10.2 (2026-03-30):**
     - LECO training for SD/SDXL is now supported. Many thanks to umisetokikaze for [PR #2285](https://github.com/kohya-ss/sd-scripts/pull/2285) and [PR #2294](https://github.com/kohya-ss/sd-scripts/pull/2294).
         - Please refer to the [documentation](./docs/train_leco.md) for details.

library/anima_models.py

@@ -738,9 +738,9 @@ class FinalLayer(nn.Module):
         x_B_T_H_W_D: torch.Tensor,
         emb_B_T_D: torch.Tensor,
         adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+        use_fp32: bool = False,
     ):
         # Compute AdaLN modulation parameters (in float32 when fp16 to avoid overflow in Linear layers)
-        use_fp32 = x_B_T_H_W_D.dtype == torch.float16
         with torch.autocast(device_type=x_B_T_H_W_D.device.type, dtype=torch.float32, enabled=use_fp32):
             if self.use_adaln_lora:
                 assert adaln_lora_B_T_3D is not None
@@ -863,11 +863,11 @@ class Block(nn.Module):
         emb_B_T_D: torch.Tensor,
         crossattn_emb: torch.Tensor,
         attn_params: attention.AttentionParams,
+        use_fp32: bool = False,
         rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
         adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
         extra_per_block_pos_emb: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        use_fp32 = x_B_T_H_W_D.dtype == torch.float16
         if use_fp32:
             # Cast to float32 for better numerical stability in residual connections. Each module will cast back to float16 by enclosing autocast context.
             x_B_T_H_W_D = x_B_T_H_W_D.float()
@@ -959,6 +959,7 @@ class Block(nn.Module):
         emb_B_T_D: torch.Tensor,
         crossattn_emb: torch.Tensor,
         attn_params: attention.AttentionParams,
+        use_fp32: bool = False,
         rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
         adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
         extra_per_block_pos_emb: Optional[torch.Tensor] = None,
@@ -972,6 +973,7 @@ class Block(nn.Module):
                     emb_B_T_D,
                     crossattn_emb,
                     attn_params,
+                    use_fp32,
                     rope_emb_L_1_1_D,
                     adaln_lora_B_T_3D,
                     extra_per_block_pos_emb,
@@ -994,6 +996,7 @@ class Block(nn.Module):
                     emb_B_T_D,
                     crossattn_emb,
                     attn_params,
+                    use_fp32,
                     rope_emb_L_1_1_D,
                     adaln_lora_B_T_3D,
                     extra_per_block_pos_emb,
@@ -1007,6 +1010,7 @@ class Block(nn.Module):
                     emb_B_T_D,
                     crossattn_emb,
                     attn_params,
+                    use_fp32,
                     rope_emb_L_1_1_D,
                     adaln_lora_B_T_3D,
                     extra_per_block_pos_emb,
@@ -1018,6 +1022,7 @@ class Block(nn.Module):
                 emb_B_T_D,
                 crossattn_emb,
                 attn_params,
+                use_fp32,
                 rope_emb_L_1_1_D,
                 adaln_lora_B_T_3D,
                 extra_per_block_pos_emb,
@@ -1338,16 +1343,19 @@ class Anima(nn.Module):
 
         attn_params = attention.AttentionParams.create_attention_params(self.attn_mode, self.split_attn)
 
+        # Determine whether to use float32 for block computations based on input dtype (use float32 for better stability when input is float16)
+        use_fp32 = x_B_T_H_W_D.dtype == torch.float16
+
         for block_idx, block in enumerate(self.blocks):
             if self.blocks_to_swap:
                 self.offloader.wait_for_block(block_idx)
 
-            x_B_T_H_W_D = block(x_B_T_H_W_D, t_embedding_B_T_D, crossattn_emb, attn_params, **block_kwargs)
+            x_B_T_H_W_D = block(x_B_T_H_W_D, t_embedding_B_T_D, crossattn_emb, attn_params, use_fp32, **block_kwargs)
 
             if self.blocks_to_swap:
                 self.offloader.submit_move_blocks(self.blocks, block_idx)
 
-        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D, use_fp32=use_fp32)
         x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)
         return x_B_C_Tt_Hp_Wp
 

library/train_util.py

@@ -4816,6 +4816,10 @@ def read_config_from_file(args: argparse.Namespace, parser: argparse.ArgumentPar
             ignore_nesting_dict[section_name] = section_dict
             continue
 
+        if section_name == "scale_weight_norms_map":
+            ignore_nesting_dict[section_name] = section_dict
+            continue
+
         # if value is dict, save all key and value into one dict
         for key, value in section_dict.items():
             ignore_nesting_dict[key] = value

networks/lora.py

@@ -5,6 +5,7 @@
 
 import math
 import os
+from fnmatch import fnmatch
 from typing import Dict, List, Optional, Tuple, Type, Union
 from diffusers import AutoencoderKL
 from transformers import CLIPTextModel
@@ -1366,7 +1367,8 @@ class LoRANetwork(torch.nn.Module):
             org_module._lora_restored = False
             lora.enabled = False
 
-    def apply_max_norm_regularization(self, max_norm_value, device):
+    @torch.no_grad()
+    def apply_max_norm_regularization(self, max_norm, device, scale_map: dict[str, float]={}):
         downkeys = []
         upkeys = []
         alphakeys = []
@@ -1381,6 +1383,11 @@ class LoRANetwork(torch.nn.Module):
                 alphakeys.append(key.replace("lora_down.weight", "alpha"))
 
         for i in range(len(downkeys)):
+            max_norm_value = max_norm
+            for key in scale_map.keys():
+                if fnmatch(downkeys[i], key):
+                    max_norm_value = scale_map[key]
+
             down = state_dict[downkeys[i]].to(device)
             up = state_dict[upkeys[i]].to(device)
             alpha = state_dict[alphakeys[i]].to(device)
@@ -1404,7 +1411,7 @@ class LoRANetwork(torch.nn.Module):
                 keys_scaled += 1
                 state_dict[upkeys[i]] *= sqrt_ratio
                 state_dict[downkeys[i]] *= sqrt_ratio
-            scalednorm = updown.norm() * ratio
+            scalednorm: torch.Tensor = updown.norm() * ratio
             norms.append(scalednorm.item())
 
         return keys_scaled, sum(norms) / len(norms), max(norms)

train_network.py

@@ -12,6 +12,8 @@ import json
 from multiprocessing import Value
 import numpy as np
 
+import ast
+
 from tqdm import tqdm
 
 import torch
@@ -1444,8 +1446,9 @@ class NetworkTrainer:
                     optimizer.zero_grad(set_to_none=True)
 
                 if args.scale_weight_norms:
+                    scale_map = args.scale_weight_norms_map if args.scale_weight_norms_map else {}
                     keys_scaled, mean_norm, maximum_norm = accelerator.unwrap_model(network).apply_max_norm_regularization(
-                        args.scale_weight_norms, accelerator.device
+                        args.scale_weight_norms, accelerator.device, scale_map=scale_map
                     )
                     mean_grad_norm = None
                     mean_combined_norm = None
@@ -1713,6 +1716,14 @@ class NetworkTrainer:
 
             logger.info("model saved.")
 
+def parse_dict(input_str):
+    """Convert string input into a dictionary."""
+    try:
+        # Use ast.literal_eval to safely evaluate the string as a Python literal (dict)
+        return ast.literal_eval(input_str)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"Invalid dictionary format: {input_str}")
+
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
@@ -1816,6 +1827,12 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="Scale the weight of each key pair to help prevent overtraing via exploding gradients. (1 is a good starting point) / 重みの値をスケーリングして勾配爆発を防ぐ（1が初期値としては適当）",
     )
+    parser.add_argument(
+        "--scale_weight_norms_map",
+        type=parse_dict,
+        default="{}",
+        help="Scale the weight of each key pair to help prevent overtraing via exploding gradients. (1 is a good starting point) / 重みの値をスケーリングして勾配爆発を防ぐ（1が初期値としては適当）",
+    )
     parser.add_argument(
         "--base_weights",
         type=str,