[WIP] GemNet Implementation

[AdeeshKolluru]

This PR adds back GemNet-OC with a few additional features.

• Adapting GemNet-OC model code from v1 to v2
• Removes torch_scatter dependency
• Attempted to remove torch_sparse dependency as well with an optional flag for now - not sure of the implementation yet - needs more tests
• Direct stress block
• Gradient based stress block
• Works across dtypes

Results check

• Reproducing training results on a small run on v1 vs v2
• Consistency of results between v1 GemNet checkpoint and v2

I've added sample configs for debugging purpose and to better understand this codebase - will delete before merge.
I have also made a few hacky changes in the codebase to get direct stress predictions working with this model.

[mshuaibii]

If there's anything we can do to help with this on our end let us know! Glad to see this coming!

[AdeeshKolluru]

Thanks @mshuaibii! My bandwidth has been pretty low in the past few weeks but I've been able to do a basic sanity check of this working with the config I pushed - the total loss is decreasing.

It'd be helpful if you could launch the OC20 run for confirmation. I also didn't get a chance to convert older checkpoint to this format but if you guys have a script already to do that, it'd be great to test v1 vs v2 results with same weights.
Also, it seems like the v2 codebase hasn't been tested for direct stress predictions (?). I've added some hacky changes here just for it to get working but you might have suggestions for cleaner solution before merging.

[mshuaibii]

Thanks @mshuaibii! My bandwidth has been pretty low in the past few weeks but I've been able to do a basic sanity check of this working with the config I pushed - the total loss is decreasing.

It'd be helpful if you could launch the OC20 run for confirmation. I also didn't get a chance to convert older checkpoint to this format but if you guys have a script already to do that, it'd be great to test v1 vs v2 results with same weights. Also, it seems like the v2 codebase hasn't been tested for direct stress predictions (?). I've added some hacky changes here just for it to get working but you might have suggestions for cleaner solution before merging.

I very much don't intend to convert old gemnet checkpoints for this. It's much easier for us to train new OC20 checkpoints and release those instead. I'll try to carve out some time to go through this, thank you!

[github-actions]

This PR has been marked as stale because it has been open for 30 days with no activity.

[samblau]

I just want to bump this, particularly given the impressive results of GemNet-OC on some of the OMol evals. Clearly this model deserves to be in fairchemV2. @mshuaibii - does this depend on you? Or is there someone else on the team who can work on it while you're on leave? Thanks!

[github-actions]

This PR has been marked as stale because it has been open for 30 days with no activity.