Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DYNAMO RNG seed update optimization #7884

Merged
merged 1 commit into from
Aug 20, 2024
Merged

Conversation

JackCaoG
Copy link
Collaborator

RNG seed had 2 forms in torch_xla, a xla tensor and a scalar. In the dynamo bridge case we want the tensor form for the execution but for the seed update we don't want to do tensor.item() since it will cause a transferFromDevice. We can directly access the scalar form of the rng seed

@JackCaoG JackCaoG requested review from lsy323 and alanwaketan August 20, 2024 17:07
@@ -54,12 +54,18 @@ def random_op(self, a):
return torch.randn(5, 5, device=a.device) + a

def test_random_op_different_result_each_run(self):
xm.wait_device_ops()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious why do we need a barrier only for this test

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is just to be safe for unit tests, there might be async operations that update the metrics. We want to wait for all async ops to finish and then clear the metrics

@JackCaoG JackCaoG marked this pull request as ready for review August 20, 2024 17:14
@JackCaoG JackCaoG merged commit 3776ccc into master Aug 20, 2024
23 checks passed
@JackCaoG JackCaoG deleted the JackCaoG/dynamo_rng_bug_fix branch August 20, 2024 17:14
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Oct 11, 2024
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Dec 11, 2024
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Dec 11, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants