When FSDP Doesn't Actually Parallelize: Diagnosing a 35-Hour EAGLE-3 Training Bottleneck

Introduction

In the middle of a complex EAGLE-3 speculative decoding training pipeline for a Kimi-K2.5 language model, a single monitoring message reveals a subtle but critical performance bug. The message, sent by an AI assistant monitoring a distributed training run on 4 NVIDIA RTX PRO 6000 Blackwell GPUs, uncovers that the training is projected to take approximately 35 hours—far longer than expected. Through careful analysis of step rates, batch counts, and GPU utilization metrics, the assistant discovers that the speculators library's training loop is not using a DistributedSampler, causing each GPU to redundantly process the entire dataset rather than splitting the work. This article examines this message in depth, exploring the reasoning, assumptions, mistakes, and insights it contains.

Context: The Road to Training

To understand this message, we must first understand what led to it. The conversation spans a massive ML engineering effort: deploying a GLM-5-NVFP4 / Kimi-K2.5 model on an 8-GPU server, developing a custom hidden state extraction patch for SGLang, extracting 37,312 training samples (87.8 million tokens, ~4.6 TB of data), and finally launching EAGLE-3 drafter training. The training itself was launched in message 4236 with the command:

torchrun --nproc_per_node=4 /tmp/04_train.py --verifier-path /shared/kimi-k2.5-int4 --data-dir /data/eagle3/synth_100k/hidden_states ... --epochs 5 --lr 3e-5 --max-seq-len 4096 --num-workers 4

The training script uses the speculators library, which implements the EAGLE-3 architecture—a speculative decoding framework where a small "draft" model predicts multiple tokens ahead that a large "verifier" model then validates. The drafter has approximately 2.6 billion total parameters, with 1.2 billion trainable (the embedding layer and verifier LM head are frozen).

By message 4242, the assistant has determined that there are 35,446 batches per epoch and 177,230 total steps across 5 epochs. The training is logging metrics at a rate of approximately 1.4 steps per second across all 4 GPUs combined. This is where the critical analysis begins.

The Message: A Deep Dive into Training Efficiency

The subject message (index 4243) opens with a precise calculation:

748 lines with epoch 0 / 4 ranks = ~187 actual steps done. At 1.4 steps/s and 35,446 steps/epoch, epoch 0 alone would take 35,446 / 1.4 = ~25,318s = ~7 hours. Total 5 epochs = ~35 hours.

This is the moment of discovery. The assistant has been watching the training metrics come in, and now it does the math to project completion time. The result—35 hours—is a red flag. The assistant immediately identifies the root cause:

That's because speculators isn't using a DistributedSampler — each GPU processes the FULL dataset. The FSDP shards the model (unnecessary for 1.2B), but doesn't split data. Each GPU redundantly processes all 35K files, 4x the work with no benefit.

This is the central insight of the message. FSDP (Fully Sharded Data Parallelism) is being used to shard the model parameters across 4 GPUs, but the data loader is not using a DistributedSampler. In PyTorch's distributed training paradigm, DistributedSampler is the mechanism that ensures each GPU sees a different subset of the data. Without it, every GPU iterates over the exact same files, computes the same gradients, and then averages them via allreduce—wasting 75% of the compute.

The Reasoning Process: A Window into Performance Debugging

What makes this message particularly valuable is the visible thinking process. The assistant doesn't just state the problem; it works through multiple hypotheses, corrects its own math, and refines its understanding in real time.

First hypothesis: The assistant initially considers whether switching to 1 GPU would be better. If FSDP isn't splitting data, then 4 GPUs are doing redundant work. Perhaps 1 GPU would be equally fast?

This is wasteful. For the 2.6B drafter, FSDP is overkill and the data-parallelism isn't actually happening. We'd be better off on 1 GPU which would take the same ~7 hours per epoch = ~35 hours total, or we need to fix the data loading.

Self-correction: The assistant then realizes its math is incomplete. It recalculates:

Actually, looking again at 1.4 steps/s — that's the GLOBAL rate (all 4 GPUs doing independent steps). If each GPU does ~0.35 steps/s, then 1 GPU at 0.35 steps/s × 35,446 steps = ~28 hours/epoch. So 4 GPUs IS faster (each does 35,446 steps at 0.35 steps/s = ~28h/epoch, but they're all doing it in parallel, so wall-clock per epoch = ~28h).

