Issues with pretraining the model in multiple GPU multiple nodes

[silabrata]

I am trying to pretraining the model using the train_pubchem_light.py on multiple GPUs across multiple nodes. When I run the code (with some modification for LSF to SLURM). For a multinode mullgpu run I am getting the following log:

fused_weight_gradient_mlp_cuda module not found. gradient accumulation fusion with weight gradient computation disabled.
Using custom data configuration default-aa670d5a2f66e342
Reusing dataset pub_chem (/tmp/silabrata.93/pubchem/pub_chem/default-aa670d5a2f66e342/0.0.0/a6cf7324273df2f2ba223646cb9ecbb36c120aab97c28e8378a3afe9dca94289)
['g073.cluster', 'g074.cluster']
g073.cluster
['g073.cluster']
g073.cluster MASTER_ADDR: g073.cluster
g073.cluster MASTER_PORT: 54966
g073.cluster NODE_RANK 0
g073.cluster NCCL_SOCKET_IFNAME: ib0
g073.cluster NCCL_DEBUG: INFO
g073.cluster NCCL_IB_CUDA_SUPPORT: 1
Using 2 Nodes---------------------------------------------------------------------
Using 2 GPUs---------------------------------------------------------------------
{'batch_size': 500, 'num_workers': 4, 'pin_memory': True}
../data/ZINC/CCEA.smi
../data/ZINC/BIBD.smi
../data/ZINC/DDCD.smi
../data/ZINC/DAEA.smi
../data/ZINC/CBCB.smi
../data/ZINC/BICA.smi
../data/ZINC/CKCA.smi
../data/ZINC/CACA.smi
../data/ZINC/AGAB.smi
../data/ZINC/BAEA.smi
../data/ZINC/BECB.smi
../data/ZINC/CEAA.smi
../data/ZINC/BJEA.smi
../data/ZINC/BKEB.smi
../data/ZINC/AKEA.smi
a configuration default-8780d85d721ae806
Reusing dataset zinc (/tmp/silabrata.93/zinc/zinc/default-8780d85d721ae806/0.0.0/81866737128724210ed5dd6e7ce0e6a97d5dca6bdcee21f0d8ab6d3e1fa68ac5)
Global seed set to 12345
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
Using environment variable NODE_RANK for node rank (0).
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
fused_weight_gradient_mlp_cuda module not found. gradient accumulation fusion with weight gradient computation disabled.
Using custom data configuration default-aa670d5a2f66e342
Reusing dataset pub_chem (/tmp/silabrata.93/pubchem/pub_chem/default-aa670d5a2f66e342/0.0.0/a6cf7324273df2f2ba223646cb9ecbb36c120aab97c28e8378a3afe9dca94289)
['g073.cluster', 'g074.cluster']
g073.cluster
['g073.cluster']
g073.cluster MASTER_ADDR: g073.cluster
g073.cluster MASTER_PORT: 54966
g073.cluster NODE_RANK 0
g073.cluster NCCL_SOCKET_IFNAME: ib0
g073.cluster NCCL_DEBUG: INFO
g073.cluster NCCL_IB_CUDA_SUPPORT: 1
Using 2 Nodes---------------------------------------------------------------------
Using 2 GPUs---------------------------------------------------------------------
{'batch_size': 500, 'num_workers': 4, 'pin_memory': True}
../data/ZINC/CCEA.smi
../data/ZINC/BIBD.smi
../data/ZINC/DDCD.smi
../data/ZINC/DAEA.smi
../data/ZINC/CBCB.smi
../data/ZINC/BICA.smi
Using custom data configuration default-8780d85d721ae806
Reusing dataset zinc (/tmp/silabrata.93/zinc/zinc/default-8780d85d721ae806/0.0.0/81866737128724210ed5dd6e7ce0e6a97d5dca6bdcee21f0d8ab6d3e1fa68ac5)
Global seed set to 12345
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1]
Global seed set to 12345
initializing ddp: GLOBAL_RANK: 1, MEMBER: 2/4
Global seed set to 12345
initializing ddp: GLOBAL_RANK: 0, MEMBER: 1/4
g073:190341:190341 [0] NCCL INFO Bootstrap : Using [0]ib0:10.73.133.90<0>
g073:190341:190341 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
g073:190341:190341 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/IB ; OOB ib0:10.73.133.90<0>
g073:190341:190341 [0] NCCL INFO Using network IB
NCCL version 2.7.8+cuda11.0
g073:190541:190541 [1] NCCL INFO Bootstrap : Using [0]ib0:10.73.133.90<0>
g073:190541:190541 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
g073:190541:190541 [1] NCCL INFO NET/IB : Using [0]mlx5_0:1/IB ; OOB ib0:10.73.133.90<0>
073:190541:190541 [1] NCCL INFO Using network IB

However, the training doe does not start and the code kind of hangs here, Although the code continues to run but fails to propagate and show the losses. Can someone help me out with why it is getting stuck. It seems that master address, nodes and gpus are getting detected properly.