Fix: torch.compile errors in threaded multi-GPU runs

[carshadi]

Resolves #20

Summary

The inference pipeline (run.py) parallelizes across GPUs by launching one GpuWorker thread per device, each holding a deepcopy of the model and optionally wrapping it with torch.compile. This combination — threads × multiple GPUs × torch.compile — exposed several failure modes, ranging from hard crashes to a model that simply cannot be compiled.

This branch:

1. Eliminates the torch.compile recompile/threading crash by keeping model input shapes constant and warming the compiled model synchronously.
2. Downgrades CUDA-graph compile modes that are unsafe under threaded multi-GPU execution.
3. Falls back to eager execution (instead of aborting the run) when a model cannot be compiled.
4. Reworks the dynamic-shape defaults around torch.compile.

---

Problems addressed and how they're fixed

1. Crash: FX symbolic trace of a dynamo-optimized function

Running the multi-GPU denoising pipeline with torch.compile and static shapes produced runtime recompile warnings followed by a hard crash in one of the GPU threads:

W ... torch._dynamo hit config.recompile_limit (8)
W ...    function: 'forward' (.../unet3d.py:76)
W ...    last reason: 0/1: tensor 'x' size mismatch at index 0. expected 32, actual 4

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

Root cause. With dynamic shapes disabled, every distinct input shape forces a fresh compilation. Because each block's patch count rarely divides evenly by batch_size, the last batch of every block is smaller (expected 32, actual 4), so dynamo recompiled at runtime — and once past recompile_limit (8) it fell back to eager for new shapes. Critically, those recompiles happened inside the concurrent gpu-* threads. Dynamo's compile path runs an FX/proxy-tensor trace that monkeypatches torch.nn.Module.__call__ process-wide, so while one thread was mid-compile, another thread's ordinary forward call was intercepted by that tracer — producing the FX symbolically trace a dynamo-optimized function" error, killing the thread, and stalling the blocks whose batches it was holding.

Fixes (complementary):

• Constant input shapes via tail-batch padding. PrepWorker now pads the final (short) batch of each block up to batch_size with zeros when compiling, so the model only ever sees one input shape. With shapes constant, no runtime recompiles happen at all, which removes the recompile stalls and the concurrent-compile race. Padded rows are ignored downstream: WriterWorker only reads rows present in starts_in_block.
• Synchronous warmup on the main thread. GpuWorker._warmup_compiled_model runs one forward pass at the canonical shape during construction (the workers are built sequentially on the main thread). This forces the initial compilation to happen serially rather than racing lazily across worker threads on their first batch.

Together these guarantee that by the time the worker threads start, every model is already compiled, and no further compilation is ever triggered.

2. CUDA-graph compile modes are unsafe under threaded multi-GPU

reduce-overhead and max-autotune enable CUDA graphs, which are fragile when capture/replay happens from multiple threads, and which return outputs backed by static buffers that the next replay overwrites — hazardous given that GpuWorker.run hands the output off to an asynchronous D2H copy on a separate stream while the next batch begins.

Fix. InferenceConfig now downgrades CUDA-graph-enabling modes to their non-graph equivalents when more than one CUDA device is configured, with a RuntimeWarning:

• reduce-overhead → default
• max-autotune → max-autotune-no-cudagraphs

3. Crash: a model torch.compile cannot trace (proteomics MAE encoder)

The proteomics pipeline uses an MAE ViT encoder with RoPE positional embeddings. Compiling it broke first with a (non-fatal) graph break and then a fatal guard assertion:

W ... Graph break from `Tensor.item()` ...
  File ".../rope_embedding.py", line 649, in forward
    float(self.resolution_z),

AssertionError: Guard failed on the same frame it was created. This is a bug - please create an issue. Guard fail reason: 1/0: tensor '___from_numpy(self.grid_size)' dispatch key set mismatch. expected DispatchKeySet(CPU, BackendSelect, ADInplaceOrView), actual DispatchKeySet(CPU, BackendSelect)

Root cause. The encoder does host-side numpy/Python work inside forward that dynamo cannot trace as tensor ops: float(self.resolution_z) converts a tensor to a Python scalar (the graph break), and self.grid_size is a numpy array (mae_encoder.py:80) that gets indexed and fed into the RoPE module. Dynamo auto-wraps the numpy array via torch.from_numpy(...) and installs a guard on it; that guard trips a known dynamo bug (dispatch-key-set mismatch on
the same frame that created it), aborting compilation. This is intrinsic to the model and independent of the threading/shape fixes above — the model is simply not traceable as written.

