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

colo init context add device attr. (#866)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-25 14:24:26 +08:00
committed by GitHub
parent 2238758c2e
commit d01d3b8cb0
3 changed files with 36 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
colossalai/tensor/colo_tensor.py
View File

@@ -58,6 +58,10 @@ class ColoTensor(object):
def shape(self):
return torch.Size(self._size)
@property
def device(self):
return self._torch_tensor.device
def size(self, dim=None):
if dim is None:
return self.shape
@@ -105,14 +109,14 @@ class ColoTensor(object):
device=self._device)
return self._torch_tensor
def set_spec(self, spec: str, lazy_shard: bool=False) -> None:
def set_spec(self, spec: str, lazy_shard: bool = False) -> None:
self._shard_spec = spec
if lazy_shard == False:
self._shard()
def _shard(self):
assert self._shard_spec is not None, 'You should call set_spec() before _shard() ColoTensor.'
if self._shard_spec == "1Drow": # TODO It actually represents the sharding layout for Linear-1Drow-weight, but we make it simpler now.
if self._shard_spec == "1Drow": # TODO It actually represents the sharding layout for Linear-1Drow-weight, but we make it simpler now.
num_partition = gpc.get_world_size(ParallelMode.TENSOR)
local_rank = gpc.get_local_rank(ParallelMode.TENSOR)
dim = -1
@@ -121,11 +125,11 @@ class ColoTensor(object):
# Reshape to get shard for this rank and we don't want autograd
# recording here for the narrow op and 'local_shard' should be a
# leaf variable in the autograd graph.
self._torch_tensor = self._torch_tensor.narrow(dim,
local_rank * chunk_size, chunk_size).detach().contiguous() # TODO Shall we clone() here since detach() will point to the old tensor?
self._torch_tensor = self._torch_tensor.narrow(dim, local_rank * chunk_size, chunk_size).detach(
).contiguous() # TODO Shall we clone() here since detach() will point to the old tensor?
self._torch_tensor.requires_grad = self._requires_grad
self._size = self._torch_tensor.size()
self._device = device # TODO A `fake` device now because torch_tensor.device always = cpu
self._device = device # TODO A `fake` device now because torch_tensor.device always = cpu
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):