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

[test] fix test: test_zero1_2

Browse Source
This commit is contained in:
hxwang
2024-07-22 05:36:20 +00:00
committed by Hongxin Liu
parent 74b03de3f9
commit 067e18f7e9
3 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/low_level/low_level_optim.py
View File

@@ -880,7 +880,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
return None
grad_flat = torch.empty((grad_store.world_size, *grad.shape), dtype=grad.dtype, device=grad.device)
dist.all_gather_into_tensor(grad_flat, grad, group=grad_store.torch_pg)
return grad_flat[: working_param.numel()].reshape_as(working_param)
return grad_flat.view(-1)[: working_param.numel()].view_as(working_param)
def get_working_grads_by_group_id(self, group_id: int) -> List[Tensor]:
working_grads = []