[zhangxin81]

Summary

• Expose data.dataloader.persistent_workers and data.dataloader.in_order for the native DataLoader path.
• Forward the knobs into DistributedDataloader and the DiT dataloader construction path.
• Guard persistent_workers when num_workers=0, and warn when in_order=False is used because StatefulDataLoader does not guarantee exact checkpoint/resume ordering in that mode.

Performance Rationale

Multimodal and video pipelines often have uneven per-sample decode/preprocess costs. persistent_workers=True avoids worker teardown/restart overhead when iterators are recreated, and in_order=False can reduce head-of-line blocking from slow workers. These are standard PyTorch DataLoader scheduling knobs, but VeOmni did not expose them through its config.

Defaults preserve current behavior.

Validation

• PYTHONPATH=/Users/bytedance/Documents/VeOmni-dataloader-scheduling PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH pytest tests/data/test_dataloader.py::test_build_dataloader_suppresses_persistent_workers_with_zero_workers tests/data/test_dataloader.py::test_build_dataloader_forwards_worker_scheduling_kwargs -q -> 2 passed
• PYTHONPATH=/Users/bytedance/Documents/VeOmni-dataloader-scheduling PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH pytest tests/data/test_dataloader.py -q -> 20 passed, 2 skipped
• PYTHONPATH=/Users/bytedance/Documents/VeOmni-dataloader-scheduling PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH make quality -> ruff check and format passed
• PYTHONPATH=/Users/bytedance/Documents/VeOmni-dataloader-scheduling PATH=/Users/bytedance/Documents/VeOmni/.venv/bin:$PATH python tests/special_sanity/check_device_api_usage.py --directory ./veomni -> passed
• /veomni-review verdict: safe after adding the checkpoint/resume warning and docs caveat

Notes: in_order=False is a throughput knob, not an exact-resume knob. It should be enabled only when the user accepts non-FIFO worker return order.

[gemini-code-assist]

Code Review

This pull request introduces two new dataloader configuration options, persistent_workers and in_order, to control worker process lifetime and batch ordering. These options are integrated into the native and diffusion dataloaders, along with corresponding documentation, tests, and warning logs when out-of-order loading is enabled. The review feedback highlights a potential runtime TypeError in build_dit_dataloader because it does not accept **kwargs or the extra parameters passed by the trainer wrapper, suggesting adding **kwargs to its signature to safely absorb these parameters.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The build_dit_dataloader function is called via the generic build_dataloader wrapper in dit_trainer.py (line 380), which passes several extra keyword arguments (such as max_seq_len, bsz_warmup_ratio, dyn_bsz, collate_fn, save_steps, etc.). Since build_dit_dataloader does not accept **kwargs or these specific parameters, configuring data.dataloader.type: "dit" will result in a runtime TypeError. To prevent this, please add **kwargs to the signature of build_dit_dataloader to safely absorb any extra configuration parameters passed by the trainer.

[cursor]

🤖 Thanks for the contribution! Cursor has started reviewing this PR and will post the full review shortly.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

Review summary — needs changes (documentation / guard, not the plumbing)

Baseline: main@e2c97423f4cb, ahead 1, behind 19. Reviewed read-only (no checkout of the PR branch; PR file versions were fetched over tracked paths, tested, then reverted).

1. Does it solve a real problem? Yes, and the plumbing is correct

Both knobs are genuinely useful — persistent_workers removes per-epoch worker respawn cost, and in_order=False removes head-of-line blocking when sample decode cost is uneven (a real issue for multimodal data). I verified the mechanics:

• StatefulDataLoader.__init__ accepts both persistent_workers and in_order on the pinned torchdata 0.11.0, so nothing raises.
• BaseTrainer._build_dataloader splats asdict(args.data.dataloader), so the new fields reach build_native_dataloader automatically; native is the only entry in DATALOADER_REGISTRY, and build_dit_dataloader enumerates its arguments explicitly, which is why the dit_trainer change is needed. No builder is left with an unexpected keyword.
• The persistent_workers and num_workers > 0 guard is necessary, not decorative: torchdata raises ValueError("persistent_workers option needs num_workers > 0") at stateful_dataloader.py:234-235.

