JittorLLMs icon indicating copy to clipboard operation
JittorLLMs copied to clipboard

显存只增不减,最终申请不到,降级到内存

Open YingchaoX opened this issue 1 year ago • 0 comments

无法使用 torch.cuda.empty_cache()来清理缓存,使用的site-packages/jtorch/cuda.py中没有清理显存相关的操作。

# cuda.py

import jittor as jt
import jtorch

def is_available():
    return jt.has_cuda

def device_count():
    return int(jt.has_cuda)

def set_device(device=None):
    pass

def get_rng_state(device=None):
    pass

def current_device():
    return jtorch.device("cuda")

class Generator:
    def __init__(self):
        pass

    def set_state(self, state):
        self.state = state

default_generators = [Generator()]
_lazy_call = lambda func: func()
device = None

LongTensor = jt.int64
FloatTensor = jt.float
HalfTensor = jt.float16
BoolTensor = jt.bool

manual_seed = jt.set_global_seed

YingchaoX avatar May 08 '23 13:05 YingchaoX