[Feature] Add row-decomposition of adj. matrix to reduce graph partitioning overhead

[chang-l]

Modulus Pull Request

Description

This PR introduces a new graph partitioning scheme for Modulus distributed GNN models (tested only with MeshGraphNet), which evenly decomposes the adjacency matrix row-wise, effectively eliminating most of the graph partitioning overhead during training.

In the MeshGraphNet model, the 1-D decomposition evenly splits the global_offsets across all ranks (i.e., distributing the nodes evenly among ranks), followed by the corresponding global_indices (which represent all incoming edges for the local nodes). Both the node feature store (node emb matrix) and edge feature store (edge emb matrix) follow this 1D decomposition scheme, and no need to distinguish between src or dst node features store. Implicitly, we assume the adjacency matrix is square, meaning the source and destination node domains are identical, i.e., the graph is not bipartite.

To update local edge store from local node store, all-to-all communication is needed for each message-passing layer to gather updated non-local node (but neighbor node) features, which are then used to update the edge feature store.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

To enable this matrix decomp scheme, developers need to pass matrix_decomp=True to partition_graph_nodewise() function

@mnabian @stadlmax @Alexey-Kamenev Can you please help review this PR?

[mnabian]

Please update the CHANGELOG

[chang-l]

@stadlmax @mnabian ready for the next review round.