The Two-GPU Drafter: Halving Training Time Through Distributed Speculative Decoding

Introduction

In the high-stakes world of large language model training, every day of compute time represents a significant cost in both resources and opportunity. When the assistant at message [msg 9336] reports "Still loading models. The 2 drafter models take extra time. Let me wait longer," it marks a pivotal moment in a multi-day journey to scale speculative decoding training across multiple GPUs. This seemingly mundane status check—a bash command followed by training logs—is actually the culmination of a complex debugging saga, a validation of distributed training infrastructure, and the beginning of a new phase in the project's evolution. The message captures the moment when a carefully engineered 2-GPU distributed drafter pipeline springs to life, delivering the promised doubling of throughput from 6.5 Ktok/s to 13.1 Ktok/s and cutting the estimated training time from fourteen days to just under seven.

Context: The DDTree Experiment and the Single-GPU Bottleneck

To understand the significance of this message, we must trace the arc of the preceding work. The project had been developing a DFlash drafter—a lightweight autoregressive model trained to predict the next several tokens in parallel, enabling speculative decoding to accelerate inference of a large target model (Qwen3.6-27B). The training pipeline was sophisticated: six target GPUs (indices 0–5) ran the large model to produce hidden states, which were fed into a queue consumed by a single drafter GPU (index 7). This architecture, while functional, had a critical bottleneck: the drafter GPU was the limiting factor.

The DDTree experiment (experiment-ddtree) introduced several architectural innovations over the v6 baseline: a larger block size of 32 (versus 16), more anchors at 1024 (versus 512), a higher gamma of 10.0 (versus 4.0) to weight deeper positions in the tree, sliding window attention on the first four layers, uniform noise scheduling, a 15% soft KL distillation loss blended with cross-entropy, and a CAP auxiliary confidence loss. These changes quadrupled the number of training tokens per batch from 8,192 to 32,768, but they also introduced a severe throughput penalty. The single drafter GPU could only process 6.5 Ktok/s, yielding a 14-day estimated time to completion for six epochs.

The user's directive was clear: "Can we distribute training to 2 GPUs?" This set in motion a chain of infrastructure work documented in the preceding messages ([msg 9316] through [msg 9335]), including implementing weight averaging across drafters (replacing a one-way copy that discarded half the gradient updates), committing the changes, deploying the updated code to the remote training machine, and—most critically—fixing a subtle torch.compile conflict that crashed multi-GPU training.

The Message: A Status Check That Tells a Story

The message itself is deceptively simple. The assistant writes:

Still loading models. The 2 drafter models take extra time. Let me wait longer:

This is followed by a bash command that sleeps for 300 seconds (five minutes) and then captures the last twelve lines of the training tmux pane. The output reveals the training is alive and running:

