Skip to content

Transfer shards from server using replication - #4508

Merged
yeounoh merged 5 commits into
masterfrom
spmd_preserve_sharding
Feb 1, 2023
Merged

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

Conversation

@yeounoh

@yeounoh yeounoh commented Jan 25, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

yeounoh commented Jan 25, 2023

Copy link
Copy Markdown
Contributor Author

The failure is due to a newly disabled test,

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

rebasing now.
@JackCaoG

JackCaoG commented Jan 26, 2023

Copy link
Copy Markdown
Collaborator

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

JackCaoG commented Jan 27, 2023

Copy link
Copy Markdown
Collaborator

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
std::move(status_or.value()));
}

ComputationClient::DataPtr PjRtComputationClient::ReplicateShardedData(

@JackCaoG JackCaoG Jan 28, 2023

Copy link
Copy Markdown
Collaborator

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?

@yeounoh yeounoh Jan 28, 2023

Copy link
Copy Markdown
Contributor Author

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

@JackCaoG JackCaoG left a comment

Copy link
Copy Markdown
Collaborator

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

@JackCaoG JackCaoG left a comment

Copy link
Copy Markdown
Collaborator

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

Labels

distributed SPMD and other distributed things.

2 participants

@yeounoh @JackCaoG