Wait—this second calculation seems contradictory. Let me re-read carefully. The assistant initially says 1.4 steps/s is the global rate. If each GPU does ~0.35 steps/s, then 1 GPU would take 35,446 / 0.35 = ~101,274 seconds = ~28 hours per epoch. But with 4 GPUs running in parallel at 1.4 steps/s combined, the wall time is 35,446 / 1.4 = ~7 hours per epoch. So 4 GPUs is actually 4x faster in wall time, but with 4x the effective batch size.

Refined understanding: The assistant then arrives at the correct interpretation:

Wait, no — if all 4 GPUs process the FULL dataset in parallel and then average gradients via FSDP allreduce, the effective training is still 35,446 steps/epoch with 4x effective batch size. The wall time per epoch = 35,446 / 1.4 = 25,318s ≈ 7h/epoch, 35h total.
With 1 GPU, it would be 35,446 steps at probably ~1.4 steps/s too (no allreduce overhead, similar compute) = same 7h/epoch but with 1x effective batch. So 4 GPUs gives us 4x effective batch size in the same time. That's legitimate scaling — but 35 hours is still long.

This is a nuanced point. The 4 GPUs are indeed providing value—they increase the effective batch size by 4x, which can improve training stability and convergence. But the 35-hour wall-clock time remains a concern.

Assumptions and Their Consequences

Several assumptions underpin this analysis, and the message reveals where they break down.

Assumption 1: FSDP implies data parallelism. The most significant incorrect assumption is that using torchrun with 4 GPUs and FSDP would automatically distribute the data across GPUs. In reality, FSDP only handles model sharding; data distribution requires explicit use of DistributedSampler in the DataLoader. The speculators library, as the assistant discovers, does not implement this. This is a common pitfall in distributed training—engineers often assume that distributed launchers handle data splitting automatically.

Assumption 2: The step rate would scale with GPUs. The assistant initially expects that 4 GPUs would process data 4x faster than 1 GPU. When it discovers each GPU processes the full dataset, it briefly considers 1 GPU as equally good. The refined analysis shows this assumption was partially wrong—4 GPUs do provide value through larger effective batch size, even if wall-clock time per epoch is similar.

Assumption 3: Data loading is the bottleneck. The assistant's final action in this message is to check disk I/O:

Let me check if we can speed it up. The bottleneck is likely data loading (reading .pt files from disk)

It runs iostat -x 1 3 to measure disk performance. This reveals another assumption: that the 35-hour training time is primarily limited by I/O rather than compute. The iostat output shows very low disk utilization (~0.04% iowait, single-digit rkB/s), suggesting that data loading is NOT the bottleneck—the GPUs are spending most of their time in computation or synchronization.

Input Knowledge Required

To fully understand this message, the reader needs familiarity with several concepts:

Distributed training fundamentals: Understanding of FSDP, DistributedSampler, data parallelism vs. model parallelism, and how torchrun coordinates multiple processes across GPUs.

Speculative decoding architecture: Knowledge of EAGLE-3, draft models, verifier models, and why a small drafter (1.2B trainable parameters) is trained separately from the large verifier.

Performance analysis metrics: Step rate (steps/s), batch size, epoch time calculation, GPU utilization percentages, power draw (watts), and memory usage (MiB).

The speculators library: Understanding that this is a third-party library implementing EAGLE-style architectures, and that its training loop may have limitations.

Linux system monitoring: Familiarity with iostat, nvidia-smi, grep, and command-line log analysis.

Output Knowledge Created

This message produces several valuable insights:

1. Identification of the DistributedSampler gap: The primary finding is that speculators does not use DistributedSampler, causing redundant data processing across GPUs. This is actionable knowledge—the training script could be patched to add a DistributedSampler, potentially reducing training time by 4x.

2. Training time projection: The assistant establishes that 5 epochs at 1.4 steps/s will take ~35 hours. This sets expectations for the user and informs decisions about whether to optimize or accept the timeline.

