AWS Deep Learning AMI not compatible with multi-GPU training of MolMIM

[xinyu-dev]

Issue

When running multi-GPU pretraining of MolMIM on AWS EC2, the following issues were observed:

1. Pretrain on 1x A10 GPU instance on AWS works OK.
2. Pretrain 1x L40S GPU instance + set devices:1 also works OK.
3. Pretrain 4x A10 GPU instance + set devices:1 also works OK.
4. Pretrain on 4x A 10 GPU instance + set devices: 4 generates error similar to below.
5. Pretrain on 4x L40S GPU instance + set devices: 4 generates similar error

Error log:
log.txt

How to replicate:

Specific AWS EC2 configurations:

1. Launch a 4xA10 GPU instance or a 4xL40S GPU instance on AWS EC2.
2. Choose Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 22.04) 20241115 (ami-0f2ad13ff5f6b6f7c) as the AMI.
3. Run

docker run --rm -it --gpus all -p 8888:8888 nvcr.io/nvidia/clara/bionemo-framework:1.10.1 "/bin/bash"

4. Try running MolMIM pertaining

Workaround

Switch to ami-075a0f15f2d44a65e (NVIDIA GPU-Optimized AMI). Note that some users might not have the liberty to switch AMIs therefore it might be great if we can figure out what is going on here. Thanks very much!