tests/data/test_dataloader.py passes (20 passed, 2 skipped) and ruff check / ruff format --check are clean.

2. Comments / docs / tests — this is where it falls short

The two tests cover exactly the two mechanical behaviors (the zero-worker downgrade, and kwarg forwarding) and do so without over-mocking. That part is fine.

The problem is the documentation of what these flags actually cost. docs/usage/arguments.md warns about in_order but says nothing at all about persistent_workers, and the persistent_workers risk is the more damaging of the two — it can silently degrade training quality with no error. Details in the inline comments. For a knob that users will reach for as a free throughput win, "opt-in" is not sufficient protection when the failure mode is invisible.

3. Cleanliness / conflicts

19 commits behind main; please rebase. No overlap with the author's other open PRs (#941, #942, #943).

Two smaller notes:

• The added in_order warning duplicates one torchdata already emits at stateful_dataloader.py:237-240 under the same condition. Keeping it is defensible since warning_rank0 is quieter and it names the VeOmni config path, but it is worth knowing users will see the message twice.
• Silently downgrading an explicit persistent_workers=True to False when num_workers == 0 avoids the crash but hides a misconfiguration. A warning_rank0 on that branch would cost one line.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

persistent_workers=True silently breaks per-epoch reshuffling for iterable datasets.

Persistent workers keep their dataset replica alive across iterator recreations, so a set_epoch call in the main process never reaches them. That matters here because DistributedDataloader.set_epoch (line 50-54) forwards to self.dataset.set_epoch(...) for iterable datasets, and every trainer calls it once per epoch (base.py:863, text_trainer.py:169, vlm_trainer.py:342, dit_trainer.py:562).

The epoch actually feeds the shuffle seed, and it is read on the worker side:

worker_info = get_worker_info()
worker_id = worker_info.id if worker_info is not None else 0
if not self._just_resumed:
    seed_seq = np.random.SeedSequence([self._seed, self._epoch, self.dp_rank, worker_id])

(veomni/data/dataset.py:541-544)

With non-persistent workers the pool is re-forked each epoch and inherits the updated _epoch. With persistent_workers=True the workers keep _epoch frozen at its first value, so every epoch replays the same sample order. No exception, no warning — just a quietly worse run, which is the hardest kind of regression to attribute.

Since the whole point of the flag is to avoid respawning workers, this is inherent rather than a bug in your wiring. But it needs to be surfaced. Suggestions, in order of preference:

1. Warn (or refuse) when persistent_workers=True is combined with num_train_epochs > 1 and a dataset that exposes set_epoch.
2. At minimum, document it in the arguments.md row — the in_order row already carries its caveat, and this one is more consequential.

A smaller point on the same line: when num_workers == 0 an explicit persistent_workers=True is silently rewritten to False. That correctly dodges torchdata's ValueError, but a one-line warning_rank0 would tell the user their setting was ignored instead of leaving them to wonder why nothing changed.

[cursor]

The in_order caveat here mentions only checkpoint/resume ordering, but upstream documents a second consequence that matters more for training quality:

Setting in_order to False can harm reproducibility and may lead to a skewed data distribution being fed to the trainer in cases with imbalanced data.

(torchdata/stateful_dataloader/stateful_dataloader.py:182)

That is the same imbalance that motivates turning the flag on — uneven per-sample decode cost — so for VeOmni's multimodal mixes the skew is likely, not hypothetical: cheap-to-decode samples get over-represented. Please add reproducibility and data-distribution skew to this row alongside the resume caveat, and to the runtime warning at data_loader.py:231-234.

The persistent_workers row above also has no caveat at all — see my comment on veomni/data/data_loader.py:250 for why it needs one.