Specifying id_tap dependency without data copy

[jaywonchung]

Hi,

I'm trying to implement a host callback that records the current time while running computation on GPUs.

import time

import jax
import jax.numpy as jnp
from jax import random
from jax.experimental import host_callback


TIME_RECORDS = []  # <- Q1. Race?

def record_time(args, transforms):
    """Record the current time and return."""
    TIME_RECORDS.append(time.time())

@jax.jit
def step(x):
    """Run computation on x and record time points in between."""
    x = host_callback.id_tap(record_time, None, result=x)
    y = x @ x @ x @ x @ x @ x @ x @ x @ x @ x @ x
    y = host_callback.id_tap(record_time, y, result=y)
    #                                     ^
    #                                     |
    #  Q2. Extra GPU -> CPU copy? --------+
    return y


if __name__ == "__main__":
    step(random.normal(random.PRNGKey(0), (10240, 10240)))
    print(TIME_RECORDS)

Output:

[1677898267.4230554, 1677898268.2604709]

I have two questions in this context:

Q1) Is this potentially a race condition on the global variable TIME_RECORDS, or are callback functions guaranteed to be executed one after another (instead of concurrently running with each other)?

Q2) I would like the second id_tap to run after computing y, and while the current code works, it feels like I'm incurring an extra data copy of y from the GPU to the CPU. Is this really the case? If so, is there any way I can specify the data dependency from y to record_time without initiating data copy?

Thanks a lot.

[sharadmv]

Could you try using jax.debug.callback(..., ordered=True)? This is our newer API though and by saying ordered=True you are forcing callbacks to be run sequentially relative to each other. host_callback behaves like ordered debug callback.

In general, there is no way to enforce a callback runs after a certain computation without adding data dependency (and extra copying) as you've done. As a result, the callbacks aren't reliable for timing and measuring performance. I'd suggest using the profiler instead.

[jaywonchung]

@sharadmv Thanks a lot for the answer!

Could you try using jax.debug.callback(..., ordered=True)? This is our newer API though and by saying ordered=True you are forcing callbacks to be run sequentially relative to each other. host_callback behaves like ordered debug callback.

I tried jax.debug.callback(record_time, x or y, ordered=True) in place of the two id_taps and they work, too. Maybe I didn't get something, but if host_callback works like ordered debug callbacks, is there a reason to switch to jax.debug.callback?

In general, there is no way to enforce a callback runs after a certain computation without adding data dependency (and extra copying) as you've done. As a result, the callbacks aren't reliable for timing and measuring performance. I'd suggest using the profiler instead.

I see. I'm actually not trying to record time using callbacks; I'm doing something else. Recording time was just an example. Do you think doing something like y = host_callback.id_tap(record_time, y[0,0], result=y) would be of smaller overhead?

Thanks again.

[sharadmv]

but if host_callback works like ordered debug callbacks, is there a reason to switch to jax.debug.callback?

There's still a lot of work left to do but we're working on moving away from HCB towards a set of new callbacks, including jax.debug.callback so we suggest using those since they will be better supported eventually.

Do you think doing something like y = host_callback.id_tap(record_time, y[0,0], result=y) would be of smaller overhead?

It's a bit tricky to identify which will have lower overhead. If you do y[0, 0] you're adding in an additional HLO op that needs to run, but the data-copying overhead will be less. With various shapes, I can imagine either being faster than the other, so it really depends on the case.

[jaywonchung]

I see. That makes sense. Thanks!