[DTensor RNG][BC Breaking] Change DTensor Philox seed and offset from int to tensor by yiming0416 · Pull Request #173876 · pytorch/pytorch

yiming0416 · 2026-01-30T00:40:12Z

When tracing a DTensor random op, it involves tracing the RNG state in OffsetBasedRNGTracker. However the .item() call in the _PhiloxState is not fake tensor friendly.

In this PR, we remove the .item() call in _PhiloxState so that the getter and setter of seed and offset directly interact with Tensors and it is more tracing friendly.

_LocalPhiloxState in _LocalOffsetBasedRNGTracker has a similar pattern that we could potentially get rid of the .item() call. Right now we don't trace local tensor, so PR keeps it unchanged but makes sure it's compatible with the new _PhiloxState

pytorch-bot · 2026-01-30T00:40:15Z

This PR needs a `release notes:` label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

pytorch-bot · 2026-01-30T00:40:18Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/173876

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a5af529 with merge base 4a6363b ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

wconstab

i think this is probably a good change. it is BC breaking though, so we need to at least document it and ensure our release notes flag it. We should also wait for some others to take a look and give feedback.

wconstab · 2026-01-30T17:52:15Z

                any_rank_philox = _PhiloxState(any_rank_device_state)
-                state.seed = any_rank_philox.seed
-                state.offset = any_rank_philox.offset
+                state.seed = int(any_rank_philox.seed.item())


you're adding a new .item() call in the per op dispatch hot path.

is this only ok because you aren't tracing localtensor? does it have perf implications to localtensor mode though?

local tensor should be fine... wouldn't worry about the overhead here.

Currently we are not tracing LocalTensor. I did see _LocalPhiloxState in _LocalOffsetBasedRNGTracker which has a similar pattern that we could potentially get rid of .item() call there. But it has a larger surface so I didn't do that.

wconstab · 2026-01-30T17:54:50Z

    @property
-    def offset(self) -> int:
-        return int(self._state[8:].view(dtype=torch.int64).item())
+    def offset(self) -> torch.Tensor:


this is technically a BC breaking change, though it is on an underscore API, it is something external users depend on (I would cc @wanchaol and @leonardo0lyj @jc-bytedance @pengyanghua to see if they mind at least).

wconstab · 2026-01-30T17:55:59Z

-            torch.uint8
-        )
-        self._state[:8] = seed_tensor
+    def seed(self, seed: torch.Tensor) -> None:


the new version of the API probably should be more strict: previosuly, it was always going to be a tensor with 1 int64 value in it viewed as uint8. now, it is not gauranteed to be a 1elt tensor anymore

added assert to make sure seed is a 1-elt tensor

wconstab · 2026-01-30T17:56:49Z


    @property
-    def seed(self) -> int:
-        return int(self._state[:8].view(dtype=torch.uint64).item())


i'm happy to get rid of this .item() call

wconstab · 2026-02-02T22:02:17Z

can you also update your PR desc with a full description of why/what you changed?

yiming0416 · 2026-02-02T22:13:37Z

can you also update your PR desc with a full description of why/what you changed?

@wconstab updated the PR desc.

yiming0416 · 2026-02-05T20:40:02Z

@pytorchbot merge

pytorchmergebot · 2026-02-05T20:42:06Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2026-02-05T20:42:08Z

Merge failed

Reason: <urlopen error [Errno 111] Connection refused>

Details for Dev Infra team

Raised by workflow job

yiming0416 · 2026-02-05T20:50:38Z

@pytorchbot merge

pytorchmergebot · 2026-02-05T20:52:37Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2026-02-05T22:59:10Z

Merge failed

Reason: New commits were pushed while merging. Please rerun the merge command.

Details for Dev Infra team

Raised by workflow job

yiming0416 · 2026-02-06T01:30:22Z

@pytorchbot merge

pytorchmergebot · 2026-02-06T01:32:17Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

