Optimizer-in-backward Fusion Implementation

[tinuademargaret]

This PR implements optimizer-in-backward fusion (suggested in this issue), which integrates optimizer updates directly into the backward pass to decrease peak memory usage by eliminating the need to store gradients. It uses pytorch's _apply_optimizer_in_backward to fuse optimizer steps with backpropagation.

This feature only works with the sft trainer.

Optimizer-in-backward fusion is incompatible with gradient accumulation and is controlled by a bwd_hook flag in the optimizer config.

There is added validation to ensure that training configurations that rely on gradient accumulation does not enable the optimizer fusion option.

[CLAassistant]

All committers have signed the CLA.

[eric-haibin-lin]

thanks! Just to double check, this supports gradient accumulation?

[tinuademargaret]

thanks! Just to double check, this supports gradient accumulation?

No, It uses pytorch's apply_optimizer_in_backward which does not support gradient accumulation.

[vermouth1992]

thanks! Just to double check, this supports gradient accumulation?

No, It uses pytorch's apply_optimizer_in_backward which does not support gradient accumulation.

I guess this constrain limit the use of apply_optimizer_in_backward because most of the cases, we need gradient accumulation to avoid OOM :(

[eric-haibin-lin]

Thanks for the contribution! For RL i think it's very likely that we need gradient accumulation to train with large batch sizes. But this feature would still be useful for SFT. Do you think it make sense if you only keep the changes for SFT trainer, while keeping the RL trainer simple for now? Thx

[tinuademargaret]

Thanks for the contribution! For RL i think it's very likely that we need gradient accumulation to train with large batch sizes. But this feature would still be useful for SFT. Do you think it make sense if you only keep the changes for SFT trainer, while keeping the RL trainer simple for now? Thx

I agree RL almost always needs gradient accumulation, so the immediate win for RL is limited. I've scoped the PR down to the SFT trainer for now and left the RL workers unchanged. PyTorch mentioned here that they are working on a more flexible post backward hook, I'm not sure of the status of this, but once that lands we can revisit a gradient accumulation friendly version for RL.