From 66ac901632a55b3edf00962f96f23e32a4c5d587 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 11 Jun 2026 12:21:20 +0200 Subject: [PATCH] fix(streaming): do not prepare the dataloader with accelerate The dataset is already rank-disjoint via split_dataset_by_node; accelerate's IterableDatasetShard wrapper kept only every Nth batch of each rank's stream, silently training on 1/N of the data per pass while decoding all of it. The --dummy benchmark path never prepared the loader, so benchmarks were unaffected. Co-Authored-By: Claude Fable 5 --- examples/scaling/train_streaming_multinode.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/scaling/train_streaming_multinode.py b/examples/scaling/train_streaming_multinode.py index f9e49b02d..ed74a40b3 100644 --- a/examples/scaling/train_streaming_multinode.py +++ b/examples/scaling/train_streaming_multinode.py @@ -115,7 +115,11 @@ def main() -> None: cfg = ACTConfig(input_features=input_features, output_features=output_features) model = ACTPolicy(cfg) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) - model, optimizer, loader = accelerator.prepare(model, optimizer, loader) + # Do NOT prepare the dataloader: the dataset is already rank-disjoint via + # split_dataset_by_node, and accelerate's IterableDatasetShard would keep only every + # world_size-th batch of it (silently training on 1/N of the data while decoding all + # of it). Batches are moved to the device manually in the loop. + model, optimizer = accelerator.prepare(model, optimizer) # Resume: restore the dataset's stream position so we don't replay already-seen data. The state holds # plain HF stream dicts + RNG state (not tensors), so weights_only=False is required; the file is a