github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Init Conext supports lazy allocate model memory (#842)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-22 18:03:35 +08:00
committed by GitHub
parent 4575a3298b
commit 8789850eea
5 changed files with 112 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/tensor/colo_tensor.py
View File

@@ -43,9 +43,10 @@ class ColoTensor(object):
torch_tensor=tensor if save_payload else torch.empty(0))
return colo_t
def del_torch_tensor(self) -> None:
self._size = (0,)
self._torch_tensor = torch.empty(self._size)
def del_torch_tensor(self, save_shape=False) -> None:
if save_shape:
self._size = (0,)
self._torch_tensor = torch.empty((0,))
def torch_tensor(self) -> torch.Tensor:
if self._torch_tensor.numel() == 0: