Edge-parallel segment sum support for GNN

[nathan-az]

Motivation:

The most memory-expensive part of the GNN is the edge-wise operations. When using sigmoid gates (as with orbmol v2), if sharding over the edge dim, many operations can be completed independently with only two relatively cheap collectives required after the segment sums.

This basic PoC is designed to be non-obtrusive and easily testable. It:

• makes the segment sum and softmax configurable, add a version that expects pre-sharded (but local) inputs, handles the collective and the grad step. I have not wired this through to full model constructors.
• allows chunked computing of edge-wise operations for local operation chunking, trading recompute for much lower memory usage

I recommend only enabling one of these at a time. Chunking if you want data-parallel inference on large systems, parallelism (probably intra-node) for multi-GPU. Effectiveness of each will depend on hardware.

Distributed edge-wise operations

Since these collectives are relatively small and the edge operations are the most runtime expensive, I expect near embarrassingly parallel runtime benefits to the GNN as well, but unfortunately I don't have a multi-GPU setup to test this. All testing was done locally on CPU with gloo backend.

Chunked edge-wise operations

This incurs additional recompute in the backward step in conservative force inference but trades massively reduced memory usage by checkpointing around each chunk.

Further work required:

This is not complete as-is. From my testing this requires gradient scaling according to the mesh size (imagine it works similarly to DP) which is an inference-time concern with the conservative force estimates. For a mature implementation this needs:

• sharding in the wider regressor's forward path for the distributed path
• gradient correction in the forcefield utils for the distributed path
• equivalency tests
• testing around the viability of selective activation checkpointing to avoid matmul recompute (needs hardware-specific testing)

---

Notes:

• Barely tested for higher order derivatives
• Currently not tested or accounting for more than 1D parallelism, not intended for use with 2D parallelism (e.g. training with FSDP). This should be a relatively trivial unlock, but too hard for me to test locally
• Not tested for using both distributed and chunked compute - I recommend treating these as mutually exclusive
• I implemented segment softmax largely to get a better idea of how online softmax works and for completion, although this gate is not used for any core models