

TITLE
Multi-node training with Deepspeed hangs when `full_determinism = True`

COMMENTS
5

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
Hey, as I've described below, I think there are problems training Deepspeed in a multi-node setting when `full_determinism = True` in the `TrainingArguments`. I've replicated this on multiple hardware configurations (i.e. different nodes and GPU types — specifically A6000, V100, RTX 3090 — on the same large cluster system). Please take a look, thank you very much!

### System Info

### `transformers-cli env`
- `transformers` version: 4.27.3
- Platform: Linux-3.10.0-1160.76.1.el7.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.1
- Huggingface_hub version: 0.13.3
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Additional info
- `deepspeed` version: 0.8.3
- gcc: 10.2
- cuda: 11.7.1
- pdsh: 2.34

### Who can help?

@sgugger @stas00 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

### Shell

Please set the following environment variables appropriately:

```bash
export NODELIST="gpu1504 gpu1505"
export NUM_NODES=2
export GPUS_PER_NODE=1
export MASTER_ADDR=gpu1504
export MASTER_PORT=9901
```

Create `train.py` from the snippet below, then run with the following commands:

```bash
conda create -n ds-trainer python==3.8.1
conda activate ds-trainer
pip install transformers[deepspeed]

echo "PATH=$PATH" > .deepspeed_env

cat /dev/null >| hostfile
for i in $NODELIST; do
    echo "$i slots=$GPUS_PER_NODE" >> hostfile;
done

deepspeed --num_gpus $GPUS_PER_NODE --num_nodes $NUM_NODES --master_addr $MASTER_ADDR --master_port $MASTER_PORT --hostfile hostfile train.py
```

### `train.py`

```python
import torch
from torch.utils.data import Dataset
from transformers import BertForMaskedLM, Trainer, TrainingArguments
import copy


## Model

model = BertForMaskedLM.from_pretrained("bert-base-uncased")


## Dataset

class DummyDataset(Dataset):
    def __init__(self, max_text_length=16, num_samples=20000) -> None:
        super().__init__()
        self.input_ids = torch.randint(0, 30522, (num_samples, max_text_length))
        self.labels = copy.deepcopy(self.input_ids)

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, index):
        return {
            "input_ids": self.input_ids[index],
            "labels": self.labels[index],
        }

train_dataset = DummyDataset()

## Training

deepspeed_config = {
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto",
        },
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}

training_arguments = TrainingArguments(
    full_determinism = True,
    output_dir = "output",
    do_train = True,
    per_device_train_batch_size = 16,
    max_steps = 100,
    deepspeed = deepspeed_config
)

trainer = Trainer(
    model=model,
    args=training_arguments,
    train_dataset=train_dataset
)

trainer.train()
```

### Expected behavior

**When I run the above code (a minimal example for DeepSpeed training) in a multi-node setting, training seems to hang after the following output:**

<details>
<summary>Output (not working)</summary>

