Investigation: faster training throughput without accuracy regression — starting with `torch.compile`

[Borda]

Motivation

We want lower wall-clock time per epoch at unchanged mAP. compile=True exists as an opt-in (config.py:502) and is documented as being "for faster throughput" (config.py:467), but there is no benchmark in the repo establishing what it actually buys, on which hardware, or for which configs. That number is unknown, and it should not be.

This issue is scoped to measuring and understanding, not to landing a fix. Any change to defaults or the compile path should follow from numbers produced here.

Hard constraint throughout: no accuracy regression. Any candidate change gets validated on mAP, not just ms/step. A throughput win that moves the metric is not a win.

What we don't know

Open questions, in no particular order. Answering any one of them is useful on its own.

Baseline

• What is the steady-state ms/step for eager vs compile=True, per model size, on common GPUs (A100, L4, 4090, T4)? With warm-up excluded and validation excluded?
• Does compilation succeed at all under each config we ship? torch._dynamo.explain(model)(samples) gives graph_count / graph_break_count / break_reasons without needing a full run.
• How much wall time does the first-step compile cost, and at what epoch count does it pay back?

Where the time actually goes

• What fraction of a training step is the model forward/backward, versus the criterion + matcher, the dataloader, GPU-side augmentation, and logging? A torch.profiler trace with schedule(wait=20, warmup=5, active=10) answers this directly.
• How often is the GPU idle waiting on input? num_workers defaults to 2 (config.py:1094) — is that the binding constraint on typical hardware before anything else matters?
• Where are the host syncs, and how much do they cost? Candidates to look at in a trace: matcher.py:726, _assignment.py:83.

Compile configuration

• How do dynamic=True (current, module_model.py:440), dynamic=None, and dynamic=False compare on identical workloads?
• What do mode="reduce-overhead" and mode="max-autotune-no-cudagraphs" do here? Does reduce-overhead work at all under Lightning, or does it need cudagraph_mark_step_begin() per step?
• What are the torch._dynamo.config settings at module_model.py:436-437 doing in practice? suppress_errors=True in particular means a failed compile is indistinguishable from a working one at the log level — worth checking what happens with it off.
• Results likely differ across torch versions; pyproject.toml allows >=2.2.0, and dynamic-shape handling and inductor coverage moved a lot over that range. Which versions were tested?

Shapes and the gate

• Under which configs do input shapes actually vary per batch? The gate at module_model.py:416 keys on multi_scale, while the scale selection also depends on do_random_resize_via_padding (coco.py:1189) — worth tracing what shapes reach the model under each combination, empirically, rather than reasoning about it.
• If shapes do vary, how many distinct ones per epoch, and how many recompiles does that trigger? (TORCH_LOGS="recompiles", and torch._dynamo.config.cache_size_limit defaults to 8.)
• What is the pos-embed interpolation path (dinov2_with_windowed_attn.py:369-401) doing per step under each config, and does it interact with compilation?

Beyond compile — same motivation, different levers, all open:

• What does augmentation_backend="gpu" cost or save, and on which hardware profile?
• Would compiling the dense loss math separately from the matcher help, or not be worth the complexity?
• Multi-GPU: how does DDPOptimizer's graph splitting change any of the above?
• Anything else a profile turns up. If the trace says the bottleneck is somewhere nobody here guessed, that's the most useful outcome.

How to contribute

Pick any question. Partial answers are welcome — a single profiler trace on one GPU and one model is a real contribution.

When posting results, please include:

• torch version, GPU, driver
• model size, batch_size, num_workers, precision, DDP or single-GPU
• the exact config flags that differ from defaults
• how steps were timed (warm-up excluded? validation excluded? how many steps?)
• for any candidate speedup: mAP on the same eval set, before and after

Raw traces and explain() output are more useful than summaries — post them if you can.

Out of scope here

