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 (https://github.com/ml-energy/zeus/blob/9381ec8d3793ac4550fe9f4ff8f53e3fca014b0f/zeus/util/framework.py#L45), 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.
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(https://github.com/ml-energy/zeus/blob/9381ec8d3793ac4550fe9f4ff8f53e3fca014b0f/zeus/util/framework.py#L45), without JAX. My current idea is to allow users to pass in a hook (Callable[[], None]) that does synchronization and default it totorch.cuda.synchronize. Then JAX users can pass inloss.block_until_readyas the synchronization hook, for example.