github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 22:52:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Tensor] init a tp network training unittest (#849)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-24 16:43:44 +08:00
committed by GitHub
parent 0dea140760
commit 62f059251b
6 changed files with 113 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/tensor/colo_tensor.py
View File

@@ -1,7 +1,9 @@
from numpy import product
from .op_wrapper import _COLOSSAL_OPS
import torch
from typing import Tuple, Optional
from .op_wrapper import _COLOSSAL_OPS
from numpy import product
class ColoTensor(object):
""" Data Structure for Tensor in Colossal-AI
@@ -52,7 +54,6 @@ class ColoTensor(object):
return product(self._size)
@staticmethod
def init_from_torch_tensor(tensor: torch.Tensor, save_payload=True) -> 'ColoTensor':
colo_t = ColoTensor(*tensor.size(),
dtype=tensor.dtype,

2
colossalai/utils/model/colo_init_context.py
View File

@@ -26,4 +26,4 @@ class ColoInitContext(InsertPostInitMethodToModuleSubClasses):
save_torch_payload = True if not self._lazy_memory_allocate else False
for name, param in name_list:
delattr(module, name)
setattr(module, name, ColoTensor.init_from_torch_tensor(tensor=param.data, save_payload=save_torch_payload))
setattr(module, name, ColoTensor.init_from_torch_tensor(tensor=param, save_payload=save_torch_payload))