Add dataloader_multiprocessing_context to TrainingArguments + Ensure dataloader_prefetch_factor fully supports IterableDataset (follow-up to #34925)
Author: muyihao
Pull request: no
Labels: name: Feature request

### Feature request

When using streaming datasets from remote storage (e.g. HDFS via pyarrow.fs.HadoopFileSystem / HdfsFileSystem), enabling `dataloader_num_workers > 0` often fails due to pickling errors when the default `'fork'` start method is used. Switching to `multiprocessing_context='spawn'` solves this reliably, but `transformers.Trainer` currently hard-codes the default context and offers no way to override it.

Additionally, for high-throughput scenarios (lightweight models + large global batch size via gradient accumulation), IO can still become a bottleneck even with workers. Higher `dataloader_prefetch_factor` (e.g. 4–10) helps a lot by prefetching more batches into the buffer ahead of time — but this only works for map-style datasets right now.

**Current status in the repo**

- There is already an open PR #34925 ("Allowing dataloader prefetching to work with iterable datasets") that aims to remove the `if not isinstance(..., IterableDataset)` guard and enable prefetch_factor for iterable-style datasets. (Opened Dec 2024, currently awaiting review.)
- Related older issues: #32169, #34867 (both pointing out the same limitation).


### Motivation

- In 'spawn' mode, users need to make sure any custom collate_fn / preprocessing functions are defined at module level (not nested closures / lambdas), which is a standard multiprocessing requirement.
- My typical use-case: streaming terabyte-scale pre-tokenized data from HDFS, using IterableDataset + packing + small models → very high token/sec demand → IO prefetch becomes critical.

### Your contribution

Yes — if the maintainers think both directions (multiprocessing_context + full prefetch support) make sense, I'm happy to open a PR implementing them (with tests + docs), possibly building on top of #34925.
Thanks for considering!

**My solution**
1. Add a new argument in `TrainingArguments`:

   ```python
   dataloader_multiprocessing_context: Optional[Literal["fork", "spawn", "forkserver"]] = None```
Then pass it to every DataLoader(..., multiprocessing_context=...) call inside Trainer (train/eval/prediction).
This would allow users to explicitly choose 'spawn' when needed without monkey-patching.

2. Once #34925 is merged (or as a complementary change), make sure dataloader_prefetch_factor is passed unconditionally to IterableDataset DataLoaders as well.In practice, I've seen 2–4× speedup / much less GPU idle time when prefetch_factor is increased in IO-heavy streaming + lightweight model settings.
3. Add other params. Even when prefetch works, the default `in_order=True` (PyTorch DataLoader behavior) enforces that batches are yielded in roughly the order they were requested. In many streaming / shuffled-buffer IterableDataset use-cases (especially when the underlying iterable already applies random buffering), strict order is not necessary — allowing `in_order=False` can enable workers to yield whatever is ready first, reducing tail latency and GPU idle time under heavy IO pressure.

Comments:
cc @sunmarc