lax.cond increases GPU memory usage a lot and causes OOM

[jjyyxx]

My model cannot fit in GPU memory after applying a "layer dropout" trick.
Inside a haiku module, before:

# Variant 0, works
for layer_type in self._block:
    x = layer(layer_type, x)

after

# Variant 1, not work (OOM)
for layer_type in self._block:
  should_drop_layer = jax.random.bernoulli(hk.next_rng_key(), self.layer_dropout)
  x = hk.cond(
      should_drop_layer,
      lambda x: x,
      lambda x: layer(layer_type, x),
      x,
  )

this can be workarounded by changing to unconditionally compute layer(layer_type, x)

# Variant 2, works
for layer_type in self._block:
    should_drop_layer = jax.random.bernoulli(hk.next_rng_key(), self.layer_dropout)
    x = jax.lax.cond(should_drop_layer, lambda x, _: x, lambda _, x: x, x, layer(layer_type, x))

But I wonder why jax needs to use more GPU memory (more than 1.6 times) for variant 1?

[cgarciae]

Hey @jjyyxx, I don't know the exact technical details but I've also experienced increased memory and slower runtime when using lax.cond. My intuition is that the compiler cannot cross optimize each branch and has to treat this type of code very differently. In many cases you can use jnp.where instead of cond and get better results.

[jjyyxx]

But in my opinion, when the condition is a scalar, lax.cond should be strictly better than jnp.where which broadcasts the condition then uses lax.select internally?