matmul consumes a lot of video memory

[980202006]

I found that when doing matmul calculations, it will consume a lot of video memory. Is there any way to do matmul to consume less or try not to consume more video memory?

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
# os.system('export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.2')
# os.system('export CUDA_VERSION=122')
import numpy as np
import time
import jax
from jax import device_put
from tqdm import tqdm
import jax.numpy as jnp
from utils import load_pkl, save_pkl
# torch.set_default_dtype(torch.float16
if __name__ == "__main__":
    cache_path = 'tens.pt'
    device = 'cuda:0'
    # ee = np.load('tens_ee.npy')
    # conn_matrix = np.load('tens_mat.npy')
    ee, conn_matrix = load_pkl('tens.pkl')
    ee = ee.astype(np.int8)
    conn_matrix = conn_matrix.astype(np.int8)
    print('done load')
    ee = device_put(ee, jax.devices()[0])
    1==1
    conn_matrix = device_put(conn_matrix, jax.devices()[0])
    print('done convert')
    start = time.time()
    num = 100
    for _ in tqdm(range(num)):
        ee = jnp.matmul(ee, conn_matrix) - ee
        1==1

[hawkinsp]

Take a look at https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

Does that answer your question?