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.

Sign up for GitHub

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

Jump to bottom

[SPMD] Use d2d replication via CopyToDevice #4336

Merged
merged 1 commit into from
Dec 15, 2022
Merged

[SPMD] Use d2d replication via CopyToDevice #4336

merged 1 commit into from
Dec 15, 2022

Conversation

yeounoh
Copy link
Contributor

@yeounoh yeounoh commented Dec 15, 2022

Introduce CopyToDevice in PjRtComputationClient to support d2d data transfer. The feature is tested as part of TEST_F(XLAShardingTest, InputHandler) for correctness.

@yeounoh yeounoh self-assigned this Dec 15, 2022
jonb377
jonb377 approved these changes Dec 15, 2022
View reviewed changes
Copy link
Collaborator

@jonb377 jonb377 left a comment

Choose a reason for hiding this comment

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

LGTM, thanks Yeounoh! 🔥

will-cromar
will-cromar approved these changes Dec 15, 2022
View reviewed changes
Copy link
Collaborator

@will-cromar will-cromar left a comment

Choose a reason for hiding this comment

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

LGTM. Thanks!

@JackCaoG
Copy link
Collaborator

Is this mainly to avoid the device->host->device overhead? Does this only works within the addressable device?

@yeounoh
Copy link
Contributor Author

yeounoh commented Dec 15, 2022

Is this mainly to avoid the device->host->device overhead? Does this only works within the addressable device?

Yes, these are highlighted with the two XLA_CHECK

@yeounoh yeounoh merged commit 5eaea4c into master Dec 15, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jonb377 jonb377 jonb377 approved these changes

@will-cromar will-cromar will-cromar approved these changes

@steventk-g steventk-g Awaiting requested review from steventk-g

@JackCaoG JackCaoG Awaiting requested review from JackCaoG

@alanwaketan alanwaketan Awaiting requested review from alanwaketan

Assignees

@yeounoh yeounoh

Labels
runtime SPMD / Distributed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@yeounoh @JackCaoG @jonb377 @will-cromar