Stop DataLoaderShard taking a device mesh it never reads

[vineethsaivs]

What breaks

DataLoaderShard(dataset, torch_device_mesh=mesh) accepts the mesh and ignores it.

Its sibling does not:

class DataLoaderDispatcher(...):
    def __init__(self, ..., torch_device_mesh=None, ...):
        self.torch_device_mesh = torch_device_mesh
        ...
        if self.torch_device_mesh and "tp" in self.torch_device_mesh.mesh_dim_names:
            self.submesh_tp = self.torch_device_mesh["tp"]

DataLoaderShard.__init__ takes the same argument, never stores it, and never reads it. Its docstring documents every other parameter and not this one.

Scope

prepare_data_loader only hands the mesh to the dispatcher, so nothing inside accelerate is affected. The trap is for anyone building a DataLoaderShard directly, which is a public, exported class: they get silence rather than either an effect or an error.

What changed

Drop the parameter. It now falls into **kwargs, where DataLoader rejects it, so an unknown argument fails the same loud way any other one does.

Tests

test_dataloader_shard_does_not_take_a_device_mesh asserts the shard rejects it and that the dispatcher still accepts one.

1 passed in 1.69s

On main:

with self.assertRaises(TypeError):
E  AssertionError: TypeError not raised

ruff 0.13.1, the pin in setup.py, is clean on both files.

Related

Same class of thing as #4251, which removes a split_batches argument Accelerator.__init__ never reads. Separate file and separate call path, so they are separate PRs, but happy to fold them together if you would rather review one.