System memory leak in Colab

[andylolu2]

System information (Google Colab free tier with T4 GPU)

• OS Platform and Distribution: Ubuntu 18.04.5 LTS
• Flax version: 0.5.2
• Jax version: 0.3.14
• Jaxlib version: 0.3.14+cuda11.cudnn805
• Python version: 3.7.13
• GPU/TPU model and memory: Tesla T4, 16 GB
• CUDA version (if applicable): 11.1

Problem you have encountered:

I have been trying to train some models on Google Colab (Free tier). There seems to be a system memory leak (Not GPU memory). More and more system memory gets used up on each call to a jit-ed apply function of a nn.Module. It increases indefinitely until Google Colab crashes.

What I have tried:

1. Just JAX

• I have tried to reproduce this error with just JAX and have referenced solutions mentioned in Memory leak in a simple loop with no side effects jax-ml/jax#2072 with no luck. I only managed to reproduce this error with Flax which is why I posted this issue here. This could totally be a JAX issue though.

2. Network structure

• When trying to create a minimal reproducible example, a one-layer network

  class Net(nn.Module):
      @nn.compact
      def __call__(self, x):
          x = nn.Conv(32, (3, 3))(x)
          return x

  does not show a memory leak.

  However, a two-layer network does:

  class Net(nn.Module):
      @nn.compact
      def __call__(self, x):
          x = nn.Conv(32, (3, 3))(x)
          x = nn.Conv(32, (3, 3))(x)
          return x

  I have not tried with other layers (e.g. using nn.Dense instead of nn.Conv) but I think the above example is minimal enough.

3. Other platforms

• I have not encountered the memory leak in other platforms. I have tried it on my local machine and on Kaggle, both ran without memory leak. So this could be Colab's issue as well, but I don't know how to debug this.

Screenshot of error:

The system memory usage growth is also reflected in Colab's resource monitor.

Steps to reproduce:

Colab notebook to reproduce: https://colab.research.google.com/drive/1oWSDpYIDUgFfAe26XqsX9Ufa-Pe_nF4d?usp=sharing

[jheek]

Inside the loop you don't actually touch any Flax code because you are compiling fwd with jax.jit. I think this is an issue with either JAX or XLA

[andylolu2]

Thanks for the reply! Unfortunately I really could not manage to reproduce the error with just JAX. I have further minimalised the example to the following:

conv_dim_nums = lax.ConvDimensionNumbers((0,3,1,2), (3,2,0,1), (0,3,1,2))

x = jnp.zeros((32,64,64,1))

def zeros(key, shape):
    return jnp.zeros(shape)

class Net(nn.Module):
    def setup(self):
        self.k1 = self.param('k1', zeros, (3,3,1,32))
        self.k2 = self.param('k2', zeros, (3,3,32,32))
        self.bias = self.param('bias', zeros, (1,1,1,32))

    def __call__(self, x):
        x = lax.conv_general_dilated(x, self.k1, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x = lax.conv_general_dilated(x, self.k2, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x += self.bias
        return x

def flax_apply():
    model = Net()
    params = model.init(jax.random.PRNGKey(0), x)
    return functools.partial(jax.jit(model.apply), params)

def jax_apply():
    k1 = jnp.zeros((3,3,1,32))
    k2 = jnp.zeros((3,3,32,32))
    bias = jnp.zeros((1,1,1,32))
    
    def apply(x):
        x = lax.conv_general_dilated(x, k1, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x = lax.conv_general_dilated(x, k2, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x += bias
        return x

    return jax.jit(apply)

Yet, the Flax version shows a memory leak, while the JAX version doesn't. It seems to suggest the issue is related to self.param but I am really out of ideas on how to further break it down.

Both flax_apply and jax_apply show the same output when passed to jax.make_jaxpr (apart from different variable namings).

I have also tried to further reduce the number of layers but it seems like 2 convolutions + 1 bias layer is the minimum example for the memory leak to show up.

Screenshot

Reproduce

Updated colab notebook: https://colab.research.google.com/drive/1oWSDpYIDUgFfAe26XqsX9Ufa-Pe_nF4d?usp=sharing

[jheek]

I was able to reproduce the issue with pure JAX as follows:

def jax_apply():
    k1 = jnp.zeros((3,3,1,32))
    k2 = jnp.zeros((3,3,32,32))
    bias = jnp.zeros((1,1,1,32))
    
    def apply(k1, k2, bias, x):
        x = lax.conv_general_dilated(x, k1, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x = lax.conv_general_dilated(x, k2, (1, 1), "SAME", dimension_numbers=conv_dim_nums)
        x += bias
        return x

    return functools.partial(jax.jit(apply), k1, k2, bias)

so passing the params as input instead of closure

[andylolu2]

Awesome, thank you! I'll move this issue to JAX then.