TPU RESOURCE_EXHAUSTED

[zhenlan0426]

Do you have any tips on how to debug "RESOURCE_EXHAUSTED: Error loading program: Attempting to reserve XXX at the bottom of memory." I reduced the batch_size to one and would still get this error. I am leveraging pmap to do data parallel training. With GPU on pytorch, I can train a whisper-large (without data parallel and 24G VRAM). But on TPU, I can only train a whisper-small. Any pointer will be helpful!

@partial(jax.pmap,axis_name='data')
def train_one_step(params,audio,input_ids,attention_mask,opt_states):
    def loss_fn(params,audio,input_ids,attention_mask):
        out = model(audio,input_ids,decoder_attention_mask=attention_mask,params=params,train=True).logits # (B, L, d)
        return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(out, input_ids)*attention_mask)
    grad_fn = jax.value_and_grad(loss_fn,has_aux=False)
    out = grad_fn(params,audio,input_ids,attention_mask)
    l,grads = jax.lax.pmean(out,'data')
    updates, opt_states = opt.update(grads, opt_states,params=params)
    params = optax.apply_updates(params, updates)
    return params,opt_states,l

devices = jax.devices()
replicated_params = jax.device_put_replicated(params, devices)
replicated_opt_states = jax.device_put_replicated(opt_states, devices)
for j,(audio,input_ids,attention_mask) in enumerate(train_loader):
        # data is of shape (n_device, batch, ...)
        replicated_params,replicated_opt_states,l = train_one_step(replicated_params,audio,input_ids,attention_mask,replicated_opt_states)

[ayaka14732]

I am also experiencing similar issues. I think you can open an issue in the JAX repository and mention this issue in the new one.