```Shell
[2023-03-24 10:26:48,202] [INFO] [multinode_runner.py:67:get_cmd] Running on the following workers: gpu1504,gpu1505
[2023-03-24 10:26:48,202] [INFO] [runner.py:550:main] cmd = pdsh -S -f 1024 -w gpu1504,gpu1505 export PYTHONPATH=/gpfs/data/csun45/akhand10/projects/test; exp
ort PATH=/gpfs/runtime/opt/pdsh/2.34/bin:/gpfs/runtime/opt/cuda/11.7.1/cuda/bin:/gpfs/runtime/opt/gcc/10.2/bin:/users/akhand10/.local/machine/bin:/users/akhan
d10/.local/machine/bin:/users/akhand10/.local/scripts:/users/akhand10/.local/bin:/users/akhand10/.local/machine/bin:/users/akhand10/palm.h/.local/miniconda3/e
nvs/ds-trainer/bin:/users/akhand10/.local/miniconda3/condabin:/users/akhand10/.local/scripts:/users/akhand10/.local/bin:/usr/lib64/qt-3.3/bin:/usr/local/bin:/
usr/bin:/usr/local/sbin:/usr/sbin:/usr/lpp/mmfs/bin:/usr/lpp/mmfs/sbin:/opt/ibutils/bin:/gpfs/runtime/bin:/opt/singularity/2.5.2/bin:/users/akhand10/bin;  cd 
/gpfs/data/csun45/akhand10/projects/test; /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/bin/python3.8 -u -m deepspeed.launcher.launch --world_info=
eyJncHUxNTA0IjogWzBdLCAiZ3B1MTUwNSI6IFswXX0= --node_rank=%n --master_addr=gpu1504 --master_port=9901 train.py
gpu1504: [2023-03-24 10:26:51,118] [INFO] [launch.py:142:main] WORLD INFO DICT: {'gpu1504': [0], 'gpu1505': [0]}
gpu1504: [2023-03-24 10:26:51,118] [INFO] [launch.py:148:main] nnodes=2, num_local_procs=1, node_rank=0
gpu1504: [2023-03-24 10:26:51,119] [INFO] [launch.py:161:main] global_rank_mapping=defaultdict(<class 'list'>, {'gpu1504': [0], 'gpu1505': [1]})
gpu1504: [2023-03-24 10:26:51,119] [INFO] [launch.py:162:main] dist_world_size=2
gpu1504: [2023-03-24 10:26:51,119] [INFO] [launch.py:164:main] Setting CUDA_VISIBLE_DEVICES=0
gpu1505: [2023-03-24 10:26:53,517] [INFO] [launch.py:142:main] WORLD INFO DICT: {'gpu1504': [0], 'gpu1505': [0]}
gpu1505: [2023-03-24 10:26:53,517] [INFO] [launch.py:148:main] nnodes=2, num_local_procs=1, node_rank=1
gpu1505: [2023-03-24 10:26:53,517] [INFO] [launch.py:161:main] global_rank_mapping=defaultdict(<class 'list'>, {'gpu1504': [0], 'gpu1505': [1]})
gpu1505: [2023-03-24 10:26:53,517] [INFO] [launch.py:162:main] dist_world_size=2
gpu1505: [2023-03-24 10:26:53,517] [INFO] [launch.py:164:main] Setting CUDA_VISIBLE_DEVICES=0
gpu1504: Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.bias', 'cls.seq_rel
ationship.weight']
gpu1504: - This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. 
initializing a BertForSequenceClassification model from a BertForPreTraining model).
gpu1504: - This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a
 BertForSequenceClassification model from a BertForSequenceClassification model).
gpu1504: [2023-03-24 10:26:55,478] [INFO] [comm.py:652:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
gpu1505: Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_r
elationship.bias']
gpu1505: - This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. 
initializing a BertForSequenceClassification model from a BertForPreTraining model).
gpu1505: - This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a
 BertForSequenceClassification model from a BertForSequenceClassification model).
```
</details>

In particular, the last line of relevance is: `[INFO] [comm.py:652:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl`. 

<details>
<summary>Extra NCCL output</summary>

If I provide the vars: `NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=ALL CUDA_LAUNCH_BLOCKING=1`

