github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[TP] change the check assert in split batch 2d (#772)

pull/781/head
Ziyue Jiang 2022-04-16 21:29:57 +08:00 committed by GitHub
parent 846406a07a
commit 4b01da24cd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/nn/layer/parallel_2d/_operation.py
View File

@ -739,11 +739,13 @@ def split_batch_2d(input_: Tensor, dim: int = 0) -> Tensor:
"""
dim_size = input_.size(dim)
world_size = gpc.get_world_size(ParallelMode.PARALLEL_2D_COL)
if world_size <= 1:
return input_
assert dim_size % world_size == 0, \
f'The batch size ({dim_size}) is not a multiple of 2D size ({world_size}).'
if input_.size(dim) <= 1:
return input_
return torch.chunk(input_, gpc.get_world_size(ParallelMode.PARALLEL_2D_COL),
dim=dim)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)].contiguous()

6
colossalai/nn/layer/parallel_2p5d/_operation.py
View File

@ -770,11 +770,13 @@ def split_batch_2p5d(input_: Tensor, dim: int = 0) -> Tensor:
"""
dim_size = input_.size(dim)
world_size = gpc.get_world_size(ParallelMode.PARALLEL_2P5D_COL)
if world_size <= 1:
return input_
assert dim_size % world_size == 0, \
f'The batch size ({dim_size}) is not a multiple of 2.5D size * depth ({world_size}).'
if input_.size(dim) <= 1:
return input_
return torch.chunk(input_, gpc.get_world_size(ParallelMode.PARALLEL_2P5D_COL),
dim=dim)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_COL)].contiguous()