|
|
@ -58,6 +58,10 @@ class ColoTensor(object):
|
|
|
|
def shape(self):
|
|
|
|
def shape(self):
|
|
|
|
return torch.Size(self._size)
|
|
|
|
return torch.Size(self._size)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
|
|
|
def device(self):
|
|
|
|
|
|
|
|
return self._torch_tensor.device
|
|
|
|
|
|
|
|
|
|
|
|
def size(self, dim=None):
|
|
|
|
def size(self, dim=None):
|
|
|
|
if dim is None:
|
|
|
|
if dim is None:
|
|
|
|
return self.shape
|
|
|
|
return self.shape
|
|
|
@ -105,14 +109,14 @@ class ColoTensor(object):
|
|
|
|
device=self._device)
|
|
|
|
device=self._device)
|
|
|
|
return self._torch_tensor
|
|
|
|
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
|
|
|
|
self._shard_spec = spec
|
|
|
|
if lazy_shard == False:
|
|
|
|
if lazy_shard == False:
|
|
|
|
self._shard()
|
|
|
|
self._shard()
|
|
|
|
|
|
|
|
|
|
|
|
def _shard(self):
|
|
|
|
def _shard(self):
|
|
|
|
assert self._shard_spec is not None, 'You should call set_spec() before _shard() ColoTensor.'
|
|
|
|
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)
|
|
|
|
num_partition = gpc.get_world_size(ParallelMode.TENSOR)
|
|
|
|
local_rank = gpc.get_local_rank(ParallelMode.TENSOR)
|
|
|
|
local_rank = gpc.get_local_rank(ParallelMode.TENSOR)
|
|
|
|
dim = -1
|
|
|
|
dim = -1
|
|
|
@ -121,11 +125,11 @@ class ColoTensor(object):
|
|
|
|
# Reshape to get shard for this rank and we don't want autograd
|
|
|
|
# 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
|
|
|
|
# recording here for the narrow op and 'local_shard' should be a
|
|
|
|
# leaf variable in the autograd graph.
|
|
|
|
# leaf variable in the autograd graph.
|
|
|
|
self._torch_tensor = self._torch_tensor.narrow(dim,
|
|
|
|
self._torch_tensor = self._torch_tensor.narrow(dim, local_rank * chunk_size, chunk_size).detach(
|
|
|
|
local_rank * chunk_size, chunk_size).detach().contiguous() # TODO Shall we clone() here since detach() will point to the old tensor?
|
|
|
|
).contiguous() # TODO Shall we clone() here since detach() will point to the old tensor?
|
|
|
|
self._torch_tensor.requires_grad = self._requires_grad
|
|
|
|
self._torch_tensor.requires_grad = self._requires_grad
|
|
|
|
self._size = self._torch_tensor.size()
|
|
|
|
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
|
|
|
|
@classmethod
|
|
|
|
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
|
|
|
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
|
|
|