Add on_optimizer_step to callback options
Author: dhruvbpai
Pull request: yes

# 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.
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #31033  (issue)


## Reviewers
As tagged in initial issue - @muellerzr @younesbelkada

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31095). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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?

> 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?

It is working fine for me. 

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

> > 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~