Landing changes to defaults or to the compile path. Those follow once we know what the numbers say; if this investigation surfaces something concrete, it gets its own PR with the measurement attached.

[JESUSROYETH]

Picking up a few of the open questions here, from two machines. Everything below is on develop@6674d858, detection-only, group_detr=13, bf16 autocast, fused AdamW, model and criterion built through the repo's own build_model_from_config / build_criterion_from_config. Steps timed with torch.cuda.synchronize() on both ends, 8-15 warm-up steps discarded, median of 20-25, OMP_NUM_THREADS=2. Validation and the dataloader are out of the step time and measured separately.

	GPU	CPU	torch
A	RTX 4060 Laptop 8 GB	laptop i9, 32 threads	2.13.0+cu130
B	NVIDIA L4 23 GB, GCE g2-standard-8	Xeon @ 2.20 GHz, 8 vCPU	2.9.1+cu129

compile=True never compiles, and when it is allowed to it raises

Two separate things, both verified end to end.

With the shipped defaults it is a no-op. module_model.py:417-418 gates on not train_config.multi_scale, and TrainConfig.multi_scale defaults to True (config.py:1066). So RFDETRNano(compile=True).train(...) runs entirely eager, with one logger.info as the only trace. A 1-epoch train() on a real 100-image COCO dataset completes normally and never enters torch.compile.

With multi_scale=False it is reached, and it raises. Both torch._dynamo.explain(model) and torch.compile(model, dynamic=True) fail at transformer.py:346:

TypeError: expected Tensor as element 0 in argument 0, but got torch.Size

Minimal repro, no rf-detr involved:

def f(x):
    return torch.stack([torch._shape_as_tensor(x)[2:4]]).to(device=x.device, dtype=torch.long)

x = torch.randn(1, 3, 8, 8, device="cuda")
f(x)                               # eager: tensor([[8, 8]])
torch.compile(f, dynamic=True)(x)  # TypeError: ... but got torch.Size

Dynamo polyfills torch._shape_as_tensor to return a torch.Size rather than a tensor, so torch.stack rejects it. Reproduced on torch 2.9.1 and 2.13.0 and at dynamic=True, None and False alike, so it is not a recent regression and not a dynamic-shapes issue.

You asked specifically what suppress_errors=True does here: it does not help. That flag only swallows errors Dynamo itself raises, and this one comes from user code, so Dynamo propagates it. Measured both ways, same failure.

The reason this survived is that test_compile_runs_when_enabled_and_static_shapes (tests/training/test_module_model.py:416) patches torch.compile with side_effect=lambda m, **_: m, it only asserts the call, never the compilation.

The fix is one condition, and the branch that works is already two lines above: torch.export cannot trace torch._shape_as_tensor either, and #1142 handled that by building spatial_shapes from the Python-int (H, W) pairs the function already collects at transformer.py:292-295. Adding torch.compiler.is_compiling() to that guard leaves eager, torch.jit.trace and the TorchScript-ONNX/TensorRT path from #1155 untouched. That is checked rather than assumed: all three report is_compiling() == False. I opened it as a separate PR since it is a crash fix on its own.

One caveat that came out of checking the PyTorch tags instead of assuming: torch.compiler.is_compiling is absent in 2.2.0 and present from 2.3.0, and torch.compiler.is_exporting only lands in 2.7.0. So on torch 2.2, the floor in pyproject.toml, both getattr guards fall through and compile=True still raises. Covering 2.2 too means using torch._dynamo.external_utils.is_compiling, a private path, for one old minor, so it's up to you whether that is worth it.

With that in, the same train(compile=True, multi_scale=False) call that raises on develop runs to the end. On the 4060 (torch 2.13.0) at nano bs=2, steady state after 8 warm-up steps, median of 25: eager 58.35 ms/step, compiled 55.00 with suppress_errors=True and 54.57 with it off, so -5.7% / -6.5%, graph_break_count = 0. First-step compile cost 375.8 s cold and 38.7 s on a warm inductor cache, but those are single samples, so treat those two as an order of magnitude rather than a measurement.

