pmap for a function containing large, static arrays

[lukepfister]

I'm curious what the best way to parallelize evaluation of a function containing large, static arrays.

I need to evaluate

def f(Y, W, X, theta): 
    return norm(Y - W * g(X, theta))

where Y, W, X are all large, but static, arrays. Each is of size 512x512. Only the parameter theta needs to be changed. I need to evaluate many thetas-- roughly 1000^2.

Option 1: Run on a single device.

f_theta = jit(lambda theta: f(Y, W, X, theta))
for t in theta:
    out.append(f_theta(t))

This results in GPU utilization ~60%.

Option 2: Parallelize over multiple devices.
I wrote a simple generator to return thetas in batches of 8, which is the number of GPUs on this machine. Then

f_theta_p = pmap(lambda theta: f(Y, W, X, theta)))
out = [f_theta_p(t) for t in my_gen(theta)]

This gives the same answer as Option 1, but at roughly the same speed. Average GPU utilization is closer to ~5% on each of the GPUs.

I'm guessing that the utilization is low on the pmap version as the arrays Y, W, X are constantly being transfered to the respective devices. I'm not sure how to test that, though.

Could there be a way to avoid this using device_put_replicated? Or, more generally, what's the best way to evaluate a function like this using multiple devices?

`

[seanmor5]

Have you tried static_broadcasted_arguments in pmap:

static_broadcasted_argnums (Union[int, Iterable[int]]) – An int or collection of ints specifying which positional arguments to treat as static (compile-time constant). Operations that only depend on static arguments will be constant-folded. Calling the pmapped function with different values for these constants will trigger recompilation. If the pmapped function is called with fewer positional arguments than indicated by static_argnums then an error is raised. Each of the static arguments will be broadcasted to all devices. Arguments that are not arrays or containers thereof must be marked as static. Defaults to ().

[lukepfister]

Ah, good catch. Unfortunately it doesn't solve my problem-- I'm still seeing ~60% utilization on one GPU, ~8% in the pmap-ed version.