Run parallel distributed tests in CI #832
Open
+144
−0
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
…feature/test_distributed
elynnwu
approved these changes
Feb 12, 2026
Comment on lines
+195
to
+197
| gathered_global = torch.zeros( | ||
| *global_shape, dtype=tensor.dtype, device=tensor.device | ||
| ) |
Contributor
There was a problem hiding this comment.
Might be a spatial parallelism PR question, do we know how often we are calling gather_global? If allocating 2x the global tensor becomes a memory issue, we may want to consider passing a pre-allocated buffer to the gather call and use it in-place, that way we don't hold both the temporary gathered and gathered_global at the same time. I don't know if this is an issue, but something we might want to consider in the future.
| pytest --durations 40 . | ||
|
|
||
| test_parallel: | ||
| torchrun --nproc-per-node 2 -m pytest ./fme/core/distributed/parallel_tests |
Contributor
There was a problem hiding this comment.
[nit] make the number of processor a Make variable so you can use a different number when calling make test_parallel
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
This PR adds some non-trivial tests of data and model parallelism.
The implementations are trivial for the non-model-parallel case, but these tests will provide good coverage of the upcoming model-parallel code. This PR makes it so that the infrastructure is in place to use these tests and add to them.
Changes:
Added
get_local_slices,data_parallel_rankandtotal_data_parallel_ranksattributes to Distributed and DistributedBackend classes, which are currently used only in unit tests but will be needed for spatial parallelism.Tests added