gemma 4 fixes

core/finetuning.py

@@ -29,7 +29,12 @@ import math
 import json
 from datetime import datetime, timezone
 from utils.dataset_helpers import load_ftpo_multi_dataset
-from utils.model_helpers import fix_gemma3_checkpoint, detie_lm_head, prepare_gemma3_for_save
+from utils.model_helpers import (
+    fix_gemma3_checkpoint,
+    detie_lm_head,
+    prepare_gemma3_for_save,
+    unwrap_clippable_linears,
+)
 # Import the new dataloader function
 from utils.trainer_dataloaders import load_and_prepare_dataset
 logger = logging.getLogger(__name__)
@@ -999,6 +1004,9 @@ def run_dpo_finetune(config: dict, experiment_run_dir: Path):
                 device_map  = {"": "cpu"},
                 trust_remote_code = True,
             )
+        unwrapped = unwrap_clippable_linears(base_fp16)
+        if unwrapped:
+            logger.info("Unwrapped %d clippable linear modules before PEFT merge.", unwrapped)
         model_fp16 = PeftModel.from_pretrained(
             base_fp16,
             lora_dir,                           # plug in the saved adapter
@@ -1051,4 +1059,4 @@ def run_dpo_finetune(config: dict, experiment_run_dir: Path):
         except Exception as e:
             logger.error(f"Error saving model: {e}", exc_info=True)
 
-    logger.info("Finetuning process completed.")
+    logger.info("Finetuning process completed.")

core/ftpo_trainer.py

@@ -309,9 +309,23 @@ class FTPOTrainer(DPOTrainer):
         frac_win  = wins_tok.float().sum(-1) / ch_mask.sum(-1).clamp(min=1e-8)  
         chosen_win = frac_win.mean().detach()                 
 
+        active_delta = delta_tok[ch_mask]
+        active_weights = weights[ch_mask]
+        margin_win = (
+            ((delta_tok >= clip_epsilon_logits) & ch_mask).float().sum()
+            / ch_mask.float().sum().clamp(min=1e-8)
+        ).detach()
+        mean_delta = active_delta.mean().detach()
+        median_delta = active_delta.median().detach()
+        active_weight = active_weights.mean().detach()
+
         metrics = {
             "pref_loss":  pref_loss.detach(),
             "chosen_win": chosen_win,
+            "margin_win": margin_win,
+            "mean_delta": mean_delta,
+            "median_delta": median_delta,
+            "active_weight": active_weight,
             **extra_metrics,
         }
         self.store_metrics(metrics, train_eval="train")

utils/model_helpers.py

@@ -7,6 +7,30 @@ from safetensors.torch import safe_open, save_file
 
 log = logging.getLogger(__name__)
 
+
+def unwrap_clippable_linears(model) -> int:
+    """
+    Replace remote-code clippable linear wrappers with their underlying
+    torch.nn.Linear modules so vanilla PEFT can inject adapters.
+
+    Some Gemma loaders expose projection modules as Gemma*ClippableLinear
+    wrappers with the real Linear stored on `.linear`. PEFT's LoRA injection
+    only accepts the inner Linear type.
+    """
+    count = 0
+
+    for child_name, child in list(model.named_children()):
+        count += unwrap_clippable_linears(child)
+
+        inner = getattr(child, "linear", None)
+        if inner is None:
+            continue
+        if child.__class__.__name__.endswith("ClippableLinear"):
+            setattr(model, child_name, inner)
+            count += 1
+
+    return count
+
 def fix_gemma3_checkpoint(ckpt_dir: str | Path) -> None:
     """
     If `ckpt_dir` is a Gemma-3 checkpoint whose tensor keys look like