Add on_optimizer_step to callback options

[dhruvbpai]

Add on_optimizer_step callback option in TrainerCallbacks

Aside: This is my first open source pull request, so any feedback would be much appreciated!

The test tests/trainer/test_trainer_callback.py has been modified appropriately to invoke the new callback method.

Fixes #31033 (issue)

Reviewers

As tagged in initial issue - @muellerzr @younesbelkada

[younesbelkada]

This looks great to me thanks !

[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.

[muellerzr]

Makes sense to me as well! Nice job. cc @amyeroberts for final review

[LysandreJik]

Awesome, LGTM! Thanks a lot 🤗

[PangziZhang523]

Hello, thank you for your contribution. I used on_optimizer_step to print the gradient, and all the printed values ​​were None, but the final grad_norm had a value. Why is that?

[colmon46]

Hello, thank you for your contribution. I used on_optimizer_step to print the gradient, and all the printed values ​​were None, but the final grad_norm had a value. Why is that?

Same question with @PangziZhang523 . Do you know how to fix this problem plz?

[Arunprakash-A]

It is working fine for me.

[YeLuoSuiYou]

Hello, thank you for your contribution. I used on_optimizer_step to print the gradient, and all the printed values ​​were None, but the final grad_norm had a value. Why is that?

Same question with @PangziZhang523 . Do you know how to fix this problem plz?

Hi @colmon46, if we use deepspeed or FSDP, we should use self.model_wrapped to get gradient of every layers, but do you know how can we get self.model_wrap in callbacks, 3ku

[Gaoyg]

Hello, thank you for your contribution. I used on_optimizer_step to print the gradient, and all the printed values ​​were None, but the final grad_norm had a value. Why is that?

Same question with @PangziZhang523 . Do you know how to fix this problem plz?

Same question. Is there a solution already? thx~