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 bdfb456fe85..b9f6b96ae0d 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,173 @@ 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.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) + + 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 "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(".") + 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_with_adapter_named_keys_is_not_paired(self): + from peft import LoraConfig, get_peft_model + + # 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), "ref": 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) + 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) + + 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 403bb866a06..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 @@ -667,6 +669,158 @@ 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 + 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." + # 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_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 + # 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." + + @pytest.mark.parametrize("sync_ref_model", [True, False]) + @require_peft + 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`. + 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=sync_ref_model, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(ValueError, 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 0bdc74f4141..36d1ff3bbb6 100644 --- a/tests/test_grpo_trainer.py +++ b/tests/test_grpo_trainer.py @@ -1552,6 +1552,127 @@ 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 = 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, + 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", + use_cpu=True, + ) + 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 + 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 + 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): + # `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." + + @pytest.mark.parametrize("sync_ref_model", [True, False]) + @require_peft + 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`. + 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=sync_ref_model, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(ValueError, 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 67c746876f2..d2b159bb3d8 100644 --- a/tests/test_kto_trainer.py +++ b/tests/test_kto_trainer.py @@ -749,6 +749,112 @@ 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 + 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." + # 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" + ) + + 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 + # 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." + + @pytest.mark.parametrize("sync_ref_model", [True, False]) + @require_peft + 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`. + 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=sync_ref_model, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(ValueError, 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 27fbe347977..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, @@ -1112,6 +1112,148 @@ 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 = 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, + 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", + use_cpu=True, + ) + 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 + 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 + 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): + # `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." + + @pytest.mark.parametrize("sync_ref_model", [True, False]) + @require_peft + 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`. + 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=sync_ref_model, + ref_model_sync_steps=2, # reduce sync steps to ensure a sync happens + report_to="none", + ) + with pytest.raises(ValueError, 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/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 134670a49e8..6578c279065 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -756,10 +756,14 @@ 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."]) + # 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) @@ -825,10 +829,18 @@ 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 + # 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."]) + 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) @@ -873,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 31d21929b50..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 @@ -414,10 +427,14 @@ 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 + # 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."]) + 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) @@ -459,10 +476,18 @@ 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 + # 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."]) + 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 f80e3208fde..e87c81f08e1 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, @@ -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 @@ -133,13 +138,56 @@ 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(): + parts = name.split(".") + # 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 + 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 + + @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_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 51f2274f6c7..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__) @@ -641,8 +642,36 @@ 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. + 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 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( + 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." + ) + + 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 @@ -652,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` " @@ -663,13 +692,32 @@ 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(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + parts = name.split(".") + # 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 + 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 # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. @@ -953,14 +1001,15 @@ 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. 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 de126792d2f..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(): @@ -445,8 +446,36 @@ 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. + 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 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( + 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." + ) + + 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 @@ -456,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` " @@ -467,13 +496,32 @@ 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(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + parts = name.split(".") + # 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 + 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 # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. @@ -1162,14 +1210,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. 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_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 632cea2423d..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__) @@ -698,8 +699,36 @@ 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. + 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 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( + 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." + ) + + 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 @@ -709,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` " @@ -720,13 +749,32 @@ 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(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + parts = name.split(".") + # 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 + 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 # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. @@ -965,14 +1013,15 @@ 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. 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 3c9de79d254..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(): @@ -358,8 +360,36 @@ 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. + 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 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( + 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." + ) + + 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 @@ -369,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` " @@ -380,13 +410,32 @@ 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(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) + parts = name.split(".") + # 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 + 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 # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. @@ -730,14 +779,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. 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))