gradient checkpointing for GPT-NeoX
# What does this PR do?

Add gradient checkpointing to GPT-NeoX model, in the style of GPT-J
_The documentation is not available anymore as the PR was closed or merged._<|||||>The current failing tests look like infra issues.<|||||>Yes, looks like a flaky failure!