[BUG] Execution hangs after 1 epoch for LLaMa-2-7B SFT.

I am noticing errors running DeepSpeed-Chat with LLaMa-2-7B for SFT where the execution simply hangs after evaluating perplexity for one epoch of training. Looks like some error with barrier sync. I am not sure about the cause of the error.

The only change I had made to the code is the support for logging with wandb. I seemed to have added code with checks such as args.local_rank == 0 or args.local_rank == 1 (check below). Is that the right way?

Another error that seems to arise is that the dataloader crashes with some NCCL error. The error is mentioned below:

Creating dataset split: 96%|█████████▋| 1927287/2000000 [29:56<01:08, 1060.83it/s][E ProcessGroupNCCL.cpp:474] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=3, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800790 milliseconds before timing out.
Traceback (most recent call last):
File "main.py", line 426, in
main()
File "main.py", line 256, in main
train_dataset, eval_dataset = create_prompt_dataset(
File "/data/train-serve-frameworks-eval/SFT/DeepSpeed-Chat/utils/data/data_utils.py", line 397, in create_prompt_dataset
return torch.load(train_fname), torch.load(eval_fname)
return torch.load(train_fname), torch.load(eval_fname)
File "/.local/lib/python3.8/site-packages/torch/serialization.py", line 986, in load
with _open_file_like(f, 'rb') as opened_file:
File "/.local/lib/python3.8/site-packages/torch/serialization.py", line 435, in _open_file_like
return _open_file(name_or_buffer, mode)
File "/.local/lib/python3.8/site-packages/torch/serialization.py", line 416, in init
super().init(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/data/data_files//traindata_ad45fc064c70bb743457d780f037709ad23971a2e16fd7cedb3404982b059fff.pt'
Traceback (most recent call last):

Some threads go ahead and check for torch.load where it doesn't exist.

Start with main.py around main() and the dataset-loading path in utils/data/data_utils.py, especially create_prompt_dataset and its torch.load calls. Reproduce the LLaMA-2-7B SFT run while examining the reported barrier timeout, missing dataset file, and rank checks. Done means the hang and dataloader failure are understood and the run completes without the reported NCCL or FileNotFoundError errors.