add torchdata version check to avoid "in_order" error

[faaany]

What does this PR do?

The test below

pytest -rA tests/test_accelerator.py::AcceleratorTester::test_save_model_with_stateful_dataloader_use_safetensors_num_workers_0_dispatch_batches

gives the following error:

self = <accelerate.data_loader.DataLoaderDispatcher object at 0x7efdf9399ed0>
dataset = <torch.utils.data.dataset.TensorDataset object at 0x7efdf9398a90>, use_stateful_dataloader = True
batch_sampler = <torch.utils.data.sampler.BatchSampler object at 0x7efdf9398760>
kwargs = {'collate_fn': <function default_collate at 0x7efecb76bd00>, 'in_order': True, 'multiprocessing_context': None, 'num_workers': 0, ...}
StatefulDataLoader = <class 'torchdata.stateful_dataloader.stateful_dataloader.StatefulDataLoader'>

    def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, **kwargs):
        self.use_stateful_dataloader = use_stateful_dataloader
        if is_torchdata_stateful_dataloader_available():
            from torchdata.stateful_dataloader import StatefulDataLoader
    
        if use_stateful_dataloader and not is_torchdata_stateful_dataloader_available():
            raise ImportError(
                "StatefulDataLoader is not available. Please install torchdata version 0.8.0 or higher to use it."
            )
        if use_stateful_dataloader:
>           self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs)
E           TypeError: StatefulDataLoader.__init__() got an unexpected keyword argument 'in_order'

src/accelerate/data_loader.py:418: TypeError

The reason is that the torchdata version doesn't support pytorch version 2.6 yet. So it doesn't accept the keyword argument "in_order". To fix it, we need to add a version check.

cc @muellerzr @sywangyi @yao-matrix

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[muellerzr]

Thanks, makes sense to me

[dvrogozh]

This PR fixes tests/test_accelerator.py::AcceleratorTester::test_save_model_with_stateful_dataloader_use_safetensors_num_workers_0_dispatch_batches, but other tests in tests/test_data_loader.py start to fail on the newly added code:

FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_dataloader_dispatcher_state_dict_num_workers_0 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_dataloader_dispatcher_state_dict_num_workers_2 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_dataloader_inheritance - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_dataloader_state_dict_num_workers_0 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_dataloader_state_dict_num_workers_2 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_decoupled_stateful_dataloader_adapter_equivalent_to_torchdata_stateful_dataloader_num_workers_0 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_decoupled_stateful_dataloader_adapter_equivalent_to_torchdata_stateful_dataloader_num_workers_2 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_end_of_dataloader - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_end_of_dataloader_dispatcher - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_skip_data_loader - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_stateful_dataloader_adapter_equivalent_to_torchdata_stateful_dataloader_num_workers_0 - KeyError: 'in_order'
FAILED tests/test_data_loader.py::StatefulDataLoaderTester::test_stateful_dataloader_adapter_equivalent_to_torchdata_stateful_dataloader_num_workers_2 - KeyError: 'in_order'

See fix for these issues in:

• dataloader: check that in_order is in kwargs before trying to drop it #3346.