The Moment Compilation Succeeds but Training Fails: Diagnosing a Multi-Threaded CUDA Graph Race

In the high-stakes world of custom multi-GPU training pipelines, few moments are more deceptive than a successful smoke test followed by a production crash. Message [msg 10350] captures exactly such a moment: the assistant has just deployed a carefully engineered fix—padding all inputs to fixed shapes and enabling torch.compile(mode="reduce-overhead", dynamic=False) on the DFlash drafter model—and is now checking whether the full multi-threaded training run has stabilized. The results are simultaneously encouraging and devastating.

The Message in Full

The assistant executes a bash command after a 240-second sleep, checking the log file of the freshly launched training run:

sleep 240 && ssh -o ConnectTimeout=10 root@10.1.2.6 "pct exec 200 -- /bin/bash -c \
  'grep -c Exception /workspace/train_cudagraph.log; \
   grep -E \"compiled|tok/s|step=\" /workspace/train_cudagraph.log | tail -20; \
   echo ===; \
   nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader; \
   echo ===CACHE===; \
   du -sh /tmp/torchinductor_root 2>/dev/null'"

The output reveals a deeply ambiguous picture:

6
  Drafter 0: forward compiled
  Drafter 1: forward compiled
  Drafter 2: forward compiled
[0m] step=0 loss=--- acc=--- streak=--- lr=--- noise=0.0000 | tgt=0.00b/s dft=0.00b/s (0.0Ktok/s) | q_pre=[21] q_hs=[0] q_hsb=[0, 0, 0, 0, 0, 0] | epoch~0.00 ETA=???
[0m] step=0 loss=--- acc=--- streak=--- lr=--- noise=0.0000 | tgt=0.20b/s dft=0.00b/s (0.0Ktok/s) | q_pre=[34] q_hs=[4] q_hsb=[0, 0, 0, 0, 1, 3] | epoch~0.00 ETA=20.8d
[1m] step=0 loss=--- acc=--- streak=--- lr=--- noise=0.0000 | tgt=0.17b/s ...

Six exceptions have been logged. Yet all three drafter threads report "forward compiled"—the CUDA graphs were successfully captured. The training metrics, however, tell a different story: throughput is stuck at 0.0Ktok/s, the estimated time to completion is ???, and the queue depths (q_pre, q_hs, q_hsb) are nearly empty, indicating that the pipeline has stalled before producing a single meaningful training step.

The Context: A Long Road to Fixed-Shape Compilation

To understand why this message matters, one must appreciate the engineering effort that preceded it. The training pipeline under development is a custom DFlash (Drafting with Flash Attention) system that trains a small "drafter" model to predict hidden states from a much larger "target" model across 8 GPUs. The pipeline uses a producer-consumer architecture: target model threads run forward passes on GPUs 0–5, producing hidden states that are consumed by drafter threads on GPUs 6–7.

Earlier in the session (see [msg 10341]), the assistant had diagnosed that the training was running at only ~12K tok/s with volatile GPU memory utilization. The root cause was a fundamental architectural mismatch: the single-process, multi-threaded pipeline produced variable-length sequences, which prevented CUDA graph replay, caused allocator churn, and created GIL contention across 12+ threads.

The assistant's response was a comprehensive redesign toward fixed-shape inputs. Over several messages ([msg 10342] through [msg 10348]), the assistant:

1. Padded all hidden state batches to the token_budget (49152 tokens), ensuring every drafter forward pass receives tensors of identical shape.
2. Preallocated persistent GPU buffers so that drafter input copies no longer allocate new memory each batch.
3. Replaced dynamic operations like nonzero and randperm in anchor selection with fixed-shape equivalents using rand + topk.
4. Replaced dynamic repeat_interleave document-id construction with fixed-shape vectorized masks.
5. Added torch.compiler.cudagraph_mark_step_begin() calls to demarcate CUDA graph capture boundaries. The smoke test in [msg 10348] was a triumph: the compiled forward+backward ran in 34 seconds on the first iteration (compilation) and 3.6 seconds on the second (replay), with stable peak memory of ~49 GB. The assistant, confident in the fix, cleared the inductor cache and launched the full training run in [msg 10349].

What the Message Reveals: Compilation vs. Execution

Message [msg 10350] is the first reality check after that launch. The output contains a critical tension: the compilation succeeded, but the training did not.

The line "Drafter 0: forward compiled, Drafter 1: forward compiled, Drafter 2: forward compiled" confirms that torch.compile ran to completion in all three drafter worker threads. This is non-trivial—the assistant had previously battled a multi-threaded FX tracing race condition (see [chunk 56.0]) where concurrent compilation attempts in different threads would corrupt each other's tracing state. The fact that all three threads compiled suggests that the per-thread execution lock and use_reentrant=False gradient checkpoint fix had resolved that particular issue.

Yet the training metrics show a pipeline that has essentially halted. The tgt=0.00b/s dft=0.00b/s (0.0Ktok/s) indicates that neither target nor drafter threads are producing work. The queue depths are near zero—q_pre=[21] (only 21 prefetch batches ready), q_hs=[0] (zero hidden state batches in the queue), and all six drafter sub-queues empty. In a healthy pipeline, these queues should be deep, with target threads continuously producing hidden states and drafter threads consuming them.

The six exceptions are the smoking gun. The assistant does not display their full tracebacks in this message—that investigation comes in the following message ([msg 10351])—but their mere presence signals that the training loop crashed before completing a single step.

The Assumptions That Proved Wrong

