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

[hotfix] fixed memory usage of shardformer module replacement (#5122)

Browse Source
This commit is contained in:
アマデウス
2023-11-28 15:38:26 +08:00
committed by GitHub
parent 7b789f4dd2
commit 126cf180bc
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/tensor/d_tensor/comm_spec.py
View File

@@ -112,7 +112,7 @@ def _split(tensor: torch.Tensor, comm_spec: CommSpec):
dim = comm_spec.shard_dim
length = tensor.shape[comm_spec.shard_dim] // dist.get_world_size(process_group)
start = length * dist.get_rank(process_group)
output = torch.narrow(tensor, dim, start, length).contiguous()
output = torch.narrow(tensor, dim, start, length).clone().contiguous()
return output