Skip to content

Transfer shards from server using replication#4508

Merged
yeounoh merged 5 commits intomasterfrom
spmd_preserve_sharding
Feb 1, 2023
Merged

Transfer shards from server using replication#4508
yeounoh merged 5 commits intomasterfrom
spmd_preserve_sharding

Conversation

@yeounoh
Copy link
Copy Markdown
Contributor

@yeounoh yeounoh commented Jan 25, 2023 •
edited
Loading

Allow shards transfer to host, so we can do

xt1 = torch.ones(16, 16).to(xm.xla_device())
xs.mark_sharding(xt1, self._get_mesh((1, self.n_devices)), (0, 1))
t1 = xt1.cpu()

Currently, xt1.cpu() would crash. This is not exposing the shards, which will be addressed in the future.

@yeounoh yeounoh added the distributed SPMD and other distributed things. label Jan 25, 2023
@yeounoh yeounoh self-assigned this Jan 25, 2023
@yeounoh yeounoh force-pushed the spmd_preserve_sharding branch 2 times, most recently from 49f21e1 to ed77f6a Compare January 25, 2023 23:07
@yeounoh
Copy link
Copy Markdown
Contributor Author

yeounoh commented Jan 25, 2023

The failure is due to a newly disabled test,

[ FAILED  ] 1 test, listed below:
[  FAILED  ] AtenXlaTensorTest.TestNonzero

rebasing now.

@JackCaoG
Copy link
Copy Markdown
Collaborator

JackCaoG commented Jan 26, 2023

@yeounoh can you add a comment in the description regarding what problem are you trying to solve?

@JackCaoG
Copy link
Copy Markdown
Collaborator

JackCaoG commented Jan 27, 2023

Thanks, I will take a look today

JackCaoG
JackCaoG reviewed Jan 28, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc Outdated
JackCaoG
JackCaoG reviewed Jan 28, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc
std::move(status_or.value()));
}

ComputationClient::DataPtr PjRtComputationClient::ReplicateShardedData(
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Jan 28, 2023

Choose a reason for hiding this comment

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

If I understand correctly, this functions takes a sharded data and then replicate it on all device and then return? This is useful when you want to all_gather a sharded data and return to the user?

Copy link
Copy Markdown
Contributor Author

@yeounoh yeounoh Jan 28, 2023

Choose a reason for hiding this comment

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

Correct, basically we perform a simple/light identity calculation, so that the compiler replicates the output (reassembled input).

JackCaoG
JackCaoG reviewed Jan 28, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

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

Mostly lgtm, minor comments

@yeounoh yeounoh force-pushed the spmd_preserve_sharding branch 3 times, most recently from eeb88a0 to a286d6e Compare January 31, 2023 04:52
@yeounoh yeounoh requested a review from will-cromar January 31, 2023 21:40
@yeounoh yeounoh force-pushed the spmd_preserve_sharding branch from a286d6e to 556cc22 Compare January 31, 2023 21:56
@yeounoh yeounoh force-pushed the spmd_preserve_sharding branch from 556cc22 to 4873623 Compare January 31, 2023 23:24
JackCaoG
JackCaoG approved these changes Feb 1, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

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

Thanks!

@yeounoh yeounoh merged commit 780a823 into master Feb 1, 2023
mateuszlewko pushed a commit that referenced this pull request Mar 15, 2023
@yeounoh @mateuszlewko
Transfer shards from server using replication (#4508)
5b9837e 
* [SPMD] Transfer shards back to host with replication
* [SPMD] Add test_transfer_sharded_data_to_host test case.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@JackCaoG JackCaoG JackCaoG approved these changes

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

@will-cromar will-cromar Awaiting requested review from will-cromar

Assignees

@yeounoh yeounoh

Labels

distributed SPMD and other distributed things.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@yeounoh @JackCaoG