3. Effective batch size analysis: The message clarifies that 4 GPUs provide 4x effective batch size even without data splitting, which is a legitimate scaling benefit for training quality.

4. Bottleneck investigation methodology: The assistant demonstrates a systematic approach to performance debugging: measure step rate, calculate projected time, identify the root cause (missing DistributedSampler), verify understanding through recalculations, and then investigate secondary bottlenecks (disk I/O).

5. I/O characterization: The iostat output shows near-zero disk activity, ruling out data loading as the primary bottleneck. This shifts the optimization focus to either compute efficiency or the DistributedSampler fix.

The Thinking Process: A Masterclass in Diagnostic Reasoning

The most striking aspect of this message is the visible chain of reasoning. The assistant works through the problem in real time, with all the false starts and self-corrections preserved. This is not a polished final answer; it's a live debugging session.

The reasoning follows a clear pattern:

  1. Observe: "748 lines with epoch 0 / 4 ranks = ~187 actual steps done"
  2. Calculate: "At 1.4 steps/s and 35,446 steps/epoch... ~7 hours per epoch"
  3. Project: "Total 5 epochs = ~35 hours"
  4. Diagnose: "speculators isn't using a DistributedSampler"
  5. Evaluate alternatives: "We'd be better off on 1 GPU"
  6. Self-correct: "Actually, looking again... 4 GPUs IS faster"
  7. Refine: "4 GPUs gives us 4x effective batch size in the same time"
  8. Investigate further: "Let me check if we can speed it up... bottleneck is likely data loading"
  9. Gather data: Runs iostat This pattern—observe, calculate, project, diagnose, evaluate, self-correct, refine, investigate—is a textbook example of systematic performance debugging. The self-correction step is particularly important. The assistant initially concludes that 4 GPUs provide no benefit over 1 GPU, then realizes that the effective batch size scaling is valuable even if wall-clock time per epoch is similar. This kind of nuanced understanding separates surface-level analysis from deep insight.

Broader Implications

This message illustrates a recurring challenge in large-scale ML engineering: the gap between infrastructure assumptions and library behavior. The assistant correctly set up torchrun with 4 GPUs and FSDP, which should have provided both model sharding and data parallelism. But the speculators library's training loop didn't implement the data parallelism side, silently wasting 75% of the compute.

This is not a bug in the traditional sense—the training works correctly, converges properly, and produces valid gradients. It's an efficiency issue that only becomes visible when someone carefully monitors step rates and projects completion times. Many ML engineers have experienced this exact scenario: a training run that works but takes much longer than expected, only to discover that the data loader isn't distributed.

The message also highlights the importance of monitoring at multiple levels. The assistant checks step rates (application-level metrics), GPU utilization and power draw (hardware-level metrics), and disk I/O (system-level metrics). Each level provides a different piece of the puzzle. The step rate reveals the overall throughput problem. The GPU utilization (100%, ~250W each from earlier messages) shows the GPUs are fully occupied—they're not waiting on data. The iostat confirms that disk I/O is negligible, ruling out the data loading hypothesis.

Conclusion

Message 4243 captures a pivotal moment in a complex ML training pipeline: the discovery that a distributed training run is not actually distributing its data. Through careful analysis of step rates, batch counts, and GPU utilization, the assistant identifies that the speculators library's training loop lacks a DistributedSampler, causing each of the 4 GPUs to redundantly process the full 35,446-batch dataset. The projected 35-hour training time becomes a call to action—either accept the timeline, patch the library to add proper data parallelism, or investigate other optimizations.

The message is a masterclass in performance debugging methodology. It demonstrates how to combine application-level metrics (step rate, batch count), system-level metrics (disk I/O, GPU utilization), and domain knowledge (FSDP, DistributedSampler, speculative decoding) to diagnose training inefficiency. The visible reasoning process, complete with false starts and self-corrections, provides an authentic window into how experienced ML engineers think about distributed training performance.

For anyone building large-scale ML training pipelines, this message serves as both a cautionary tale and a diagnostic template. The caution: never assume that distributed launchers automatically distribute data—always verify that your data loader uses DistributedSampler or equivalent. The template: when training seems slow, measure step rates, project completion times, identify the bottleneck layer by layer, and verify your assumptions at each level of the system.