diff --git a/core/finetuning.py b/core/finetuning.py index 19f8256..32a114a 100644 --- a/core/finetuning.py +++ b/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.") \ No newline at end of file + logger.info("Finetuning process completed.") diff --git a/core/ftpo_trainer.py b/core/ftpo_trainer.py index 7c981cd..f5cfe4c 100644 --- a/core/ftpo_trainer.py +++ b/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") diff --git a/utils/model_helpers.py b/utils/model_helpers.py index c74d53d..410985a 100644 --- a/utils/model_helpers.py +++ b/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