EquiformerV2 + DeNS model and trainer

[kyonofx]

No description provided.

[codecov]

Codecov Report

Attention: Patch coverage is 14.07129% with 458 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...core/models/equiformer_v2/trainers/dens_trainer.py	11.53%	322 Missing ⚠️
...em/core/models/equiformer_v2/equiformer_v2_dens.py	19.52%	136 Missing ⚠️

Files with missing lines	Coverage Δ
src/fairchem/core/trainers/ocp_trainer.py	69.66% <ø> (ø)
...em/core/models/equiformer_v2/equiformer_v2_dens.py	19.52% <19.52%> (ø)
...core/models/equiformer_v2/trainers/dens_trainer.py	11.53% <11.53%> (ø)

[IliasChair]

Hey,
Looks like we were tackling the same issue, but you beat me to it! Great solution! Our implementations are pretty similar, but I didn’t add a separate noise head since I didn’t think it looks cleaner without one. While I'm not a acossiated with the OCP repo, I fully support the changes made. I plan to test this implementation comming week if I find the time. I also have some minor suggestions - please see below:

[IliasChair]

Please see my suggestions below.

[IliasChair]

Hello again,
I finally got around to testing your implementation. I have noticed that running it as-is requires a lot of GPU memory. I am running this on one Enterprise NVIDIA H100 GPU, and I can't even get it to run with a batch size of 4. Whereas with the original version, I could easily get away with a batch size of 96. I don't think this is expected behavior; I might have missed a bug somewhere in my review. Is this a known problem? I have also checked that I am running the correct dependency versions.

I am very eager to try out DeNS with the recent OCP additions, If there is anything I can to do assist, let me know!
Since I don't encounter this issue on my Fork I will compare our implementations again and see if I have missed anything after all.

Best
Ilias

[lbluque]

Hello again, I finally got around to testing your implementation. I have noticed that running it as-is requires a lot of GPU memory. I am running this on one Enterprise NVIDIA H100 GPU, and I can't even get it to run with a batch size of 4. Whereas with the original version, I could easily get away with a batch size of 96. I don't think this is expected behavior; I might have missed a bug somewhere in my review. Is this a known problem? I have also checked that I am running the correct dependency versions.

I am very eager to try out DeNS with the recent OCP additions, If there is anything I can to do assist, let me know! Since I don't encounter this issue on my Fork I will compare our implementations again and see if I have missed anything after all.

Best Ilias

Hi @IliasChair,

Thanks for looking into this PR and your suggestions. We did not run into issue using a batch size of 8 on A100 GPUs.

If you get a chance to compare with your implementation and flag any difference that may be causing additional memory use please let us know!

[IliasChair]

I have found another bug in the code. Pleae see below.
Best
Ilias

[IliasChair]

Hello,
I've got a few more comments for you. Sorry about the back-and-forth. I'll try to put everything into one bigger review if anything else comes up.