github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 23:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

[colotensor] add Tensor.view op and its unit test (#1343)

Browse Source 
[colotensor] add megatron initialization for gpt2
This commit is contained in:
HELSON
2022-07-21 10:53:15 +08:00
committed by GitHub
parent 6160a1d6a7
commit 7a8702c06d
16 changed files with 309 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/utils/checkpoint/utils.py
View File

@@ -1,7 +1,7 @@
import torch
import torch.distributed as dist
from colossalai.tensor import ColoTensor, ColoTensorSpec
from colossalai.tensor.distspec import _DistSpec
from colossalai.tensor.distspec import _DistSpec, DistPlacementPattern
def gather_tensor(colo_tensor: ColoTensor) -> None:
@@ -26,7 +26,7 @@ def gather_tensor(colo_tensor: ColoTensor) -> None:
def scatter_tensor(colo_tensor: ColoTensor, dist_spec: _DistSpec) -> None:
"""Reversal operation of `gather_tensor`.
"""
if dist_spec.placement == 'r':
if dist_spec.placement == DistPlacementPattern.REPLICATE:
dist.broadcast(colo_tensor.data, 0)
else:
global_size = colo_tensor.size_global()