The Per-Epoch Batch Rebuild: A Surgical Fix for Gradient Diversity in DFlash Training

The Message

The plan: replace the fixed build_batches + shuffle-order approach with per-epoch batch rebuilding that shuffles within length buckets. The _feed_loop will call build_batches_shuffled each epoch instead of reordering fixed batches.

[edit] /data/dflash/scripts/train_dflash_pipeline.py Edit applied successfully.

This short message, index 8690 in the conversation, represents the implementation of a critical fix to the DFlash training pipeline. On its surface, it is a simple code edit — a few lines changed in a Python training script. But beneath that surface lies a carefully reasoned intervention into the fundamental data ordering strategy of a large-scale machine learning training run, one that balances computational efficiency against gradient diversity, and that required the assistant to diagnose a subtle flaw in the batching logic, design a solution that preserved throughput, and implement it with surgical precision.

The Problem: Static Batch Composition

To understand why this message was written, one must trace back through the preceding conversation. The DFlash training pipeline, running on 8× RTX PRO 6000 Blackwell GPUs, had been achieving impressive throughput — around 30 Ktok/s with a 4.2-day estimated time to completion. But the user noticed something concerning in the Weights & Biases logs: the loss curve was "jumpy," with periodic spikes, and accuracy was hovering around 0.03–0.05 even after hundreds of training steps. The assistant initially dismissed this as normal early-training noise, but the user persisted with a sharper concern: "won't we run into catastrophic forgetting loops if the model doesn't reinforce previous sample distribution?"

This question forced a re-examination of the batching logic. The assistant traced the issue to the build_batches method in train_dflash_pipeline.py. That method sorted all 902,000 training samples by sequence length and then greedily packed them into batches respecting a token budget. The result was that batch 0 always contained the shortest sequences (perhaps 100–200 tokens of simple QA pairs), while the final batch always contained the longest sequences (8000+ tokens of complex coding problems). The _feed_loop did shuffle the order of batches each epoch — so the model saw batches in a different sequence — but the composition of each batch was frozen. The same samples appeared together in every epoch.

This is a subtle but important flaw. When the optimizer processes a gradient update from a batch of uniformly short samples, it adjusts weights to perform well on short, simple inputs. The next batch might contain uniformly long, complex samples, pulling the weights in a different direction. Over the course of an epoch, the optimizer oscillates between these extremes rather than converging smoothly. And critically, across six epochs, the model never sees a diverse mix of sample lengths within a single gradient accumulation window. The user's intuition about "catastrophic forgetting loops" was not about forgetting entire tasks — it was about the model failing to reinforce a consistent input distribution, potentially leading to slower convergence or unstable training dynamics.

The Design Space: Three Options

The assistant had to navigate a constrained design space. Three broad approaches existed:

Option 1: Full random shuffle. Rebuild batches each epoch by randomly shuffling all samples and then packing them. This would give maximum gradient diversity — every batch would contain a mix of short and long sequences. But it would also destroy padding efficiency. Mixing a 200-token sample with an 8000-token sample means the shorter sample is padded to 8000 tokens, wasting 97.5% of its computation. The assistant had tested this approach and watched throughput collapse from ~30 Ktok/s to ~12 Ktok/s — a catastrophic 60% drop that would push the training ETA from 4 days to over 10 days.

Option 2: Rebuild batches each epoch with the same length-sorted approach. This would vary which samples appear together, since the random shuffle before sorting would produce different groupings. But wait — the build_batches method sorted by length after any shuffling, so shuffling before sorting produces the same result: samples are always grouped by length. This option is actually identical to the original approach.

Option 3: Bucketed shuffle. Group samples into length buckets (e.g., 0–770 tokens, 770–1216 tokens, etc.), shuffle samples within each bucket, then greedily pack into batches. This preserves the key efficiency insight — similar-length samples are batched together, minimizing padding waste — while ensuring that the specific composition of each batch changes every epoch. The optimizer sees different combinations of samples within each length range, breaking the frozen batch composition while maintaining throughput.

