MACE oeq benchmarking

[malixian]

OpenEq is a great job. I try to test the mace performance by mace_driver.py, but get a error:

TypeError: MACE.init() got an unexpected keyword argument 'oeq_config'

Maybe some modifications for mace, can you give me some suggestions? The version of mace is 0.3.10.

[vbharadwaj-bk]

We appreciate the interest in the package! Yes, you need a modified version of MACE to run that driver, and we plan to make that repo public this week. We'll drop you a message when it's online.

[vbharadwaj-bk]

@malixian We have made a modified version of MACE public and have done some cursory testing - let us know if you can run it by following the instructions here: https://github.com/PASSIONLab/OpenEquivariance?tab=readme-ov-file#running-mace

[malixian]

@malixian We have made a modified version of MACE public and have done some cursory testing - let us know if you can run it by following the instructions here: https://github.com/PASSIONLab/OpenEquivariance?tab=readme-ov-file#running-mace

Exciting!! I've done some simple tests.
On the NVIDIA H800, for the carbon.xyz data, 1.38x performance improvement compared to cueq in FP32, but there is still room for optimization in FP64.

[vbharadwaj-bk]

Thanks for verifying! Good to know, cuE may be using the tensor cores for FP64. Do you mind sharing roughly how much slower we are against cuE in FP64 precision? We may attempt some additional optimizations in that case.

If you have the time, I would also appreciate pasting the output jsons that the script produces of us vs. cuE here on FP64 so we can determine if the convolution kernel is taking the most time.

ETA: A few brief scans of the internet tell me that the FP64 performance on the H800 may be orders of magnitude lower than the H100 (correct me if I'm wrong). That doesn't explain any anomalous performance, but it's not a hardware configuration that we have optimized for.

[malixian]

Thanks for verifying! Good to know, cuE may be using the tensor cores for FP64. Do you mind sharing roughly how much slower we are against cuE in FP64 precision? We may attempt some additional optimizations in that case.

If you have the time, I would also appreciate pasting the output jsons that the script produces of us vs. cuE here on FP64 so we can determine if the convolution kernel is taking the most time.

ETA: A few brief scans of the internet tell me that the FP64 performance on the H800 may be orders of magnitude lower than the H100 (correct me if I'm wrong). That doesn't explain any anomalous performance, but it's not a hardware configuration that we have optimized for.

For carbon.xyz, openEq is 415ms and cuEq is 523ms on FP64. For BOTNet, openEq is 229ms and cuEq is 86ms on FP64. Using TensorCore on the H100 is expected to bring a noticeable performance improvement. I will conduct further tests on an H100 GPU.

[vbharadwaj-bk]

Awesome! Thank you, we will take a look at the BOTNet configuration. Shall we migrate this to a Github Discussion where you can post any further results? The comments, etc. will be preserved, but we can close this issue if you are able to run MACE successfully.

[malixian]

Awesome! Thank you, we will take a look at the BOTNet configuration. Shall we migrate this to a Github Discussion where you can post any further results? The comments, etc. will be preserved, but we can close this issue if you are able to run MACE successfully.

we can go ahead and close this issue. Thanks for your efforts!