LinearMeasurement slow on GPU

[cynddl]

Instantiating a LinearMeasurement object appears to be very slow on GPU:

https://github.com/ryan112358/private-pgm/blob/57554e604733a52f6dd217bd67feef0aa7a7435b/src/mbi/marginal_loss.py#L74-L80

I tried removing the converter and the code seems to run fine, and actually faster. My guess is that many calls to send data in advance to the GPU memory are not efficient at all.

[ryan112358]

Interesting, can you send a minimum example that demonstrates this behavior? It seems surprising that this would be the bottleneck compared to fitting the model or other parts of the code, but maybe depends on your use case

[cynddl]

Apologies for the delay. Here is a small reproducible benchmark. There seems to be a big hit with small domain sizes and clique sizes.

Benchmarking code

import os
import time

import click
import pandas as pd


def setup_platform(platform: str, xla_flags: str | None = None):
    """Configure JAX platform before import."""
    os.environ["JAX_PLATFORMS"] = platform
    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

    if xla_flags:
        os.environ["XLA_FLAGS"] = xla_flags

    if platform == "cpu":
        # Enable CPU multithreading
        os.environ["XLA_FLAGS"] = os.environ.get("XLA_FLAGS", "") + " --xla_cpu_multi_thread_eigen=true"


def run_benchmark(domain_size: int, num_attrs: int, num_rows: int, clique_size: int, iters: int):
    """Run mirror_descent benchmark with given parameters."""
    import jax
    import numpy as np

    from mbi import Dataset, Domain, estimation
    from mbi.marginal_loss import LinearMeasurement

    # Create domain with specified size
    attrs = [f"A{i}" for i in range(num_attrs)]
    sizes = [domain_size] * num_attrs
    domain = Domain(attrs, sizes)

    # Create synthetic dataset
    data = np.random.randint(0, domain_size, size=(num_rows, num_attrs))
    df = pd.DataFrame(data, columns=attrs)
    dataset = Dataset(df, domain)

    # Create cliques of specified size (sliding window)
    workload = [tuple(attrs[i : i + clique_size]) for i in range(num_attrs - clique_size + 1)]
    measurements = [
        LinearMeasurement(
            noisy_measurement=dataset.project(cl).datavector(),
            clique=cl,
            stddev=1.0,
        )
        for cl in workload
    ]

    # Warmup (JIT compilation)
    _ = estimation.mirror_descent(domain, measurements, iters=10)
    jax.block_until_ready(_)

    # Timed run
    t0 = time.perf_counter()
    result = estimation.mirror_descent(domain, measurements, iters=iters)
    jax.block_until_ready(result)
    elapsed = time.perf_counter() - t0

    return elapsed


@click.command()
@click.option("--platform", type=click.Choice(["cpu", "cuda"]), default="cpu")
@click.option("--xla-flags", type=str, default=None, help="Extra XLA flags")
def main(platform: str, xla_flags: str | None):
    setup_platform(platform, xla_flags)

    import jax

    print(f"\n{'='*60}")
    print(f"Platform: {platform} (devices: {jax.devices()})")
    print(f"XLA_FLAGS: {os.environ.get('XLA_FLAGS', 'none')}")
    print(f"{'='*60}\n")

    # Test 1: Varying domain size (GPU should scale better)
    print("=== Test 1: Domain Size Scaling ===")
    print(f"{'Domain Size':<15} {'Time (s)':<10}")
    for domain_size in [10, 20, 50, 100]:
        t = run_benchmark(domain_size=domain_size, num_attrs=5, num_rows=10000, clique_size=2, iters=500)
        print(f"{domain_size:<15} {t:.3f}")

    # Test 2: Varying clique size (larger matrices)
    print("\n=== Test 2: Clique Size Scaling ===")
    print(f"{'Clique Size':<15} {'Time (s)':<10}")
    for clique_size in [2, 3, 4]:
        t = run_benchmark(domain_size=10, num_attrs=8, num_rows=10000, clique_size=clique_size, iters=500)
        print(f"{clique_size:<15} {t:.3f}")

    # Test 3: Varying number of attributes (more measurements)
    print("\n=== Test 3: Number of Attributes ===")
    print(f"{'Num Attrs':<15} {'Time (s)':<10}")
    for num_attrs in [5, 10, 15, 20]:
        t = run_benchmark(domain_size=10, num_attrs=num_attrs, num_rows=10000, clique_size=2, iters=500)
        print(f"{num_attrs:<15} {t:.3f}")

    # Test 4: Iteration scaling
    print("\n=== Test 4: Iteration Scaling ===")
    print(f"{'Iterations':<15} {'Time (s)':<10} {'Time/iter (ms)':<15}")
    for iters in [100, 500, 1000, 2000]:
        t = run_benchmark(domain_size=20, num_attrs=10, num_rows=10000, clique_size=2, iters=iters)
        print(f"{iters:<15} {t:.3f}      {t/iters*1000:.3f}")


if __name__ == "__main__":
    main()

Results on CPU

============================================================
Platform: cpu (devices: [CpuDevice(id=0)])
XLA_FLAGS:  --xla_cpu_multi_thread_eigen=true
============================================================

=== Test 1: Domain Size Scaling ===
Domain Size     Time (s)  
10              0.607
20              0.808
50              1.510
100             3.673

=== Test 2: Clique Size Scaling ===
Clique Size     Time (s)  
2               1.072
3               2.022
4               5.773

=== Test 3: Number of Attributes ===
Num Attrs       Time (s)  
5               0.575
10              1.401
15              2.029
20              2.808

=== Test 4: Iteration Scaling ===
Iterations      Time (s)   Time/iter (ms) 
100             1.056      10.561
500             1.565      3.129
1000            1.996      1.996
2000            3.066      1.533

Results on CUDA (two L40s cards)

============================================================
Platform: cuda (devices: [CudaDevice(id=0), CudaDevice(id=1)])
XLA_FLAGS: none
============================================================

=== Test 1: Domain Size Scaling ===
Domain Size     Time (s)  
10              1.624
20              1.956
50              1.741
100             1.619

=== Test 2: Clique Size Scaling ===
Clique Size     Time (s)  
2               2.235
3               2.447
4               2.121

=== Test 3: Number of Attributes ===
Num Attrs       Time (s)  
5               1.618
10              2.733
15              4.054
20              5.338

=== Test 4: Iteration Scaling ===
Iterations      Time (s)   Time/iter (ms) 
100             2.127      21.266
500             3.150      6.300
1000            4.601      4.601
2000            7.019      3.510

[cynddl]

(tested on both mbi 1.0.0 and on the latest revision from the main branch, with similar findings)

[ryan112358]

I'm not sure I see how LinearMeasurement is the bottleneck here, since the test appears to to measure performance of the mirror descent only.

The differences between CPU & GPU benchmark times here do not seem too surprising to me, I think you'll tend to benefit from GPU when the size of each attribute is larger (at least 128) and you have larger graphs.

The fact that time / iter varies so much in test 4 is interesting though, and suggests that compile time may be a dominant factor (I'll look into this).