xla icon indicating copy to clipboard operation
xla copied to clipboard

[SPMD] Use UNKNOWN sharding type for implicit replication

Open yeounoh opened this issue 1 year ago • 3 comments

This is part of auto-sharding PoC described in #6322 . This PR addresses the following:

  • Adapt xla::OpSharding::UNKNOWN for implicit replication type.
  • Refactoring of ShardingUtil::GetOutputSharding, ShardingUtil::CreateShardedData
  • Remove resolved TODOs

It is important to distinguish implicit and explicit replciation for XLA compiler to generate shardings for.

yeounoh avatar Mar 12 '24 07:03 yeounoh

I guess my question is why do we need to distinguish implicit sharding and explicit sharding?

JackCaoG avatar Mar 12 '24 20:03 JackCaoG

I guess my question is why do we need to distinguish implicit sharding and explicit sharding?

Good question, updated the description. TLDR, XLA needs to know what is implicit and expcliit so that can consider implicit for auto-sharding.

yeounoh avatar Mar 12 '24 20:03 yeounoh

Can you remind me why do we need to explicit set the output sharding to be replicated before and how is left it to unknown(implicit sharding) will do the same thing?

For the outputs, we set them to explicit replicated. Let me sync with you offline!

yeounoh avatar Mar 12 '24 21:03 yeounoh

Ok CI passed, will rebase and land

yeounoh avatar Mar 13 '24 17:03 yeounoh