Skip to content

Commit

Permalink
dataloader: check that in_order is in kwargs before trying to drop it (…
Browse files Browse the repository at this point in the history 
…#3346)

This fixes tests/test_data_loader.py::StatefulDataLoaderTester tests which
started to fail after 828aae4:
```
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'
```

The reason for the failure is that "in_order" is added only if data loader
is created with `prepare_data_loader` or `skip_first_batches()`. Tests in
`tests/test_data_loader.py::StatefulDataLoaderTester` however are creating
data loaders directly as classes and "in_order" was not added. Hence the
issue.

Fixes: 828aae4 ("add torchdata version check to avoid in_order error (#3344)")

Signed-off-by: Dmitry Rogozhkin <[email protected]>
  • Loading branch information
@dvrogozh
dvrogozh authored Jan 15, 2025
1 parent d09040d commit fbfa53b
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion src/accelerate/data_loader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,7 +419,11 @@ def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, *
)
if use_stateful_dataloader:
torchdata_version = version.parse(importlib.metadata.version("torchdata"))
if compare_versions(torchdata_version, "<", "0.11") and is_torch_version(">=", "2.6.0"):
if (
"in_order" in kwargs
and compare_versions(torchdata_version, "<", "0.11")
and is_torch_version(">=", "2.6.0")
):
kwargs.pop("in_order")
self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs)
else:
Expand Down

0 comments on commit fbfa53b

Please sign in to comment.