github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-31 16:40:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[tensor] add zero_like colo op, important for Optimizer (#1236)

Browse Source
This commit is contained in:
Jiarui Fang
2022-07-08 14:55:27 +08:00
committed by GitHub
parent 3b500984b1
commit 4a76084dc9
4 changed files with 16 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/tensor/colo_parameter.py
View File

@@ -55,7 +55,7 @@ class ColoParameter(ColoTensor, torch.nn.Parameter):
return tensor
def __repr__(self):
return f'ColoParameter: {torch.Tensor.__repr__(self)}'
return f'ColoParameter: {ColoTensor.__repr__(self)}'
@classmethod
def __torch_function__(cls, func, types, args=..., kwargs=None):

3
colossalai/tensor/colo_tensor.py
View File

@@ -271,3 +271,6 @@ class ColoTensor(torch.Tensor):
def is_shard_1drow(self):
return self.dist_spec.placement == DistPlacementPattern.SHARD \
and len(self.dist_spec.dims) == 1 and self.dist_spec.dims[0] == 0
def is_sharded(self):
return self.dist_spec.placement == DistPlacementPattern.SHARD