I got RESOURCE_EXHAUSTED but I still have memory

[kkeevin123456]

Here is the question,
I have 4 gpus but jax always allocate kernel matrix to single device.
How could I use all gpus?

Thanks for help.

[YouJiacheng]

You need pmap to leverage multiple devices.
https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html
https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html
FYI, we cannot use vmap for multiple devices because of the non-unified memory nature of multiple devices, thus inter device parallel is different from intra device parallel.
It just like in pytorch, we need to use DDP or DP for multiple devices.

[kkeevin123456]

Thanks for reply.
I want to apply pmap on empirical_ntk_fn like below.

eNTK = nt.empirical_ntk_fn(apply_fn, vmap_axes=0, trace_axes=(), implementation=2)
gram = eNTK(inputs, None, get_params(opt_state))

How can I pass the params?
My idea is like this, but it still has some bug.

eNTK = nt.empirical_ntk_fn(apply_fn, vmap_axes=0, trace_axes=(), implementation=2)
gram = pmap(eNTK)(inputs, None, get_params(opt_state))

[YouJiacheng]

Could you give a runnable code snippet?
I have no idea what nt.empirical_ntk_fn is.

[kkeevin123456]

Here is the code that can reproduce the problem~
The code below is to load data and train a toy CNN.
Here I want to calculate the kernel matrix of first layer that will cause the problem.

import numpy as np
import torch
import torchvision.datasets as datasets
import torch.nn.functional as F

import jax
from jax import random
import jax.numpy as jnp
from jax.example_libraries import optimizers
from jax import jit, grad, vmap, pmap
import functools
import neural_tangents as nt
from neural_tangents import stax

from tqdm import tqdm

mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=None)
class_5_idx = np.where(mnist_trainset.targets.numpy() == 5)
class_3_idx = np.where(mnist_trainset.targets.numpy() == 3)
class_5 = mnist_trainset.data[class_5_idx]
class_3 = mnist_trainset.data[class_3_idx]

# reproduce problem
M = 200
W = H = 28
C = 1
P = 200
eta = 0.1

inputs = np.vstack((class_5[:P//2], class_3[:P//2]))
idx = np.arange(P)
np.random.shuffle(idx)
inputs = inputs[idx].reshape(-1, 28, 28, 1).astype('float32')
inputs = jnp.array(inputs)

ys = np.vstack((*np.ones(P//2), *np.zeros(P//2)))
ys = ys[idx].astype('float32')
ys = jnp.array(ys)

# kernel for all layers
init_fn, apply_fn, kernel_fn = stax.serial(
    stax.Conv(10, (5, 5), (1, 1), 'SAME'), stax.Relu(),  # shape (output_c, filter_size, stride_size)
    stax.Flatten(), stax.Dense(1)
)

# kernel for first layers
init_fn_a, apply_fn_a, kernel_fn_a = stax.serial(
    stax.Conv(10, (5, 5), (1, 1), 'SAME'), stax.Relu(),  # shape (output_c, filter_size, stride_size)
    stax.Flatten()
)

shape, params = init_fn(random.PRNGKey(np.random.randint(1e6)), inputs.shape)
eNTK = nt.empirical_ntk_fn(apply_fn, vmap_axes=0, trace_axes=(), implementation=2) # implementation 2 會比較快，1會比較慢但比較精確

opt_init, opt_update, get_params = optimizers.sgd(eta)
opt_state = opt_init(params)

eNTK_a = nt.empirical_ntk_fn(apply_fn_a, vmap_axes=0, trace_axes=(), implementation=2)

loss = jit(lambda params, x, y: 0.5 * jnp.mean((apply_fn(params, x) - y) ** 2))
grad_loss = jit(lambda state, x, y: grad(loss)(get_params(state), x, y))

nsteps = 1000

for i in range(nsteps):
    opt_state = opt_update(i, grad_loss(opt_state, inputs, ys), opt_state)
    gram_a = eNTK_a(inputs, None, [get_params(opt_state)[0], get_params(opt_state)[1], get_params(opt_state)[2]])
    print(gram_a.shape)

[YouJiacheng]

Ok, you should pmap over inputs only.

[kkeevin123456]

When I try to pmap over inputs only it cause this error.

Then I try to pmap another term which will alse cause error.

[YouJiacheng]

Emmm.

gram_a = pmap(lambda ins: eNTK_a(ins, None, [get_params(opt_state)[0], get_params(opt_state)[1], get_params(opt_state)[2]]))(inputs)

[kkeevin123456]

I use your code but it cause this error.

So I change into this:

gram_a = pmap(lambda ins: eNTK_a(ins, None, [get_params(opt_state)[0], get_params(opt_state)[1], get_params(opt_state)[2]]))(np.array([inputs]))

But it still put all data into one device.

I'm sure the jax can see all devices.

[YouJiacheng]

What is the shape of inputs.