Generalized CUDA synchronize

[jaywonchung]

PyTorch has torch.cuda.synchronize, which syncs CPU and GPU code execution. This is essential for accurate measurement. But there isn't one for JAX, which we do hope to support as first class in the future.

Right now we have some infrastructure for generalized CUDA synchronize in zeus.util.framework (

zeus/zeus/util/framework.py

Line 45 in 9381ec8

def cuda_sync(device: int | None = None) -> None:

), without JAX.
My current idea is to allow users to pass in a hook (Callable[[], None]) that does synchronization and default it to torch.cuda.synchronize. Then JAX users can pass in loss.block_until_ready as the synchronization hook, for example.