Skip to content

[SPMD] Preserve parameter sharding with output data sharding#4721

Merged
yeounoh merged 4 commits intomasterfrom
new_param_sharding_fix
Mar 8, 2023
Merged

[SPMD] Preserve parameter sharding with output data sharding#4721
yeounoh merged 4 commits intomasterfrom
new_param_sharding_fix

Conversation

@yeounoh
Copy link
Copy Markdown
Contributor

@yeounoh yeounoh commented Mar 3, 2023

This addresses the same problem as in #4696 with an alternative solution. We shard the replicated output while handling the computation results. This avoids post traversal pass to replace original data node with a sharded one, thus more efficient. Key changes include:

  • Introduce ShardingUtil::OutputHandler
  • Add XLAShardingTest.OutputHandler test for unit testing, test_optimizer_step_with_sharding checks the validity of the change with a simple e2e example already.
  • Add std::optional<xla::Shape> to ShardingSpec
  • Add std::optional<xla::OpSharding> to PjRtShardedData
  • Pass an additional std::vector<XLATensor::ShardingSpecPtr> param to XLAGraphExecutor::ScheduleSyncTensorsGraph, since the async function now calls ShardingUtil::OutputHandler
  • Introduce & call XLAGraphExecutor::CollectShardingSpecs before calling ScheduleSyncTensorsGraph
  • Introduce WrapDataShards and GetDataSharding APIs in ComputationClient.

@yeounoh yeounoh added the distributed SPMD and other distributed things. label Mar 3, 2023
@yeounoh yeounoh requested review from JackCaoG and alanwaketan March 3, 2023 23:52
@yeounoh yeounoh self-assigned this Mar 3, 2023
@yeounoh yeounoh marked this pull request as draft March 3, 2023 23:52
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch from 54c51f1 to 420d701 Compare March 3, 2023 23:55
JackCaoG
JackCaoG reviewed Mar 4, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc
JackCaoG
JackCaoG reviewed Mar 4, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc
JackCaoG
JackCaoG reviewed Mar 4, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.h Outdated
JackCaoG
JackCaoG reviewed Mar 4, 2023
View reviewed changes
Comment thread torch_xla/csrc/init_python_bindings.cpp Outdated
JackCaoG
JackCaoG reviewed Mar 4, 2023
View reviewed changes
Comment thread torch_xla/csrc/tensor.h
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch 15 times, most recently from 3eac5e6 to f26b305 Compare March 4, 2023 04:02
@yeounoh yeounoh marked this pull request as ready for review March 4, 2023 04:20
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch 2 times, most recently from 5c3e631 to 0ddee73 Compare March 6, 2023 22:59
JackCaoG
JackCaoG reviewed Mar 7, 2023
View reviewed changes
Comment thread torch_xla/csrc/xla_graph_executor.cpp Outdated
Comment thread torch_xla/csrc/xla_sharding_util.cpp Outdated
Comment thread test/cpp/test_xla_sharding.cpp Outdated
Comment thread third_party/xla_client/pjrt_computation_client.cc
JackCaoG
JackCaoG reviewed Mar 7, 2023
View reviewed changes
Comment thread third_party/xla_client/xrt_computation_client.h Outdated
JackCaoG
JackCaoG reviewed Mar 7, 2023
View reviewed changes
Comment thread torch_xla/csrc/xla_graph_executor.cpp Outdated
JackCaoG
JackCaoG reviewed Mar 7, 2023
View reviewed changes
Comment thread torch_xla/csrc/xla_sharding_util.h
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch from a90760e to e45ab94 Compare March 7, 2023 06:25
@yeounoh yeounoh mentioned this pull request Mar 7, 2023
Closed
@JackCaoG
Copy link
Copy Markdown
Collaborator

JackCaoG commented Mar 7, 2023

OutputHandler seems to crash

@yeounoh yeounoh force-pushed the new_param_sharding_fix branch from e45ab94 to 5ba829f Compare March 7, 2023 21:04
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch from 5ba829f to 26279e3 Compare March 7, 2023 21:04
@yeounoh
Copy link
Copy Markdown
Contributor Author

yeounoh commented Mar 7, 2023

OutputHandler seems to crash

Yea, we need at least 2 devices to create Hlo sharding.

2023-03-07 21:01:25.297118: F external/org_tensorflow/tensorflow/compiler/xla/hlo/ir/hlo_sharding.cc:54] Check failed: num_tiles > 1 (1 vs. 1)

Added the safeguard.

@yeounoh yeounoh requested review from jonb377 and steventk-g March 7, 2023 22:03
@yeounoh yeounoh changed the title Preserve parameter sharding with output data sharding [SPMD] Preserve parameter sharding with output data sharding Mar 7, 2023
JackCaoG
JackCaoG reviewed Mar 7, 2023
View reviewed changes
Comment thread torch_xla/csrc/xla_graph_executor.cpp Outdated
@yeounoh yeounoh force-pushed the new_param_sharding_fix branch from 26279e3 to 8d83ef4 Compare March 7, 2023 22:55
JackCaoG
JackCaoG approved these changes Mar 7, 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!

jonb377
jonb377 approved these changes Mar 8, 2023
View reviewed changes
Copy link
Copy Markdown
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!

Comment thread third_party/xla_client/pjrt_computation_client.h
}

void Assign(const Data& data) override {
XLA_ERROR() << __FUNCTION__ << " not supported.";
Copy link
Copy Markdown
Collaborator

@jonb377 jonb377 Mar 7, 2023

Choose a reason for hiding this comment

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

Nice! We can retry the simple MpDeviceLoader hack for SPMD once this lands, this was the blocker.

@yeounoh yeounoh merged commit e2abcaf into master Mar 8, 2023
mateuszlewko pushed a commit that referenced this pull request Mar 15, 2023
@yeounoh @mateuszlewko
[SPMD] Preserve parameter sharding with output data sharding (#4721)
609d06c 
[SPMD] Persist tensor sharding with XLA sharding propagation
@yeounoh yeounoh mentioned this pull request Mar 22, 2023
Merged
ManfeiBai pushed a commit to ManfeiBai/PyTorchXLA that referenced this pull request Mar 29, 2023
@yeounoh @ManfeiBai
[SPMD] Preserve parameter sharding with output data sharding (pytorch…
73567b2 
…#4721)

[SPMD] Persist tensor sharding with XLA sharding propagation
ManfeiBai pushed a commit to ManfeiBai/PyTorchXLA that referenced this pull request Mar 29, 2023
@yeounoh @ManfeiBai
[SPMD] Preserve parameter sharding with output data sharding (pytorch…
ebb3c1b 
…#4721)

[SPMD] Persist tensor sharding with XLA sharding propagation
@yeounoh yeounoh mentioned this pull request May 16, 2023
Merged
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

@JackCaoG JackCaoG JackCaoG approved these changes

@alanwaketan alanwaketan Awaiting requested review from alanwaketan

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

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.

3 participants

@yeounoh @JackCaoG @jonb377