```Shell
...
...
gpu1504: [2023-03-24 10:48:06,695] [INFO] [comm.py:652:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
gpu1504: gpu1504:30024:30024 [0] NCCL INFO Bootstrap : Using ib0:172.25.211.4<0>
gpu1504: gpu1504:30024:30024 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
gpu1504: gpu1504:30024:30024 [0] NCCL INFO cudaDriverVersion 11070
gpu1504: NCCL version 2.14.3+cuda11.7
gpu1505: gpu1505:22060:22060 [0] NCCL INFO cudaDriverVersion 11070
gpu1504: gpu1504:30024:30024 [0] NCCL INFO init.cc:1147 Cuda Host Alloc Size 4 pointer 0x7f62ffe00000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO NET/IB : Using [0]mlx5_2:1/IB [1]mlx5_0:1/IB ; OOB ib0:172.25.211.4<0>
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Using network IB
gpu1505: gpu1505:22060:22060 [0] NCCL INFO Bootstrap : Using ib0:172.25.211.5<0>
gpu1505: gpu1505:22060:22060 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
gpu1505: gpu1505:22060:22060 [0] NCCL INFO init.cc:1147 Cuda Host Alloc Size 4 pointer 0x7f97a3e00000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO NET/IB : Using [0]mlx5_2:1/IB [1]mlx5_0:1/IB ; OOB ib0:172.25.211.5<0>
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Using network IB
gpu1505: gpu1505:22060:22171 [0] NCCL INFO NET/IB : GPU Direct RDMA Disabled for HCA 0 'mlx5_2'
gpu1505: gpu1505:22060:22171 [0] NCCL INFO NET/IB : GPU Direct RDMA Disabled for HCA 1 'mlx5_0'
gpu1505: gpu1505:22060:22171 [0] NCCL INFO === System : maxBw 12.5 totalBw 24.0 ===
gpu1505: gpu1505:22060:22171 [0] NCCL INFO CPU/0 (1/2/-1)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO + SYS[5000.0] - CPU/1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO + PCI[24.0] - GPU/1000 (1)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO + PCI[12.0] - NIC/23000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO               + NET[12.5] - NET/1 (3ad0a20003723f04/1/12.500000)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO CPU/1 (1/2/-1)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO + SYS[5000.0] - CPU/0
gpu1505: gpu1505:22060:22171 [0] NCCL INFO + PCI[24.0] - NIC/C2000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO               + NET[12.5] - NET/0 (82d0a20003723f04/1/12.500000)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO ==========================================
gpu1505: gpu1505:22060:22171 [0] NCCL INFO GPU/1000 :GPU/1000 (0/5000.000000/LOC) CPU/0 (1/24.000000/PHB) CPU/1 (2/24.000000/SYS) NET/1 (3/12.000000/PHB) NET/0 (4/12.500000/SYS) 
gpu1505: gpu1505:22060:22171 [0] NCCL INFO NET/1 :GPU/1000 (3/12.000000/PHB) CPU/0 (2/12.000000/PHB) CPU/1 (3/12.000000/SYS) NET/1 (0/5000.000000/LOC) NET/0 (5/12.000000/SYS) 
gpu1505: gpu1505:22060:22171 [0] NCCL INFO NET/0 :GPU/1000 (4/12.500000/SYS) CPU/0 (3/12.500000/SYS) CPU/1 (2/12.500000/PHB) NET/1 (5/12.000000/SYS) NET/0 (0/5000.000000/LOC) 
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Setting affinity for GPU 0 to 04
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Pattern 4, crossNic 0, nChannels 1, bw 12.000000/12.000000, type LOC/PHB, sameChannels 1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO  0 : NET/1 GPU/1 NET/1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Pattern 3, crossNic 0, nChannels 1, bw 24.000000/12.000000, type LOC/PHB, sameChannels 1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO  0 : NET/1 GPU/1 NET/1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Pattern 3, crossNic 0, nChannels 0, bw 0.000000/0.000000, type LOC/PIX, sameChannels 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO NET/IB : GPU Direct RDMA Disabled for HCA 0 'mlx5_2'
gpu1504: gpu1504:30024:30241 [0] NCCL INFO NET/IB : GPU Direct RDMA Disabled for HCA 1 'mlx5_0'
gpu1504: gpu1504:30024:30241 [0] NCCL INFO === System : maxBw 12.5 totalBw 24.0 ===
gpu1504: gpu1504:30024:30241 [0] NCCL INFO CPU/0 (1/2/-1)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO + SYS[5000.0] - CPU/1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO + PCI[24.0] - GPU/41000 (0)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO + PCI[12.0] - NIC/23000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO               + NET[12.5] - NET/1 (b4e3ff0003a1420c/1/12.500000)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO CPU/1 (1/2/-1)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO + SYS[5000.0] - CPU/0
gpu1504: gpu1504:30024:30241 [0] NCCL INFO + PCI[24.0] - NIC/C2000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO               + NET[12.5] - NET/0 (d2cfa20003723f04/1/12.500000)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO ==========================================
gpu1504: gpu1504:30024:30241 [0] NCCL INFO GPU/41000 :GPU/41000 (0/5000.000000/LOC) CPU/0 (1/24.000000/PHB) CPU/1 (2/24.000000/SYS) NET/1 (3/12.000000/PHB) NET/0 (4/12.500000/SYS) 
gpu1504: gpu1504:30024:30241 [0] NCCL INFO NET/1 :GPU/41000 (3/12.000000/PHB) CPU/0 (2/12.000000/PHB) CPU/1 (3/12.000000/SYS) NET/1 (0/5000.000000/LOC) NET/0 (5/12.000000/SYS) 
gpu1504: gpu1504:30024:30241 [0] NCCL INFO NET/0 :GPU/41000 (4/12.500000/SYS) CPU/0 (3/12.500000/SYS) CPU/1 (2/12.500000/PHB) NET/1 (5/12.000000/SYS) NET/0 (0/5000.000000/LOC) 
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Setting affinity for GPU 0 to 10
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Pattern 4, crossNic 0, nChannels 1, bw 12.000000/12.000000, type LOC/PHB, sameChannels 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO  0 : NET/1 GPU/0 NET/1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Pattern 3, crossNic 0, nChannels 1, bw 24.000000/12.000000, type LOC/PHB, sameChannels 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO  0 : NET/1 GPU/0 NET/1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Pattern 3, crossNic 0, nChannels 0, bw 0.000000/0.000000, type LOC/PIX, sameChannels 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Tree 0 : -1 -> 0 -> 1/-1/-1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Tree 1 : 1 -> 0 -> -1/-1/-1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 00/02 :    0   1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 01/02 :    0   1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Ring 00 : 1 -> 0 -> 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Ring 01 : 1 -> 0 -> 1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] -1/-1/-1->0->1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO misc/utils.cc:235 memory stack hunk malloc(65536)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO channel.cc:23 Cuda Alloc Size 1152 pointer 0x7f6301c00000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO channel.cc:27 Cuda Alloc Size 8 pointer 0x7f6301c00600
gpu1504: gpu1504:30024:30241 [0] NCCL INFO channel.cc:23 Cuda Alloc Size 1152 pointer 0x7f6301c00800
gpu1504: gpu1504:30024:30241 [0] NCCL INFO channel.cc:27 Cuda Alloc Size 8 pointer 0x7f6301c00e00
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Tree 0 : 0 -> 1 -> -1/-1/-1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Tree 1 : -1 -> 1 -> 0/-1/-1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Ring 00 : 0 -> 1 -> 0
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Ring 01 : 0 -> 1 -> 0
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Trees [0] -1/-1/-1->1->0 [1] 0/-1/-1->1->-1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO misc/utils.cc:235 memory stack hunk malloc(65536)
gpu1505: gpu1505:22060:22171 [0] NCCL INFO channel.cc:23 Cuda Alloc Size 1152 pointer 0x7f97a5c00000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO channel.cc:27 Cuda Alloc Size 8 pointer 0x7f97a5c00600
gpu1505: gpu1505:22060:22171 [0] NCCL INFO channel.cc:23 Cuda Alloc Size 1152 pointer 0x7f97a5c00800
gpu1505: gpu1505:22060:22171 [0] NCCL INFO channel.cc:27 Cuda Alloc Size 8 pointer 0x7f97a5c00e00
gpu1505: gpu1505:22060:22174 [0] NCCL INFO Mem Realloc old size 0, new size 8 pointer 0x7f96f00009c0
gpu1505: gpu1505:22060:22174 [0] NCCL INFO Allocated 4194656 bytes of shared memory in /dev/shm/nccl-3elXjI
gpu1505: 
gpu1505: gpu1505:22060:22174 [0] NCCL INFO New proxy recv connection 0 from local rank 0, transport 2
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f96f0004010
gpu1504: gpu1504:30024:30248 [0] NCCL INFO Mem Realloc old size 0, new size 8 pointer 0x7f62840008c0
gpu1504: gpu1504:30024:30248 [0] NCCL INFO Allocated 4194656 bytes of shared memory in /dev/shm/nccl-HznK2t
gpu1504: 
gpu1504: gpu1504:30024:30248 [0] NCCL INFO New proxy recv connection 0 from local rank 0, transport 2
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f6284004010
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Channel 00/0 : 0[41000] -> 1[1000] [receive] via NET/IB/1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 00/0 : 1[1000] -> 0[41000] [receive] via NET/IB/1
gpu1505: gpu1505:22060:22174 [0] NCCL INFO New proxy recv connection 1 from local rank 0, transport 2
gpu1504: gpu1504:30024:30248 [0] NCCL INFO New proxy recv connection 1 from local rank 0, transport 2
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f96f0004050
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f6284004050
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Channel 01/0 : 0[41000] -> 1[1000] [receive] via NET/IB/1
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 01/0 : 1[1000] -> 0[41000] [receive] via NET/IB/1
gpu1504: gpu1504:30024:30248 [0] NCCL INFO New proxy send connection 2 from local rank 0, transport 2
gpu1505: gpu1505:22060:22174 [0] NCCL INFO New proxy send connection 2 from local rank 0, transport 2
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f6284004090
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 00/0 : 0[41000] -> 1[1000] [send] via NET/IB/1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f96f0004090
gpu1504: gpu1504:30024:30248 [0] NCCL INFO New proxy send connection 3 from local rank 0, transport 2
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Channel 00/0 : 1[1000] -> 0[41000] [send] via NET/IB/1
gpu1505: gpu1505:22060:22174 [0] NCCL INFO New proxy send connection 3 from local rank 0, transport 2
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f62840040d0
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Channel 01/0 : 0[41000] -> 1[1000] [send] via NET/IB/1
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f96f00040d0
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Channel 01/0 : 1[1000] -> 0[41000] [send] via NET/IB/1
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:596 Ib Alloc Size 26560 pointer 0x7f6284020000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:596 Ib Alloc Size 26560 pointer 0x7f96f0020000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO NET/IB: Dev 1 Port 1 qpn 173 mtu 5 LID 1038
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:653 Ib Alloc Size 552 pointer 0x7f96f0034000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net.cc:571 Cuda Host Alloc Size 9641984 pointer 0x7f97a7400000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO Mem Realloc old size 0, new size 768 pointer 0x7f96f0033ff0
gpu1504: gpu1504:30024:30248 [0] NCCL INFO NET/IB: Dev 1 Port 1 qpn 153 mtu 5 LID 1036
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:596 Ib Alloc Size 26560 pointer 0x7f96f0035000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:653 Ib Alloc Size 552 pointer 0x7f6284034000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO NET/IB: Dev 1 Port 1 qpn 174 mtu 5 LID 1038
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:653 Ib Alloc Size 552 pointer 0x7f96f003f000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net.cc:571 Cuda Host Alloc Size 9641984 pointer 0x7f97a9400000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net.cc:571 Cuda Host Alloc Size 9641984 pointer 0x7f6303400000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO Mem Realloc old size 0, new size 768 pointer 0x7f6284034050
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:683 Ib Alloc Size 21688 pointer 0x7f96f003f000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:696 Ib Alloc Size 552 pointer 0x7f96f0046000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:596 Ib Alloc Size 26560 pointer 0x7f6284035000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:771 Ib Alloc Size 552 pointer 0x7f96f0049000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO NET/IB: Dev 1 Port 1 qpn 154 mtu 5 LID 1036
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:653 Ib Alloc Size 552 pointer 0x7f628403f000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net.cc:698 Cuda Host Alloc Size 9641984 pointer 0x7f97ab400000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net.cc:571 Cuda Host Alloc Size 9641984 pointer 0x7f6305400000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:683 Ib Alloc Size 21688 pointer 0x7f628403f000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:696 Ib Alloc Size 552 pointer 0x7f6284046000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:771 Ib Alloc Size 552 pointer 0x7f6284049000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net.cc:698 Cuda Host Alloc Size 9641984 pointer 0x7f6307400000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:683 Ib Alloc Size 21688 pointer 0x7f96f0049000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:696 Ib Alloc Size 552 pointer 0x7f96f0050000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net_ib.cc:771 Ib Alloc Size 552 pointer 0x7f96f0053000
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net.cc:698 Cuda Host Alloc Size 9641984 pointer 0x7f97ad400000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connected all rings
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connected all trees
gpu1505: gpu1505:22060:22171 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
gpu1505: gpu1505:22060:22171 [0] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
gpu1505: gpu1505:22060:22174 [0] NCCL INFO New proxy send connection 4 from local rank 0, transport 2
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:683 Ib Alloc Size 21688 pointer 0x7f6284049000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:696 Ib Alloc Size 552 pointer 0x7f6284050000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net_ib.cc:771 Ib Alloc Size 552 pointer 0x7f6284053000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net.cc:698 Cuda Host Alloc Size 9641984 pointer 0x7f6309400000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connected all rings
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connected all trees
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Latency/AlgBw |    Tree/    LL |    Tree/ LL128 |    Tree/Simple |    Ring/    LL |    Ring/ LL128 |    Ring/Simple | CollNetDirect/    LL | CollNetDirect/ LL128 | CollNetDirect/Simple | CollNetChain/    LL | CollNetChain/ LL128 | CollNetChain/Simple |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO  Max NThreads |            512 |            640 |            512 |            512 |            640 |            256 |              0 |              0 |            512 |              0 |              0 |            512 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO     Broadcast |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     6.3/   3.0 |    14.0/   0.0 |    18.0/  12.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO        Reduce |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     6.3/   3.0 |    14.0/   0.0 |    18.0/  12.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO     AllGather |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     6.3/   6.0 |    14.0/   0.0 |    18.0/  24.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO ReduceScatter |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     6.3/   6.0 |    14.0/   0.0 |    18.0/  24.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |     0.0/   0.0 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO     AllReduce |    14.4/   2.4 |    21.4/   0.0 |    56.0/   9.2 |    10.8/   3.0 |    21.0/   0.0 |    35.4/  12.0 |     4.4/   0.0 |     4.4/   0.0 |    10.7/   0.0 |     4.4/   0.0 |     4.4/   0.0 |     0.0/   0.0 |
gpu1504: gpu1504:30024:30241 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
gpu1504: gpu1504:30024:30241 [0] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
gpu1505: gpu1505:22060:22171 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f96f0004110
gpu1505: gpu1505:22060:22171 [0] NCCL INFO init.cc:367 Cuda Alloc Size 5168 pointer 0x7f97a5c01000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO New proxy send connection 4 from local rank 0, transport 2
gpu1505: gpu1505:22060:22174 [0] NCCL INFO transport/net.cc:376 Cuda Alloc Size 4194304 pointer 0x7f97af400000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO init.cc:392 Cuda Host Alloc Size 33554432 pointer 0x7f96ca000000
gpu1505: gpu1505:22060:22171 [0] NCCL INFO init.cc:398 Cuda Host Alloc Size 128 pointer 0x7f97a3e00200
gpu1505: gpu1505:22060:22171 [0] NCCL INFO comm 0x560589dc07a0 rank 1 nranks 2 cudaDev 0 busId 1000 - Init COMPLETE
gpu1505: gpu1505:22060:22060 [0] NCCL INFO Broadcast: opCount 0 sendbuff 0x7f97bc000000 recvbuff 0x7f97bc000000 count 93763584 datatype 0 op 0 root 0 comm 0x560589dc07a0 [nranks=2] stream 0x560589bf2f70
gpu1505: gpu1505:22060:22060 [0] NCCL INFO misc/utils.cc:235 memory stack hunk malloc(65536)
gpu1504: gpu1504:30024:30241 [0] NCCL INFO Connection to proxy localRank 0 -> connection 0x7f6284004110
gpu1504: gpu1504:30024:30241 [0] NCCL INFO init.cc:367 Cuda Alloc Size 5168 pointer 0x7f6301c01000
gpu1504: gpu1504:30024:30248 [0] NCCL INFO transport/net.cc:376 Cuda Alloc Size 4194304 pointer 0x7f630b400000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO init.cc:392 Cuda Host Alloc Size 33554432 pointer 0x7f6226000000
gpu1504: gpu1504:30024:30241 [0] NCCL INFO init.cc:398 Cuda Host Alloc Size 128 pointer 0x7f62ffe00200
gpu1504: gpu1504:30024:30241 [0] NCCL INFO comm 0x55a62826f2d0 rank 0 nranks 2 cudaDev 0 busId 41000 - Init COMPLETE
gpu1504: gpu1504:30024:30024 [0] NCCL INFO Broadcast: opCount 0 sendbuff 0x7f6335e00000 recvbuff 0x7f6335e00000 count 93763584 datatype 0 op 0 root 0 comm 0x55a62826f2d0 [nranks=2] stream 0x55a628085b40
gpu1504: gpu1504:30024:30024 [0] NCCL INFO misc/utils.cc:235 memory stack hunk malloc(65536)
```
</details>

