Fix MoE router LoRA target + attention-only fallback

train_securecoder.py

@@ -571,10 +571,14 @@ def print_stats(stats: list[dict], records: list[dict], tokenizer=None) -> None:
 # --------------------------------------------------------------------------
 # Model + training
 # --------------------------------------------------------------------------
-# Attention + router only by default: on a 128-expert MoE, adapting every expert
-# MLP means ~800M trainable parameters, which dominates VRAM and step time.
-# Pass --target-modules all-linear when you want the MLP/expert capacity too.
-ATTENTION_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate"]
+# Attention projections only by default, for two reasons:
+#   * Qwen3's MoE router is a custom `Qwen3MoeTopKRouter` module, not nn.Linear, so
+#     listing it as a LoRA target dies with "Target module ... is not supported".
+#   * adapting all 128 experts x 32 layers is ~800M trainable parameters, which
+#     dominates VRAM and step time.
+# Use --target-modules all-linear to include the expert MLPs (PEFT then skips the
+# modules it cannot adapt instead of failing).
+ATTENTION_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj"]
 
 
 def load_model_and_tokenizer(args):
@@ -591,10 +595,8 @@ def load_model_and_tokenizer(args):
     if isinstance(targets, str) and targets != "all-linear":
         targets = [t.strip() for t in targets.split(",") if t.strip()]
 
-    model = FastLanguageModel.get_peft_model(
-        model,
+    peft_kwargs: dict[str, Any] = dict(
         r=args.lora_r,
-        target_modules=targets,
         lora_alpha=args.lora_alpha,
         lora_dropout=0.0,
         bias="none",
@@ -602,9 +604,22 @@ def load_model_and_tokenizer(args):
         random_state=args.seed,
         use_rslora=False,
     )
+
+    try:
+        model = FastLanguageModel.get_peft_model(model, target_modules=targets, **peft_kwargs)
+    except ValueError as exc:
+        if "is not supported" not in str(exc) or targets == ATTENTION_TARGETS:
+            raise
+        log.warning("LoRA targets rejected (%s); retrying with attention projections only",
+                    str(exc).splitlines()[0][:180])
+        model = FastLanguageModel.get_peft_model(
+            model, target_modules=ATTENTION_TARGETS, **peft_kwargs
+        )
+
     trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    log.info("trainable parameters: %s (%.2f%% of the model)",
-             f"{trainable:,}", 100 * trainable / max(sum(p.numel() for p in model.parameters()), 1))
+    total = sum(p.numel() for p in model.parameters())
+    log.info("trainable parameters: %s (%.2f%% of the model)", f"{trainable:,}",
+             100 * trainable / max(total, 1))
     return model, tokenizer