JittorLLMs
JittorLLMs copied to clipboard
显存只增不减,最终申请不到,降级到内存
无法使用 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