使用 GPU 在 Google Colab 上出现 Cupy memory 错误 - 但只是我第二次运行代码

Question

I am trying to do a matrix multiplication on two large arrays using Cupy since it is significantly faster (about 100x) than using the CPU. My problem is that it works the first time I run it, but then the second time and so on it gives me a memory error. It is a step in a loop so this is a problem, I can't be restarting the runtime each time.

Here is the reproducible code with same array size and data type:

import cupy as cp
import datetime

cp.get_default_memory_pool().free_all_blocks()
cp.get_default_pinned_memory_pool().free_all_blocks()

x = cp.random.uniform(-1,1,size = (3000,300000))
w = cp.random.uniform(-1,1,size= (300000,1000))

start = datetime.now()
ans = cp.matmul(x,w)
stop = datetime.now()
print(stop-start)

Here is the error I get when I run it for the second time in the same runtime:

---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-5-43db33b58bc8> in <module>()
      2 cp.get_default_pinned_memory_pool().free_all_blocks()
      3 
----> 4 x = cp.random.uniform(-1,1,size = (3000,300000))
      5 w = cp.random.uniform(-1,1,size= (300000,1000))
      6 

4 frames
/usr/local/lib/python3.6/dist-packages/cupy/creation/basic.py in empty(shape, dtype, order)
     20 
     21     """
---> 22     return cupy.ndarray(shape, dtype, order=order)
     23 
     24 

cupy/core/core.pyx in cupy.core.core.ndarray.__init__()

cupy/cuda/memory.pyx in cupy.cuda.memory.alloc()

cupy/cuda/memory.pyx in cupy.cuda.memory.MemoryPool.malloc()

cupy/cuda/memory.pyx in cupy.cuda.memory.MemoryPool.malloc()

cupy/cuda/memory.pyx in cupy.cuda.memory.SingleDeviceMemoryPool.malloc()

cupy/cuda/memory.pyx in cupy.cuda.memory.SingleDeviceMemoryPool._malloc()

cupy/cuda/memory.pyx in cupy.cuda.memory._try_malloc()

OutOfMemoryError: Out of memory allocating 7,200,000,000 bytes (allocated so far: 9,624,000,000 bytes).

Can this be fixed? I'm trying to clear the GPU memory in the first two lines but not sure if thats correct. Maybe using a dask array would work instead? But can that be done while still using the GPU for speed?

Answer 1

Yes, using CuPy-backed Dask arrays would probably be fine here. You would want to make sure to use the single-threaded scheduler ( .compute(scheduler="single-threaded") . Assuming that you're making your arrays in such a way that Dask can load in chunks at a time then Dask will probably be able to load in a few chunks, do some of the computation, throw out intermediates, and then load in other chunks.