Fix BEMACallback parameter alignment when some parameters are frozen

tests/test_callbacks.py

@@ -16,6 +16,7 @@
 import os
 from unittest.mock import call, patch
 
+import torch
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, Trainer, TrainingArguments
 
@@ -236,3 +237,35 @@ def test_no_ema(self):
             callbacks=[bema_callback],
         )
         trainer.train()
+
+    def test_frozen_parameters(self):
+        """Test that BEMACallback writes the BEMA weights to the right parameters when some are frozen."""
+        self.model.model.embed_tokens.weight.requires_grad_(False)
+        training_args = TrainingArguments(output_dir=self.tmp_dir, report_to="none")
+        bema_callback = BEMACallback(update_freq=1)
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=self.dataset["train"],
+            processing_class=self.tokenizer,
+            callbacks=[bema_callback],
+        )
+        trainer.train()
+
+        # Total 9 steps, updated at every step: the running model must hold the BEMA weights of the last update for
+        # each trainable parameter, matched by name, and the frozen parameter must be left untouched.
+        running_params = dict(bema_callback.running_model.named_parameters())
+        alpha = bema_callback._bema_alpha(9)
+        for name, thetat, theta0, ema in zip(
+            bema_callback.param_names,
+            bema_callback.thetat_params,
+            bema_callback.theta0_params,
+            bema_callback.ema_params,
+            strict=True,
+        ):
+            torch.testing.assert_close(
+                running_params[name], ema + alpha * (thetat.detach() - theta0), check_dtype=False
+            )
+        torch.testing.assert_close(
+            running_params["model.embed_tokens.weight"], self.model.model.embed_tokens.weight, check_dtype=False
+        )

trl/trainer/callbacks.py

@@ -665,6 +665,7 @@ def __init__(
         self.theta0_params = []  # θ₀ buffers (on self.device)
         self.ema_params = []  # EMA buffers (on self.device)
         self.running_model = None  # a copy of the model to run BEMA on
+        self.running_params = []  # references to running model params, matched by name with thetat_params
 
     @staticmethod
     def _unwrap_model(model):
@@ -698,12 +699,15 @@ def on_train_begin(
         self.running_model = type(model)(model.config).to(self.device)
         self.running_model.load_state_dict(model.state_dict())
 
-        # Cache trainable parameters once in a fixed order
+        # Cache trainable parameters once in a fixed order, along with their counterparts in the running model. Frozen
+        # parameters are skipped, so the two must be matched by name rather than by position.
+        running_params = dict(self.running_model.named_parameters())
         for name, param in model.named_parameters():
             if not param.requires_grad:
                 continue
             self.param_names.append(name)
             self.thetat_params.append(param)
+            self.running_params.append(running_params[name])
 
             # Clone θ₀ and EMA on the same device as model
             theta0 = param.detach().clone().to(self.device)
@@ -725,7 +729,7 @@ def _update_bema_weights(self, step: int):
 
         # Compute EMA + BEMA in-place and write directly to running_model
         for thetat, theta0, ema, run_param in zip(
-            self.thetat_params, self.theta0_params, self.ema_params, self.running_model.parameters(), strict=True
+            self.thetat_params, self.theta0_params, self.ema_params, self.running_params, strict=True
         ):
             thetat = thetat.detach().to(self.device)
             ema.mul_(1 - beta).add_(thetat, alpha=beta)  # EMA update: ema = (1 - beta) * ema + beta * θₜ