<details>
<summary>py-spy output</summary>

`py-spy dump --pid [pid]`

```Python
Thread 27321 (active): "MainThread"
    broadcast (torch/distributed/distributed_c10d.py:1555)
    wrapper (torch/distributed/distributed_c10d.py:1436)
    broadcast (deepspeed/comm/torch.py:78)
    broadcast (deepspeed/comm/comm.py:228)
    log_wrapper (deepspeed/comm/comm.py:123)
    _broadcast_model (deepspeed/runtime/engine.py:1105)
    _configure_distributed_model (deepspeed/runtime/engine.py:1182)
    __init__ (deepspeed/runtime/engine.py:297)
    initialize (deepspeed/__init__.py:125)
    deepspeed_init (transformers/deepspeed.py:378)
    _inner_training_loop (transformers/trainer.py:1702)
    train (transformers/trainer.py:1633)
    <module> (train.py:64)
```
</details>

This code works fine in a single-node setup (i.e. with `deepspeed train.py`).

<details>
<summary>Continued output (for single-node, working)</summary>

```Shell
...
...
[2023-03-24 10:36:01,318] [INFO] [comm.py:652:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117/fused_adam/build.ninja...
Building extension module fused_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module fused_adam...
Time to load fused_adam op: 0.304279088973999 seconds
Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
Emitting ninja build file /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117/utils/build.ninja...
Building extension module utils...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module utils...
Time to load utils op: 0.23477387428283691 seconds
{'train_runtime': 12.686, 'train_samples_per_second': 126.124, 'train_steps_per_second': 7.883, 'train_loss': 1.1398809814453126, 'epoch': 0.08}              
[2023-03-24 10:36:18,214] [INFO] [launch.py:350:main] Process 24746 exits successfully.
```
</details>