[5m] step=11 loss=12.2268 acc=0.011 streak=0.0 lr=5.14e-06 noise=0.0004 | tgt=0.
47b/s dft=0.32b/s (12.9Ktok/s) | q_pre=[50, 50, 50, 50, 50, 50] q_hs=[20, 20] |
epoch~0.01 ETA=6.9d
[5m] step=12 loss=11.7024 acc=0.014 streak=0.0 lr=5.57e-06 noise=0.0004 | tgt=0.
47b/s dft=0.32b/s (13.0Ktok/s) | q_pre=[50, 50, 50, 50, 50, 50] q_hs=[20, 20] |
epoch~0.01 ETA=6.9d
[5m] step=12 loss=11.7161 acc=0.008 streak=0.0 lr=5.57e-06 noise=0.0004 | tgt=0.
47b/s dft=0.32b/s (13.1Ktok/s) | q_pre=[50, 50, 50, 50, 5...

The output is truncated, but the critical information is visible. The training is running at approximately 12.9–13.1 Ktok/s—essentially double the 6.5 Ktok/s achieved with a single drafter GPU. The ETA has dropped from ~14 days to ~6.9 days. The hidden state queues (q_hs=[20, 20]) are both full at depth 20, meaning both drafter GPUs are consuming hidden states as fast as the six target GPUs can produce them. The prefetch queues (q_pre) are all at 50, indicating the data loading pipeline is keeping up. The system is balanced.

What the Training Logs Reveal

Beyond the throughput numbers, the logs reveal the early-stage dynamics of the DDTree experiment. At step 11–12, the loss is around 11.7–12.2, the accuracy is extremely low (0.008–0.014), and the streak metric (average number of correct tokens in a row) is 0.0. This is expected behavior at the very beginning of training—the learning rate is still in its warmup phase at 5.14e-06 to 5.57e-06, far below the target of 6e-4. The noise parameter is at 0.0004, which is the uniform noise being annealed from its initial value of 0.05 toward 0.01. The model has barely begun to learn.

The "epoch~0.01" indicator shows that less than 1% of the first epoch has been completed. With 46,669 batches per epoch and each batch containing 32,768 block tokens, the model is seeing a massive amount of training signal per step—31,744 predictions per batch (1024 anchors × 31 positions each, since the first position is the anchor itself). This is 4× more training signal per step than the v6 baseline, which used 512 anchors × 15 positions = 7,680 predictions per batch.

The Infrastructure Behind the Success

The successful deployment of the 2-GPU drafter pipeline required solving several non-trivial engineering challenges. The first was the weight synchronization mechanism. The original implementation used a one-way copy from drafter 0 to all other drafters every 100 steps, which discarded the gradient updates accumulated by the secondary drafters. The assistant recognized this flaw in [msg 9316] and implemented proper weight averaging: all drafters contribute equally to the averaged weights, preserving the information from their independent gradient trajectories. The sync frequency was also increased from every 100 steps to every 50 steps for tighter convergence.

The second and more vexing challenge was the torch.compile conflict. The flex_attention implementation used torch.compile to optimize the causal masking computation. When two drafter GPUs ran simultaneously in separate threads, the compiled kernel cache—which is device-specific—caused a RuntimeError: Detected that you are using FX to symbolically trace a dynamo-optimized function. The fix, implemented in [msg 9330] through [msg 9333], was to maintain a per-device cache of compiled flex_attention functions with thread-safe lazy initialization. Each GPU gets its own compiled version, eliminating the cross-device tracing conflict.

Assumptions and Decisions

Several key assumptions underpin this message. The assistant assumes that the training logs visible in the tmux pane are representative of the full training state and that no silent errors are occurring on the other drafter GPU. The decision to wait 300 seconds before checking reflects an assumption that model loading and initial warmup would complete within that window—an assumption validated by the output. The assistant also assumes that the weight averaging strategy (every 50 steps) is sufficient for convergence, accepting the known limitation that optimizer states (AdamW momentum and variance) are not averaged, only the model weights.

The choice to use two independent drafter models with periodic weight averaging rather than true data parallelism (splitting a single batch across two GPUs with gradient accumulation) reflects a pragmatic engineering decision. The independent-drafter approach is simpler to implement in the existing pipeline architecture, avoids the complexity of gradient synchronization across devices, and naturally handles the asynchronous nature of the hidden state queue consumption.

Significance and Implications

This message represents a successful infrastructure milestone. The 2-GPU drafter deployment validates the distributed training architecture and demonstrates that the DDTree experiment can complete in a reasonable timeframe. The throughput of 13.1 Ktok/s, while still far below the v6 baseline of 26 Ktok/s, represents a 2× improvement over the single-drafter configuration. More importantly, the full hidden state queues indicate that the system is not yet bottlenecked—there is headroom for further scaling if additional drafter GPUs were added.

The message also reveals the early training dynamics of the DDTree experiment. The low accuracy and zero streak at step 12 are not causes for concern but rather expected behavior at the start of a difficult learning task. The model must learn to predict up to 31 tokens ahead across a vocabulary of 248,000 tokens, with the added complexity of the CAP loss and soft KL distillation. The real test will come hundreds of thousands of steps later, when the accuracy and streak metrics begin to climb.

For the broader project, this message marks the transition from infrastructure debugging to sustained training. The assistant's next steps will involve monitoring the training trajectory, comparing against the v6 baseline, and eventually evaluating the trained drafter's performance in actual speculative decoding. The 6.9-day ETA provides a concrete timeline for when those evaluations can begin.