Settings

Theme

Forget, then remember: checkpointing the adjoint sweep on a GPU

nablatensor.com

1 points by petrpravda · 1 comment

Reader

1 thread
petrpravdaOP

Author here. Started because of one number: our Vulkan kernel does 125M paths/sec pricing this tape, but drops to 15M once you also ask for the five adjoint Greeks -- 8x slower, vs. 1.7x on CPU for the same request. Turned out to be a register-spill problem, not an arithmetic one: the reverse sweep needs ~760 forward values per path and a GPU lane doesn't have room.

Fix is checkpointing (Griewank's 1992 "revolve," same idea modern LLM training uses under "gradient checkpointing"): keep ~38 bookmarks instead of 760 values, recompute what's needed when the reverse sweep asks for it. 2.9x, or 6.8x batching four markets per dispatch, bit-exact vs. the original.

One fun problem along the way: the compiler kept noticing the recomputation was identical to the original and merging them back together (undoing the whole point), so I had to XOR every bookmarked value with a runtime-zero constant just to stop it from proving that.

Keyboard Shortcuts

j
Next item
k
Previous item
o / Enter
Open selected item
?
Show this help
Esc
Close modal / clear selection