SageMaker OOM on 11_train_gptj_smp_tensor_parallel_notebook.ipynb

[Link to the notebook.](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt-j/11_train_gptj_smp_tensor_parallel_notebook.ipynb)

I am trying to run the example notebook above on SageMaker using the gpt-j-xl model but keep running into an out of memory error on a ml.p3.16xlarge Ec2 instance. I have tried other suggested instances (eg ml.g4dn.12xlarge) as well but get the same error. I've attached the latest error below. I've tried to set the train and val batch sizes to as low as 2 and still run into OOM issues. Any guidance would be appreciated.


Start by running 11_train_gptj_smp_tensor_parallel_notebook.ipynb with gpt-j-xl on the reported SageMaker instances and compare the displayed OOM with the train and validation batch-size settings. Done would require a confirmed notebook configuration or documented guidance that allows the example to run without running out of memory.

Written by the indexing model from the issue text.