Commit 7a963ff
LocalTensor for random_ops tests (pytorch#166540)
adds support for randomness in localtensor. tl;dr it needs to be able to handle RNG the same way (i.e., rng tracking/syncing across shards, user-defined seeds, user-defined generators, etc.
we extend the existing OffsetBasedRNGTracker to play nicely with localtensor's setup, creating a few small subclasses and patching the core RNG logic to manage the per-rank seeds and offsets correctly.
i still haven't done the per-rank generator support (since the existing tests imply a globally-seeded generator), but that it something that should be done.
Pull Request resolved: pytorch#166540
Approved by: https://github.com/dzmitry-huba1 parent 9f94c7b commit 7a963ff
File tree
6 files changed
+765
-119
lines changed- test/distributed/tensor
- torch
- distributed
- _local_tensor
- tensor
- testing/_internal/distributed/_tensor
6 files changed
+765
-119
lines changed
0 commit comments