jax.jit makes traing very slow

[davidshen84]

Hi,

I have a simple dummy module like this.

class Dummy(nn.Module):
  ...

  @compact
  def __call__(self, x: jax.Array, v: Optional[float] = None):
    if v is None:
      return x, None

    # some Dense layers

    return x, y

If I use v=None during training for all samples, I can ignore this argument. So I can apply the jax.jit decorator to the train_step directly, and the training runs fair fast, at about 40 it/s.

But if I want to use a real float value, I have to use the partial decorator with jax.jit, like this:

@partial(jax.jit, static_argnames=["v"])
def train_step(state, batch, v, rngs):
  ...

The training speed is at about 13 s/it... that is 13 seconds per iteration.

I found no caveats about slowness when using the static_argnames argument. So I wonder if it is caused by jax.jit or maybe by the if branch.

[chiamp]

What sort of values are you passing in as v during the training loop?
Since you defined static_argnames=["v"], jax.jit will trigger a re-compile every time you pass in a different value for v in train_step. Read more about it here.

[davidshen84]

v is a random float number...so, it triggers a recompile every time... lesson learnt. Thanks

…

On Thu, 25 Jan 2024, 15:22 Marcus Chiam, ***@***.***> wrote: What sort of values are you passing in as v during the training loop? Since you defined static_argnames=["v"], jax.jit will trigger a re-compile every time you pass in a different value for v in train_step. Read more about it here <https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html#caching>. — Reply to this email directly, view it on GitHub <#3636 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAAQBTIPPALWLPM4P44GJODYQHMXRAVCNFSM6AAAAABCDWMKK2VHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4DENBQGU3DO> . You are receiving this because you authored the thread.Message ID: ***@***.***>