… int to tensor (pytorch#173876) When tracing a DTensor random op, it involves tracing the RNG state in `OffsetBasedRNGTracker`. However the `.item()` call in the `_PhiloxState` is not fake tensor friendly. In this PR, we remove the `.item()` call in `_PhiloxState` so that the getter and setter of `seed` and `offset` directly interact with Tensors and it is more tracing friendly. `_LocalPhiloxState` in `_LocalOffsetBasedRNGTracker` has a similar pattern that we could potentially get rid of the `.item()` call. Right now we don't trace local tensor, so PR keeps it unchanged but makes sure it's compatible with the new `_PhiloxState` Pull Request resolved: pytorch#173876 Approved by: https://github.com/dolpm, https://github.com/wconstab

pytorch-bot Bot added ciflow/inductor release notes: distributed (dtensor) release notes category labels Jan 30, 2026

yiming0416 added the topic: not user facing topic category label Jan 30, 2026

yiming0416 force-pushed the yiming/dtensor_philox branch from 72a5f5e to 9faa3fd Compare January 30, 2026 00:47

yiming0416 marked this pull request as ready for review January 30, 2026 01:38

yiming0416 requested review from bdhirsh, dolpm and wconstab January 30, 2026 01:38

dolpm approved these changes Jan 30, 2026

View reviewed changes

wconstab reviewed Jan 30, 2026

View reviewed changes

yiming0416 requested a review from wanchaol February 2, 2026 21:16

yiming0416 force-pushed the yiming/dtensor_philox branch 2 times, most recently from f551b79 to a8e94d3 Compare February 2, 2026 22:01

wconstab changed the title ~~[DTensor RNG] Change DTensor Philox seed and offset from int to tensor~~ [DTensor RNG][BC Breaking] Change DTensor Philox seed and offset from int to tensor Feb 2, 2026

yiming0416 force-pushed the yiming/dtensor_philox branch from a8e94d3 to 1a139d7 Compare February 2, 2026 22:04

yiming0416 force-pushed the yiming/dtensor_philox branch 2 times, most recently from b8a3c6c to 8b49c6b Compare February 4, 2026 18:00

yiming0416 requested a review from wconstab February 4, 2026 18:00

yiming0416 removed the topic: not user facing topic category label Feb 4, 2026

wconstab approved these changes Feb 5, 2026

View reviewed changes

yiming0416 force-pushed the yiming/dtensor_philox branch from 8b49c6b to 97e2528 Compare February 5, 2026 18:16

yiming0416 added the ciflow/trunk Trigger trunk jobs on your pull request label Feb 5, 2026

pytorchmergebot added the merging label Feb 5, 2026

pytorchmergebot removed the merging label Feb 5, 2026

pytorchmergebot added the merging label Feb 5, 2026

yiming0416 removed the merging label Feb 5, 2026

yiming0416 force-pushed the yiming/dtensor_philox branch from 97e2528 to 9161613 Compare February 5, 2026 22:57

change dtensor philox from int to tensor

a5af529

yiming0416 force-pushed the yiming/dtensor_philox branch from 9161613 to a5af529 Compare February 5, 2026 22:58

pytorchmergebot added the merging label Feb 6, 2026

pytorchmergebot added the Merged label Feb 6, 2026

pytorchmergebot closed this in 67112c4 Feb 6, 2026

pytorchmergebot removed the merging label Feb 6, 2026

github-actions Bot deleted the yiming/dtensor_philox branch March 9, 2026 02:23

Uh oh!

Conversation

yiming0416 commented Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jan 30, 2026

This PR needs a release notes: label

Uh oh!

pytorch-bot Bot commented Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/173876

✅ No Failures

Uh oh!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

wconstab Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

dolpm Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

yiming0416 Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

yiming0416 Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab commented Feb 2, 2026

Uh oh!

yiming0416 commented Feb 2, 2026

Uh oh!

yiming0416 commented Feb 5, 2026

Uh oh!

pytorchmergebot commented Feb 5, 2026

Merge started

Uh oh!

pytorchmergebot commented Feb 5, 2026

Merge failed

Uh oh!

yiming0416 commented Feb 5, 2026

Uh oh!

pytorchmergebot commented Feb 5, 2026

Merge started

Uh oh!

pytorchmergebot commented Feb 5, 2026

Merge failed

Uh oh!

yiming0416 commented Feb 6, 2026

Uh oh!

pytorchmergebot commented Feb 6, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yiming0416 commented Jan 30, 2026 •

edited

Loading

This PR needs a `release notes:` label

pytorch-bot Bot commented Jan 30, 2026 •

edited

Loading