Skip to content

Add custom iterable dataloader support in Accelerator.prepare#3984

Open
t7r0n wants to merge 1 commit intohuggingface:mainfrom
t7r0n:issue-2975-custom-iterable-dataloader
Open

Add custom iterable dataloader support in Accelerator.prepare#3984
t7r0n wants to merge 1 commit intohuggingface:mainfrom
t7r0n:issue-2975-custom-iterable-dataloader

Conversation

@t7r0n
Copy link
Copy Markdown

@t7r0n t7r0n commented Mar 19, 2026

Summary\n- add DataLoaderConfiguration.custom_classes so users can register iterable dataloader-like classes\n- add a lightweight CustomIterableDataLoader wrapper that:\n - iterates the underlying object\n - optionally moves each batch to device\n - preserves accelerate dataloader state tracking for gradient synchronization\n- update Accelerator.prepare / prepare_data_loader flow to treat configured custom classes as dataloaders\n\n## Tests\n- add test_prepare_data_loader_with_custom_iterable_loader in tests/test_data_loader.py\n- add test_prepare_with_custom_iterable_dataloader in tests/test_accelerator.py\n\n## Local verification\n- python3 -m compileall src/accelerate/accelerator.py src/accelerate/data_loader.py src/accelerate/utils/dataclasses.py tests/test_data_loader.py tests/test_accelerator.py\n- direct runtime smoke check with PYTHONPATH=src for both prepare_data_loader and Accelerator.prepare custom iterable paths\n- full pytest run in this environment is blocked by missing test dependency: parameterized\n\nCloses #2975

@t7r0n
Copy link
Copy Markdown
Author

t7r0n commented Mar 19, 2026

Additional local validation run in a uv-managed virtual environment: pytest tests/test_data_loader.py tests/test_accelerator.py -k not setting_cpu_affinity. Result: 81 passed, 5 skipped, 1 deselected. The deselected CPU affinity test is hardware-specific on this machine (NVML CPU affinity not supported) and unrelated to this PR changes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant