The Per-Device Compilation Fix: Diagnosing a Multi-GPU torch.compile Cache Conflict

[assistant] [bash] cd /data/dflash/scripts && git add -A && git commit -m "fix: per-device flex_attention compile for multi-GPU drafter

>

torch.compile caches are device-specific. Sharing a single compiled flex_attention across GPUs in multi-drafter training causes: 'RuntimeError: Detected that you are using FX to symbolically trace a dynamo-optimized function'

>

Fix: compile per device key with thread-safe lazy init." 2>&1 [experiment-ddtree 9e45dc1] fix: per-device flex_attention compile for multi-GPU drafter 1 file changed, 15 insertions(+), 10 deletions(-)

In distributed deep learning, the most insidious bugs are those that only manifest under concurrency—race conditions, cache invalidation, and device-specific state leakage that single-GPU development never exposes. Message [msg 9333] captures exactly such a moment: a single git commit that fixes a runtime crash in a multi-GPU speculative decoding training pipeline, caused by a shared torch.compile cache being reused across different GPU devices. The commit message itself tells the story concisely:

fix: per-device flex_attention compile for multi-GPU drafter

>

torch.compile caches are device-specific. Sharing a single compiled flex_attention across GPUs in multi-drafter training causes: 'RuntimeError: Detected that you are using FX to symbolically trace a dynamo-optimized function'

>

Fix: compile per device key with thread-safe lazy init.

But behind this terse summary lies a deep diagnostic journey—one that reveals fundamental assumptions about how PyTorch's torch.compile interacts with multi-GPU training, and the kind of thinking required to debug distributed systems where the same code path executes simultaneously on different devices.

The Context: Scaling the Drafter to Two GPUs

The broader project involves training a DFlash speculative decoding drafter—a small language model that learns to predict the next several tokens a large target model would generate, enabling faster inference through speculative execution. The training pipeline uses a novel architecture with 6 target GPUs (producing hidden states) feeding into a drafter model that learns to predict those hidden states. Initially, the drafter ran on a single GPU (GPU 7), achieving 6.5 Ktok/s throughput with a projected 14-day training time.

The user requested distributing the drafter across 2 GPUs to double throughput ([msg 9315]). The assistant implemented this by using both GPU 6 and GPU 7 as independent drafter instances, each consuming from 3 target GPUs in a round-robin fashion, with weight averaging every 50 steps (<msg id=9316-9328>). The updated code was deployed and the training run restarted.

The Crash: A Cryptic Runtime Error

When the assistant checked the training output after deployment ([msg 9329]), it found a crash:

RuntimeError: Detected that you are using FX to symbolically trace
a dynamo-optimized function. This is not supported at the moment.

The stack trace pointed to the flex_attention call inside the drafter model's forward pass. Flex attention is a custom attention mechanism used in the DFlash architecture, and the code had been wrapping it with torch.compile(flex_attention) for performance. The error occurred during step 8 of training—early enough to suggest a systematic initialization issue rather than a data-dependent edge case.

The Diagnostic Process: Three False Starts

The assistant's reasoning in [msg 9330] reveals a meticulous, iterative debugging process that cycled through three distinct hypotheses before landing on the correct diagnosis.

Hypothesis 1: Gradient checkpointing tracing through compiled code. The assistant first suspected that torch.utils.checkpoint.checkpoint(use_reentrant=False) was using FX tracing internally, which conflicted with the already-compiled flex_attention. This is a plausible explanation: gradient checkpointing saves memory by recomputing intermediate activations during the backward pass, and the reentrant=False variant uses FX tracing to capture the computation graph. If that tracing encountered a function already compiled by torch.compile, the two compilation systems could collide.

Hypothesis 2: Simultaneous compilation on two GPUs. The assistant then reconsidered: the gradient checkpoint wraps only the _chunk_fwd function (containing lm_head and loss computation), not the flex attention layers. The error occurred during the forward pass itself, not the backward. This led to a second hypothesis: with two drafter GPUs running simultaneously, the singleton _compiled_flex_attention was being compiled on both GPUs at the same time, causing a device-level conflict.

