Make training reproducible: seeded per-loader shuffles, no global RNG side effects

[solarsys]

Why

A downstream EHR project trained PyHealth models (MLP, RNN, Transformer, RETAIN) twice on identical data, identical fitted processors and fixed seeds (torch.manual_seed, np.random.seed). Held-out PR-AUC still moved by up to 0.06 between runs, enough to flip whether a neural model was "tied" with an XGBoost baseline. Model initialization and CPU kernels were already deterministic; every difference came from random-number handling:

1. In-memory shuffle used Python's global random. InMemorySampleDataset.__iter__ called random.shuffle, which torch and NumPy seeds don't touch.
2. litdata reseeds the caller's RNGs on cache writes. DataProcessor.run calls random.seed(42), np.random.seed(42) and torch.manual_seed(42) in the calling process. This runs through create_sample_dataset(in_memory=False), and set_task's write also consumed Python's random. Whether a user's seed survived therefore depended on whether a cache already existed. Filed upstream as DataProcessor.run reseeds the caller's global random, NumPy and torch generators Lightning-AI/litData#944.
3. The disk-backed shuffle ignored the user's seed.
4. Loaders interfered. get_dataloader called dataset.set_shuffle(...) on the shared dataset, so creating an unshuffled evaluation loader turned off shuffling for an existing training loader on the same dataset.
5. Best epoch not restored without logging. With Trainer(enable_logging=False), early stopping returned the weights from patience epochs after the best one, with no warning.
6. Import-time seeding. Importing pyhealth.models reseeded NumPy (gnn.py, the same fix as Remove module-level RNG reset from gnn.py #1227).

What changes

• get_dataloader(dataset, batch_size, shuffle=False, seed=None). The returned loader (a small DataLoader subclass) keeps its own shuffle flag, seed and epoch, and applies them to the dataset with set_shuffle(shuffle, seed) and set_epoch(epoch) at the start of every pass, in the main process. loader.dataset is the dataset itself. Each pass is a new, deterministic order derived from (seed, epoch), also with num_workers > 0:

  • in memory, NumPy default_rng([seed, epoch]);
  • on disk, litdata's own per-epoch shuffle with the seed passed in.

  With seed=None, the seed is drawn from torch's generator when the loader is created, so torch.manual_seed(...) alone, which is what users already do, makes runs reproducible. Each loader also gets a private torch.Generator; otherwise DataLoader takes a draw from torch's global generator on every pass. set_shuffle(shuffle, seed=None) on datasets is seeded the same way. Nothing shuffles with Python's global random anymore.

• pyhealth.utils.preserve_rng_state() saves and restores the Python, NumPy, torch and (if initialized) CUDA RNG states. BaseDataset.set_task and create_sample_dataset are wrapped in it.

• Trainer.train(..., keep_best_in_memory=True). When monitor is set and there's no output directory, the best weights are kept as a CPU copy and restored at the end. keep_best_in_memory=False opts out with a warning (the docs cover the memory cost). The on-disk behavior with logging is unchanged.

• No library code seeds global RNGs.

  • Removed gnn.py's import-time seeding (from Remove module-level RNG reset from gnn.py #1227).
  • SHAP, LIME and RandomBaseline, SCRIB search, train_lstm_model and the KG splitter now seed only for their own call.
  • pyhealth.utils.set_seed, which users call explicitly, stays. Its docs now cover the GPU flags and note that PYTHONHASHSEED only affects child processes.

• Docs and example. "Reproducible Runs" in docs/api/trainer.rst (including what GPU users additionally need, without changing those flags globally), "Shuffling and Reproducibility" in docs/api/datasets.rst, and examples/reproducible_training.py.

Validation (tests/core/test_reproducibility.py, 17 tests)

Check	Result
Same seed gives an identical batch order over 2 epochs; the epochs differ; a recreated loader restarts the sequence (in-memory and disk)	pass
Different seeds give different orders (both backends)	pass
Iterating a loader leaves Python, NumPy and torch global RNG states unchanged (both backends)	pass
torch.manual_seed(s) alone, with seed=None, gives reproducible orders	pass
A shuffled training loader still shuffles after an unshuffled evaluation loader is created on the same dataset	pass
loader.dataset is the dataset itself (both backends)	pass
With num_workers=2: a new order each epoch, every sample exactly once per epoch, identical across two loaders with the same seed (both backends)	pass
set_task (fresh and existing cache) and create_sample_dataset(in_memory=False) leave global RNG states unchanged; later shuffles follow the user's seed	pass
import pyhealth.models leaves NumPy and torch unchanged (subprocess)	pass
Best epoch restored with enable_logging=False (final validation PR-AUC equals the best, which isn't the last epoch); identical results with logging on; opt-out warns	pass
End to end: MLP and RNN in two separate processes give bit-identical predictions with the same seed and different ones with another seed (in memory; on disk with fresh and existing cache)	pass

• Full core suite: 1,507 tests; the only errors are 4 that also occur on master in this environment (test_processor_schemas reading a stale shared test cache, and test_sdoh's working-directory path, fixed in Find test resources relative to the test file, not the working directory #1292).
• tools/check_pr_rules.py passes.

Changes after review

The first version wrapped the dataset in a per-loader _LoaderView. Review pointed out that it hid the dataset behind loader.dataset, relied on litdata private fields, and its epoch counter didn't advance with num_workers > 0. I confirmed the last point: with 2 workers, every epoch repeated the same order, because __iter__ runs on worker copies of the wrapper.

• _LoaderView is gone. The loader now applies (shuffle, seed, epoch) to the dataset itself at the start of each pass, using only public litdata calls (set_shuffle, set_epoch, seed).
• The epoch stays on the loader and advances in the main process. An epoch counter on the dataset bumped in __iter__ would have the same worker problem.
• Fixed a duplicate-sample bug with workers. With workers, the in-memory dataset yielded every sample once per worker: 32 samples from 16 with 2 workers. This bug is already on master. Each worker now gets a disjoint share.
• Caveat: two loaders on one dataset still can't be iterated at the same time, since the settings live on the dataset. The datasets docs say so; sequential use, such as training then evaluation, is fine.

Notes

• After Store non-tensor sample fields as one opaque value per sample (fix variable-length raw fields on the disk cache) #1277: create_sample_dataset no longer calls litdata.optimize, so its preserve_rng_state guard is mostly defensive now; the set_task guard is still needed.
• Not covered:
  • GPU determinism is documented but untested (no GPU here).
  • Importing pyhealth.models still advances (doesn't reseed) Python's random once, through the third-party rich package; seeding after import is unaffected.
  • pyhealth/medcode/pretrained_embeddings/kg_emb already fails to import on master, so its splitter change is untested at runtime.

🤖 Generated with Claude Code

[jhnwu3]

merge conflict.

[jhnwu3]

What the PR is doing

The PR is trying to make two training runs with the same seeds produce the same results. The author says held-out PR-AUC moved by up to 0.06 between identical runs. The six causes it lists boil down to:

1. InMemorySampleDataset.__iter__ used random.shuffle. That's Python's global random, which torch.manual_seed and np.random.seed don't touch.
2. litdata calls random.seed(42), np.random.seed(42) and torch.manual_seed(42) in the calling process when it writes a cache. So whether your seed survived depended on whether the cache already existed.
3. The disk-backed (litdata) shuffle ignored the user's seed.
4. Shuffle state lives on the dataset object. get_dataloader calls dataset.set_shuffle(shuffle), which changes the shared dataset.
5. Early stopping without logging didn't restore the best epoch.
6. gnn.py reseeded NumPy when it was imported.

Why the custom IterableDataset (_LoaderView)

_LoaderView only addresses #4.

Today, SampleDataset is a litdata StreamingDataset, which is an IterableDataset. DataLoader can't apply a sampler or shuffle=True to an iterable dataset, so the shuffle has to happen inside the dataset's __iter__. That's why PyHealth uses a set_shuffle flag on the dataset. But the flag belongs to the dataset object, so:

train_loader = get_dataloader(ds, 32, shuffle=True)
eval_loader  = get_dataloader(ds, 32, shuffle=False)   # turns off shuffling for train_loader too

_LoaderView is a thin wrapper that gives each loader its own (shuffle, seed, epoch). Its __iter__ calls dataset._iter_for_loader(shuffle, seed, epoch), which works on a copy and leaves the real dataset alone. It has to subclass IterableDataset so that DataLoader keeps calling __iter__ instead of trying index-based access. The __getattr__ passes through attributes like processors and input_schema. The epoch counter is there so each pass gives a new order that is still deterministic.

Is it worth it?

I'd argue mostly no:

• #4 rarely happens in practice. The normal workflow is split_by_patient → three separate dataset objects → three loaders. Two loaders on the same dataset object, one shuffled and one not, is unusual.
• The reproducibility fix doesn't need the wrapper. Fixes #1–3 come down to replacing random.shuffle with np.random.default_rng([seed, epoch]), passing the seed to litdata, and wrapping set_task/create_sample_dataset in preserve_rng_state(). All of that fits in the existing set_shuffle(shuffle, seed) plus an epoch counter on the dataset.
• It adds hidden complexity:
  • loader.dataset is no longer a SampleDataset, so isinstance checks will fail. Attribute access only works through the __getattr__ pass-through. I checked our tree and found no loader.dataset usages today, but user code might have some.
  • The epoch counter likely breaks with num_workers > 0. I reasoned this from how DataLoader works; I didn't run it. Without persistent workers, DataLoader sends a copy of the dataset to the worker processes on every pass and calls __iter__ there. So self.epoch += 1 updates the workers' copies, the main process stays at epoch 0, and every epoch gets the same order. get_dataloader doesn't pass num_workers today, so it's dormant, but nothing stops it from being hit later.
  • _iter_for_loader shallow-copies a litdata StreamingDataset and sets private fields (shuffler, _state_dict, current_epoch). That relies on litdata internals.

Suggestion

• Keep: the seeded shuffles (#1–3), preserve_rng_state, removing the import-time and library-level global seeding, and the Trainer best-weights fix. These are real bugs with clear fixes.
• Drop _LoaderView: put the seed and epoch on set_shuffle/the dataset instead. If the shared-dataset problem matters, document that one dataset object shouldn't back two loaders with different shuffle settings, or raise a warning when it happens.

That would remove most of the utils.py diff you linked and keep the actual reproducibility gains.

Tldr; It's unclear what the _LoaderView here does that improves reproducibility here.

[jhnwu3]

What the PR is doing

The PR is trying to make two training runs with the same seeds produce the same results. The author says held-out PR-AUC moved by up to 0.06 between identical runs. The six causes it lists boil down to:

1. InMemorySampleDataset.__iter__ used random.shuffle. That's Python's global random, which torch.manual_seed and np.random.seed don't touch.
2. litdata calls random.seed(42), np.random.seed(42) and torch.manual_seed(42) in the calling process when it writes a cache. So whether your seed survived depended on whether the cache already existed.
3. The disk-backed (litdata) shuffle ignored the user's seed.
4. Shuffle state lives on the dataset object. get_dataloader calls dataset.set_shuffle(shuffle), which changes the shared dataset.
5. Early stopping without logging didn't restore the best epoch.
6. gnn.py reseeded NumPy when it was imported.

Why the custom IterableDataset (_LoaderView)

_LoaderView only addresses #4.

Today, SampleDataset is a litdata StreamingDataset, which is an IterableDataset. DataLoader can't apply a sampler or shuffle=True to an iterable dataset, so the shuffle has to happen inside the dataset's __iter__. That's why PyHealth uses a set_shuffle flag on the dataset. But the flag belongs to the dataset object, so:

train_loader = get_dataloader(ds, 32, shuffle=True)
eval_loader  = get_dataloader(ds, 32, shuffle=False)   # turns off shuffling for train_loader too

_LoaderView is a thin wrapper that gives each loader its own (shuffle, seed, epoch). Its __iter__ calls dataset._iter_for_loader(shuffle, seed, epoch), which works on a copy and leaves the real dataset alone. It has to subclass IterableDataset so that DataLoader keeps calling __iter__ instead of trying index-based access. The __getattr__ passes through attributes like processors and input_schema. The epoch counter is there so each pass gives a new order that is still deterministic.

Is it worth it?

I'd argue mostly no:

• #4 rarely happens in practice. The normal workflow is split_by_patient → three separate dataset objects → three loaders. Two loaders on the same dataset object, one shuffled and one not, is unusual.
• The reproducibility fix doesn't need the wrapper. Fixes #1–3 come down to replacing random.shuffle with np.random.default_rng([seed, epoch]), passing the seed to litdata, and wrapping set_task/create_sample_dataset in preserve_rng_state(). All of that fits in the existing set_shuffle(shuffle, seed) plus an epoch counter on the dataset.
• It adds hidden complexity:
  • loader.dataset is no longer a SampleDataset, so isinstance checks will fail. Attribute access only works through the __getattr__ pass-through. I checked our tree and found no loader.dataset usages today, but user code might have some.
  • The epoch counter likely breaks with num_workers > 0. I reasoned this from how DataLoader works; I didn't run it. Without persistent workers, DataLoader sends a copy of the dataset to the worker processes on every pass and calls __iter__ there. So self.epoch += 1 updates the workers' copies, the main process stays at epoch 0, and every epoch gets the same order. get_dataloader doesn't pass num_workers today, so it's dormant, but nothing stops it from being hit later.
  • _iter_for_loader shallow-copies a litdata StreamingDataset and sets private fields (shuffler, _state_dict, current_epoch). That relies on litdata internals.

Suggestion

Ask the author to split the PR:

• Keep: the seeded shuffles (#1–3), preserve_rng_state, removing the import-time and library-level global seeding, and the Trainer best-weights fix. These are real bugs with clear fixes.
• Drop _LoaderView: put the seed and epoch on set_shuffle/the dataset instead. If the shared-dataset problem matters, document that one dataset object shouldn't back two loaders with different shuffle settings, or raise a warning when it happens.

That would remove most of the utils.py diff you linked and keep the actual reproducibility gains.

[solarsys]

Thanks, agreed on dropping _LoaderView; it's gone in 81ad324. You were right about workers: I reproduced it (2 workers → the same order every epoch). One catch with putting the epoch counter on the dataset: it has the same problem, since __iter__ runs on worker copies. So get_dataloader now returns a thin DataLoader subclass that calls dataset.set_shuffle(shuffle, seed) and dataset.set_epoch(epoch) in the main process at the start of each pass. loader.dataset is the dataset again, it uses no litdata private fields, and it's tested with num_workers=2. It also fixes an existing bug where the in-memory dataset yielded every sample once per worker.

[jhnwu3]

lgtm