Fix MoE training: per-expert LR logging and param group splitting

This commit fixes two critical issues with Mixture of Experts (MoE) training
for dual-transformer models like WAN 2.2 14B I2V:

**Issue 1: Averaged LR logging masked expert-specific behavior**
- Previous logging averaged LR across all param groups (both experts)
- Made it impossible to verify LR was resuming correctly per expert
- Example: High Noise at 0.0005, Low Noise at 0.00001 → logged as 0.00026

**Fix:** Per-expert LR display (BaseSDTrainProcess.py lines 2198-2226)
- Detects MoE via multiple param groups
- Shows separate LR for each expert: "lr0: 5.0e-04 lr1: 3.5e-05"
- Makes expert-specific LR adaptation visible and debuggable

**Issue 2: Transformer detection bug prevented param group splitting**
- _prepare_moe_optimizer_params() checked for '.transformer_1.' (dots)
- But lora_name uses '$$' separator: "transformer$$transformer_1$$blocks..."
- Check never matched, all params went into single group → no per-expert LRs

**Fix:** Corrected substring matching (lora_special.py lines 622-630)
- Changed from '.transformer_1.' to 'transformer_1' substring check
- Now correctly creates separate param groups for transformer_1/transformer_2
- Enables per-expert lr_bump, min_lr, max_lr with automagic optimizer

**Result:**
- Visible per-expert LR adaptation: lr0 and lr1 tracked independently
- Proper LR state preservation when experts switch every N steps
- Accurate monitoring of training progress for each expert

Example output:
```
lr0: 2.8e-05 lr1: 0.0e+00 loss: 8.414e-02  # High Noise active
lr0: 5.2e-05 lr1: 1.0e-05 loss: 7.821e-02  # After switch to Low Noise
lr0: 5.2e-05 lr1: 3.4e-05 loss: 6.103e-02  # Low Noise adapting, High preserved
```

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

jobs/process/BaseSDTrainProcess.py

@@ -1791,6 +1791,8 @@ def run(self):
                     config['default_lr'] = self.train_config.lr
                 if 'learning_rate' in sig.parameters:
                     config['learning_rate'] = self.train_config.lr
+                if 'optimizer_params' in sig.parameters:
+                    config['optimizer_params'] = self.train_config.optimizer_params
                 params_net = self.network.prepare_optimizer_params(
                     **config
                 )
@@ -2203,7 +2205,13 @@ def run(self):
                     # torch.cuda.empty_cache()
                     # if optimizer has get_lrs method, then use it
                     if hasattr(optimizer, 'get_avg_learning_rate'):
-                        learning_rate = optimizer.get_avg_learning_rate()
+                        # Check if this is MoE with multiple param groups
+                        if hasattr(optimizer, 'get_learning_rates') and len(optimizer.param_groups) > 1:
+                            # Show per-expert LRs for MoE
+                            group_lrs = optimizer.get_learning_rates()
+                            learning_rate = None  # Will use group_lrs instead
+                        else:
+                            learning_rate = optimizer.get_avg_learning_rate()
                     elif hasattr(optimizer, 'get_learning_rates'):
                         learning_rate = optimizer.get_learning_rates()[0]
                     elif self.train_config.optimizer.lower().startswith('dadaptation') or \
@@ -2215,7 +2223,16 @@ def run(self):
                     else:
                         learning_rate = optimizer.param_groups[0]['lr']
 
-                    prog_bar_string = f"lr: {learning_rate:.1e}"
+                    # Format LR string (per-expert for MoE, single value otherwise)
+                    if hasattr(optimizer, 'get_avg_learning_rate') and learning_rate is None:
+                        # MoE: show each expert's LR
+                        lr_strings = []
+                        for i, lr in enumerate(group_lrs):
+                            lr_val = lr.item() if hasattr(lr, 'item') else lr
+                            lr_strings.append(f"lr{i}: {lr_val:.1e}")
+                        prog_bar_string = " ".join(lr_strings)
+                    else:
+                        prog_bar_string = f"lr: {learning_rate:.1e}"
                     for key, value in loss_dict.items():
                         prog_bar_string += f" {key}: {value:.3e}"
 

toolkit/lora_special.py

@@ -570,10 +570,110 @@ def create_modules(
                 unet.conv_in = self.unet_conv_in
                 unet.conv_out = self.unet_conv_out
 
-    def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):
-        # call Lora prepare_optimizer_params
+    def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr, optimizer_params=None):
+        # Check if we're training a WAN 2.2 14B MoE model
+        base_model = self.base_model_ref() if self.base_model_ref is not None else None
+        is_wan22_moe = base_model is not None and hasattr(base_model, 'arch') and base_model.arch in ["wan22_14b", "wan22_14b_i2v"]
+
+        # If MoE model and optimizer_params provided, split param groups for high/low noise experts
+        if is_wan22_moe and optimizer_params is not None and self.unet_loras:
+            return self._prepare_moe_optimizer_params(text_encoder_lr, unet_lr, default_lr, optimizer_params)
+
+        # Otherwise use standard param group creation
         all_params = super().prepare_optimizer_params(text_encoder_lr, unet_lr, default_lr)
 
