Metal training: per-step host waits leave the GPU idle (planning fixes) #3905

Activity

1. MrNeRF commented on Oct 29, 2026

   

   MrNeRF

   on Oct 29, 2026

   Owner

   More actions

   Thanks, this is a careful write-up. I checked the four waits against dev (50a1c5e), and they all match.

   One requirement for all four fixes: they have to stay behind the Tensor facade, so they work on every backend (Metal, Vulkan, CUDA and HIP). Fix each wait where the facade causes it, not with Metal-only paths in the trainer.

   We're also evaluating #3844, which reworks the upload, readback, completion and queue services that all four waits go through. Please build and measure on top of it, or at least check each fix against it, so the same wait isn't fixed twice.

   1. Deferred count with the bilateral grid and PPISP. When #2867 landed, exposure correction was still opt-in. #2963 made it an MRNF default a day later, so default Metal training has waited on the count every step since. That explains your 51%.

   #2867 kept those iterations synchronous, but most of what's needed is already in place:

   • The persistent updates (the grid's step_image, PPISP's optimizer_step and scheduler_step) run after the resolve, and a Retry returns before them.
   • The bilateral grid's backward zeroes its slice gradient on every call.
   • The gap is PPISP: PPISP::backward_in_place accumulates into exposure_grad_ and its other gradient buffers before the count is resolved, so a replay would add the failed attempt's gradients again.

   So I'd go with your second option: keep the updates after the resolve, zero the PPISP gradients on replay, and leave controller-phase iterations synchronous. Please extend #2867's forced-overflow test (capacity 1) with exposure correction on, and check that the replay matches a large-capacity run.

   The deferral is opt-in through the FastGS ops table, and only Metal provides it today. Keep the trainer side backend-neutral, so other backends can opt in the same way.

   2. Crop-box constants. Passing the matrix and the camera centre in the roi_masks parameter block is fine, since that kernel is single-source Slang. The underlying problem, a small synchronous upload waiting on a recycled buffer, is in the facade's upload path and will hit other callers too. Please fix it there.

   3. Loss readback. The same goes for Tensor::cat, which waits on the host once per input on Metal. Please fix that in the Tensor library instead of avoiding cat in the trainer.

   4. Denominator. No per-step item(). Either cache it per camera, keyed on the camera, its mask and the crop box (the crop box can change while training runs in the GUI), or keep it on the device.

   Benchmark: the 59.3 → 55.7 s was GUI training on garden images_4, 7,000 iterations, median of 3, on an M5 Max, before #2963. Today's MRNF defaults include exposure correction, so please measure with it both on and off. Alternate the before and after runs: thermals moved the same headless run by about a third on my machine.

   Duplicate work: nobody else is working on these. #3893 also touches trainer.cpp, so expect some rebasing.

   One PR per fix into dev works for me. Please run tools/check_backend_neutrality.py and the error-debt census before each one.

   👍React with 👍1Reacted by oli-har
2. oli-har commented on Oct 30, 2026

   

   oli-har

   on Oct 30, 2026

   Author

   More actions

   Thanks for the guidance. We followed it, and found one more wait on the way. Each fix is its own branch from dev 2849c644 and will be its own PR:

   • Deferred count with exposure correction on (item 1, option b). The grid's and PPISP's updates already run after the resolve, so a replay only has to zero PPISP's gradients first; controller-phase steps stay synchronous. Resolve the Metal FastGS instance count at the commit instead of the forward #2867's forced-overflow test also runs with both on.
   • Each bilateral grid slot gets its own storage (new). With both slots in one tensor, writing back the slot being replaced waited on Metal for the step still using the other one.
   • Synchronous Metal uploads stage instead of waiting (item 2), in MetalBackendOps::copy_host_to_device.
   • The thin-structure denominator is kept per camera (item 4), keyed on the mask and its settings and, for a camera outside the crop box, the crop box and its loss weight.
   • Tensor::cat copies without a host wait per input (item 3).

   Alone, each one hardly changes the step time on the Mac, because the next wait in the same step takes over. Together, on our tile (M3 Ultra, steps 0-3,000, mean of two runs): 89.64 → 72.06 ms per step (-19.6%), with the GPU going from about 89% to 97% busy. At the 48M cap: 297.46 → 281.91 (-5.2%). Garden at 7,000 steps with exposure correction on: 8.670 → 8.642 (-0.3%), since its GPU is already 97-99% busy. Losses and PSNR stay within the run-to-run spread. Most of what is left on the tile is the trainer resolving the deferred count at the end of the backward, while the GPU is still busy.

   #3926 (split from #3844) and #3930 are on dev now. They don't touch the code the five change, and all five branches merge cleanly with today's dev (20d82adf).

   The deferred count and the upload staging only change what Metal runs, so their PRs come first. The other three also change what CUDA and Vulkan run, so they follow once our PC (RTX 5090) has run both suites and the same benchmarks on them.
3. MrNeRF

   added 2 commits that reference this issue on Oct 30, 2026

   Stage synchronous Metal uploads into storage the GPU still uses (#3934)

   ...

   c411413

   Defer the FastGS instance count with exposure correction on (#3933)

   ...

   8092517