[Enhancement] New Callback event, before and after backward

[muellerzr]

This PR introduces a new Callback event: after_backward. This is designed so custom backward pass implementations can be used (such as accelerate).

after_backwards is the same as before_step, but I opted to keep both because there isn't a clean way to combine the two into one without introducing too many changes.

If a callback implements their own backward pass, they just need to attach a RaiseBackwardException in the implementation.

cc @jph00

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[muellerzr]

ci failed unrelated to the current notebook

[jph00]

Looks like you need to re-run this after fixing your cuda env

[warner-benjamin]

Could you update the Callback documentation to mention after_backward and what it is used for? Currently it reads:

- `before_backward`: called after the loss has been computed, but only in training mode (i.e. when the backward pass will be used)
- `before_step`: called after the backward pass, but before the update of the parameters. It can be used to do any change to the gradients before said update (gradient clipping for instance).

[muellerzr]

@warner-benjamin after_backward can be synonymous with before_step. But everyone should use before_step really, hence why it was left unmentioned.

[warner-benjamin]

That strikes me as an important thing to mention.