!deepspeed examples/summarize_rlhf/sft/train_gptj_summarize.py is failing

**Bug**
Hello ,

I am trying to run summarize_rlhf example using [this blog on wandb](https://wandb.ai/carperai/summarize_RLHF/reports/Implementing-RLHF-Learning-to-Summarize-with-trlX--VmlldzozMzAwODM2).
This script is failing with attached logs; however I am not able to locate the actual issue.
Error snapshot -
`Loading extension module utils...
Time to load utils op: 18.591187238693237 seconds
Rank: 0 partition count [1] and sizes[(6050882784, False)]
[2023-04-01 14:17:09,378] [INFO] [launch.py:318:sigkill_handler] Killing subprocess 3605
[2023-04-01 14:17:09,383] [ERROR] [launch.py:324:sigkill_handler] ['/usr/bin/python3', '-u', 'examples/summarize_rlhf/sft/train_gptj_summarize.py', '--local_rank=0'] exits with return code = -9`

Detailed error - [error.txt](https://github.com/CarperAI/trlx/files/11129818/error.txt)
My enviornment looks like this - [pip_freeze.txt](https://github.com/CarperAI/trlx/files/11129836/pip_freeze.txt)

Can someone guide me in investigating this issue?
Is there any other way to get more diagnostic information?

Note : Before this, I encountered base64 decoding issue (padding error) with ds_config_gptj.json so I converted it to dict and added in the train_gptj_summarize.py.

**Which trlX version are you using?**
Editable install with no version control (trlx==0.6.0)

Installed using pip install - e . using source