That number does not carry to the L4, and the failure there is worth more than the speedup. On the L4 (torch 2.9.1+cu129, same nano bs=2), eager is 134.74 ms/step [132.42, 158.47] and the fix does get past spatial_shapes, but compilation then fails on its own inside Inductor, with both suppress_errors settings:

eager                          : 134.74 ms/step  [min 132.42, max 158.47]
compile(dynamic=True) sup=True : RAISED InductorError: SubprocException
compile(dynamic=True) sup=False: RAISED InductorError: SubprocException
  Name=triton_red_fused__to_copy_add_native_layer_norm_backward_view_193

So suppress_errors=True does not absorb this one either, same as the TypeError above but for a different reason. And it is not simply a broken toolchain on that box: a trivial layer_norm forward+backward under torch.compile(dynamic=True) compiles and runs there without complaint. I have not reduced it past that, so I am not claiming a cause, only the observation: on this pair, transformer.py:346 is necessary but not sufficient, and the next blocker is a fused layer-norm-backward kernel from your own graph.

So the answer to "what is the steady-state ms/step for eager vs compile=True" is: there is no single figure yet. One of my two machines compiles cleanly and gains 6.5%. The other one cannot compile at all. That is worth knowing before the flag's default is discussed.

Where the time actually goes, and the part I did not expect

"GPU busy" below is the sum of leaf CUDA kernel self time from torch.profiler (the ATen-op rows are excluded so nothing is counted twice), over the same steps whose wall time is reported. One config per process.

	wall ms/step	GPU busy	utilisation	kernels/step
4060 nano bs=1	41.4	29.8	72.0%	4214
4060 nano bs=4	105.9	96.4	91.0%	4379
4060 small bs=1	54.8	40.9	74.6%	4781
L4 nano bs=1	132.9	25.8	19.4%	4228
L4 nano bs=4	141.0	64.4	45.7%	4472
L4 small bs=1	151.1	32.2	21.3%	4810
L4 small bs=2	155.5	51.5	33.1%	4958

The L4 does the same step in less GPU time than the laptop 4060 (64.4 ms against 96.4 at nano bs=4) and is still 1.3x slower on the clock. The step issues 4200-4800 CUDA kernels and that count does not move with batch size (4228 at bs=1, 4379 at bs=4), and at bs=1 about three quarters of them are under 5 µs of device time (3205 of 4228 on the L4). So on a stock cloud L4 shape this workload is host-bound, not GPU-bound. Taking utilisation as a rough ceiling, that gives 2.2x at nano bs=4 and about 5x at bs=1, against the -6% I measured on a laptop where the GPU is already the bottleneck. That is the regime torch.compile and CUDA graphs exist for.

Phase split on the 4060, sync-delimited:

	nano bs=4	small bs=2
forward	31.1 ms (29.0%)	25.3 ms (29.0%)
criterion	9.7 ms (9.1%)	9.0 ms (10.4%)
backward	61.2 ms (57.1%)	47.5 ms (54.6%)
optimizer	5.0 ms (4.7%)	5.3 ms (6.0%)

Inside the criterion, HungarianMatcher._match_many is 58.6%, loss_labels 17.4%, loss_boxes 11.0%. GEMMs are not the problem: aten::mm + aten::addmm is 30.75 ms of the 95.5 ms of device time at nano bs=4 and every shape runs at 16-26 TFLOP/s.

Host syncs: the two candidates in the issue cost nothing

torch.cuda.set_sync_debug_mode("warn") with the Python stack captured at each warning gives 78 CUDA host syncs in one step (nano bs=4). Almost all of them are in criterion.py: 12 at :837, 12 at :839, 6 each at :496, :510, :513, :514, :656, 3 at :636.