The assistant and user had already collaborated on Option 3 in the preceding chunk, running an analytical optimization script against the actual sequence length distribution of the 902K samples. That script determined the optimal six bucket boundaries: [0, 770, 1216, 1728, 2432, 3296, 8192], achieving an estimated ~87% padding efficiency.

The Implementation Decision

Message 8690 is the moment this plan becomes code. The assistant's edit replaces the single-call build_batches + shuffle-order approach with a per-epoch build_batches_shuffled that the _feed_loop calls fresh each epoch. The reasoning is precise:

Move batch construction into the epoch loop. Instead of calling build_batches once at startup and storing the result in self.all_batches, the _feed_loop now calls a new method each epoch. This ensures that the sample-to-batch assignment is regenerated, not reused.
Shuffle within length buckets before packing. The new build_batches_shuffled method first groups samples by length bucket, shuffles each bucket's indices, then concatenates and packs greedily. This is the key algorithmic change: samples of similar length still end up together (preserving efficiency), but which specific samples co-occur changes each epoch.
Preserve the batch-order shuffle. The existing random.shuffle(order) on line 319 remains, so the sequence of batches within an epoch is also randomized. This is a secondary randomization on top of the primary fix. The assistant does not show the full diff in the message — it simply states "Edit applied successfully" — but the reasoning is clear from the plan statement. The edit is surgical: it modifies only the batching logic, leaving the rest of the pipeline (the async CSP architecture, the GPU topology, the loss computation) untouched.

Assumptions and Trade-offs

The implementation makes several assumptions. First, that shuffling within length buckets provides sufficient gradient diversity. This is a weaker form of randomization than full shuffle — samples from different buckets never mix within a batch. The assistant implicitly assumes that the model benefits more from diverse combinations of similar-length samples than from mixing lengths directly. This is plausible because the gradient signal is dominated by the actual tokens, not the padding, and samples of similar length tend to have similar padding ratios.

Second, the assistant assumes that the analytical bucket boundaries remain optimal across all six epochs. In reality, the optimal boundaries depend on the sample length distribution, which is static for this dataset. So this assumption holds.

Third, the assistant assumes that the ~87% padding efficiency estimate translates to a similar throughput recovery. In practice, the final run achieved 25.1 Ktok/s — about 78% of the original sorted throughput, slightly below the estimate. The gap is attributed to the overhead of variable batch sizes on model execution (since bucketed shuffling produces batches with slightly different size distributions than strict sorting).

Input Knowledge Required

To understand this message, one needs knowledge of:

The DFlash training architecture, particularly the async CSP pipeline with decoupled drafter, target, and verifier GPUs
The build_batches method's length-sorting logic and its implications for batch composition
The concept of padding efficiency in transformer training — how mixing sequence lengths wastes computation on padding tokens
The analytical optimization approach used to determine bucket boundaries
The training run's throughput metrics (30 Ktok/s original, 12 Ktok/s with full shuffle, 25.1 Ktok/s target)

Output Knowledge Created

This message produces:

A concrete implementation of the bucketed shuffle strategy in the production training script
The knowledge that the fix was successfully applied and the training run was restarted
A documented design decision that balances gradient diversity against computational efficiency
A template for future batching strategies in similar training pipelines

The Thinking Process

The assistant's reasoning, visible in the preceding messages, shows a progression from initial dismissal ("This isn't actually a problem for convergence") to deep engagement with the user's concern. The key insight came when the assistant traced through the data flow: build_batches is called once, _feed_loop shuffles batch order but not composition, therefore the same samples always appear together. The assistant then systematically evaluated the three options, tested the full random shuffle empirically (observing the 12 Ktok/s collapse), and collaborated with the user on the analytical optimization of bucket boundaries.

Message 8690 is the culmination of this reasoning chain. It is not a speculative proposal or a discussion — it is an action. The edit is applied, the training run is stopped and restarted, and the new strategy is live. The message's brevity belies the depth of analysis that preceded it, and the precision of the fix reflects a clear understanding of both the computational constraints and the training dynamics at play.