search

atomicarchitects / equiformer_v2

[ICLR 2024] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
https://arxiv.org/abs/2306.12059
MIT License
218 stars 28 forks source link

Multi-node multi-gpu training #4

Open gxglxy opened 1 year ago

gxglxy commented 1 year ago

Could you provide instructions on how to run experiments under the multi-node multi-gpu setting without using the submitit? For example, I have 2 nodes, each of which contains 16 gpus. How should I modify the scripts you provide to reproduce the reported results?

Thanks!

yilunliao commented 1 year ago

Hi @gxglxy

You can take a look at here: https://github.com/atomicarchitects/equiformer_v2/blob/main/oc20/trainer/dist_setup.py#L14-L89

This is the part where we set up distributed training. If you use MPI to launch multi-node training, you can replace --submit with --summit in *.sh.

Depending on how you launch multi-node training, you might need to use this part or this part. They differ in how they set up environment variables. It might be hard for me to tell what is exactly the modification you need to have given the information you have.

Best