decouple backward and step in accelerator/deepspeed

[naomili0924]

What does this PR do?

Addressing this issue: https://github.com/huggingface/accelerate/issues/2951

This PR enables user to debug model gradient after backward+gradient_clipping and before optimizer.step().

To achieve this functionality, we move the gradient clipping logic from deepspeed engine (another public repo) to accelerator. As a result, self.engine.set_gradient_clipping(0.0) must be set in the user's config. Or add an extra line:

self.engine.set_gradient_clipping(0.0) in the step() function.

We didn't replace amp logic with normal pytorch.cuda.amp for clipping fp16 gradient. As a result, you might need to manually clone the apex repo to use amp.

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir ./

Fixes # (issue) https://github.com/huggingface/accelerate/issues/2951

Before submitting

• [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• [x] Did you read the contributor guideline, Pull Request section?
• [x] Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
• [ ] Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
• [x] Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.