Converting np.ndarrays to jnp.arrays causes memory leak

[jzhao321]

I've been playing around on Google Colab with GPU instances.
When I call jax.device_put() on a converted numpy array, there seems to be memory allocated that cannot be recovered even after running jax.clear_backends() and gc.collect()

not sure if this is intended behavior, I wasn't able to find anything about this in the docs.

Jax version is one prepackaged with Google Colab

import gc
from psutil import Process
import numpy as np
import jax.numpy as jnp
import jax

os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = "false"
os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = "platform"

def mem_profile():
  return print_mem(Process().memory_info().rss)

def print_mem(mem):
  mb = mem / (1024 * 1024)
  return "{:.2f} MB".format(mb)

def move_np_to_jnp(state):
  flattened, structure = jax.tree_util.tree_flatten(state)

  unsharded = []
  for flattened_array in flattened:
    jnp_array = jnp.asarray(flattened_array)
    unsharded.append(jax.device_put(jnp_array, device=jax.devices()[0]))

  state = jax.tree_util.tree_unflatten(structure, unsharded)
  return state
  
for index in range(10):
  print("Before Loop: ", index, mem_profile())
  state_1 = move_np_to_jnp(state = {
    "test_key_1": np.ones((1024,512,256)),
    "test_key_2": np.ones((1024,512,256))
  })

  state_2 = move_np_to_jnp(state = {
    "test_key_3": np.ones((1024,1024,128)),
    "test_key_4": np.ones((1024,1024,128))
  })

  del state_1, state_2
  jax.clear_backends()
  gc.collect()
  print("After Loop: ", index, mem_profile())

print("After All: ", mem_profile())

[jakevdp]

I think this is working as intended. The device does not grow its memory pool until it is needed, but it will never shrink its memory pool once it's grown. After the initial allocations, the memory pool never expands again, because new allocations fit in the existing pool.

Does that make sense?

Side note, your move_to_jnp function could be written more efficiently like this:

def move_np_to_jnp(state):
  return jax.device_put(state)

(device_put natively handles pytrees)

[jzhao321]

@jakevdp I have a follow up question, not sure if topical.
When running just this:

state_1 = jax.device_put({
    "test_key_1": np.ones((1024,512,256)),
    "test_key_2": np.ones((1024,512,256))
  })

I see in colab the CPU and GPU memory both go up, suggesting that JAX is storing the data on both. Is this intended behavior?
If so, I don't see how I would be able to load OPT 6.7B (13.31 gb fp16) in colab's limited 12.7 gb CPU buffer.

[jakevdp]

Do you have X64 mode enabled? If not, then pushing 64-bit arrays to JAX will result in a 32-bit copy on host-side, along with a 32-bit copy on device. To avoid this you can either enable X64 mode, or you can use 32-bit source arrays.

[jzhao321]

@jakevdp It seems to be allocating regardless of X64 mode and dtype of the numpy array.
It also seems to be happening even with data defined by jnp.ones(...)

EDIT:
I can't seem to reproduce CPU memory allocation with jnp.ones(...) after Deleting / Starting a new runtime in Colab.

EDIT 2:
There's definitely some weird behavior going on. JAX memory allocation seems to be inconsistant with jnp.ones(...). Sometimes it takes up CPU ram, sometimes it doesn't. It changes between runs. Haven't figured out what causes it yet.
Running jax.device_put(...) with numpy arrays always allocate CPU memory, even if declared inline.

EDIT 3:
When declaring using jnp.ones(...), it "consumes" cpu memory because I had the Colab Variable inspector open. The variable inspector forces jax to fetch the values to cpu ram, thus explaining the memory consumption.

However, when declaring a numpy array then calling device.put() on the value, it still consumes CPU memory even after removing the reference.

import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = "false"
os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = "platform"
os.environ['JAX_ENABLE_X64'] = "True"
import gc
from psutil import Process
import jax.numpy as jnp
import jax
from jax._src.lib import xla_client
import numpy as np
import time

def mem_profile():
  return print_mem(Process().memory_info().rss)

def print_mem(mem):
  mb = mem / (1024 * 1024)
  return "{:.2f} MB".format(mb)

def clear_mem():
  jax.clear_backends()
  xla_client._xla.collect_garbage()
  gc.collect()

# Test 1
print("Baseline", mem_profile())

state_1 = jax.device_put(jnp.ones((1024,1024,1024), dtype=jnp.float16), device=jax.devices()[0])
time.sleep(5) # Ensure Async ops are executed
print("After Init JNP", mem_profile())

del state_1
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared JNP", mem_profile())

# Test 2
print("Baseline 2", mem_profile())

state_2_numpy = {
  "test_key_1": np.ones((1024,1024,512), dtype=np.float16),
  "test_key_2": np.ones((1024,1024,512), dtype=np.float16)
}
state_2 = jax.device_put(state_2_numpy, device=jax.devices()[0])
time.sleep(5) # Ensure Async ops are executed
print("After Init Numpy Dict", mem_profile())

