introduce batch sharding strategy

[zpcore]

(Split out the large PR from #46)

Introduce the batch sharding strategy:

from torch.distributed.tensor._op_schema import RuntimeSchemaInfo
from autoparallel.dtensor_util.utils import batch_shard_strategy
from torch.distributed.tensor._ops.utils import register_op_strategy
# create strategy with input tensor 1 replicated, input tensor 2 shard on dim 0. Output tensor shard on dim 0:
custom_shard_strategy = functools.partial(batch_shard_strategy, input_shard_dim=[None, 0], output_shard_dim=[0])
# register the strategy:
register_op_strategy(new_op)(custom_shard_strategy)

For details, check func description in autoparallel/dtensor_util/utils.py and example usage in tests/test_dtensor.py.

Stack from ghstack (oldest at bottom):

• -> introduce batch sharding strategy #50
• Support of implicit fallback #49

[ezyang]

The terminology is confusing here. Is the shard dim the batch dim? Or is it something else? If it is the batch dim, the comments below seem to imply there is only one batch dim, why is this a list?

[zpcore]

It's batch dim. There can be multiple tensor input. Each element in input_shard_dim maps to one input tensor. Same to output tensor(s).

[ezyang]

Is there a reason this takes an OpSchema and not an OpStrategy?

[zpcore]

The type definition of OpSchema is very unclear in DTensor. Here the OpSchema is a collection of OpStrategy for all input args. While for the output, most of the time we just have one output tensor, so OpStrategy is sufficient for output.

[ezyang]

Running str on the type, very suspicious!

[zpcore]

What is the recommended way to check the tensor output type?

[zpcore]

Check to see if there's any concerns for this PR. Or should we merge and do a try?

[ezyang]

This is pretty reversible so I don't mind landing it and deciding what to do with it later. @zpcore I do hope this gets obsoleted by whatever we end up deciding to do with DTensor sharding though!

[wconstab]

let's land this and get some experience using the API for deepseek enablement.

@zpcore have you rebased? I'd like to at least kick off a job on llama3 mast to make sure nothing got broken