matcher.py:726 and _assignment.py:83 are one sync each, and timed individually they are 0.02 ms/step. The one blocking call that shows up in a timer is a different line, transformer.py:346, the same one above, whose 16-byte .to(device=...) with non_blocking=False drains the whole CUDA queue mid-forward and blocks the host for 24.57 ms/step, 23.4% of the step. Before anyone spends a day on that: caching the device tensor per (shape, device) is bit-identical and measures +0.11% to +0.27%, i.e. slightly worse, across nano bs=4/8 and small bs=2/4. The step is GPU-bound on that machine and the drained work has to happen anyway.

The 60-odd criterion syncs come from one place. _assignment._solve_to_indices ends with pair = torch.stack((rows, cols)).cpu(), documented as "the single device-to-host transfer for the whole step", and _assign_padded then does its bookkeeping on the host. Every downstream use in criterion.py indexes a CUDA tensor with those host indices, which pushes each one back to the device. One deliberate D2H traded for ~60 implicit H2D round trips. Keeping the tail of _assign_padded on the device is about six lines, takes the count from 77 to 11 and the full 13-key loss dict stays bitwise identical, but it is only -0.81% (4060) and -0.85% (L4) at nano bs=4, so I am reporting it rather than proposing it.

num_workers=2 is not the binding constraint

Real images, the shipped torchvision train pipeline, __getitem__ timed directly, median of 120 calls: 6.24 ms at 384, 7.90 at 512, 8.41 at 576, 9.81 at 560. 45% of that is PIL.Image.resize, 176 resizes for 120 samples, because Option B of the train resize does resize → crop → resize. Two workers sustain ~320 samples/s at 384 and ~200/s at 560, so a bs=4 batch costs 12.5 ms of worker time against a 105-155 ms step. Not binding on either machine. Worth re-checking where the GPU is fast and the vCPU count is high, but it should not be the first lever.

Things I measured and dropped

• Batching the 13-iteration two-stage group_detr loop: the block alone goes 7.57 → 6.63 ms and 1101 → 352 kernels, but only -0.9% end to end, because the batched form adds ~1.0 ms of real GPU work (materialising (G,B,S,C), the permute copy, the LayerNorm affine outside the fused kernel) and eats half the bubble it recovers. The variant with concatenated enc_output weights is faster (-1.7%) but not bit-identical, and topk turns a 1.6e-2 bf16 difference into a different set of selected proposals.
• Routing deformable attention through the repo's own gather path on CUDA: 9-33x slower (33x forward, 9x forward+backward). The docstring's claim that F.grid_sample is faster there is correct.
• Running the deformable-attention core in bf16 instead of the fp32 it inherits from autocast's promote policy: forward alone drops 1.06 → 0.66 ms, but forward+backward goes 3.38 → 6.87 ms, because of bf16 atomicAdd in grid_sampler_2d_backward. fp32 is the faster choice.

What this does not cover

Worth being clear about the limits before anyone builds on these numbers.

The step timings run on a synthetic batch: real model and real criterion, built through the repo's own build_model_from_config / build_criterion_from_config, but the images are torch.randn and the boxes are random. That is fine for kernel counts, GPU utilisation and host syncs, which depend on shapes rather than pixel values, and the box count per image is set to a realistic 12. It is not fine for anything that depends on the data, so treat it that way. The dataloader and validation sections are the exception, those ran on real COCO images.

No mAP anywhere. Nothing here was trained to convergence, so I cannot speak to your "no accuracy regression" bar for any of it. The one change I am proposing does not touch a value in eager, the new branch is unreachable outside torch.compile, but that is an argument from the code, not a measurement.

Two machines only, a laptop 4060 and one L4. The issue asks about A100, 4090 and T4 as well and I have none of them. No DDP, and I did not touch augmentation_backend="gpu".

I'm happy to post raw traces or the scripts for any of these if it helps. The one thing I would do first is the transformer.py:346 fix, since until that lands there is no compile number to measure at all ..