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

[Tensor] add embedding tp1d row (#904)

Browse Source
This commit is contained in:
Ziyue Jiang
2022-04-29 14:10:05 +08:00
committed by GitHub
parent 16122d5fac
commit f593a5637e
5 changed files with 108 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/tensor/colo_tensor.py
View File

@@ -166,6 +166,7 @@ class ColoTensor(object):
dim = -1
self._torch_tensor = gather_forward_split_backward(self._torch_tensor, parallel_action.parallel_mode, dim=dim)
self._shard_pattern = ShardPattern.NA
self._size = self._torch_tensor.size()
def is_gathered(self) -> bool:
return self._shard_pattern == ShardPattern.NA