Benchmark runtime of varying size of attention matrix

[ksmdnl]

Hi all, I'm quite new to Jax and am trying to find out how to correctly setup my small investigation regarding the runtime of my model.
Specifically, I have a graph transformer model that takes nodes features and edge features as input. I would like to see how the runtime changes when I increase the graph size, our theoretical study shows it should be a cubic complexity but we want to make sure.
Since I'm not very familiar with the standard/correct approach of doing this in and generally with Jax, I'd like a second of third opinion from the community.

def model(x):
    net = factory(hidden_dim, num_layers=1)
    return net(*x)

def inference(fn, rng, params, x, mode="runtime"):
    start = time.time()
    _ = jax.block_until_ready(fn)(params, x)
    end = time.time()
    print(f"{mode}: {end - start} s")
    return end - start

def main():
    if args.single == 0:
        nb_nodes = np.arange(args.max_nb_node) + 1
    else:
        nb_nodes = [args.max_nb_node]
    runtimes = []
    rng = jax.random.PRNGKey(42)
    for nb_node in nb_nodes:
        print(f"Number of node: {nb_node}")
        net = hk.without_apply_rng(hk.transform(model))
        node_fts = jax.random.normal(rng, (batch_size, nb_node, hidden_dim))
        edge_fts = jax.random.normal(rng, (batch_size, nb_node, nb_node, hidden_dim))
        x = (node_fts, edge_fts)
        params = net.init(rng, x)
        apply_fn = jax.jit(net.apply)

        # compile time
        _ = inference(apply_fn, rng, params, x, mode="compile time")

        # execution time
        runtime = inference(apply_fn, rng, params, x, mode="execution time")
        runtimes.append(runtime)

In the main function it iterates over a range of graph size. In inference the runtime is measured, which is why jax.block_until_ready() is used to synchronize all results. My question is, if this is the correct way to do it using such loop, or does the measurement corresponding to each graph size have to be done in each separate process?
Looking forward for a discussion!

[jakevdp]

That looks like a reasonable approach, though you may consider using the tools in the timeit module (or in IPython, the %timeit magic command), either of which are better suited to getting accurate benchmarks than straight use of time.time.

Also, see JAX FAQ: Benchmarking JAX code for more JAX-specific tips on getting accurate benchmarks.

[ksmdnl]

Thank you for the feedback and the inputs!

I wanted to see if this would work up to for example nb_node = 1000.
For nb_node = 414 I'm getting an OOM error in the initialization (when performing the Einstein summation), which looks as follows:

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 54101757696 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:  249.86MiB
              constant allocation:         0B
        maybe_live_out allocation:   50.39GiB
     preallocated temp allocation:         0B
                 total allocation:   50.63GiB
              total fragmentation:         0B (0.00%)

This is quite strange, since I'm using one A100 80GB.
However, when I jit the init funtion, there is no OOM error even after 500 node size.
My question is, would this be a correct workaround given in this case?

[jakevdp]

It's hard to answer without knowing more about what code you are running.

Are you literally looping from nb_node = 1 up to 1000 and finding when the device runs out of memory? If so, I suspect the issue may have to do with arrays leftover from previous iterations still taking up space on the device.

If that's not what you're doing, it would be more helpful if you could edit your question to show exactly what you're running.