Fix BEMACallback parameter alignment when some parameters are frozen

[yurekami]

What does this PR do?

BEMACallback caches only the trainable parameters of the trained model in on_train_begin (it skips requires_grad=False), but _bema_update writes the BEMA weights by zipping that list against self.running_model.parameters(), which yields every parameter. As soon as any parameter is frozen the two sequences are positionally misaligned: zip(..., strict=True) raises on the length mismatch, and without strict each BEMA tensor would be written into the wrong parameter of the running model.

Repro: freeze the embedding matrix (model.model.embed_tokens.weight.requires_grad_(False)) and train with BEMACallback(update_freq=1):

RuntimeError: The size of tensor a (151665) must match the size of tensor b (8) at non-singleton dimension 0

This PR caches the matching running-model parameters by name in on_train_begin (self.running_params) and zips against that list, so frozen parameters are skipped consistently on both sides.

Test: TestBEMACallback::test_frozen_parameters freezes the embedding matrix, trains for 9 steps, and checks every trainable parameter of the running model against ema + alpha * (theta_t - theta_0) matched by name, while the frozen embedding is left equal to the trained model's. Fails on main with the error above, passes here. check_dtype=False because running_model is built from the config and may carry a different dtype than the trained model (pre-existing, unrelated to this fix).

Fixes # (none filed; found while reading trl/trainer/callbacks.py)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? (not needed: behavior fix only)
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

🤖 Generated with Claude Code

https://claude.ai/code/session_011vZKTbRaDh8FuXSbVpQomD

---

Note

Medium Risk
Changes training callback weight-update logic; incorrect alignment could corrupt saved BEMA checkpoints, though the fix is narrow and covered by a regression test.

Overview
Fixes BEMACallback when some weights are frozen (requires_grad=False). Trainable tensors were tracked in a filtered list, but BEMA updates were applied by zipping against every parameter in running_model, which misaligned or crashed (strict=True length mismatch / wrong tensor shapes).

on_train_begin now builds running_params by parameter name for each trainable weight, and _update_bema_weights writes only into those tensors. Frozen weights in the running copy are no longer touched by the BEMA loop.

Adds test_frozen_parameters: freezes embed_tokens, trains with update_freq=1, checks trainable running weights match ema + alpha * (θₜ - θ₀) and the frozen embedding stays equal to the training model.

^{Reviewed by Cursor Bugbot for commit 71dc9f0. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 71dc9f0. Configure here.}

[cursor]

Test compares tensors across devices

Medium Severity

The new frozen-parameter test builds the expected BEMA weights from thetat on the training device and compares them to running_model tensors, which live on the callback device (default cpu). On GPU that subtraction and assert_close fail with a device mismatch, so the regression test does not run in this project's GPU CI.

 

^{Reviewed by Cursor Bugbot for commit 71dc9f0. Configure here.}