Description

Part of the iter_batches consumer pipeline cleanup (#63660, #63682). Gates restore_original_order behind DataContext.execution_options.preserve_order (default off). When one format/collate worker lags, the reorder buffer grows with the other workers' completed batches, and ready batches aren't allowed to be yielded; this PR skips the reorder step when ordering isn't required. Recovers next-batch latency from PR1+2's regressed 113 ms steady back to 23 ms (lower than master's 32 ms), with no other regressions.

Why does restore_original_order reduce throughput?

Here's the chain that leads to lower average throughput:

1. Straggler in one format/collate worker → other workers' batches dam up in the reorder buffer.
2. Consumer's next(loader) blocks waiting for the buffer to flush → long next_batch latency, consumer idle.
3. When the straggler finally finishes, the buffer releases the whole wave in batch_idx order → consumer pulls the straggler's batch first.
4. That batch's H2D was just kicked off (finalize ran right before it entered the buffer) — so forward() queues kernels that depend on a tensor whose copy is still in flight → CUDA stream stalls on the H2D.
5. The training step time is inflated for these straggler batches, dragging up the average.

Removing restore_original_order in the default case of preserve_order=False means that we never have this straggler issue. Fastest-producer collate thread wins.

Profiling the training worker

Before: loss.item() takes up ~25% of the profiling duration, since loss.item() is the cuda/CPU sync point. The slow straggler batch time manifests in loss.item() timing outliers.

After: loss.item() takes <2% of the profiling duration.

Release test results

PR3 closes the temporary next-batch latency regression that PR1+2 alone introduced (32 → 113 ms in the throughput variant), bringing it down to 23 ms — lower than master's 32 ms — while preserving PR1+2's −44% peak object store memory and +21% throughput wins.

Follow-up work

Extra GPU memory cost:

• The reorder buffer holds post-finalize GPU tensors while it waits for the next batch_idx to come in, and it can accumulate up to N items (one per format/collate worker that ran ahead) on top of the outer make_async_gen output queue of size prefetch_batches. ([data] Fix iter_batches spilling (1/n): Remove outer make_async_gen to reduce untracked buffered batches + reduce prefetch onto GPU #63660 is removing that output queue.)
• This means that the reorder buffer holds up some GPU memory, on top of the prefetch_batches × batch_bytes outer queue, the in-flight finalize, the current step's forward activations, gradients, etc. Gating the reorder buffer behind preserve_order=False reclaims this GPU memory entirely — batches flow through immediately rather than waiting on the slowest worker.
• We should update the iter_batches pipeline to finalize after the restore_original_order instead to remove the unexpected buffer in GPU memory.

This PR mitigates the issue for the default case because the restore_original_order buffer isn't used.