How do I debug a memory leak in the training loop?

[NightMachinery]

How do I debug a memory leak in the training loop?

I am using JAX on Colab with GPU.

I have tried stopping the training, and then running the following, but it doesn't show any big global variables:

import sys
def sizeof_fmt(num, suffix='B'):
    ''' by Fred Cirera,  https://stackoverflow.com/a/1094933/1870254, modified'''
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f %s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f %s%s" % (num, 'Yi', suffix)

for name, size in sorted(((name, sys.getsizeof(value)) for name, value in locals().items()),
                         key= lambda x: -x[1])[:10]:
    print("{:>30}: {:>8}".format(name, sizeof_fmt(size)))

                        losses: 173.9 KiB
                     my_locals:  9.1 KiB
                          _i35:  3.2 KiB
                          _i46:  2.3 KiB
              accuracies_valid:  1.9 KiB
              accuracies_train:  1.9 KiB
                          _i48:  1.6 KiB
                          _i31:  1.5 KiB
                   Classifier1:  1.4 KiB
                          _i44:  1.2 KiB

---

top -b -n 1 -o %MEM

top - 14:23:42 up  4:04,  0 users,  load average: 0.00, 0.09, 0.31
Tasks:  19 total,   1 running,  15 sleeping,   0 stopped,   3 zombie
%Cpu(s): 14.0 us,  2.0 sy,  0.0 ni, 83.2 id,  0.7 wa,  0.0 hi,  0.1 si,  0.0 st
KiB Mem : 13298572 total,   221188 free, 12250040 used,   827344 buff/cache
KiB Swap:        0 total,        0 free,        0 used.   780940 avail Mem 

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND
  12490 root      20   0 52.983g 0.011t 199800 S   0.0 88.0   9:25.17 python3
      8 root      20   0  401904  82200  14104 S   0.0  0.6   0:44.35 node
     56 root      20   0  237532  54748   7192 S   0.0  0.4   0:37.86 jupyter-+
     31 root      20   0  156140  34204   6160 S   0.0  0.3   0:00.39 colab-fi+
  12510 root      20   0  129744  17324   4620 S   0.0  0.1   0:04.17 python3
  10971 root      20   0  129492  17016   4564 S   0.0  0.1   0:08.16 python3
     57 root      20   0  716504   8508   2284 S   0.0  0.1   0:13.89 dap_mult+
  13278 root      20   0   63184   6628   4968 R   0.0  0.0   0:00.00 top
  13277 root      20   0   39196   6424   4940 S   0.0  0.0   0:00.15 bash
     19 root      20   0   35888   4412   3320 S   0.0  0.0   0:01.07 tail
  12605 root      25   5   51620   3752   2044 S   0.0  0.0   0:00.03 zsh
  12675 root      25   5   51620   3708   2012 S   0.0  0.0   0:00.01 zsh
  12606 root      25   5   49512   2088    440 S   0.0  0.0   0:00.00 zsh
  12607 root      25   5   49512   2088    440 S   0.0  0.0   0:00.00 zsh
  12608 root      25   5   49512   2088    440 S   0.0  0.0   0:00.00 zsh
      1 root      20   0     992      4      0 S   0.0  0.0   0:00.76 docker-i+
     30 root      20   0       0      0      0 Z   0.0  0.0   0:04.49 python3
  12604 root      20   0       0      0      0 Z   0.0  0.0   0:00.00 zsh
  12674 root      20   0       0      0      0 Z   0.0  0.0   0:00.00 zsh

---

I have attached a copy of the whole notebook I run. You need to run the headings Bootstrap and Fanfic classification. The memory usage will continue to grow as the training progresses until the machine crashes.

And here is my training loop:

rng = jax.random.PRNGKey(43)

num_training_steps = 3 * 10**6

#: @todo1 schedule this
batch_size = 2**6 
# batch_size = 2**12 

losses = []
accuracies_valid = []
accuracies_train = []

log_every_n = 100
save_every_n = 5000

my_locals = locals()

for step in range(0, num_training_steps + 1):
  rng, k1 = jrnd.split(rng)
  input_batch, target_batch = books2batch_jit(k1, books_bytes_train, batch_size)
  # target_batch = jax.nn.one_hot(jnp.zeros(target_batch.shape[0]), 2) #: @obviousCase

  if step == 0:
    k1, k2 = jrnd.split(k1)
    params, state = forward.init(k1, input_batch, k2, is_training=True)
    opt_state = optimizer.init(params)
    continue

  params, state, opt_state, model_output, loss = (
    train_step(params, state, opt_state, input_batch, target_batch, k1))

  if True:
    loss = jax.device_get(loss)
    losses.append(loss)
    if step % log_every_n == 0:
      print(f'[{step}/{num_training_steps}] ' + 
            (f'train loss: {loss}'), end='')
      
      if (step % log_every_n == 0) or (step % save_every_n == 0):
        model_output = jax.device_get(model_output)
        out = model_output
        probs = jax.nn.softmax(out)

        accu = accuracy_fn(probs, target_batch)
        print(f'; accu={accu}; p_m: {jnp.mean(probs, axis=0)}; p[0]: {probs[0]}; ' 
              + f'out_m: {jnp.mean(out, axis=0)}; out[0]: {out[0]}')

        # preds = probs2pred(probs)
        # print(preds)

        # print(target_batch)
        for ds_name in ('books_bytes', 'books_bytes_train', 'books_bytes_valid'):
          ds = my_locals[ds_name]
          accuracies = []
          for i in range(len(ds)):
            rng, rng_sub = jrnd.split(rng)
            accu = validate_jit(ds[i], 
                            i, rng_sub, params, state, forward.apply)
            print(f"{ds_name}, t={i}: {accu}")
            accuracies.append(accu)
          
          accu_mean = np.mean(accuracies)
          print(f"{ds_name}, total: {accu_mean}\n")
          if ds_name.endswith('_valid'):
            accuracies_valid.append(accu_mean)
          elif ds_name.endswith('_train'):
            accuracies_train.append(accu_mean)
        print("################")
        if (step % save_every_n == 0): #: step > 1 or 
          model_save(f'{step}_{accuracies_valid[-1]}')

[hawkinsp]

There are known memory leaks in JAX under the CUDA/CuDNN version in Colab. There isn't much we can do about them except wait for Colab to update. That may or may not be your problem.

It's not going to be possible for us to debug this without a complete self-contained runnable example.

JAX device objects have a .live_buffers() method that enumerates the list of GPU buffers, one thing you might try is looking to see if the list is growing over time.

[NightMachinery]

1. Just to confirm, I am talking about a main memory leak, not a GPU memory leak. Is .live_buffers() also useful for debugging that?

2. Can't I manually install the latest JAX version on Colab?

[hawkinsp]

1. yes, .live_buffers() is for GPU buffers.
2. in principle sure. Note you don't need to update JAX, you need to update CUDA and CuDNN.

Does this reproduce outside Colab?

[NightMachinery]

I don't have a machine with a GPU now. I'll try to report back when I have access to such a machine.

But this problem is not present when using Colab with a TPU.

---

Again, to emphasize, the leak happens in the main memory (RAM). The GPU memory is just fine.

[hawkinsp]

If you share a complete runnable repro, we could try it on a GPU VM in cloud easily enough. You haven't shared code I can run, so that's not an option at the moment...

[NightMachinery]

@hawkinsp I have actually attached my whole notebook!

https://github.com/google/jax/files/9269481/nb.zip

You need to run the headings Bootstrap and Fanfic classification. The memory usage will continue to grow as the training progresses until the machine crashes.