del state_2_numpy
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Dict On CPU", mem_profile())

del state_2
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Dict On GPU", mem_profile())

# Test 3
print("Baseline 3", mem_profile())

state_3 = jax.device_put({
  "test_key_1": np.ones((1024,1024,512), dtype=np.float16),
  "test_key_2": np.ones((1024,1024,512), dtype=np.float16)
}, device=jax.devices()[0])
time.sleep(5) # Ensure Async ops are executed
print("After Init Numpy Dict Inline", mem_profile())

del state_3
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Dict Inline", mem_profile())

# Test 4
print("Baseline 4", mem_profile())

state_4_numpy = np.ones((1024,1024,1024), dtype=np.float16)
state_4 = jax.device_put(state_4_numpy, device=jax.devices()[0])
time.sleep(5) # Ensure Async ops are executed
print("After Init Numpy Array", mem_profile())

del state_4_numpy
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Array On CPU", mem_profile())

del state_4
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Array On GPU", mem_profile())

# Test 5
print("Baseline 5", mem_profile())

state_5 = jax.device_put(np.ones((1024,1024,1024), dtype=np.float16), device=jax.devices()[0])
time.sleep(5) # Ensure Async ops are executed
print("After Init Numpy Array Inline", mem_profile())

del state_5
clear_mem()
time.sleep(5) # Ensure Async ops are executed
print("Cleared Numpy Array Inline", mem_profile())

Output:

Baseline 1126.88 MB
After Init JNP 1130.17 MB
Cleared JNP 1130.19 MB
Baseline 2 1130.19 MB
After Init Numpy Dict 5290.35 MB
Cleared Numpy Dict On CPU 3242.23 MB
Cleared Numpy Dict On GPU 1130.21 MB
Baseline 3 1130.21 MB
After Init Numpy Dict Inline 5290.37 MB
Cleared Numpy Dict Inline 1130.21 MB
Baseline 4 1130.21 MB
After Init Numpy Array 5290.14 MB
Cleared Numpy Array On CPU 3242.23 MB
Cleared Numpy Array On GPU 1130.21 MB
Baseline 5 1130.21 MB
After Init Numpy Array Inline 5290.15 MB
Cleared Numpy Array Inline 1130.22 MB

From this, something is definetly taking up memory on the CPU as a result of JAX ops behind the scenes.

[jzhao321]

Output from nvidia-smi:

Fri Mar 10 00:36:11 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   78C    P0    34W /  70W |    361MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

Output from jax.devices()

[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]

[jzhao321]

@jakevdp Problem can be illustrated more though this snippet:

import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = "false"
os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = "platform"
os.environ['JAX_ENABLE_X64'] = "True"
import gc
from psutil import Process
import jax.numpy as jnp
import jax
from jax._src.lib import xla_client
import numpy as np
import time

print("Jax Version:", jax.__version__)

def mem_profile():
  return print_mem(Process().memory_info().rss)

def print_mem(mem):
  mb = mem / (1024 * 1024)
  return "{:.2f} MB".format(mb)

def clear_mem():
  jax.clear_backends()
  xla_client._xla.collect_garbage()
  gc.collect()

for index in range(10):
  print("Baseline", index, mem_profile())
  state_4_numpy = np.ones((512,512,512), dtype=np.float16)
  state_4_jax_numpy = jnp.asarray(state_4_numpy, dtype=jnp.float16)
  state_4 = jax.device_put(state_4_jax_numpy, device=jax.devices()[0])
  print("After Init", mem_profile())

  del state_4_numpy, state_4_jax_numpy, state_4
  clear_mem()
  time.sleep(10) # Ensure Async ops are executed
  print("After Delete", mem_profile())

Output:

Jax Version: 0.4.6
Baseline 0 150.53 MB
After Init 772.03 MB
After Delete 772.26 MB
Baseline 1 772.26 MB
After Init 1036.50 MB
After Delete 1036.45 MB
Baseline 2 1036.45 MB
After Init 1300.70 MB
After Delete 1300.61 MB
Baseline 3 1300.61 MB
After Init 1564.86 MB
After Delete 1564.76 MB
Baseline 4 1564.76 MB
After Init 1831.33 MB
After Delete 1831.24 MB
Baseline 5 1831.24 MB
After Init 2095.52 MB
After Delete 2095.43 MB
Baseline 6 2095.43 MB
After Init 2359.67 MB
After Delete 2359.62 MB
Baseline 7 2359.62 MB
After Init 2623.86 MB
After Delete 2623.77 MB
Baseline 8 2623.77 MB
After Init 2888.05 MB
After Delete 2887.96 MB
Baseline 9 2887.96 MB
After Init 3154.54 MB
After Delete 3154.47 MB