Option to disable unwrapping model for generation in PPO/RLOO/OnlineDPO

[dawidm]

Feature request

Add trainer option to disable unwrapping model (unwrap_model_for_generation()) in online methods.

Motivation

As discussed in #2250, when deepspeed stage 3 is used and policy model is bigger than single GPU VRAM, unwrap_model_for_generation() will cause OOM. The option to disable unwrapping will make such training scenarios possible (but slow).

Your contribution

I will likely create PR after it's discussed and approved.

[yiyepiaoling0715]

same question,how to resolve thie?

[qgallouedec]

The solution that you're suggesting sounds good to me, feel free to open a PR

[Superskyyy]

@dawidm @qgallouedec The merged code leads to a runtime error upon trying on my RLOO train script, disabling this option the trainer works fine. Can you help to check if there's something wrong? Thanks.

0]:   File "/shared_workspace/xx/xx-rl-pipeline/trainer.py", line 520, in <module>
[rank0]:     main(
[rank0]:   File "/shared_workspace/xx/xx-rl-pipeline/trainer.py", line 435, in main
[rank0]:     trainer.train()
[rank0]:   File "/shared_workspace/xx/xx-rl-pipeline/venv/lib/python3.12/site-packages/trl/trainer/rloo_trainer.py", line 320, in train
[rank0]:     with unwrap_model_for_generation(
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/lib/python3.12/contextlib.py", line 149, in __exit__
[rank0]:     raise RuntimeError("generator didn't stop")
[rank0]: RuntimeError: generator didn't stop

[Superskyyy]

I think we are missing an else check before Line 189

else:  # missing else here
    with deepspeed.zero.GatheredParameters(model.parameters())

When the context manager exits it will again yield a unwrap model leading to more than one yield -> runtime error

If we can confirm this indeed needs an else I will open a PR to fix.

[qgallouedec]

Yes feel free to open a PR. Thanks