Support sync_ref_model with PEFT by syncing the reference adapter

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).
 

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.

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))

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")