Hypothesis 3: Thread-safety and GPU-specific caching. The assistant refined further, noting that torch.compile caches GPU-specific state. When a thread running on GPU 6 calls a function that was compiled on GPU 7, the cached kernel may reference device-specific resources (e.g., CUDA streams, workspace allocations) that are invalid on the other device. The error message about "FX to symbolically trace a dynamo-optimized function" was a red herring—the actual root cause was a device mismatch in the compiled kernel cache.

The assistant initially considered the simplest fix: just disable torch.compile for flex_attention during training, since compilation primarily benefits inference throughput. But then arrived at a more elegant solution: compile per device with a thread-safe lazy initialization cache.

The Fix: Per-Device Compilation with Thread-Safe Lazy Init

The commit in [msg 9333] implements exactly this. Instead of a single global _compiled_flex_attention singleton, the code now maintains a dictionary keyed by device (e.g., "cuda:6", "cuda:7"), with lazy initialization protected against concurrent access. Each GPU gets its own compiled version of flex_attention, compiled on its own device, eliminating the cross-device cache conflict.

The fix required two changes: modifying the _get_compiled_flex_attention() function to accept and use a device key, and updating the caller to pass the device. The syntax check in [msg 9332] confirmed the code compiled cleanly.

Assumptions and Lessons Learned

Several assumptions were implicit in the original code that proved incorrect:

1. torch.compile produces device-agnostic kernels. The original code assumed a single compiled function could be shared across devices. In reality, torch.compile may generate device-specific code or cache device-specific state, making the compiled function non-portable across GPUs.
2. Singleton patterns are safe in multi-threaded training. The original _compiled_flex_attention was a module-level singleton, assuming single-device usage. Multi-GPU training with independent threads broke this assumption.
3. FX tracing and torch.compile compose cleanly. The error message suggested a conflict between two PyTorch compilation subsystems, but this was a symptom of the device mismatch rather than the root cause. The key insight—that torch.compile caches are device-specific—is not widely documented and represents a subtle but critical piece of knowledge for anyone building multi-GPU training pipelines with compiled models.

Output Knowledge Created

This message creates several pieces of valuable output knowledge:

• A reproducible fix pattern for multi-GPU torch.compile conflicts: per-device compilation with thread-safe lazy initialization.
• A diagnostic heuristic for FX/dynamo tracing errors: when the error occurs across devices, suspect device-specific compilation caches before assuming a tracing incompatibility.
• A version-controlled record (commit 9e45dc1 on branch experiment-ddtree) of the exact code change, enabling rollback and review. The fix enabled the 2-GPU drafter training to proceed, doubling throughput from 6.5 Ktok/s to an expected ~13 Ktok/s and halving the training ETA from 14 days to approximately 7 days. More importantly, it established a pattern for scaling to even more drafter GPUs in the future—each additional GPU would simply get its own compiled attention kernel from the device-keyed cache.

The Broader Significance

This message exemplifies a class of bugs that are becoming increasingly common as distributed training pipelines grow in complexity. The interaction between torch.compile (a graph-level optimization), torch.utils.checkpoint (a memory-saving technique), and multi-GPU parallelism creates a three-way tension that PyTorch's compilation stack is still evolving to handle gracefully. The assistant's diagnostic process—cycling through hypotheses, checking the stack trace against code structure, and ultimately identifying a device-scoping issue—is a model for debugging in this space.

The fix itself is elegant in its simplicity: a dictionary keyed by device string, with lazy initialization. No changes to the model architecture, no disabling of compilation entirely, no complex synchronization primitives. Just the recognition that a compiled function is not a pure mathematical object but a device-bound computational artifact, and that distributed training must respect these physical boundaries.