The assistant made a reasonable but ultimately incorrect assumption: that a successful single-threaded compilation smoke test would generalize to the multi-threaded pipeline. The smoke test in [msg 10348] ran on a single GPU (cuda:7) with a single thread, sequentially calling model.forward() twice. The full pipeline, by contrast, spawns multiple Python threads that each call torch.compile on their own drafter model instance, while simultaneously target threads are launching CUDA kernels on other GPUs.

The critical oversight was timing. In the smoke test, compilation happened before any other CUDA work was in flight. In the full pipeline, the target threads begin processing batches immediately, launching CUDA kernels on GPUs 0–5, while the drafter threads are still compiling. As the assistant would later realize in [msg 10352]: "CUDA graph capture happened lazily inside drafter threads while target threads were already launching CUDA work in the same process. That is not safe: capture has to happen before the async pipeline starts."

This is a subtle but important constraint of PyTorch's CUDA graph capture: the process of recording CUDA operations into a graph requires exclusive access to the CUDA stream state. If other threads are concurrently launching kernels (even on different GPUs), the CUDA runtime's internal state can become inconsistent, leading to assertion failures, segfaults, or silent corruption.

The Input Knowledge Required

To fully understand this message, one needs familiarity with several layers of the PyTorch ecosystem:

• torch.compile and Inductor: PyTorch's compiler infrastructure that traces Python operations into an intermediate representation (FX graph), then compiles them to efficient CUDA kernels via Triton or CUDA graphs. The mode="reduce-overhead" flag enables CUDA graph capture, which records the sequence of GPU kernel launches into a replayable graph for maximum throughput.
• CUDA Graphs (CUDAGraph Trees): A CUDA runtime feature that captures a sequence of GPU operations into a graph object that can be replayed with minimal CPU overhead. In PyTorch, this is exposed through torch.cuda.CUDAGraph and used internally by torch.compile in reduce-overhead mode.
• Multi-threaded PyTorch: PyTorch's Python threading model allows multiple threads to use different GPUs within the same process, but the CUDA runtime and PyTorch's internal caches (including the inductor compilation cache) have complex thread-safety requirements.
• FX Tracing Race Condition: When multiple threads attempt to trace Python code into an FX graph simultaneously, they can interfere with each other's tracing state, causing corrupted graphs or crashes. This was the subject of extensive debugging in earlier chunks ([chunk 56.0]).

The Output Knowledge Created

This message produces several critical pieces of knowledge:

1. Compilation is no longer the bottleneck: All three drafter threads successfully compiled their forward passes, proving that the per-thread execution lock and use_reentrant=False fixes were effective against the FX tracing race.
2. The crash occurs post-compilation: Since the compilation succeeded but the pipeline stalled, the failure mode must be in the interaction between compiled execution and the concurrent target threads, not in the compilation itself.
3. The pipeline topology is the suspect: The queue depths being near-zero suggests that either the target threads crashed (stopping production of hidden states) or the drafter threads crashed during their first compiled forward pass (stopping consumption). The six exceptions point to the latter.
4. CUDA graph capture has a thread-safety dimension beyond compilation: Even after the FX tracing race is resolved, the actual replay of captured graphs in a multi-threaded context introduces new failure modes. This is a deeper and more fundamental constraint than the assistant initially appreciated.

The Thinking Process Visible

The assistant's reasoning is implicit in the structure of the diagnostic command. The choice of what to check reveals a sophisticated mental model of the failure modes:

• Check for exceptions first (grep -c Exception): The assistant knows the run may have crashed and wants a quick binary signal. Six exceptions is a clear "something went wrong" indicator.
• Check compilation status (grep -E "compiled|..."): The assistant needs to distinguish between "compilation failed" and "compilation succeeded but execution failed." The output showing all three drafter threads compiled narrows the problem space dramatically.
• Check training metrics (step=, tok/s): The assistant looks for evidence of actual training progress. Zero throughput confirms the pipeline is not producing useful work.
• Check GPU state (nvidia-smi): Though not shown in the truncated output, the command includes GPU memory and utilization checks, which would reveal whether GPUs are idle (confirming a stall) or busy (suggesting a hang or infinite loop).
• Check compile cache (du -sh /tmp/torchinductor_root): The assistant checks whether the inductor cache is growing, which would indicate ongoing compilation attempts rather than cached replay. This diagnostic structure—checking for errors, then narrowing to compilation vs. execution, then checking pipeline health, then checking hardware state—is a textbook example of systematic debugging in distributed ML systems.

The Broader Significance

Message [msg 10350] sits at a pivotal moment in the session. It represents the boundary between two phases of debugging: the first phase focused on making torch.compile work in a multi-threaded context (solved by the per-thread lock and fixed-shape inputs), and the second phase focused on making compiled execution safe alongside concurrent CUDA work (solved by sequential graph warmup before thread start, as implemented in [msg 10352]).

The message also illustrates a recurring theme in ML engineering: the gap between "works in isolation" and "works in production." The smoke test was a necessary but insufficient validation. The real system introduced emergent complexity—thread concurrency, queue timing, GPU stream synchronization—that the simplified test could not capture. Every layer of the stack, from Python threading to the CUDA runtime to PyTorch's inductor compiler, introduced potential failure modes, and the assistant had to iterate through them one by one.

For the reader, this message is a case study in the importance of production-like testing environments, the subtle dangers of CUDA graph capture in multi-threaded processes, and the value of structured diagnostic commands that can quickly distinguish between different failure classes. It is also a reminder that in ML systems engineering, a successful compilation is not the same as a successful training run—and the distance between them can be measured in exceptions.