Support list of ShardingSpec in MpDeviceLoader #5789
Conversation
jonb377
force-pushed
the
jonbolin/multi-spec
branch
from
8cd4e2e to
ab617f8
Compare
| for sharding in input_sharding: | ||
| if sharding.can_apply(tensor): | ||
| shardings[i] = sharding.xla_spec(tensor) | ||
| break |
There was a problem hiding this comment.
Maybe we also add some comments about the first match is applied?
There was a problem hiding this comment.
Done, thanks Yeounoh!
There was a problem hiding this comment.
Is first match something common practice? Or will it be better:
- either one sharding spec for everything;
- or one sharding spec for each input and in order.
That the behavior is well defined. Otherwise, this "compatible check" is totally a black box for the user.
There was a problem hiding this comment.
It's definitely not ideal... Probably better would be to have the user provide sharding specs in a structure matching that of the inputs and avoid the black-box compatibility check, e.g.:
# Inputs of the form:
{"a": torch.randn(16), "b": torch.randn(16, 16), "c": torch.randn(16, 16, 16)}
# Then input_sharding would be:
{"a": ShardingSpec(mesh, (0,)), "b": ShardingSpec(mesh, (0, None)), "c": ShardingSpec(mesh, (0, None, None))}This is a bit of a refactor of the TensorToXlaArena code though, so I took the easy way out. I'll go ahead and take a stab at the cleaner approach now.
There was a problem hiding this comment.
No worries. Take your time.
jonb377
added
the
DO_NOT_MERGE_YET
In #5768, there is a use case to shard different-rank input tensors. Currently, the MpDeviceLoader sharding only accepts a single spec, which is applied to compatible tensors. With this change, multiple specs can be provided in construction, and each input tensor will have the first compatible spec applied to it.