Upload train_sft.py with huggingface_hub

train_sft.py

@@ -31,11 +31,35 @@ Jobs box is ephemeral so push_to_hub must be on.
 import argparse
 
 import torch
+import torch.nn as nn
 from datasets import load_dataset
 from peft import LoraConfig
 from transformers import AutoModelForImageTextToText, AutoProcessor
 from trl import SFTConfig, SFTTrainer
 
+PROJ_SUFFIXES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+
+
+def find_lora_targets(model):
+    """FULL module names of the language-tower projection Linears.
+
+    Gemma 4's vision tower wraps its projections in `Gemma4ClippableLinear`
+    (not `nn.Linear`), which PEFT can't adapt, and matching by bare suffix (e.g.
+    "q_proj") would also grab those and crash. So we walk the graph and return
+    the *fully-qualified* names of the real `nn.Linear` projections under
+    `language_model`. PEFT matches those exactly, leaving the vision encoder
+    frozen and skipping MoE expert Parameters (which aren't Modules).
+    """
+    targets = []
+    for name, module in model.named_modules():
+        if not isinstance(module, nn.Linear):
+            continue
+        if "language_model" not in name:
+            continue
+        if name.split(".")[-1] in PROJ_SUFFIXES:
+            targets.append(name)
+    return targets
+
 
 def build_collator(processor, model_config):
     """Render our (system/user/assistant + image) rows into a batch.
@@ -133,11 +157,15 @@ def main():
 
     # Adapt attention + MLP projections across the language tower; leave the
     # vision encoder frozen (grounding signal comes from the LM reading tokens).
+    targets = find_lora_targets(model)
+    print(f"[train] LoRA targets: {len(targets)} language-model Linear projections "
+          f"(e.g. {targets[:2]})")
+    if not targets:
+        raise RuntimeError("no language_model projection Linears found for LoRA")
     peft_config = LoraConfig(
         r=16, lora_alpha=32, lora_dropout=0.05, bias="none",
         task_type="CAUSAL_LM",
-        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
-                        "gate_proj", "up_proj", "down_proj"],
+        target_modules=targets,
         modules_to_save=None,
     )