+        if self.full_train_in_out:
+            if self.is_pixart or self.is_auraflow or self.is_flux or (base_model is not None and base_model.arch == "wan21"):
+                all_params.append({"lr": unet_lr, "params": list(self.transformer_pos_embed.parameters())})
+                all_params.append({"lr": unet_lr, "params": list(self.transformer_proj_out.parameters())})
+            else:
+                all_params.append({"lr": unet_lr, "params": list(self.unet_conv_in.parameters())})
+                all_params.append({"lr": unet_lr, "params": list(self.unet_conv_out.parameters())})
+
+        return all_params
+
+    def _prepare_moe_optimizer_params(self, text_encoder_lr, unet_lr, default_lr, optimizer_params):
+        """
+        Prepare optimizer params with separate groups for High Noise and Low Noise experts.
+        Allows per-expert lr_bump, min_lr, max_lr configuration for automagic optimizer.
+        """
+        self.requires_grad_(True)
+        all_params = []
+
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+        # Handle text encoder loras (standard, no splitting)
+        if self.text_encoder_loras:
+            param_data = {"params": enumerate_params(self.text_encoder_loras)}
+            if text_encoder_lr is not None:
+                param_data["lr"] = text_encoder_lr
+            all_params.append(param_data)
+
+        # Split unet_loras by transformer (High Noise = transformer_1, Low Noise = transformer_2)
+        if self.unet_loras:
+            high_noise_loras = []
+            low_noise_loras = []
+            other_loras = []
+
+            for lora in self.unet_loras:
+                # Note: lora_name uses $$ as separator, so check for 'transformer_1' substring
+                # This correctly matches names like "transformer$$transformer_1$$blocks$$0$$attn1$$to_q"
+                if 'transformer_1' in lora.lora_name:
+                    high_noise_loras.append(lora)
+                elif 'transformer_2' in lora.lora_name:
+                    low_noise_loras.append(lora)
+                else:
+                    other_loras.append(lora)
+
+            # Extract per-expert optimizer params with fallback to defaults
+            default_lr_bump = optimizer_params.get('lr_bump')
+            default_min_lr = optimizer_params.get('min_lr')
+            default_max_lr = optimizer_params.get('max_lr')
+
+            # High Noise Expert param group
+            if high_noise_loras:
+                high_noise_params = {"params": enumerate_params(high_noise_loras)}
+                if unet_lr is not None:
+                    high_noise_params["lr"] = unet_lr
+
+                # Add per-expert optimizer params if using automagic
+                if default_lr_bump is not None:
+                    high_noise_params["lr_bump"] = optimizer_params.get('high_noise_lr_bump', default_lr_bump)
+                if default_min_lr is not None:
+                    high_noise_params["min_lr"] = optimizer_params.get('high_noise_min_lr', default_min_lr)
+                if default_max_lr is not None:
+                    high_noise_params["max_lr"] = optimizer_params.get('high_noise_max_lr', default_max_lr)
+
+                all_params.append(high_noise_params)
+
+            # Low Noise Expert param group
+            if low_noise_loras:
+                low_noise_params = {"params": enumerate_params(low_noise_loras)}
+                if unet_lr is not None:
+                    low_noise_params["lr"] = unet_lr
+
+                # Add per-expert optimizer params if using automagic
+                if default_lr_bump is not None:
+                    low_noise_params["lr_bump"] = optimizer_params.get('low_noise_lr_bump', default_lr_bump)
+                if default_min_lr is not None:
+                    low_noise_params["min_lr"] = optimizer_params.get('low_noise_min_lr', default_min_lr)
+                if default_max_lr is not None:
+                    low_noise_params["max_lr"] = optimizer_params.get('low_noise_max_lr', default_max_lr)
+
+                all_params.append(low_noise_params)
+
+            # Other loras (not transformer-specific) - use defaults
+            if other_loras:
+                other_params = {"params": enumerate_params(other_loras)}
+                if unet_lr is not None:
+                    other_params["lr"] = unet_lr
+                all_params.append(other_params)
+
+        # Add full_train_in_out params if needed
         if self.full_train_in_out:
             base_model = self.base_model_ref() if self.base_model_ref is not None else None
             if self.is_pixart or self.is_auraflow or self.is_flux or (base_model is not None and base_model.arch == "wan21"):