Fix. GpuWorker._compile_model now wraps both the torch.compile() call and the warmup forward (where lazy compilation actually executes) in a try/except. On any failure it logs a warning and restores the original eager
module, so the run proceeds in eager mode instead of crashing:

WARNING - torch.compile failed on cuda:0 (AssertionError); falling back to
eager execution.

Because warmup runs on the main thread, the failure and fallback both happen there, before any worker thread executes — so a failed compile never reintroduces the threading hazard.

---

Dynamic-shape defaults

The handling of torch.compile(dynamic=...) was reworked:

• Default changed from dynamic=True to dynamic=None (automatic). With inputs now shape-constant, automatic dynamic compiles a single static, shape-specialized graph at warmup (fastest kernels) and would only promote to a dynamic graph if an unexpected shape ever appeared — the most graceful degradation of the three options.
• compile_dynamic is now Optional[bool] (None = auto / True = force dynamic / False = recompile per shape).
• Default compile_mode changed from reduce-overhead to default (the former silently enabled the unsafe CUDA-graph path under multi-GPU).

---

Behavioral / breaking changes

• CLI (breaking): --no-compile-dynamic is removed. Use
  --compile-dynamic false for the old behavior, or the new default
  --compile-dynamic auto. Any caller still passing --no-compile-dynamic
  will get an argparse error.
• Default compile_mode is now default (was reduce-overhead).
• Default compile_dynamic is now auto/None (was True).
• A model that fails to compile no longer aborts the run; it logs a warning and
  runs eager (one warning + traceback per GPU worker).

---

Benchmarks

Run on a g6.12xlarge (4X L4 GPUs)
Test script - baseline with no compile

#!/usr/bin/env bash

: "${AWS_MAX_ATTEMPTS:=100}"
: "${OMP_NUM_THREADS:=1}"
: "${MKL_NUM_THREADS:=1}"
: "${OPENBLAS_NUM_THREADS:=1}"
: "${NUMEXPR_NUM_THREADS:=1}"
: "${TENSORSTORE_HTTP_THREADS:=32}"
: "${TENSORSTORE_HTTP2_MAX_CONCURRENT_STREAMS:=64}"
: "${TENSORSTORE_S3_REQUEST_CONCURRENCY:=256}"

export AWS_REGION=us-west-2
export AWS_RETRY_MODE=standard
export AWS_MAX_ATTEMPTS

export OMP_NUM_THREADS
export MKL_NUM_THREADS
export OPENBLAS_NUM_THREADS
export NUMEXPR_NUM_THREADS

export TENSORSTORE_HTTP_THREADS
export TENSORSTORE_HTTP2_MAX_CONCURRENT_STREAMS
export TENSORSTORE_S3_REQUEST_CONCURRENCY

python -m aind_torch_utils.run \
  --in-spec /root/capsule/scratch/specs/in_spec.json \
  --out-spec /root/capsule/scratch/specs/out_spec.json \
  --model-type denoise-net \
  --weights /root/capsule/scratch/BM4DNet-20250905-169-0.0073.pth \
  --t 0 --c 0 \
  --patch 64 64 64 \
  --overlap 10 \
  --block 256 256 256 \
  --batch 64 \
  --devices cuda:0 cuda:1 cuda:2 cuda:3 \
  --seam-mode blend \
  --halo 5 \
  --max-inflight-batches 64 \
  --norm-lower 0.5 --norm-upper 99.9 \
  --clip-norm 0.0 5.0 \
  --prep-workers 10 \
  --writer-workers 10 \
  --metrics-json /results/metrics.json \
  --metrics-interval 1
  # --compile \
  # --compile-mode "default" \
  # --compile-dynamic "auto"

2026-06-24 23:45:47,258 - main - INFO - Total time: 608.65s
2026-06-24 23:45:47,258 - main - INFO - Throughput: 127.02MB/s

With compile (uncomment above options):
2026-06-24 22:01:54,000 - main - INFO - Total time: 498.90s
2026-06-24 22:01:54,000 - main - INFO - Throughput: 154.97MB/s

So a ~22% boost using the default options.

[carshadi]

Closing to reopen against dev as per #23