From f27f21b8c0eb9c7160b62dadbbf4b202769217c1 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Sat, 29 Aug 2026 12:59:26 -0700 Subject: [PATCH 1/6] feat(peft): let sync_ref_model track the reference adapter `sync_ref_model=True` raised NotImplementedError on every PEFT model in GRPO, DPO, KTO and RLOO, on the grounds that PEFT keeps no separate `ref_model` to synchronize. All four trainers already build what a sync needs: a frozen "ref" adapter copied from "default", created when an already-PEFT model is passed in. The guard rejected that case too, so it was broader than the limitation it described. SyncRefModelCallback now handles it. The existing `_sync_target_model` zips the parameters of two modules, which cannot express a reference that lives inside the policy model, so the PEFT branch pairs each "default" parameter with its "ref" counterpart by name, reusing the mapping that initializes the adapter. The ZeRO-3 gather wrapper mirrors the one already beside it. The trainers also create the "ref" adapter when the caller passes `peft_config` and asks for a synced reference. Without that, the same config would work or raise depending only on how the model was constructed. A caller who leaves `sync_ref_model=False` still gets no extra adapter, so nobody pays for one they did not ask for. The guard now fires only when no reference adapter could be created at all, which is `peft<0.20.0` with a LoRA config using `target_parameters`. Resolves #3108 --- tests/test_dpo_trainer.py | 35 ++++++++++++++++++++++++++++++++ tests/test_grpo_trainer.py | 40 +++++++++++++++++++++++++++++++++++++ tests/test_kto_trainer.py | 35 ++++++++++++++++++++++++++++++++ tests/test_rloo_trainer.py | 40 +++++++++++++++++++++++++++++++++++++ trl/trainer/callbacks.py | 37 ++++++++++++++++++++++++++++++---- trl/trainer/dpo_trainer.py | 24 +++++++++++++++------- trl/trainer/grpo_trainer.py | 24 +++++++++++++++------- trl/trainer/kto_trainer.py | 24 +++++++++++++++------- trl/trainer/rloo_trainer.py | 24 +++++++++++++++------- 9 files changed, 251 insertions(+), 32 deletions(-) diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index 403bb866a06..adefae46b54 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -667,6 +667,41 @@ def test_train_with_sync_ref_model(self): new_ref_param = trainer.ref_model.get_parameter(n) assert not torch.equal(previous_ref_params[n], new_ref_param), f"Ref Parameter {n} has not changed." + @require_peft + def test_train_with_sync_ref_model_and_peft(self): + # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the + # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. + dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") + + training_args = DPOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = DPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(), + ) + + assert trainer.ref_model is None # PEFT keeps the reference as an adapter, not a separate model + model = trainer.accelerator.unwrap_model(trainer.model) + assert "ref" in model.peft_config # the EMA target the callback syncs into + previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} + assert previous_ref_params # guard against the loop below vacuously passing + + trainer.train() + + assert trainer.state.log_history[-1]["train_loss"] is not None + + # Check that the reference adapter has tracked the policy + for n, param in previous_ref_params.items(): + new_param = model.get_parameter(n) + assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_model_dtype(self): dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index 9901ebc4553..aa2e6b039cf 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1483,6 +1483,46 @@ def test_train_with_sync_ref_model(self): new_ref_param = trainer.ref_model.get_parameter(n) assert not torch.equal(previous_ref_params[n], new_ref_param), f"Ref Parameter {n} has not changed." + @require_peft + def test_train_with_sync_ref_model_and_peft(self): + # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the + # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = GRPOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = GRPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(), + ) + + assert trainer.ref_model is None # PEFT keeps the reference as an adapter, not a separate model + model = trainer.accelerator.unwrap_model(trainer.model) + assert "ref" in model.peft_config # the EMA target the callback syncs into + previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} + assert previous_ref_params # guard against the loop below vacuously passing + + trainer.train() + + assert trainer.state.log_history[-1]["train_loss"] is not None + + # Check that the reference adapter has tracked the policy + for n, param in previous_ref_params.items(): + new_param = model.get_parameter(n) + assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_beta_non_zero(self): dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = GRPOConfig( diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 67c746876f2..3df199edc46 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -749,6 +749,41 @@ def test_train_with_sync_ref_model(self): new_ref_param = trainer.ref_model.get_parameter(n) assert not torch.equal(previous_ref_params[n], new_ref_param), f"Ref Parameter {n} has not changed." + @require_peft + def test_train_with_sync_ref_model_and_peft(self): + # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the + # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. + dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") + + training_args = KTOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = KTOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(), + ) + + assert trainer.ref_model is None # PEFT keeps the reference as an adapter, not a separate model + model = trainer.accelerator.unwrap_model(trainer.model) + assert "ref" in model.peft_config # the EMA target the callback syncs into + previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} + assert previous_ref_params # guard against the loop below vacuously passing + + trainer.train() + + assert trainer.state.log_history[-1]["train_loss"] is not None + + # Check that the reference adapter has tracked the policy + for n, param in previous_ref_params.items(): + new_param = model.get_parameter(n) + assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_model_dtype(self): dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index 27fbe347977..ffeb05ba78a 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -1112,6 +1112,46 @@ def test_train_with_sync_ref_model(self): new_ref_param = trainer.ref_model.get_parameter(n) assert not torch.equal(previous_ref_params[n], new_ref_param), f"Ref Parameter {n} has not changed." + @require_peft + def test_train_with_sync_ref_model_and_peft(self): + # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the + # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = RLOOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = RLOOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(), + ) + + assert trainer.ref_model is None # PEFT keeps the reference as an adapter, not a separate model + model = trainer.accelerator.unwrap_model(trainer.model) + assert "ref" in model.peft_config # the EMA target the callback syncs into + previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} + assert previous_ref_params # guard against the loop below vacuously passing + + trainer.train() + + assert trainer.state.log_history[-1]["train_loss"] is not None + + # Check that the reference adapter has tracked the policy + for n, param in previous_ref_params.items(): + new_param = model.get_parameter(n) + assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_beta_zero(self): dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = RLOOConfig( diff --git a/trl/trainer/callbacks.py b/trl/trainer/callbacks.py index f80e3208fde..7af8fde6a37 100644 --- a/trl/trainer/callbacks.py +++ b/trl/trainer/callbacks.py @@ -18,7 +18,7 @@ import torch from accelerate import Accelerator from accelerate.state import AcceleratorState -from accelerate.utils import gather_object, is_wandb_available +from accelerate.utils import gather_object, is_peft_model, is_wandb_available from transformers import ( GenerationConfig, PreTrainedModel, @@ -133,13 +133,42 @@ def sync_target_model(model, target_model, alpha): else: SyncRefModelCallback._sync_target_model(model, target_model, alpha) + @staticmethod + def _sync_ref_adapter(model, alpha): + # With PEFT the reference is not a separate module but a second adapter inside the policy model, so + # `_sync_target_model`'s parameter-wise zip of two modules does not apply. Pair each `"default"` parameter with + # its `"ref"` counterpart by name instead; this is the same mapping used to initialize the `"ref"` adapter. + for name, param in model.named_parameters(): + if ".default." in name: + ref_param = model.get_parameter(name.replace(".default.", ".ref.")) + ref_param.data.mul_(1.0 - alpha).add_(param.data, alpha=alpha) + + @staticmethod + def sync_ref_adapter(model, alpha): + deepspeed_plugin = AcceleratorState().deepspeed_plugin + if deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3: + import deepspeed + + with deepspeed.zero.GatheredParameters(list(model.parameters()), modifier_rank=0): + if deepspeed.comm.get_rank() == 0: + SyncRefModelCallback._sync_ref_adapter(model, alpha) + else: + SyncRefModelCallback._sync_ref_adapter(model, alpha) + def on_step_end(self, args, state, control, **kwargs): model: PreTrainedModel = kwargs["model"] - if self.ref_model is not None and state.global_step % args.ref_model_sync_steps == 0: - if self.accelerator: - model = self.accelerator.unwrap_model(model) + if state.global_step % args.ref_model_sync_steps != 0: + return + + if self.accelerator: + model = self.accelerator.unwrap_model(model) + + if self.ref_model is not None: self.sync_target_model(model, self.ref_model, args.ref_model_mixup_alpha) + elif is_peft_model(model) and "ref" in model.peft_config: + # PEFT keeps the reference as a frozen `"ref"` adapter rather than a standalone `ref_model`. + self.sync_ref_adapter(model, args.ref_model_mixup_alpha) class RichProgressCallback(TrainerCallback): diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index a43f3764163..c565d785be5 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -650,8 +650,19 @@ def __init__( ): get_peft_model_kwargs["autocast_adapter_dtype"] = False model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + # A freshly created adapter is zero-initialized, so disabling it recovers the base model, and that base + # model is the reference. That equivalence only holds while the reference stays fixed: with + # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to + # move. So in that case create the "ref" adapter here as well. + needs_ref_adapter = args.sync_ref_model and ref_model is None elif is_peft_model(model) and ref_model is None: + needs_ref_adapter = True + + else: + needs_ref_adapter = False + + if needs_ref_adapter: # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy # of the "default" adapter, so that we can use it as the reference model during DPO training. Before PEFT # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see @@ -962,14 +973,13 @@ def __init__( self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) if args.sync_ref_model: - if is_peft_model(self.model): + if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, DPOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." + "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " + "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " + "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0`, or use `sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index de126792d2f..b5cbc74a782 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -445,8 +445,19 @@ def __init__( ): get_peft_model_kwargs["autocast_adapter_dtype"] = False model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + # A freshly created adapter is zero-initialized, so disabling it recovers the base model, and that base + # model is the reference. That equivalence only holds while the reference stays fixed: with + # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to + # move. So in that case create the "ref" adapter here as well. + needs_ref_adapter = args.sync_ref_model and args.beta != 0.0 elif is_peft_model(model) and args.beta != 0.0: + needs_ref_adapter = True + + else: + needs_ref_adapter = False + + if needs_ref_adapter: # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy # of the "default" adapter, so that we can use it as the reference model during GRPO training. Before PEFT # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see @@ -1162,14 +1173,13 @@ def cast_outputs_to_original_dtype(module, args, output): "during training. Consequently, GRPOTrainer does not create a `ref_model` instance, and there is " "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." ) - if is_peft_model(model): + if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, GRPOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." + "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " + "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " + "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0`, or use `sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index 632cea2423d..9bb6af7aabe 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -698,8 +698,19 @@ def __init__( ): get_peft_model_kwargs["autocast_adapter_dtype"] = False model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + # A freshly created adapter is zero-initialized, so disabling it recovers the base model, and that base + # model is the reference. That equivalence only holds while the reference stays fixed: with + # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to + # move. So in that case create the "ref" adapter here as well. + needs_ref_adapter = args.sync_ref_model and ref_model is None elif is_peft_model(model) and ref_model is None: + needs_ref_adapter = True + + else: + needs_ref_adapter = False + + if needs_ref_adapter: # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy # of the "default" adapter, so that we can use it as the reference model during KTO training. Before PEFT # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see @@ -965,14 +976,13 @@ def __init__( self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) if args.sync_ref_model: - if is_peft_model(self.model): + if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, KTOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." + "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " + "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " + "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0`, or use `sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 3c9de79d254..c39ad2f0377 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -358,8 +358,19 @@ def __init__( ): get_peft_model_kwargs["autocast_adapter_dtype"] = False model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + # A freshly created adapter is zero-initialized, so disabling it recovers the base model, and that base + # model is the reference. That equivalence only holds while the reference stays fixed: with + # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to + # move. So in that case create the "ref" adapter here as well. + needs_ref_adapter = args.sync_ref_model elif is_peft_model(model): + needs_ref_adapter = True + + else: + needs_ref_adapter = False + + if needs_ref_adapter: # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy # of the "default" adapter, so that we can use it as the reference model during the training. Before PEFT # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see @@ -730,14 +741,13 @@ def __init__( "during training. Consequently, RLOOTrainer does not create a `ref_model` instance, and there is " "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." ) - if is_peft_model(model): + if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, RLOOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." + "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " + "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " + "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0`, or use `sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) From debdd5e5520e176f97dd7ea970038633c85cee9b Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Sat, 29 Aug 2026 23:16:12 -0700 Subject: [PATCH 2/6] fix(peft): sync every reference-adapter parameter and reject bias-bearing LoRA The "default" to "ref" pairing matched the adapter name as the substring ".default.", which requires a path component after it. `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names end in ".default" with nothing following. The predicate skipped them, and the reference copy of the token deltas stayed frozen while the policy moved. Matching "default" as a path component covers both placements: on a tiny Qwen2 with `trainable_token_indices=[0, 1]`, one of five adapter parameter pairs went unsynced before, none after. A LoRA config with `bias` other than "none" trains bias terms shared with the base model, and PEFT permits one such adapter per model: `add_adapter("ref", ...)` raises `LoraModel supports only 1 adapter with bias` for both "all" and "lora_only", while "none" succeeds. That call sat after the guard, so those configurations failed during trainer construction. The guard now skips the reference adapter for them, which routes into the existing `sync_ref_model` rejection rather than a PEFT stack trace. That rejection named `target_parameters` and `peft<0.20.0` as the only cause, which is wrong once bias reaches it, so the message now names both. All five call sites move together, as the trainers require. --- tests/test_dpo_trainer.py | 54 ++++++++++++++++++++++++++++++++ tests/test_grpo_trainer.py | 62 +++++++++++++++++++++++++++++++++++++ tests/test_kto_trainer.py | 52 +++++++++++++++++++++++++++++++ tests/test_rloo_trainer.py | 62 +++++++++++++++++++++++++++++++++++++ trl/trainer/callbacks.py | 9 ++++-- trl/trainer/dpo_trainer.py | 27 +++++++++++----- trl/trainer/grpo_trainer.py | 27 +++++++++++----- trl/trainer/kto_trainer.py | 27 +++++++++++----- trl/trainer/rloo_trainer.py | 27 +++++++++++----- 9 files changed, 317 insertions(+), 30 deletions(-) diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index adefae46b54..eadb54ce1b1 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -702,6 +702,60 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): + # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names + # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and + # leave the reference copy of the token deltas frozen at its initial value while the policy moves. + dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") + + training_args = DPOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = DPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(trainable_token_indices=[0, 1]), + ) + + model = trainer.accelerator.unwrap_model(trainer.model) + token_ref_params = { + n: param.clone() for n, param in model.named_parameters() if "trainable_tokens_delta" in n and "ref" in n + } + assert token_ref_params # the config must actually produce token deltas, or the loop below passes vacuously + + trainer.train() + + for n, param in token_ref_params.items(): + assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + + def test_train_with_sync_ref_model_and_peft_bias(self): + # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one + # such adapter per model. Creating a "ref" adapter would raise, so the trainer falls back to the base model as + # the reference instead of failing to construct. + dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") + + training_args = DPOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + # No "ref" adapter can be created, and the base model is a fixed reference that cannot track the policy, so + # asking for both is rejected rather than silently training against a frozen reference. + with pytest.raises(NotImplementedError, match="bias"): + DPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(bias="all"), + ) + def test_train_model_dtype(self): dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index aa2e6b039cf..adaa168a941 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1523,6 +1523,68 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): + # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names + # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and + # leave the reference copy of the token deltas frozen at its initial value while the policy moves. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = GRPOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = GRPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(trainable_token_indices=[0, 1]), + ) + + model = trainer.accelerator.unwrap_model(trainer.model) + token_ref_params = { + n: param.clone() for n, param in model.named_parameters() if "trainable_tokens_delta" in n and "ref" in n + } + assert token_ref_params # the config must actually produce token deltas, or the loop below passes vacuously + + trainer.train() + + for n, param in token_ref_params.items(): + assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + + def test_train_with_sync_ref_model_and_peft_bias(self): + # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one + # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot + # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = GRPOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(NotImplementedError, match="bias"): + GRPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(bias="all"), + ) + def test_train_beta_non_zero(self): dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = GRPOConfig( diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 3df199edc46..8a3a0ecf524 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -784,6 +784,58 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): + # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names + # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and + # leave the reference copy of the token deltas frozen at its initial value while the policy moves. + dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") + + training_args = KTOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = KTOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(trainable_token_indices=[0, 1]), + ) + + model = trainer.accelerator.unwrap_model(trainer.model) + token_ref_params = { + n: param.clone() for n, param in model.named_parameters() if "trainable_tokens_delta" in n and "ref" in n + } + assert token_ref_params # the config must actually produce token deltas, or the loop below passes vacuously + + trainer.train() + + for n, param in token_ref_params.items(): + assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + + def test_train_with_sync_ref_model_and_peft_bias(self): + # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one + # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot + # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") + + training_args = KTOConfig( + output_dir=self.tmp_dir, + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(NotImplementedError, match="bias"): + KTOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(bias="all"), + ) + def test_train_model_dtype(self): dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index ffeb05ba78a..29bb8824cf7 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -1152,6 +1152,68 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): + # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names + # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and + # leave the reference copy of the token deltas frozen at its initial value while the policy moves. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = RLOOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + trainer = RLOOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(trainable_token_indices=[0, 1]), + ) + + model = trainer.accelerator.unwrap_model(trainer.model) + token_ref_params = { + n: param.clone() for n, param in model.named_parameters() if "trainable_tokens_delta" in n and "ref" in n + } + assert token_ref_params # the config must actually produce token deltas, or the loop below passes vacuously + + trainer.train() + + for n, param in token_ref_params.items(): + assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + + def test_train_with_sync_ref_model_and_peft_bias(self): + # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one + # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot + # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + + training_args = RLOOConfig( + output_dir=self.tmp_dir, + beta=0.1, # ensure the reference is used so sync has an effect + per_device_train_batch_size=3, # reduce the batch size to reduce memory usage + num_generations=3, # reduce the number of generations to reduce memory usage + max_completion_length=8, # reduce the completion length to reduce memory usage + learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates + sync_ref_model=True, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(NotImplementedError, match="bias"): + RLOOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", + args=training_args, + train_dataset=dataset, + peft_config=LoraConfig(bias="all"), + ) + def test_train_beta_zero(self): dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = RLOOConfig( diff --git a/trl/trainer/callbacks.py b/trl/trainer/callbacks.py index 7af8fde6a37..c430598aa1c 100644 --- a/trl/trainer/callbacks.py +++ b/trl/trainer/callbacks.py @@ -138,9 +138,14 @@ def _sync_ref_adapter(model, alpha): # With PEFT the reference is not a separate module but a second adapter inside the policy model, so # `_sync_target_model`'s parameter-wise zip of two modules does not apply. Pair each `"default"` parameter with # its `"ref"` counterpart by name instead; this is the same mapping used to initialize the `"ref"` adapter. + # The adapter name is a path component, and it is not always followed by one: `trainable_token_indices` stores + # its deltas in a `ParameterDict` keyed by adapter, so those names END in `.default`. Splitting on `.` matches + # both placements, where a `".default." in name` substring test silently skips the terminal ones. for name, param in model.named_parameters(): - if ".default." in name: - ref_param = model.get_parameter(name.replace(".default.", ".ref.")) + parts = name.split(".") + if "default" in parts: + ref_name = ".".join("ref" if part == "default" else part for part in parts) + ref_param = model.get_parameter(ref_name) ref_param.data.mul_(1.0 - alpha).add_(param.data, alpha=alpha) @staticmethod diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index c565d785be5..ed589f4ce6c 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -670,7 +670,15 @@ def __init__( # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] - if ( + if isinstance(default_config, LoraConfig) and default_config.bias != "none": + logger.warning( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " + "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " + "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " + "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " + "adapter instead." + ) + elif ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") @@ -685,9 +693,12 @@ def __init__( ) else: model.add_adapter("ref", default_config) + # The adapter name is a path component that is not always followed by one: `trainable_token_indices` + # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") + parts = name.split(".") + if "default" in parts: + ref_name = ".".join("ref" if part == "default" else part for part in parts) ref_param = model.get_parameter(ref_name) ref_param.data.copy_(param.data) @@ -976,10 +987,12 @@ def __init__( if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " - "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " - "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " - "`peft>=0.20.0`, or use `sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " + "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " + "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " + "cases the reference log probs come from the base model with adapters disabled, which is fixed " + "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " + "`sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index b5cbc74a782..d97005f11ec 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -465,7 +465,15 @@ def __init__( # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] - if ( + if isinstance(default_config, LoraConfig) and default_config.bias != "none": + logger.warning( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " + "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " + "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " + "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " + "adapter instead." + ) + elif ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") @@ -480,9 +488,12 @@ def __init__( ) else: model.add_adapter("ref", default_config) + # The adapter name is a path component that is not always followed by one: `trainable_token_indices` + # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") + parts = name.split(".") + if "default" in parts: + ref_name = ".".join("ref" if part == "default" else part for part in parts) ref_param = model.get_parameter(ref_name) ref_param.data.copy_(param.data) @@ -1176,10 +1187,12 @@ def cast_outputs_to_original_dtype(module, args, output): if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " - "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " - "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " - "`peft>=0.20.0`, or use `sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " + "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " + "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " + "cases the reference log probs come from the base model with adapters disabled, which is fixed " + "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " + "`sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index 9bb6af7aabe..d8812eb464f 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -718,7 +718,15 @@ def __init__( # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] - if ( + if isinstance(default_config, LoraConfig) and default_config.bias != "none": + logger.warning( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " + "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " + "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " + "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " + "adapter instead." + ) + elif ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") @@ -733,9 +741,12 @@ def __init__( ) else: model.add_adapter("ref", default_config) + # The adapter name is a path component that is not always followed by one: `trainable_token_indices` + # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") + parts = name.split(".") + if "default" in parts: + ref_name = ".".join("ref" if part == "default" else part for part in parts) ref_param = model.get_parameter(ref_name) ref_param.data.copy_(param.data) @@ -979,10 +990,12 @@ def __init__( if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " - "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " - "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " - "`peft>=0.20.0`, or use `sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " + "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " + "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " + "cases the reference log probs come from the base model with adapters disabled, which is fixed " + "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " + "`sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index c39ad2f0377..93d5f9ed261 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -378,7 +378,15 @@ def __init__( # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] - if ( + if isinstance(default_config, LoraConfig) and default_config.bias != "none": + logger.warning( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " + "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " + "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " + "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " + "adapter instead." + ) + elif ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") @@ -393,9 +401,12 @@ def __init__( ) else: model.add_adapter("ref", default_config) + # The adapter name is a path component that is not always followed by one: `trainable_token_indices` + # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") + parts = name.split(".") + if "default" in parts: + ref_name = ".".join("ref" if part == "default" else part for part in parts) ref_param = model.get_parameter(ref_name) ref_param.data.copy_(param.data) @@ -744,10 +755,12 @@ def __init__( if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. This happens with `peft<0.20.0` when the LoRA " - "config uses `target_parameters` (peft#3340): the reference log probs then come from the base " - "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " - "`peft>=0.20.0`, or use `sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " + "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " + "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " + "cases the reference log probs come from the base model with adapters disabled, which is fixed " + "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " + "`sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) From be926d57f5a11839a0d5be78e5435ac5d3f08d5b Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Sun, 30 Aug 2026 03:30:32 -0700 Subject: [PATCH 3/6] fix(peft): gate the new tests on PEFT and pair only adapter-owned parameters The eight tests added in the previous commit build a `LoraConfig`, which is imported only when PEFT is installed, and they carried no `@require_peft` while every neighbouring PEFT test does. All eight failed the no-optional-dependency CI lane with `NameError: name 'LoraConfig' is not defined`. A local run cannot show this, because PEFT is installed there. Pairing "default" with "ref" on the path component alone was also wrong in both directions. PEFT reserves neither name, so a base model owning a parameter called `default` produced a reference name nothing provides: `base_model.model.default -> base_model.model.ref` raised `AttributeError` during construction. Taking the last such component instead was wrong the other way: `modules_to_save` wraps the saved module in a `ModuleDict` keyed by adapter, so a saved module with its own `default` child reads `...modules_to_save.default.default.weight`, whose adapter key is the first component. The last one belongs to the base module, so the parameter was skipped and its reference never moved, which is the defect this branch set out to fix. Adapter parameters are the ones PEFT keys by adapter name inside a `ModuleDict` or a `ParameterDict`, so the scan now walks the candidates from the end and takes the first whose container also holds a "ref" key. That covers the LoRA matrices, the `trainable_token_indices` deltas and `modules_to_save` alike, and it declines a base container that happens to hold a `default` key with no counterpart. Five callback tests cover the cases: a base module and parameter named `default`, a `modules_to_save` child of the same name, a base `ModuleDict` without a "ref" key, the terminal token delta, and the EMA equation itself. Each of the four ways to get the rule wrong fails one of them. --- tests/test_callbacks.py | 163 +++++++++++++++++++++++++++++++++++- tests/test_dpo_trainer.py | 2 + tests/test_grpo_trainer.py | 2 + tests/test_kto_trainer.py | 2 + tests/test_rloo_trainer.py | 2 + trl/trainer/callbacks.py | 18 ++-- trl/trainer/dpo_trainer.py | 17 ++-- trl/trainer/grpo_trainer.py | 17 ++-- trl/trainer/kto_trainer.py | 17 ++-- trl/trainer/rloo_trainer.py | 17 ++-- 10 files changed, 225 insertions(+), 32 deletions(-) diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index bdfb456fe85..7356338194c 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -16,12 +16,14 @@ import os from unittest.mock import call, patch +import torch from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, Trainer, TrainingArguments from trl import BEMACallback, LogCompletionsCallback +from trl.trainer.callbacks import SyncRefModelCallback -from .testing_utils import TrlTestCase, require_comet, require_wandb +from .testing_utils import TrlTestCase, require_comet, require_peft, require_wandb class TestLogCompletionsCallback(TrlTestCase): @@ -236,3 +238,162 @@ def test_no_ema(self): callbacks=[bema_callback], ) trainer.train() + + +@require_peft +class TestSyncRefModelCallbackAdapterPairing(TrlTestCase): + """`_sync_ref_adapter` pairs a "default" parameter with its "ref" counterpart by name.""" + + @staticmethod + def _peft_model_with_base_parameter_named_default(): + from peft import LoraConfig, get_peft_model + + # A base model owning both a parameter and a submodule literally called "default". PEFT reserves neither, so + # the adapter parameters here read `...default.proj.lora_A.default.weight`, with a base component and an + # adapter component of the same name, and the base parameter reads `...default` with no component after it. + class Inner(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(2, 2) + + def forward(self, x): + return self.proj(x) + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.default = Inner() + self.register_parameter("default_bias", torch.nn.Parameter(torch.ones(2))) + + def forward(self, x): + return self.default(x) + + config = LoraConfig(target_modules=["proj"]) + model = get_peft_model(Model(), config) + model.add_adapter("ref", config) + return model + + def test_base_module_and_parameter_named_default_are_not_paired(self): + model = self._peft_model_with_base_parameter_named_default() + adapter_names = [n for n, _ in model.named_parameters() if n.split(".").count("default") == 2] + assert adapter_names, "the fixture must produce a path with both a base and an adapter 'default' component" + + SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) # must not raise + + # Only the adapter component may be rewritten; the base module keeps its name. + for name in adapter_names: + parts = name.split(".") + index = len(parts) - 1 - parts[::-1].index("default") + model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + + def test_adapter_parameters_follow_the_ema_equation(self): + from peft import LoraConfig, get_peft_model + from transformers import AutoModelForCausalLM + + config = LoraConfig(r=4, target_modules=["q_proj"], trainable_token_indices=[0, 1]) + model = get_peft_model( + AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"), config + ) + model.add_adapter("ref", config) + + alpha = 0.25 + with torch.no_grad(): + for name, param in model.named_parameters(): + if "default" in name.split("."): + param.fill_(4.0) + elif "ref" in name.split("."): + param.fill_(8.0) + + SyncRefModelCallback._sync_ref_adapter(model, alpha=alpha) + + expected = (1.0 - alpha) * 8.0 + alpha * 4.0 + checked = 0 + for name, param in model.named_parameters(): + if "ref" in name.split("."): + assert torch.allclose(param, torch.full_like(param, expected)), name + checked += 1 + assert checked, "no reference parameter was examined" + + def test_terminal_adapter_parameter_is_paired(self): + from peft import LoraConfig, get_peft_model + from transformers import AutoModelForCausalLM + + # `trainable_token_indices` deltas live in a `ParameterDict` keyed by adapter, so the adapter name is the + # final path component with nothing after it. + config = LoraConfig(r=4, target_modules=["q_proj"], trainable_token_indices=[0, 1]) + model = get_peft_model( + AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"), config + ) + model.add_adapter("ref", config) + + delta = "base_model.model.model.embed_tokens.token_adapter.trainable_tokens_delta" + with torch.no_grad(): + model.get_parameter(f"{delta}.default").fill_(7.0) + model.get_parameter(f"{delta}.ref").fill_(0.0) + + SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) + + assert torch.equal( + model.get_parameter(f"{delta}.ref"), torch.full_like(model.get_parameter(f"{delta}.ref"), 7.0) + ) + + def test_modules_to_save_with_a_child_named_default_is_paired(self): + from peft import LoraConfig, get_peft_model + + # `modules_to_save` wraps the saved module in a `ModuleDict` keyed by adapter, so a saved module that itself + # contains a child called "default" produces two "default" components and the adapter key is the FIRST one. + class Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.default = torch.nn.Linear(4, 4) + + def forward(self, x): + return self.default(x) + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(4, 4) + self.block = Block() + + def forward(self, x): + return self.block(self.proj(x)) + + config = LoraConfig(target_modules=["proj"], modules_to_save=["block"]) + model = get_peft_model(Model(), config) + model.add_adapter("ref", config) + + policy = "base_model.model.block.modules_to_save.default.default.weight" + reference = "base_model.model.block.modules_to_save.ref.default.weight" + with torch.no_grad(): + model.get_parameter(policy).fill_(5.0) + model.get_parameter(reference).fill_(0.0) + + SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) + + assert torch.equal(model.get_parameter(reference), torch.full_like(model.get_parameter(reference), 5.0)), ( + "a modules_to_save parameter was skipped because the adapter key was not the last 'default' component" + ) + + def test_base_module_dict_without_a_ref_key_is_not_paired(self): + from peft import LoraConfig, get_peft_model + + # A `ModuleDict` in the base model may hold a "default" key that PEFT knows nothing about, so there is no + # "ref" counterpart to pair it with. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(4, 4) + self.heads = torch.nn.ModuleDict({"default": torch.nn.Linear(4, 4)}) + + def forward(self, x): + return self.heads["default"](self.proj(x)) + + config = LoraConfig(target_modules=["proj"]) + model = get_peft_model(Model(), config) + model.add_adapter("ref", config) + head = model.get_parameter("base_model.model.heads.default.weight").clone() + + SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) # must not raise + + assert torch.equal(model.get_parameter("base_model.model.heads.default.weight"), head) diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index eadb54ce1b1..b2edda175b2 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -702,6 +702,7 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and @@ -733,6 +734,7 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one # such adapter per model. Creating a "ref" adapter would raise, so the trainer falls back to the base model as diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index adaa168a941..7a17b414c32 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1523,6 +1523,7 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and @@ -1559,6 +1560,7 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 8a3a0ecf524..37951d0656a 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -784,6 +784,7 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and @@ -815,6 +816,7 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index 29bb8824cf7..236f7b40787 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -1152,6 +1152,7 @@ def test_train_with_sync_ref_model_and_peft(self): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names # END in ".default" with no component after it. Pairing "default" with "ref" by substring would skip them and @@ -1188,6 +1189,7 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot diff --git a/trl/trainer/callbacks.py b/trl/trainer/callbacks.py index c430598aa1c..4fec4b3ec4b 100644 --- a/trl/trainer/callbacks.py +++ b/trl/trainer/callbacks.py @@ -138,15 +138,19 @@ def _sync_ref_adapter(model, alpha): # With PEFT the reference is not a separate module but a second adapter inside the policy model, so # `_sync_target_model`'s parameter-wise zip of two modules does not apply. Pair each `"default"` parameter with # its `"ref"` counterpart by name instead; this is the same mapping used to initialize the `"ref"` adapter. - # The adapter name is a path component, and it is not always followed by one: `trainable_token_indices` stores - # its deltas in a `ParameterDict` keyed by adapter, so those names END in `.default`. Splitting on `.` matches - # both placements, where a `".default." in name` substring test silently skips the terminal ones. for name, param in model.named_parameters(): parts = name.split(".") - if "default" in parts: - ref_name = ".".join("ref" if part == "default" else part for part in parts) - ref_param = model.get_parameter(ref_name) - ref_param.data.mul_(1.0 - alpha).add_(param.data, alpha=alpha) + # PEFT keys adapter parameters by adapter name inside a `ModuleDict` (LoRA matrices, `modules_to_save`) or a + # `ParameterDict` (`trainable_token_indices` deltas), and the key is not always the last "default" component: + # `modules_to_save` wraps a module that may itself contain one. Scan the candidates from the end and take the + # first whose container also holds a "ref" key. Names where none qualifies belong to the base model, even when + # a module or parameter there happens to be called "default". + for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): + parent = model.get_submodule(".".join(parts[:index])) if index else model + if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + ref_param.data.mul_(1.0 - alpha).add_(param.data, alpha=alpha) + break @staticmethod def sync_ref_adapter(model, alpha): diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index ed589f4ce6c..ee13e87b8fc 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -693,14 +693,19 @@ def __init__( ) else: model.add_adapter("ref", default_config) - # The adapter name is a path component that is not always followed by one: `trainable_token_indices` - # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): parts = name.split(".") - if "default" in parts: - ref_name = ".".join("ref" if part == "default" else part for part in parts) - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + # PEFT keys adapter parameters by adapter name inside a `ModuleDict` (LoRA matrices, `modules_to_save`) or a + # `ParameterDict` (`trainable_token_indices` deltas), and the key is not always the last "default" component: + # `modules_to_save` wraps a module that may itself contain one. Scan the candidates from the end and take the + # first whose container also holds a "ref" key. Names where none qualifies belong to the base model, even when + # a module or parameter there happens to be called "default". + for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): + parent = model.get_submodule(".".join(parts[:index])) if index else model + if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + ref_param.data.copy_(param.data) + break # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index d97005f11ec..1c2fdddbbb8 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -488,14 +488,19 @@ def __init__( ) else: model.add_adapter("ref", default_config) - # The adapter name is a path component that is not always followed by one: `trainable_token_indices` - # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): parts = name.split(".") - if "default" in parts: - ref_name = ".".join("ref" if part == "default" else part for part in parts) - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + # PEFT keys adapter parameters by adapter name inside a `ModuleDict` (LoRA matrices, `modules_to_save`) or a + # `ParameterDict` (`trainable_token_indices` deltas), and the key is not always the last "default" component: + # `modules_to_save` wraps a module that may itself contain one. Scan the candidates from the end and take the + # first whose container also holds a "ref" key. Names where none qualifies belong to the base model, even when + # a module or parameter there happens to be called "default". + for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): + parent = model.get_submodule(".".join(parts[:index])) if index else model + if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + ref_param.data.copy_(param.data) + break # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index d8812eb464f..e3f05333991 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -741,14 +741,19 @@ def __init__( ) else: model.add_adapter("ref", default_config) - # The adapter name is a path component that is not always followed by one: `trainable_token_indices` - # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): parts = name.split(".") - if "default" in parts: - ref_name = ".".join("ref" if part == "default" else part for part in parts) - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + # PEFT keys adapter parameters by adapter name inside a `ModuleDict` (LoRA matrices, `modules_to_save`) or a + # `ParameterDict` (`trainable_token_indices` deltas), and the key is not always the last "default" component: + # `modules_to_save` wraps a module that may itself contain one. Scan the candidates from the end and take the + # first whose container also holds a "ref" key. Names where none qualifies belong to the base model, even when + # a module or parameter there happens to be called "default". + for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): + parent = model.get_submodule(".".join(parts[:index])) if index else model + if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + ref_param.data.copy_(param.data) + break # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 93d5f9ed261..3ef109b829c 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -401,14 +401,19 @@ def __init__( ) else: model.add_adapter("ref", default_config) - # The adapter name is a path component that is not always followed by one: `trainable_token_indices` - # keys its deltas by adapter, so those parameter names end in `.default`. Split on `.` to match both. for name, param in model.named_parameters(): parts = name.split(".") - if "default" in parts: - ref_name = ".".join("ref" if part == "default" else part for part in parts) - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + # PEFT keys adapter parameters by adapter name inside a `ModuleDict` (LoRA matrices, `modules_to_save`) or a + # `ParameterDict` (`trainable_token_indices` deltas), and the key is not always the last "default" component: + # `modules_to_save` wraps a module that may itself contain one. Scan the candidates from the end and take the + # first whose container also holds a "ref" key. Names where none qualifies belong to the base model, even when + # a module or parameter there happens to be called "default". + for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): + parent = model.get_submodule(".".join(parts[:index])) if index else model + if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) + ref_param.data.copy_(param.data) + break # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. From bd05e294855eda314802686549a73428768a3417 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Wed, 2 Sep 2026 19:47:30 -0700 Subject: [PATCH 4/6] fix(peft): reject bias-bearing LoRA outright, keep "ref" copies out of the vLLM sync, and correct the docs A LoRA config with `bias` other than `"none"` trains bias terms that live in the base model. The previous commit turned PEFT's crash on the second adapter into a warning saying the reference would come from the base model with adapters disabled. That reference is not fixed: with `bias="all"` six base biases train, and the adapter-disabled logits move when they move (PEFT warns about the same thing). All four trainers now raise a `ValueError` at construction, with or without `sync_ref_model`, and the `NotImplementedError` for a missing "ref" adapter no longer lists bias as a cause. DPO and KTO also name the case it had left out: a standalone `ref_model` passed alongside a PEFT policy, whose parameters do not pair with the adapter's. The vLLM weight sync skipped the "ref" adapter's LoRA matrices through the tuner prefix but not its `modules_to_save` copy: `lm_head.modules_to_save.ref.weight` was streamed to vLLM under that name. The four sync loops (server and colocate, in `VLLMGeneration` and OnlineDPO) now skip every `modules_to_save` copy except the "default" one. Reachable since the previous commit through GRPO and RLOO with `peft_config`, `sync_ref_model=True`, `use_vllm=True` and `modules_to_save`. RLOO created the "ref" adapter for `beta=0.0` and then raised that there is nothing to synchronize; its guard now carries GRPO's `beta != 0.0` clause. The DPO and KTO sync tests checked only that the adapter's parameters moved. They now also compute the reference log probs through the trainer's own path after training and require them to differ from the base model's, taken with adapters disabled at the same moment: a fresh "ref" adapter reproduces the base model exactly, so only a synced one can differ. A copy of the tree whose reference path ignores the "ref" adapter fails that assertion in both trainers. (A before/after comparison was tried first and rejected: enabling gradient checkpointing during `train()` shifts the reference forward by 1e-3 on identical weights, which a same-moment comparison does not see.) The callbacks fixture said it owned a base parameter named "default" but registered "default_bias"; it now owns one, and the test asserts the sync leaves it alone. The DPO and KTO docs still said `sync_ref_model=True` is unsupported with PEFT models; they now describe the "ref" adapter copy and the three cases that are rejected. --- docs/source/dpo_trainer.md | 2 +- docs/source/kto_trainer.md | 2 +- tests/test_callbacks.py | 11 +++++++-- tests/test_dpo_trainer.py | 20 +++++++++++----- tests/test_grpo_trainer.py | 8 +++---- tests/test_kto_trainer.py | 18 ++++++++++---- tests/test_rloo_trainer.py | 8 +++---- .../online_dpo/online_dpo_trainer.py | 12 ++++++---- trl/generation/vllm_generation.py | 12 ++++++---- trl/trainer/dpo_trainer.py | 24 +++++++++---------- trl/trainer/grpo_trainer.py | 22 ++++++++--------- trl/trainer/kto_trainer.py | 24 +++++++++---------- trl/trainer/rloo_trainer.py | 24 +++++++++---------- 13 files changed, 108 insertions(+), 79 deletions(-) diff --git a/docs/source/dpo_trainer.md b/docs/source/dpo_trainer.md index 854b1ce3187..2533b78cf3e 100644 --- a/docs/source/dpo_trainer.md +++ b/docs/source/dpo_trainer.md @@ -154,7 +154,7 @@ Some argument combinations are intentionally restricted in the current [`DPOTrai * only a single `loss_type` is supported, * `compute_metrics` is not supported, * `precompute_ref_log_probs=True` is not supported. -* `sync_ref_model=True` is not supported when training with PEFT models that do not keep a standalone `ref_model`. +* `sync_ref_model=True` with a PEFT model syncs a frozen `"ref"` adapter copy instead of a standalone `ref_model`. It is not supported with `peft<0.20.0` when the LoRA config uses `target_parameters`, or together with a standalone `ref_model`; a LoRA `bias` other than `"none"` is rejected with or without it, because the trained biases live in the base model and the reference cannot stay fixed. * `sync_ref_model=True` cannot be combined with `precompute_ref_log_probs=True`. * `precompute_ref_log_probs=True` is not supported with `IterableDataset` (train or eval). diff --git a/docs/source/kto_trainer.md b/docs/source/kto_trainer.md index b1200d5cb57..64b277062ef 100644 --- a/docs/source/kto_trainer.md +++ b/docs/source/kto_trainer.md @@ -157,7 +157,7 @@ Some argument combinations are intentionally restricted in the current [`KTOTrai * `compute_metrics` is not supported, * `precompute_ref_log_probs=True` is not supported, * PEFT models are not supported. -* `sync_ref_model=True` is not supported when training with PEFT models that do not keep a standalone `ref_model`. +* `sync_ref_model=True` with a PEFT model syncs a frozen `"ref"` adapter copy instead of a standalone `ref_model`. It is not supported with `peft<0.20.0` when the LoRA config uses `target_parameters`, or together with a standalone `ref_model`; a LoRA `bias` other than `"none"` is rejected with or without it, because the trained biases live in the base model and the reference cannot stay fixed. * `sync_ref_model=True` cannot be combined with `precompute_ref_log_probs=True`. * `precompute_ref_log_probs=True` is not supported with `IterableDataset` (train or eval) or with vision datasets. * Loss types that estimate the KL divergence term (all except `"apo_zero_unpaired"`) require `train_sampling_strategy="sequential"` and a per-device train batch size greater than 1. diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 7356338194c..9733fc216f7 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -250,11 +250,13 @@ def _peft_model_with_base_parameter_named_default(): # A base model owning both a parameter and a submodule literally called "default". PEFT reserves neither, so # the adapter parameters here read `...default.proj.lora_A.default.weight`, with a base component and an - # adapter component of the same name, and the base parameter reads `...default` with no component after it. + # adapter component of the same name, and the base parameter reads `...default.default` with no component + # after it: its parent is a plain module, not an adapter dict, so the predicate must leave it alone. class Inner(torch.nn.Module): def __init__(self): super().__init__() self.proj = torch.nn.Linear(2, 2) + self.register_parameter("default", torch.nn.Parameter(torch.ones(2))) def forward(self, x): return self.proj(x) @@ -275,11 +277,16 @@ def forward(self, x): def test_base_module_and_parameter_named_default_are_not_paired(self): model = self._peft_model_with_base_parameter_named_default() - adapter_names = [n for n, _ in model.named_parameters() if n.split(".").count("default") == 2] + adapter_names = [n for n, _ in model.named_parameters() if "lora_" in n and n.split(".").count("default") == 2] assert adapter_names, "the fixture must produce a path with both a base and an adapter 'default' component" + base_param = model.get_parameter("base_model.model.default.default") + base_before = base_param.clone() SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) # must not raise + # The base parameter ends in "default" too, but its parent is a plain module, so it is not an adapter slot. + assert torch.equal(base_param, base_before) + # Only the adapter component may be rewritten; the base module keeps its name. for name in adapter_names: parts = name.split(".") diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index b2edda175b2..9aceeba5fdf 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -692,6 +692,7 @@ def test_train_with_sync_ref_model_and_peft(self): assert "ref" in model.peft_config # the EMA target the callback syncs into previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} assert previous_ref_params # guard against the loop below vacuously passing + batch = next(iter(trainer.get_train_dataloader())) trainer.train() @@ -701,6 +702,15 @@ def test_train_with_sync_ref_model_and_peft(self): for n, param in previous_ref_params.items(): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + # The trainer's own reference path has to read the synced adapter: a fresh "ref" adapter is a zero-initialized + # copy of "default" and reproduces the base model exactly, so once the sync has moved it the reference log probs + # must differ from the base model's, taken with adapters disabled at the same moment. + ref_logps = [t for t in trainer.compute_ref_log_probs(model, batch) if t is not None] + with model.disable_adapter(): + base_logps = [t for t in trainer.compute_ref_log_probs(model, batch) if t is not None] + assert not all(torch.equal(r, b) for r, b in zip(ref_logps, base_logps, strict=True)), ( + "the reference log probs equal the base model's, so the reference path is not reading the synced adapter" + ) @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @@ -736,9 +746,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): - # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one - # such adapter per model. Creating a "ref" adapter would raise, so the trainer falls back to the base model as - # the reference instead of failing to construct. + # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter + # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be + # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") training_args = DPOConfig( @@ -748,9 +758,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) - # No "ref" adapter can be created, and the base model is a fixed reference that cannot track the policy, so - # asking for both is rejected rather than silently training against a frozen reference. - with pytest.raises(NotImplementedError, match="bias"): + with pytest.raises(ValueError, match="bias"): DPOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index 7a17b414c32..16e67e2d9d1 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1562,9 +1562,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): - # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one - # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot - # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter + # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be + # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = GRPOConfig( @@ -1578,7 +1578,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) - with pytest.raises(NotImplementedError, match="bias"): + with pytest.raises(ValueError, match="bias"): GRPOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 37951d0656a..12d2ea34e70 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -774,6 +774,7 @@ def test_train_with_sync_ref_model_and_peft(self): assert "ref" in model.peft_config # the EMA target the callback syncs into previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} assert previous_ref_params # guard against the loop below vacuously passing + batch = next(iter(trainer.get_train_dataloader())) trainer.train() @@ -783,6 +784,15 @@ def test_train_with_sync_ref_model_and_peft(self): for n, param in previous_ref_params.items(): new_param = model.get_parameter(n) assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + # The trainer's own reference path has to read the synced adapter: a fresh "ref" adapter is a zero-initialized + # copy of "default" and reproduces the base model exactly, so once the sync has moved it the reference log probs + # must differ from the base model's, taken with adapters disabled at the same moment. + ref_logps = [t for t in trainer.compute_ref_log_probs(model, batch) if t is not None] + with model.disable_adapter(): + base_logps = [t for t in trainer.compute_ref_log_probs(model, batch) if t is not None] + assert not all(torch.equal(r, b) for r, b in zip(ref_logps, base_logps, strict=True)), ( + "the reference log probs equal the base model's, so the reference path is not reading the synced adapter" + ) @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @@ -818,9 +828,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): - # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one - # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot - # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter + # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be + # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") training_args = KTOConfig( @@ -830,7 +840,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) - with pytest.raises(NotImplementedError, match="bias"): + with pytest.raises(ValueError, match="bias"): KTOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index 236f7b40787..78892a3f99c 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -1191,9 +1191,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): @require_peft def test_train_with_sync_ref_model_and_peft_bias(self): - # A LoRA config with `bias != "none"` trains bias terms shared with the base model, and PEFT permits only one - # such adapter per model. No "ref" adapter can be created, and the base model is a fixed reference that cannot - # track the policy, so asking for both is rejected rather than silently training against a frozen reference. + # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter + # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be + # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") training_args = RLOOConfig( @@ -1207,7 +1207,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) - with pytest.raises(NotImplementedError, match="bias"): + with pytest.raises(ValueError, match="bias"): RLOOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5", diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 134670a49e8..6fffa111f1e 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -756,8 +756,9 @@ def _sync_fsdp2_params_to_vllm(self, module: nn.Module): # Skip PEFT layers: they don’t exist in vLLM, and they are merged already. if is_peft_model(module) and module.prefix in name: continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: + # When module to save, remove its prefix and discard the original module, as well as the copies held by + # other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in name or (".modules_to_save." in name and ".modules_to_save.default." not in name): continue name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) @@ -825,8 +826,11 @@ def _move_model_to_vllm_inner(self): # Skip PEFT layers: they don’t exist in vLLM, and they are merged already. if self.model.prefix in name: continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: + # When module to save, remove its prefix and discard the original module, as well as the copies + # held by other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in name or ( + ".modules_to_save." in name and ".modules_to_save.default." not in name + ): continue name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) diff --git a/trl/generation/vllm_generation.py b/trl/generation/vllm_generation.py index 31d21929b50..5eca71bd30f 100644 --- a/trl/generation/vllm_generation.py +++ b/trl/generation/vllm_generation.py @@ -414,8 +414,9 @@ def _iter_fsdp2_params(self, module: nn.Module): # Skip PEFT layers: they don't exist in vLLM, and they are merged already. if is_peft_model(module) and module.prefix in name: continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: + # When module to save, remove its prefix and discard the original module, as well as the copies held by + # other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in name or (".modules_to_save." in name and ".modules_to_save.default." not in name): continue name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) @@ -459,8 +460,11 @@ def _iter_named_params(self): # Skip PEFT layers: they don't exist in vLLM, and they are merged already. if model.prefix in name: continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: + # When module to save, remove its prefix and discard the original module, as well as the copies + # held by other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in name or ( + ".modules_to_save." in name and ".modules_to_save.default." not in name + ): continue name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index 457a6d65c3b..32a32d225ac 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -662,12 +662,12 @@ def __init__( # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": - logger.warning( - f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " - "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " - "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " - "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " - "adapter instead." + raise ValueError( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that live in the base model " + "rather than in the adapter, so disabling the adapter does not recover a fixed reference: the " + "trained biases stay in it. PEFT also allows only one such adapter per model (`LoraModel supports " + "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " + "train against a copy of your adapter." ) elif ( isinstance(default_config, LoraConfig) @@ -983,12 +983,12 @@ def __init__( if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " - "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " - "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " - "cases the reference log probs come from the base model with adapters disabled, which is fixed " - "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " - "`sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped with `peft<0.20.0` when the " + "LoRA config uses `target_parameters` (peft#3340); the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. It is also skipped when " + "a standalone `ref_model` is passed alongside a PEFT policy, whose parameters do not pair with the " + "adapter's. Upgrade to `peft>=0.20.0`, drop `ref_model` to sync against a 'ref' adapter copy, or " + "use `sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 1c2fdddbbb8..68f625f1651 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -466,12 +466,12 @@ def __init__( # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": - logger.warning( - f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " - "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " - "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " - "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " - "adapter instead." + raise ValueError( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that live in the base model " + "rather than in the adapter, so disabling the adapter does not recover a fixed reference: the " + "trained biases stay in it. PEFT also allows only one such adapter per model (`LoraModel supports " + "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " + "train against a copy of your adapter." ) elif ( isinstance(default_config, LoraConfig) @@ -1192,12 +1192,10 @@ def cast_outputs_to_original_dtype(module, args, output): if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " - "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " - "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " - "cases the reference log probs come from the base model with adapters disabled, which is fixed " - "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " - "`sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped with `peft<0.20.0` when the " + "LoRA config uses `target_parameters` (peft#3340); the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0` or use `sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index e3f05333991..65be38ea281 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -719,12 +719,12 @@ def __init__( # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": - logger.warning( - f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " - "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " - "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " - "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " - "adapter instead." + raise ValueError( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that live in the base model " + "rather than in the adapter, so disabling the adapter does not recover a fixed reference: the " + "trained biases stay in it. PEFT also allows only one such adapter per model (`LoraModel supports " + "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " + "train against a copy of your adapter." ) elif ( isinstance(default_config, LoraConfig) @@ -995,12 +995,12 @@ def __init__( if is_peft_model(self.model) and "ref" not in self.model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " - "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " - "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " - "cases the reference log probs come from the base model with adapters disabled, which is fixed " - "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " - "`sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped with `peft<0.20.0` when the " + "LoRA config uses `target_parameters` (peft#3340); the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. It is also skipped when " + "a standalone `ref_model` is passed alongside a PEFT policy, whose parameters do not pair with the " + "adapter's. Upgrade to `peft>=0.20.0`, drop `ref_model` to sync against a 'ref' adapter copy, or " + "use `sync_ref_model=False`." ) if args.precompute_ref_log_probs: raise ValueError( diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 3ef109b829c..00acdd5d766 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -362,7 +362,7 @@ def __init__( # model is the reference. That equivalence only holds while the reference stays fixed: with # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to # move. So in that case create the "ref" adapter here as well. - needs_ref_adapter = args.sync_ref_model + needs_ref_adapter = args.sync_ref_model and args.beta != 0.0 elif is_peft_model(model): needs_ref_adapter = True @@ -379,12 +379,12 @@ def __init__( # parameters, which holds here since the "ref" adapter reuses the "default" config. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": - logger.warning( - f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that are shared with the " - "base model rather than owned by the adapter, and PEFT allows only one such adapter per model " - "(`LoraModel supports only 1 adapter with bias`). The reference log probs are therefore computed " - "from the base model (adapters disabled). Set `bias='none'` to train against a copy of your " - "adapter instead." + raise ValueError( + f"A LoRA config with `bias={default_config.bias!r}` trains bias terms that live in the base model " + "rather than in the adapter, so disabling the adapter does not recover a fixed reference: the " + "trained biases stay in it. PEFT also allows only one such adapter per model (`LoraModel supports " + "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " + "train against a copy of your adapter." ) elif ( isinstance(default_config, LoraConfig) @@ -760,12 +760,10 @@ def __init__( if is_peft_model(model) and "ref" not in model.peft_config: raise NotImplementedError( "You passed `sync_ref_model=True` while using a PEFT model whose reference adapter could not be " - "created, so there is nothing to synchronize. The adapter is skipped when the LoRA config sets " - "`bias` to anything other than `'none'`, because PEFT allows only one bias-bearing adapter per " - "model, and with `peft<0.20.0` when the config uses `target_parameters` (peft#3340). In both " - "cases the reference log probs come from the base model with adapters disabled, which is fixed " - "and cannot track the policy. Set `bias='none'`, upgrade to `peft>=0.20.0`, or use " - "`sync_ref_model=False`." + "created, so there is nothing to synchronize. The adapter is skipped with `peft<0.20.0` when the " + "LoRA config uses `target_parameters` (peft#3340); the reference log probs then come from the base " + "model with adapters disabled, which is fixed and cannot track the policy. Upgrade to " + "`peft>=0.20.0` or use `sync_ref_model=False`." ) self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) From 04fdedd757d744b81cbd2f8a58cd25064321502f Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Wed, 2 Sep 2026 21:42:36 -0700 Subject: [PATCH 5/6] fix(peft): reject bias-bearing LoRA on every PEFT reference path and keep trainable-token state out of the vLLM sync The bias check ran only when a "ref" adapter was about to be created, so with sync_ref_model off a LoRA config with bias="all" was still accepted although the docs and the error message said otherwise. The reason the message gives holds on both paths: the trained biases live in the base model, so disabling the adapter does not give a fixed reference whether or not the reference is synced. The check now runs whenever the trainer takes its reference from the PEFT model, on the peft_config path and on the pretrained-adapter path, in all four trainers. The bias tests cover sync on and off. With trainable_token_indices the vLLM sync pushed three names vLLM does not own: the wrapped embedding weight under its token_adapter prefix and the per-adapter deltas for "default" and "ref". Only the "ref" copy is new in this PR; the other two mean the feature never worked with vLLM. PEFT merges the deltas into the embedding before the sync, so the four loops now skip trainable_tokens_delta and strip the token_adapter prefix, and the merged weight lands under its base name. Measured on a tiny Qwen2 with a fake process group: three leaked names before, none after. --- tests/test_dpo_trainer.py | 5 +++-- tests/test_grpo_trainer.py | 5 +++-- tests/test_kto_trainer.py | 5 +++-- tests/test_rloo_trainer.py | 5 +++-- trl/trainer/dpo_trainer.py | 27 ++++++++++++++++++--------- trl/trainer/grpo_trainer.py | 27 ++++++++++++++++++--------- trl/trainer/kto_trainer.py | 27 ++++++++++++++++++--------- trl/trainer/rloo_trainer.py | 27 ++++++++++++++++++--------- 8 files changed, 84 insertions(+), 44 deletions(-) diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index 9aceeba5fdf..c62f533d557 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -744,8 +744,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @pytest.mark.parametrize("sync_ref_model", [True, False]) @require_peft - def test_train_with_sync_ref_model_and_peft_bias(self): + def test_train_with_sync_ref_model_and_peft_bias(self, sync_ref_model): # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. @@ -754,7 +755,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): training_args = DPOConfig( output_dir=self.tmp_dir, learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - sync_ref_model=True, + sync_ref_model=sync_ref_model, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index a6fe51300b0..9aeb17d4c62 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1629,8 +1629,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @pytest.mark.parametrize("sync_ref_model", [True, False]) @require_peft - def test_train_with_sync_ref_model_and_peft_bias(self): + def test_train_with_sync_ref_model_and_peft_bias(self, sync_ref_model): # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. @@ -1643,7 +1644,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): num_generations=3, # reduce the number of generations to reduce memory usage max_completion_length=8, # reduce the completion length to reduce memory usage learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - sync_ref_model=True, + sync_ref_model=sync_ref_model, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 12d2ea34e70..88f1a563a04 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -826,8 +826,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @pytest.mark.parametrize("sync_ref_model", [True, False]) @require_peft - def test_train_with_sync_ref_model_and_peft_bias(self): + def test_train_with_sync_ref_model_and_peft_bias(self, sync_ref_model): # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. @@ -836,7 +837,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): training_args = KTOConfig( output_dir=self.tmp_dir, learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - sync_ref_model=True, + sync_ref_model=sync_ref_model, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index 78892a3f99c..472e8ac828e 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -1189,8 +1189,9 @@ def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): for n, param in token_ref_params.items(): assert not torch.equal(param, model.get_parameter(n)), f"Ref token delta {n} has not changed." + @pytest.mark.parametrize("sync_ref_model", [True, False]) @require_peft - def test_train_with_sync_ref_model_and_peft_bias(self): + def test_train_with_sync_ref_model_and_peft_bias(self, sync_ref_model): # A LoRA config with `bias != "none"` trains bias terms that live in the base model, so disabling the adapter # does not give a fixed reference, and PEFT permits only one such adapter per model, so no "ref" copy can be # made either. The trainer rejects the config at construction, with or without `sync_ref_model`. @@ -1203,7 +1204,7 @@ def test_train_with_sync_ref_model_and_peft_bias(self): num_generations=3, # reduce the number of generations to reduce memory usage max_completion_length=8, # reduce the completion length to reduce memory usage learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - sync_ref_model=True, + sync_ref_model=sync_ref_model, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", ) diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index 32a32d225ac..8631e6443fc 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -645,21 +645,21 @@ def __init__( # model is the reference. That equivalence only holds while the reference stays fixed: with # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to # move. So in that case create the "ref" adapter here as well. - needs_ref_adapter = args.sync_ref_model and ref_model is None + uses_peft_reference = ref_model is None + needs_ref_adapter = args.sync_ref_model and uses_peft_reference elif is_peft_model(model) and ref_model is None: + uses_peft_reference = True needs_ref_adapter = True else: + uses_peft_reference = False needs_ref_adapter = False - if needs_ref_adapter: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during DPO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see - # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. + if uses_peft_reference: + # The reference is the base model with the adapter disabled, or a frozen copy of the adapter. Neither is + # fixed when the LoRA config trains bias terms: those live in the base model, so the reference moves with + # the policy whether or not it is synced. Refuse the configuration on both paths. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": raise ValueError( @@ -669,7 +669,16 @@ def __init__( "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " "train against a copy of your adapter." ) - elif ( + + if needs_ref_adapter: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during DPO training. Before PEFT + # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see + # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log + # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same + # parameters, which holds here since the "ref" adapter reuses the "default" config. + default_config = model.peft_config["default"] + if ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 68f625f1651..a2a9cd7041a 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -449,21 +449,21 @@ def __init__( # model is the reference. That equivalence only holds while the reference stays fixed: with # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to # move. So in that case create the "ref" adapter here as well. - needs_ref_adapter = args.sync_ref_model and args.beta != 0.0 + uses_peft_reference = args.beta != 0.0 + needs_ref_adapter = args.sync_ref_model and uses_peft_reference elif is_peft_model(model) and args.beta != 0.0: + uses_peft_reference = True needs_ref_adapter = True else: + uses_peft_reference = False needs_ref_adapter = False - if needs_ref_adapter: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during GRPO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see - # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. + if uses_peft_reference: + # The reference is the base model with the adapter disabled, or a frozen copy of the adapter. Neither is + # fixed when the LoRA config trains bias terms: those live in the base model, so the reference moves with + # the policy whether or not it is synced. Refuse the configuration on both paths. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": raise ValueError( @@ -473,7 +473,16 @@ def __init__( "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " "train against a copy of your adapter." ) - elif ( + + if needs_ref_adapter: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during GRPO training. Before PEFT + # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see + # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log + # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same + # parameters, which holds here since the "ref" adapter reuses the "default" config. + default_config = model.peft_config["default"] + if ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index 65be38ea281..86ed6e84a88 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -702,21 +702,21 @@ def __init__( # model is the reference. That equivalence only holds while the reference stays fixed: with # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to # move. So in that case create the "ref" adapter here as well. - needs_ref_adapter = args.sync_ref_model and ref_model is None + uses_peft_reference = ref_model is None + needs_ref_adapter = args.sync_ref_model and uses_peft_reference elif is_peft_model(model) and ref_model is None: + uses_peft_reference = True needs_ref_adapter = True else: + uses_peft_reference = False needs_ref_adapter = False - if needs_ref_adapter: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during KTO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see - # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. + if uses_peft_reference: + # The reference is the base model with the adapter disabled, or a frozen copy of the adapter. Neither is + # fixed when the LoRA config trains bias terms: those live in the base model, so the reference moves with + # the policy whether or not it is synced. Refuse the configuration on both paths. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": raise ValueError( @@ -726,7 +726,16 @@ def __init__( "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " "train against a copy of your adapter." ) - elif ( + + if needs_ref_adapter: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during KTO training. Before PEFT + # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see + # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log + # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same + # parameters, which holds here since the "ref" adapter reuses the "default" config. + default_config = model.peft_config["default"] + if ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 00acdd5d766..94a54001d8c 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -362,21 +362,21 @@ def __init__( # model is the reference. That equivalence only holds while the reference stays fixed: with # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to # move. So in that case create the "ref" adapter here as well. - needs_ref_adapter = args.sync_ref_model and args.beta != 0.0 + uses_peft_reference = args.beta != 0.0 + needs_ref_adapter = args.sync_ref_model and uses_peft_reference elif is_peft_model(model): + uses_peft_reference = True needs_ref_adapter = True else: + uses_peft_reference = False needs_ref_adapter = False - if needs_ref_adapter: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during the training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see - # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. + if uses_peft_reference: + # The reference is the base model with the adapter disabled, or a frozen copy of the adapter. Neither is + # fixed when the LoRA config trains bias terms: those live in the base model, so the reference moves with + # the policy whether or not it is synced. Refuse the configuration on both paths. default_config = model.peft_config["default"] if isinstance(default_config, LoraConfig) and default_config.bias != "none": raise ValueError( @@ -386,7 +386,16 @@ def __init__( "only 1 adapter with bias`), so no frozen 'ref' copy can be created either. Set `bias='none'` to " "train against a copy of your adapter." ) - elif ( + + if needs_ref_adapter: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during the training. Before PEFT + # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` (see + # peft#3340, fixed in peft#3350), so in that case we skip the "ref" adapter and compute the reference log + # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same + # parameters, which holds here since the "ref" adapter reuses the "default" config. + default_config = model.peft_config["default"] + if ( isinstance(default_config, LoraConfig) and default_config.target_parameters and Version(peft.__version__) < Version("0.20.0") From db949375946d1bb1a156a57952d2ff0cde72dbf8 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Thu, 3 Sep 2026 21:52:46 -0700 Subject: [PATCH 6/6] fix(peft-ref): harden the frozen reference adapter across PEFT versions and wrappers The reference-adapter path accessed LoraConfig.target_parameters unconditionally, which older PEFT releases do not define; the version window is now checked before the attribute. AdaLoRA allows a single trainable adapter, so sync_ref_model with an AdaLoraConfig raises a clear ValueError instead of failing inside PEFT. Trainable-token deltas are skipped and the token_adapter prefix stripped in every vLLM weight loop, including the FSDP1 loop and the three Online DPO loops that had not received the modules_to_save filter. The name-pairing predicate now requires the container's owner to be a PEFT tuner layer or ModulesToSaveWrapper, so a base-model ModuleDict keyed "default"/"ref" is left alone. DPO and KTO config text no longer says PEFT is unsupported. Tests: reference log-probs must differ from the adapter-disabled base (GRPO) and the logged KL must match the ref adapter within float noise and differ from the base (RLOO); target_parameters absent, AdaLoRA refused, merged embedding emitted, FSDP1 filter, base ModuleDict left alone, config help text. --- tests/test_callbacks.py | 18 ++++-- tests/test_dpo_trainer.py | 56 +++++++++++++++- tests/test_grpo_trainer.py | 24 +++++-- tests/test_kto_trainer.py | 6 ++ tests/test_rloo_trainer.py | 47 ++++++++++++-- tests/test_vllm_client_server.py | 64 ++++++++++++++++++- .../online_dpo/online_dpo_trainer.py | 25 +++++++- trl/generation/vllm_generation.py | 25 +++++++- trl/trainer/callbacks.py | 14 +++- trl/trainer/dpo_config.py | 8 +-- trl/trainer/dpo_trainer.py | 18 +++++- trl/trainer/grpo_trainer.py | 18 +++++- trl/trainer/kto_config.py | 8 +-- trl/trainer/kto_trainer.py | 18 +++++- trl/trainer/rloo_trainer.py | 19 +++++- 15 files changed, 324 insertions(+), 44 deletions(-) diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 9733fc216f7..b9f6b96ae0d 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -382,16 +382,16 @@ def forward(self, x): "a modules_to_save parameter was skipped because the adapter key was not the last 'default' component" ) - def test_base_module_dict_without_a_ref_key_is_not_paired(self): + def test_base_module_dict_with_adapter_named_keys_is_not_paired(self): from peft import LoraConfig, get_peft_model - # A `ModuleDict` in the base model may hold a "default" key that PEFT knows nothing about, so there is no - # "ref" counterpart to pair it with. + # A `ModuleDict` in the base model may use both adapter names for unrelated modules. Container membership does + # not make those modules PEFT adapter parameters. class Model(torch.nn.Module): def __init__(self): super().__init__() self.proj = torch.nn.Linear(4, 4) - self.heads = torch.nn.ModuleDict({"default": torch.nn.Linear(4, 4)}) + self.heads = torch.nn.ModuleDict({"default": torch.nn.Linear(4, 4), "ref": torch.nn.Linear(4, 4)}) def forward(self, x): return self.heads["default"](self.proj(x)) @@ -399,8 +399,12 @@ def forward(self, x): config = LoraConfig(target_modules=["proj"]) model = get_peft_model(Model(), config) model.add_adapter("ref", config) - head = model.get_parameter("base_model.model.heads.default.weight").clone() + default_head = model.get_parameter("base_model.model.heads.default.weight") + ref_head = model.get_parameter("base_model.model.heads.ref.weight") + with torch.no_grad(): + default_head.fill_(7.0) + ref_head.fill_(3.0) - SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) # must not raise + SyncRefModelCallback._sync_ref_adapter(model, alpha=1.0) - assert torch.equal(model.get_parameter("base_model.model.heads.default.weight"), head) + assert torch.equal(ref_head, torch.full_like(ref_head, 3.0)) diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index c62f533d557..474aec0b647 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import PropertyMock, patch + import pytest import torch import transformers @@ -37,7 +39,7 @@ if is_peft_available(): import peft - from peft import LoraConfig, PromptTuningConfig, get_peft_model + from peft import AdaLoraConfig, LoraConfig, PromptTuningConfig, get_peft_model from peft.utils import TaskType @@ -712,6 +714,58 @@ def test_train_with_sync_ref_model_and_peft(self): "the reference log probs equal the base model's, so the reference path is not reading the synced adapter" ) + @require_peft + def test_init_with_sync_ref_model_rejects_adalora(self): + dataset = Dataset.from_dict({"prompt": ["a"], "chosen": [" b"], "rejected": [" c"]}) + training_args = DPOConfig(output_dir=self.tmp_dir, sync_ref_model=True, report_to="none", use_cpu=True) + + # `AdaLoraModel` allows a single trainable adapter, so a frozen "ref" copy cannot be added; refuse up front rather + # than failing inside PEFT. + with pytest.raises(ValueError, match="AdaLoRA"): + DPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + peft_config=AdaLoraConfig(target_modules=["q_proj"], total_step=10), + ) + + @require_peft + def test_init_with_sync_ref_model_before_target_parameters(self): + model = get_peft_model( + AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"), + LoraConfig(target_modules=["q_proj"]), + ) + original_add_adapter = model.add_adapter + + def add_adapter(adapter_name, peft_config): + # The installed PEFT still reads this newer field internally, so expose it only after the trainer's + # compatibility check, where PEFT 0.10 did not have it. + with patch.object(LoraConfig, "target_parameters", None): + original_add_adapter(adapter_name, peft_config) + + model.add_adapter = add_adapter + dataset = Dataset.from_dict({"prompt": ["a"], "chosen": [" b"], "rejected": [" c"]}) + training_args = DPOConfig(output_dir=self.tmp_dir, sync_ref_model=True, report_to="none", use_cpu=True) + + with ( + patch.object(peft, "__version__", "0.10.0"), + patch.object( + LoraConfig, + "target_parameters", + new_callable=PropertyMock, + side_effect=AttributeError("target_parameters is unavailable"), + ), + ): + trainer = DPOTrainer(model=model, args=training_args, train_dataset=dataset) + + assert "ref" in trainer.model.peft_config + + def test_sync_ref_model_help_describes_peft_support(self): + help_text = DPOConfig.__dataclass_fields__["sync_ref_model"].metadata["help"] + + assert 'frozen `"ref"` adapter' in help_text + assert "not yet compatible with PEFT" not in help_text + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py index 9aeb17d4c62..36d1ff3bbb6 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1556,7 +1556,9 @@ def test_train_with_sync_ref_model(self): def test_train_with_sync_ref_model_and_peft(self): # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. - dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + dataset = Dataset.from_dict( + {"prompt": ["What is 2+2?", "Name a color.", "Say hello.", "What is 1+1?", "Name a pet.", "Say bye."]} + ) training_args = GRPOConfig( output_dir=self.tmp_dir, @@ -1568,6 +1570,7 @@ def test_train_with_sync_ref_model_and_peft(self): sync_ref_model=True, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", + use_cpu=True, ) trainer = GRPOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", @@ -1582,15 +1585,28 @@ def test_train_with_sync_ref_model_and_peft(self): assert "ref" in model.peft_config # the EMA target the callback syncs into previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} assert previous_ref_params # guard against the loop below vacuously passing + with torch.no_grad(): + for name, param in model.named_parameters(): + if ".lora_B.default." in name: + param.add_(0.1) + batch = next(iter(trainer.get_train_dataloader())) trainer.train() assert trainer.state.log_history[-1]["train_loss"] is not None # Check that the reference adapter has tracked the policy - for n, param in previous_ref_params.items(): - new_param = model.get_parameter(n) - assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + assert any(not torch.equal(param, model.get_parameter(name)) for name, param in previous_ref_params.items()) + # The generated reference log probs must come from the synced adapter, not from the base model with adapters + # disabled. Recompute the latter on the exact generated tokens to make the distinction observable. + generated = trainer._generate_and_score_completions(batch) + input_ids = torch.cat([generated["prompt_ids"], generated["completion_ids"]], dim=1) + attention_mask = torch.cat([generated["prompt_mask"], generated["completion_mask"]], dim=1) + with model.disable_adapter(), torch.no_grad(): + base_logps, _, _ = trainer._get_per_token_logps_and_entropies( + trainer.model, input_ids, attention_mask, generated["completion_ids"].size(1) + ) + assert not torch.equal(generated["ref_per_token_logps"], base_logps) @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): diff --git a/tests/test_kto_trainer.py b/tests/test_kto_trainer.py index 88f1a563a04..d2b159bb3d8 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -794,6 +794,12 @@ def test_train_with_sync_ref_model_and_peft(self): "the reference log probs equal the base model's, so the reference path is not reading the synced adapter" ) + def test_sync_ref_model_help_describes_peft_support(self): + help_text = KTOConfig.__dataclass_fields__["sync_ref_model"].metadata["help"] + + assert 'frozen `"ref"` adapter' in help_text + assert "not yet compatible with PEFT" not in help_text + @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): # `trainable_token_indices` stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names diff --git a/tests/test_rloo_trainer.py b/tests/test_rloo_trainer.py index 472e8ac828e..367f1c3aa5c 100644 --- a/tests/test_rloo_trainer.py +++ b/tests/test_rloo_trainer.py @@ -17,7 +17,7 @@ import pytest import torch import transformers -from datasets import DatasetDict, IterableDatasetDict, load_dataset +from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset from packaging.version import Version from transformers import ( AutoModelForCausalLM, @@ -1116,7 +1116,9 @@ def test_train_with_sync_ref_model(self): def test_train_with_sync_ref_model_and_peft(self): # With PEFT there is no standalone `ref_model`; the reference lives in a frozen "ref" adapter inside the # policy model. Check that `sync_ref_model=True` creates that adapter and that it tracks the policy. - dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + dataset = Dataset.from_dict( + {"prompt": ["What is 2+2?", "Name a color.", "Say hello.", "What is 1+1?", "Name a pet.", "Say bye."]} + ) training_args = RLOOConfig( output_dir=self.tmp_dir, @@ -1128,6 +1130,7 @@ def test_train_with_sync_ref_model_and_peft(self): sync_ref_model=True, ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens report_to="none", + use_cpu=True, ) trainer = RLOOTrainer( model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", @@ -1142,15 +1145,49 @@ def test_train_with_sync_ref_model_and_peft(self): assert "ref" in model.peft_config # the EMA target the callback syncs into previous_ref_params = {n: param.clone() for n, param in model.named_parameters() if ".ref." in n} assert previous_ref_params # guard against the loop below vacuously passing + with torch.no_grad(): + for name, param in model.named_parameters(): + if ".lora_B.default." in name: + param.add_(0.1) + batch = next(iter(trainer.get_train_dataloader())) trainer.train() assert trainer.state.log_history[-1]["train_loss"] is not None # Check that the reference adapter has tracked the policy - for n, param in previous_ref_params.items(): - new_param = model.get_parameter(n) - assert not torch.equal(param, new_param), f"Ref adapter parameter {n} has not changed." + assert any(not torch.equal(param, model.get_parameter(name)) for name, param in previous_ref_params.items()) + # The KL logged by the generation path must use the synced adapter. Recompute it from that adapter on the + # exact generated tokens and compare it with the value the trainer consumed. + generated = trainer._generate_and_score_completions(batch) + input_ids = torch.cat([generated["prompt_ids"], generated["completion_ids"]], dim=1) + attention_mask = torch.cat([generated["prompt_mask"], generated["completion_mask"]], dim=1) + previous_adapter = model.active_adapter + model.set_adapter("ref") + try: + with torch.no_grad(): + ref_logps, _, _ = trainer._get_per_token_logps_and_entropies( + trainer.model, input_ids, attention_mask, generated["completion_ids"].size(1) + ) + finally: + model.set_adapter(previous_adapter) + with model.disable_adapter(), torch.no_grad(): + base_logps, _, _ = trainer._get_per_token_logps_and_entropies( + trainer.model, input_ids, attention_mask, generated["completion_ids"].size(1) + ) + completion_mask = generated["completion_mask"] + + def sequence_kl(reference_logps): + return ( + (generated["old_logps"] - (reference_logps * completion_mask).sum(1)).sum() / completion_mask.sum() + ).item() + + logged_kl = trainer._metrics["train"]["kl"][-1] + # The trainer scores the reference in `batch_size` chunks while this recomputation runs the whole batch at once, + # so float32 reduction order alone moves a KL of order 1e-7 by a few 1e-8; the tolerance must sit above that + # noise and below the gap to the base model, which the second assertion pins. + assert logged_kl == pytest.approx(sequence_kl(ref_logps), abs=1e-6) + assert abs(logged_kl - sequence_kl(base_logps)) > 1e-6 @require_peft def test_train_with_sync_ref_model_and_peft_trainable_tokens(self): diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index d697692ca91..1546d7c9cf7 100644 --- a/tests/test_vllm_client_server.py +++ b/tests/test_vllm_client_server.py @@ -14,26 +14,35 @@ import os import subprocess +from contextlib import nullcontext from types import SimpleNamespace import pytest +import torch +from torch import nn from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer from transformers.testing_utils import torch_device +from transformers.utils import is_peft_available from trl.generation.vllm_client import VLLMClient, parse_logprobs -from trl.generation.vllm_generation import extract_logprobs +from trl.generation.vllm_generation import VLLMGeneration, extract_logprobs from trl.import_utils import is_vllm_available from .testing_utils import ( TrlTestCase, kill_process, require_3_accelerators, + require_peft, require_torch_multi_accelerator, require_vision, require_vllm, ) +if is_peft_available(): + from peft import LoraConfig, get_peft_model + + if is_vllm_available(): from vllm import LLM, SamplingParams @@ -127,6 +136,59 @@ def test_extract_logprobs_returns_none_token_ids_when_logprobs_missing(self): assert all_token_ids is None +@require_peft +class TestVLLMGenerationParameterIteration(TrlTestCase): + def test_trainable_token_wrapper_emits_merged_embedding(self): + class DistributedBackend: + is_fsdp = False + + @staticmethod + def gather_params(params): + return nullcontext() + + model = get_peft_model( + AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"), + LoraConfig(target_modules=["q_proj"], trainable_token_indices=[0, 1]), + ) + ref_config = LoraConfig(target_modules=["q_proj"], trainable_token_indices=[0, 1], inference_mode=True) + model.add_adapter("ref", ref_config) + expected = model.get_input_embeddings().token_adapter.get_merged_weights(["default"]).detach().clone() + + generation = object.__new__(VLLMGeneration) + generation.model = model + generation._dist = DistributedBackend() + parameters = dict(generation._iter_named_params()) + + torch.testing.assert_close(parameters["model.embed_tokens.weight"], expected) + assert not any("trainable_tokens_delta" in name for name in parameters) + + def test_fsdp1_filters_reference_modules_to_save(self, monkeypatch): + class FakeFSDP(nn.Module): + @staticmethod + def summon_full_params(module, recurse=False, writeback=False): + return nullcontext() + + class Model(FakeFSDP): + def __init__(self): + super().__init__() + self.base_model = nn.Module() + self.base_model.model = nn.Module() + self.base_model.model.lm_head = nn.Module() + self.base_model.model.lm_head.modules_to_save = nn.ModuleDict( + { + "default": nn.Linear(2, 2, bias=False), + "ref": nn.Linear(2, 2, bias=False), + } + ) + + monkeypatch.setattr("trl.generation.vllm_generation.FSDP", FakeFSDP) + generation = object.__new__(VLLMGeneration) + + parameters = dict(generation._iter_fsdp1_params(Model())) + + assert set(parameters) == {"lm_head.weight"} + + @pytest.mark.slow @require_torch_multi_accelerator @require_vllm diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 6fffa111f1e..6578c279065 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -760,7 +760,10 @@ def _sync_fsdp2_params_to_vllm(self, module: nn.Module): # other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy if "original_module" in name or (".modules_to_save." in name and ".modules_to_save.default." not in name): continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default.", "token_adapter."]) if param.is_cpu: param = param.to(self.accelerator.device) @@ -832,7 +835,12 @@ def _move_model_to_vllm_inner(self): ".modules_to_save." in name and ".modules_to_save.default." not in name ): continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in name: + continue + name = self._fix_param_name_to_vllm( + name, extra_prefixes=["modules_to_save.default.", "token_adapter."] + ) if self.vllm_mode == "server" and self.accelerator.is_main_process: self.vllm_client.update_named_param(name, param.data) @@ -877,6 +885,19 @@ def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visite for param_name, param in module.named_parameters(): full_name = f"{prefix}.{param_name}" if prefix else param_name full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + full_name = full_name.removeprefix("base_model.model.").replace(".base_layer", "") + # When module to save, remove its prefix and discard the original module, as well as the copies + # held by other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in full_name or ( + ".modules_to_save." in full_name and ".modules_to_save.default." not in full_name + ): + continue + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in full_name: + continue + full_name = self._fix_param_name_to_vllm( + full_name, extra_prefixes=["modules_to_save.default.", "token_adapter."] + ) if full_name in visited: continue # skip FSDP subtrees already traversed diff --git a/trl/generation/vllm_generation.py b/trl/generation/vllm_generation.py index 5eca71bd30f..4511b571ea5 100644 --- a/trl/generation/vllm_generation.py +++ b/trl/generation/vllm_generation.py @@ -398,6 +398,19 @@ def _iter_fsdp1_params(self, module: nn.Module, prefix: str = "", visited: set[s for param_name, param in module.named_parameters(): full_name = f"{prefix}.{param_name}" if prefix else param_name full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + full_name = full_name.removeprefix("base_model.model.").replace(".base_layer", "") + # When module to save, remove its prefix and discard the original module, as well as the copies + # held by other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy + if "original_module" in full_name or ( + ".modules_to_save." in full_name and ".modules_to_save.default." not in full_name + ): + continue + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in full_name: + continue + full_name = self._fix_param_name_to_vllm( + full_name, extra_prefixes=["modules_to_save.default.", "token_adapter."] + ) if full_name in visited: continue # skip FSDP subtrees already traversed @@ -418,7 +431,10 @@ def _iter_fsdp2_params(self, module: nn.Module): # other adapters (such as the frozen "ref" one); vLLM sees only the "default" copy if "original_module" in name or (".modules_to_save." in name and ".modules_to_save.default." not in name): continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default.", "token_adapter."]) if param.is_cpu: param = param.to(self.accelerator.device) @@ -466,7 +482,12 @@ def _iter_named_params(self): ".modules_to_save." in name and ".modules_to_save.default." not in name ): continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + # Trainable token deltas are already merged into the embedding weight and do not exist in vLLM. + if ".trainable_tokens_delta." in name: + continue + name = self._fix_param_name_to_vllm( + name, extra_prefixes=["modules_to_save.default.", "token_adapter."] + ) yield name, param.data # Unmerge adapters while parameters are still gathered diff --git a/trl/trainer/callbacks.py b/trl/trainer/callbacks.py index 4fec4b3ec4b..e87c81f08e1 100644 --- a/trl/trainer/callbacks.py +++ b/trl/trainer/callbacks.py @@ -30,7 +30,7 @@ TrainingArguments, ) from transformers.trainer_utils import has_length -from transformers.utils import is_rich_available +from transformers.utils import is_peft_available, is_rich_available from ..data_utils import maybe_apply_chat_template from ..import_utils import is_weave_available @@ -51,6 +51,11 @@ import wandb +if is_peft_available(): + from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils.other import ModulesToSaveWrapper + + if is_weave_available(): import weave from weave import EvaluationLogger @@ -147,7 +152,12 @@ def _sync_ref_adapter(model, alpha): # a module or parameter there happens to be called "default". for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): parent = model.get_submodule(".".join(parts[:index])) if index else model - if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + owner = model.get_submodule(".".join(parts[: index - 1])) if index > 1 else model + if ( + isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) + and isinstance(owner, (BaseTunerLayer, ModulesToSaveWrapper)) + and "ref" in parent + ): ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) ref_param.data.mul_(1.0 - alpha).add_(param.data, alpha=alpha) break diff --git a/trl/trainer/dpo_config.py b/trl/trainer/dpo_config.py index 864de6634d8..bb8aaaf9eb1 100644 --- a/trl/trainer/dpo_config.py +++ b/trl/trainer/dpo_config.py @@ -118,8 +118,8 @@ class DPOConfig(_BaseConfig): sync_ref_model (`bool`, *optional*, defaults to `False`): Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible with - PEFT or `precompute_ref_log_probs=True`. + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. With PEFT, this synchronizes a frozen `"ref"` + adapter. It is not compatible with `precompute_ref_log_probs=True`. ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): α parameter from the TR-DPO paper, which controls the mix between the current policy and the previous reference policy during updates. The reference policy is updated according to the equation: `π_ref = α * @@ -319,8 +319,8 @@ class DPOConfig(_BaseConfig): metadata={ "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " "steps, using the `ref_model_mixup_alpha` parameter. This synchronization originates from the " - "[TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible " - "with PEFT or `precompute_ref_log_probs=True`." + '[TR-DPO](https://huggingface.co/papers/2404.09656) paper. With PEFT, this synchronizes a frozen `"ref"` ' + "adapter. It is not compatible with `precompute_ref_log_probs=True`." }, ) ref_model_mixup_alpha: float = field( diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index 8631e6443fc..bb9e7bf2eef 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -73,8 +73,9 @@ if is_peft_available(): import peft - from peft import LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model + from peft import AdaLoraConfig, LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils.other import ModulesToSaveWrapper logger = get_logger(__name__) @@ -680,8 +681,8 @@ def __init__( default_config = model.peft_config["default"] if ( isinstance(default_config, LoraConfig) + and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") and default_config.target_parameters - and Version(peft.__version__) < Version("0.20.0") ): logger.warning( "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " @@ -691,6 +692,12 @@ def __init__( "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." ) + elif isinstance(default_config, AdaLoraConfig): + raise ValueError( + "`sync_ref_model=True` is not supported with an AdaLoRA adapter: `AdaLoraModel` allows a single " + "trainable adapter, so no frozen 'ref' copy can be added. Disable `sync_ref_model` to train against " + "the base model with the adapter disabled." + ) else: model.add_adapter("ref", default_config) for name, param in model.named_parameters(): @@ -702,7 +709,12 @@ def __init__( # a module or parameter there happens to be called "default". for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): parent = model.get_submodule(".".join(parts[:index])) if index else model - if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + owner = model.get_submodule(".".join(parts[: index - 1])) if index > 1 else model + if ( + isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) + and isinstance(owner, (BaseTunerLayer, ModulesToSaveWrapper)) + and "ref" in parent + ): ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) ref_param.data.copy_(param.data) break diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index a2a9cd7041a..49ee937b65d 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -107,8 +107,9 @@ if is_peft_available(): import peft - from peft import LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model + from peft import AdaLoraConfig, LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils.other import ModulesToSaveWrapper if is_trackio_available(): @@ -484,8 +485,8 @@ def __init__( default_config = model.peft_config["default"] if ( isinstance(default_config, LoraConfig) + and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") and default_config.target_parameters - and Version(peft.__version__) < Version("0.20.0") ): logger.warning( "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " @@ -495,6 +496,12 @@ def __init__( "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." ) + elif isinstance(default_config, AdaLoraConfig): + raise ValueError( + "`sync_ref_model=True` is not supported with an AdaLoRA adapter: `AdaLoraModel` allows a single " + "trainable adapter, so no frozen 'ref' copy can be added. Disable `sync_ref_model` to train against " + "the base model with the adapter disabled." + ) else: model.add_adapter("ref", default_config) for name, param in model.named_parameters(): @@ -506,7 +513,12 @@ def __init__( # a module or parameter there happens to be called "default". for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): parent = model.get_submodule(".".join(parts[:index])) if index else model - if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + owner = model.get_submodule(".".join(parts[: index - 1])) if index > 1 else model + if ( + isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) + and isinstance(owner, (BaseTunerLayer, ModulesToSaveWrapper)) + and "ref" in parent + ): ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) ref_param.data.copy_(param.data) break diff --git a/trl/trainer/kto_config.py b/trl/trainer/kto_config.py index a91ef9d1bd4..0c273f07abd 100644 --- a/trl/trainer/kto_config.py +++ b/trl/trainer/kto_config.py @@ -86,8 +86,8 @@ class KTOConfig(_BaseConfig): sync_ref_model (`bool`, *optional*, defaults to `False`): Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible with - PEFT or `precompute_ref_log_probs=True`. + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. With PEFT, this synchronizes a frozen `"ref"` + adapter. It is not compatible with `precompute_ref_log_probs=True`. ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): α parameter from the TR-DPO paper, which controls the mix between the current policy and the previous reference policy during updates. The reference policy is updated according to the equation: `π_ref = α * @@ -224,8 +224,8 @@ class KTOConfig(_BaseConfig): metadata={ "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " "steps, using the `ref_model_mixup_alpha` parameter. This synchronization originates from the " - "[TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible " - "with PEFT or `precompute_ref_log_probs=True`." + '[TR-DPO](https://huggingface.co/papers/2404.09656) paper. With PEFT, this synchronizes a frozen `"ref"` ' + "adapter. It is not compatible with `precompute_ref_log_probs=True`." }, ) ref_model_mixup_alpha: float = field( diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py index 86ed6e84a88..b2192aa9a61 100644 --- a/trl/trainer/kto_trainer.py +++ b/trl/trainer/kto_trainer.py @@ -79,8 +79,9 @@ if is_peft_available(): import peft - from peft import LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model + from peft import AdaLoraConfig, LoraConfig, PeftConfig, PeftModel, PromptLearningConfig, get_peft_model from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils.other import ModulesToSaveWrapper logger = get_logger(__name__) @@ -737,8 +738,8 @@ def __init__( default_config = model.peft_config["default"] if ( isinstance(default_config, LoraConfig) + and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") and default_config.target_parameters - and Version(peft.__version__) < Version("0.20.0") ): logger.warning( "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " @@ -748,6 +749,12 @@ def __init__( "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." ) + elif isinstance(default_config, AdaLoraConfig): + raise ValueError( + "`sync_ref_model=True` is not supported with an AdaLoRA adapter: `AdaLoraModel` allows a single " + "trainable adapter, so no frozen 'ref' copy can be added. Disable `sync_ref_model` to train against " + "the base model with the adapter disabled." + ) else: model.add_adapter("ref", default_config) for name, param in model.named_parameters(): @@ -759,7 +766,12 @@ def __init__( # a module or parameter there happens to be called "default". for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): parent = model.get_submodule(".".join(parts[:index])) if index else model - if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + owner = model.get_submodule(".".join(parts[: index - 1])) if index > 1 else model + if ( + isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) + and isinstance(owner, (BaseTunerLayer, ModulesToSaveWrapper)) + and "ref" in parent + ): ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) ref_param.data.copy_(param.data) break diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py index 94a54001d8c..f5c336da7b3 100644 --- a/trl/trainer/rloo_trainer.py +++ b/trl/trainer/rloo_trainer.py @@ -86,7 +86,9 @@ if is_peft_available(): import peft - from peft import LoraConfig, PeftConfig, PeftModel, get_peft_model + from peft import AdaLoraConfig, LoraConfig, PeftConfig, PeftModel, get_peft_model + from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils.other import ModulesToSaveWrapper if is_trackio_available(): @@ -397,8 +399,8 @@ def __init__( default_config = model.peft_config["default"] if ( isinstance(default_config, LoraConfig) + and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") and default_config.target_parameters - and Version(peft.__version__) < Version("0.20.0") ): logger.warning( "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " @@ -408,6 +410,12 @@ def __init__( "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." ) + elif isinstance(default_config, AdaLoraConfig): + raise ValueError( + "`sync_ref_model=True` is not supported with an AdaLoRA adapter: `AdaLoraModel` allows a single " + "trainable adapter, so no frozen 'ref' copy can be added. Disable `sync_ref_model` to train against " + "the base model with the adapter disabled." + ) else: model.add_adapter("ref", default_config) for name, param in model.named_parameters(): @@ -419,7 +427,12 @@ def __init__( # a module or parameter there happens to be called "default". for index in (i for i in reversed(range(len(parts))) if parts[i] == "default"): parent = model.get_submodule(".".join(parts[:index])) if index else model - if isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) and "ref" in parent: + owner = model.get_submodule(".".join(parts[: index - 1])) if index > 1 else model + if ( + isinstance(parent, (torch.nn.ModuleDict, torch.nn.ParameterDict)) + and isinstance(owner, (BaseTunerLayer, ModulesToSaveWrapper)) + and "ref" in parent + ): ref_param = model.get_parameter(".".join(parts[:index] + ["ref"] + parts[index + 1 :])) ref_param.data.copy_(param.data) break