Comments related to training scripts

[MJordahn]

Hi @yilunliao!

I have few more things that I discovered that may be helpful for future users of the repo. I am currently pretraining with direct force predictions on OMat24 and ran into the following errors when getting things up and running:

1. The training script threw errors for the me as the TORCHRUN variable is not set by default in the environment setup. I fixed this by simply adding a TORCHRUN=$(which torchrun) line in the script, but a better solution would be to just set an environment variable directly.
2. I think the CONFIG_PATH in the training script points to the incorrect config path. It currently points to direct training on mptrj when I think it should be pointing to direct training on omat24.
3. When running my_main.py I originally got a AttributeError: 'Namespace' object has no attribute 'distributed' because the distributed flag was not defined. I fixed this by adding the distributed flag https://github.com/atomicarchitects/equiformer_v3/blob/main/src/fairchem/core/common/flags.py.

I hope this is useful. If not, feel free to delete the issue.

[yilunliao]

Hi @MJordahn

Thanks for reporting these.

I have a few comments below:

1. I think first using which torchrun to get the path to the torchrun you want to use (inside some conda / uv env) and then updating the TORCHRUN variable in the training script would be great. Alternatively, you can replace TORCHRUN with torchrun if the path to torchrun is correct by default. Let me update this part soon.
2. Yes, you are correct. Let me fix that.
3. For the distributed flag, do you get that because of running submitit or using the training scripts?

Feel free to let me know if there is any other question.

[MJordahn]

No problem at all.

1. Agree, it is easy to fix, but just wanted to make you aware that the training script requires setting this environment variable which is not detailed in the manual on how to run the training scripts. Again, it is trivial to set this variable on ones own accord, but perhaps thought it was useful for you to know.
2. Nice, thanks.
3. I get the error from this line when running the training script (which runs the my_main.py script).

   equiformer_v3/my_main.py

   Line 73 in a7300c5

   tasks_per_node=(args.num_gpus if args.distributed else 1),

One other question I have (again, sorry for spam!):

In the appendix of your paper, it says that you run batch size 512 with OMat24 (Table 9 in the appendix), but in the script it is set to 16 - I assume 16 is correct? I at least get memory issues on H100s when running with batch size 512 😅

[yilunliao]

3. I see. That one is related to submitit, which I was not able to test. Let me try to fix that.

For batch size, 16 is on each GPU, and 32 GPUs are used -- the effective batch size is 16 * 32 = 512.

[MJordahn]

Ah perfect - thank you!

Let me know if there is anything I can do to assist, e.g. testing etc.