## Problem: `full_determinism = True`

**If you set `full_determinism = False` in TrainingArguments, multi-node training does work:**

<details>
<summary>Working multi-node output</summary>

```Shell
[2023-03-23 16:40:59,614] [INFO] [multinode_runner.py:67:get_cmd] Running on the following workers: gpu1504,gpu1505
[2023-03-23 16:40:59,614] [INFO] [runner.py:550:main] cmd = pdsh -S -f 1024 -w gpu1504,gpu1505 export PYTHONPATH=/gpfs/data/csun45/akhand10/projects/test_ds; export PATH=/users/akhand10/.local/machine/bin:/users/akhand10/.local/machine/bin:/users/akhand10/.local/scripts:/users/akhand10/.local/bin:/gpfs/runtime/opt/cuda/11.7.1/cuda/bin:/gpfs/runtime/opt/gcc/10.2/bin:/users/akhand10/.local/machine/bin:/users/akhand10/.local/scripts:/users/akhand10/.local/bin:/gpfs/home/akhand10/.vscode-cli/server-stable/bin/ee2b180d582a7f601fa6ecfdad8d9fd269ab1884/bin/remote-cli:/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/bin:/users/akhand10/.local/miniconda3/condabin:/usr/lib64/qt-3.3/bin:/usr/local/bin:/usr/bin:/usr/local/sbin:/usr/sbin:/usr/lpp/mmfs/bin:/usr/lpp/mmfs/sbin:/opt/ibutils/bin:/gpfs/runtime/bin:/opt/singularity/2.5.2/bin:/users/akhand10/bin;  cd /gpfs/data/csun45/akhand10/projects/test_ds; /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/bin/python3.8 -u -m deepspeed.launcher.launch --world_info=eyJncHUxNTA0IjogWzBdLCAiZ3B1MTUwNSI6IFswXX0= --node_rank=%n --master_addr=gpu1504 --master_port=9901 deepspeed_trainer_mvp.py
gpu1504: [2023-03-23 16:41:01,849] [INFO] [launch.py:142:main] WORLD INFO DICT: {'gpu1504': [0], 'gpu1505': [0]}
gpu1504: [2023-03-23 16:41:01,849] [INFO] [launch.py:148:main] nnodes=2, num_local_procs=1, node_rank=0
gpu1504: [2023-03-23 16:41:01,849] [INFO] [launch.py:161:main] global_rank_mapping=defaultdict(<class 'list'>, {'gpu1504': [0], 'gpu1505': [1]})
gpu1504: [2023-03-23 16:41:01,849] [INFO] [launch.py:162:main] dist_world_size=2
gpu1504: [2023-03-23 16:41:01,850] [INFO] [launch.py:164:main] Setting CUDA_VISIBLE_DEVICES=0
gpu1505: [2023-03-23 16:41:05,197] [INFO] [launch.py:142:main] WORLD INFO DICT: {'gpu1504': [0], 'gpu1505': [0]}
gpu1505: [2023-03-23 16:41:05,197] [INFO] [launch.py:148:main] nnodes=2, num_local_procs=1, node_rank=1
gpu1505: [2023-03-23 16:41:05,197] [INFO] [launch.py:161:main] global_rank_mapping=defaultdict(<class 'list'>, {'gpu1504': [0], 'gpu1505': [1]})
gpu1505: [2023-03-23 16:41:05,197] [INFO] [launch.py:162:main] dist_world_size=2
gpu1505: [2023-03-23 16:41:05,197] [INFO] [launch.py:164:main] Setting CUDA_VISIBLE_DEVICES=0
gpu1504: Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']
gpu1504: - This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
gpu1504: - This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
gpu1504: [2023-03-23 16:41:05,717] [INFO] [comm.py:652:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
gpu1505: Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']
gpu1505: - This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
gpu1505: - This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
gpu1505: Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
gpu1504: Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
gpu1505: Detected CUDA files, patching ldflags
gpu1505: Emitting ninja build file /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117/fused_adam/build.ninja...
gpu1505: Building extension module fused_adam...
gpu1505: Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
gpu1505: [1/3] /gpfs/runtime/opt/cuda/11.7.1/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/includes -I/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/adam -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/TH -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/THC -isystem /gpfs/runtime/opt/cuda/11.7.1/cuda/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86 -std=c++17 -c /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o 
gpu1505: [2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/includes -I/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/adam -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/TH -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/THC -isystem /gpfs/runtime/opt/cuda/11.7.1/cuda/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -O3 -std=c++14 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -c /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o 
gpu1505: [3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/gpfs/runtime/opt/cuda/11.7.1/cuda/lib64 -lcudart -o fused_adam.so
gpu1505: Loading extension module fused_adam...
gpu1505: Time to load fused_adam op: 41.71729230880737 seconds
gpu1505: Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
gpu1504: Loading extension module fused_adam...
gpu1504: Time to load fused_adam op: 41.791842222213745 seconds
gpu1504: Using /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...
gpu1505: Emitting ninja build file /gpfs/home/akhand10/.cache/torch_extensions/py38_cu117/utils/build.ninja...
gpu1505: Building extension module utils...
gpu1505: Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
gpu1505: [1/2] c++ -MMD -MF flatten_unflatten.o.d -DTORCH_EXTENSION_NAME=utils -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/TH -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/include/THC -isystem /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -c /users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/deepspeed/ops/csrc/utils/flatten_unflatten.cpp -o flatten_unflatten.o 
gpu1505: [2/2] c++ flatten_unflatten.o -shared -L/users/akhand10/palm.h/.local/miniconda3/envs/ds-trainer/lib/python3.8/site-packages/torch/lib -lc10 -ltorch_cpu -ltorch -ltorch_python -o utils.so
gpu1505: Loading extension module utils...
gpu1505: Time to load utils op: 14.304872751235962 seconds
gpu1504: Loading extension module utils...
gpu1504: Time to load utils op: 14.238191604614258 seconds
gpu1505: {'train_runtime': 31.3539, 'train_samples_per_second': 102.061, 'train_steps_per_second': 3.189, 'train_loss': 0.9621941375732422, 'epoch': 0.16}
gpu1504: {'train_runtime': 31.3186, 'train_samples_per_second': 102.176, 'train_steps_per_second': 3.193, 'train_loss': 0.9663803863525391, 'epoch': 0.16}
100%|██████████| 100/100 [00:31<00:00,  3.13it/s]
100%|██████████| 100/100 [00:31<00:00,  3.13it/s]
gpu1505: [2023-03-23 16:42:40,309] [INFO] [launch.py:350:main] Process 26888 exits successfully.
gpu1504: [2023-03-23 16:42:40,966] [INFO] [launch.py:350:main] Process 37514 exits successfully.
```
</details>