Fix RLOO checkpointing

[bartoszzuk]

This PR fixes RLOO checkpointing (in the same way as a recent fix for PPOv2 PR #2080).

This is needed after changes to _save_checkpoint method introduced in transformers v4.45.0.dev. Specifically we will get KeyError: 'TrainerControl' during saving of the trainer state (here is the exact line causing the issue). By passing stateful_callbacks to OnlineTrainerState explicitly the TrainerControl object is stored and can be properly accessed in _save_checkpoint.

[qgallouedec]

Nice, thanks @bartoszzuk. Without your fix, does it cause any error when running RLOO?

[qgallouedec]

make sure to run make precommit by the way to make the CI happy

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[bartoszzuk]

@qgallouedec Yes when using transformers v4.45.0.dev I'm getting

...
File "/usr/local/lib/python3.10/dist-packages/trl/trainer/rloo_trainer.py", line 449, in train
  self._save_checkpoint(model, trial=None, metrics=metrics)
File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 3016, in _save_checkpoint
  if isinstance(self.state.stateful_callbacks[cb_name], list):
KeyError: 'TrainerControl'

Because the self.state.stateful_callbacks is an empty dict

Sorry, totally forgot about make precommit will fix it ASAP

[qgallouedec]

Thanks. I'm not sure to understand why this failing mode doesn't break our CI

[sahandrez]

I am not sure if this is related, but I have observed a strange behaviour in RLOO checkpointing. For example, I have set it to checkpoint every 500 steps and it follows that for some time, but after a while it starts generating checkpoints every 2 steps. Is this an intended functionality?

[lewtun]

Hi @bartoszzuk thanks for the fix! Would you mind writing a regression test in test_rloo_trainer.py that fails on main but passes on your branch? That would help ensure future code changes don't accidentally introduce the bug again

[bartoszzuk]

Hey @lewtun, sorry for late response. I added a simple regression test for RLOO checkpointing. Hopefully it somewhat follows the conventions found in other tests (let me know If any improvements are required). The test should:

• Succeed without the fix for transformers<4.45.0
• Fail without the fix for transformers>=4.45.0
• Succeed with the fix for transformers>=4.45.0

I also changed the test function to ensure that tokenizer matches SFT and Reward model.

[qgallouedec]

• Succeed without the fix for transformers<4.45.0

✅

• Fail without the fix for transformers>=4.45.0

✅

• Succeed with the fix for transformers>=4.45.0

✅

[qgallouedec]

LGTM, thanks a lot @bartoszzuk, I'll merge as soon as the CI passes

[lewtun]

Thanks a lot for the fix and regression test @bartoszzuk! I think the CI is failing for some unrelated reason, so I've rerun it to be sure

[qgallouedec]

Yep, not related