Apply sigma to contrastive guidance to balance loss better. Prevent noise grads on images/non audio datasets. Prep for training adapters on MiniMax H3

toolkit/config_modules.py

@@ -590,6 +590,11 @@ class TrainConfig:
         self.do_guidance_loss = kwargs.get('do_guidance_loss', False)
         self.guidance_loss_target: Union[int, List[int, int]] = kwargs.get('guidance_loss_target', 3.0)
         self.do_guidance_loss_cfg_zero: bool = kwargs.get('do_guidance_loss_cfg_zero', False)
+        # 'constant' uses guidance_loss_target as is. 'sigma' decays the target
+        # toward 1.0 as sigma falls (effective = 1 + (target - 1) * sigma) so the
+        # extrapolation never amplifies the unpredictable fresh-noise term at low
+        # sigma. Needed for guidance-distilled models with no guidance embedding.
+        self.guidance_loss_schedule: str = kwargs.get('guidance_loss_schedule', 'sigma')
         self.unconditional_prompt: str = kwargs.get('unconditional_prompt', '')
         if isinstance(self.guidance_loss_target, tuple):
             self.guidance_loss_target = list(self.guidance_loss_target)

toolkit/network_mixins.py

@@ -451,7 +451,9 @@ class ToolkitModuleMixin:
         if is_ao_quantized:
             from toolkit.util.quantize import get_torchao_config, requantize_module_weight
             config = get_torchao_config(self._get_base_qtype())
-            if config is None:
+            if config is None and not getattr(self.org_module[0], "is_ostris_quantized", False):
+                # ostris-quantized layers re-quantize through their own backend
+                # (requantize_module_weight) and need no torchao config
                 print_once(f"Warning: merging into quantized layer {getattr(self, 'lora_name', '?')} "
                            f"without a known qtype; it will be left dequantized")
             requantize_module_weight(self.org_module[0], weight.to(weight_device), orig_dtype, config)