MeshGraphNet Performance: Automaticaly Use transformer engine for LayerNorm. #1036
Conversation
This file contains hidden or 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
… training speed up on GPU.
coreyjadams
added
! - Release
|
/blossom-ci |
coreyjadams
commented
Jul 29, 2025
coreyjadams
commented
Jul 29, 2025
coreyjadams
commented
Jul 29, 2025
coreyjadams
commented
Jul 29, 2025
|
FYI - this PR looks bigger than it is. Many tests are updated but most are not really changed: I have some env variables to force the torch version of layernorm on tests that explicitly use CPU. This prevents transformer engine from being used on CPU models. There are new additional tests for layer norm, however, that should get reviewed before merge. |
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
Update docstring to use torch layernorm (for CPU tests).
Disable TE for docstring tests.
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
Alexey-Kamenev
approved these changes
Aug 1, 2025
- remove warnings about deprecation - add warning if env variable PHYSICSNEMO_FORCE_TE is set, but to an unexpected value.
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
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.
PhysicsNeMo Pull Request
Description
This PR provides performance enhancements for mesh graph net on GPUs by automatically using transformer engine for layer norm. Here's what happens:
A note about test changes: transformer engine is not supported on CPU. So the tests are modified to let physicsnemo select the layernorm backend optimally for performance, except if the test is a CPU test: then it is forced to torch instead of TE. MGN tests therefore use both torch and TE backends during testing, and both show agreement after restoring from file.
On synthetic data, the performance improvement is pretty good on large graphs. I measured with PyG and DGL, up to 200k nodes and up to 1M+ edges (more goes out of memory). At small graphs, the performance is similar. For large graphs, transformer engine is better - especially during training.
Additionally, PyG is faster than DGL in all measurements.
Including performance measurements here in the PR for posterity.
Float 32
Training
Inference
Float16
Training
Inference
BFloat16
Training
Inference
Checklist
Dependencies