jitted function get recompiled and compile is very slow

[zhenlan0426]

Trying to fine tune a pre-trained model. train_one_step_embed function somehow get recompiled for each batch and each time trigger a warnning "Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results." Any help is appreciated!

this function do one gradient update step

@jax.jit
def train_one_step_embed(embedding,params,audio,input_ids,attention_mask,opt_states):
    def loss_fn(embedding,params,audio,input_ids,attention_mask):
        params['model']['decoder']['embed_tokens']['embedding'] = embedding
        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[:,3:-1], input_ids[:,4:])*attention_mask[:,4:])
    grad_fn = jax.value_and_grad(loss_fn,has_aux=False)
    l,grads = grad_fn(embedding,params,audio,input_ids,attention_mask)
    updates, opt_states = opt.update(grads, opt_states,params=embedding)
    embedding = optax.apply_updates(embedding, updates)
    return embedding,opt_states,l

main training loop, call train_one_step_embed for each batch.
But the jit function is outside the loop, NOT called for each batch.
How come the recompilation happens for each batch?

for i in range(epochs):
    # train
    train_loss = 0
    for j,(audio,input_ids,attention_mask) in enumerate(train_loader):
        audio,input_ids,attention_mask = jnp.array(audio,dtype=dtype),jnp.array(input_ids),jnp.array(attention_mask)
        embedding,opt_states,l = train_one_step_embed(embedding,params,audio,input_ids,attention_mask,opt_states)
        train_loss += l.item()

I would get a warning like this for EACH BATCH,

2023-07-30 00:18:32.871780: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 5m59.993335541s [Compiling module jit_train_one_step_embed] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

[zhenlan0426]

Not sure if this is related. But I would also run out of RAM (not VRAM and I have 32G of RAM) and notebook would crash. I only load a batch of 32 from disk. So the RAM is not used by the data but used by the compilation.