How to provide sharding annotation for MpDeviceLoader when data has different dimensions

[hanzhi713]

❓ Questions and Help

Let's say my dataloader yields a dict when iterating over and the members of this dict has different dimensions

{
    "input_ids": shape = (batch, seq),
    "masks": shape = (batch, seq, seq),
}

pl.MpDeviceLoader appears to only able to provide one sharding annotation. I'm currently using it like this:

data_loader = pl.MpDeviceLoader(
    data_loader,
    dev,
    input_sharding=xs.ShardingSpec(mesh, ('data', None, None)))

Obviously, ('data', None, None) is not valid for input_ids which has only 2 dimensions. But this seems to work. I wonder what's the proper way of using MpDeviceLoader in this case.

[hanzhi713]

cc @alanwaketan

[alanwaketan]

cc @yeounoh

[jonb377]

@hanzhi713 MpDeviceLoader currently only accepts a single ShardingSpec, which will only be applied to compatible tensors (i.e. tensors whose rank matches the partition spec). In your case, masks will be sharded and input_ids will not.

We are developing a general sharding-aware dataloader to replace the MpDeviceLoader approach for more advanced use cases, but for now we can generalize this to support multiple ShardingSpecs for different inputs. I've opened #5789 with the patch.

With that change, you can provide multiple ShardingSpec to the MpDeviceLoder input_sharding, which would be:

specs = [xs.ShardingSpec(mesh, spec) for spec in [('data', None), ('data', None, None)]
data_loader = pl.MpDeviceLoader(
    data_loader,
    dev,
    input_sharding=specs)