[Feature] Add row-decomposition of adj. matrix to reduce graph partitioning overhead #720
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Please update the CHANGELOG |
stadlmax
reviewed
Nov 26, 2024
| @@ -769,6 +908,12 @@ def get_src_node_features_in_partition( | |||
| ) -> torch.Tensor: # pragma: no cover | |||
| # if global features only on local rank 0 also scatter, split them | |||
| # according to the partition and scatter them to other ranks | |||
|
|
|||
| if self.graph_partition.matrix_decomp: | |||
There was a problem hiding this comment.
Ideally, one shouldn't make the code using these utilities dependent on how a graph is partitioned. Couldn't one instead of throwing this error just use get_dst_node_features underneath?
stadlmax
reviewed
Nov 26, 2024
| @@ -872,6 +1017,12 @@ def get_global_src_node_features( | |||
| if partitioned_node_features.device != self.device: | |||
| raise AssertionError(error_msg) | |||
|
|
|||
| if self.graph_partition.matrix_decomp: | |||
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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
Dependencies
To enable this matrix decomp scheme, developers need to pass
matrix_decomp=Truetopartition_graph_nodewise()function@mnabian @stadlmax @Alexey-Kamenev Can you please help review this PR?