From de3d371f65ed35ae10282074ea78412b73a06130 Mon Sep 17 00:00:00 2001 From: Hongxin Liu Date: Tue, 10 Dec 2024 16:46:15 +0800 Subject: [PATCH] [hotfix] fix zero comm buffer init (#6154) --- colossalai/zero/low_level/low_level_optim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py index 3c67299bb..8abaf8fc6 100644 --- a/colossalai/zero/low_level/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -371,7 +371,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper): for i, sz in enumerate(bucket_store.sizes): grp = bucket_store.torch_pg if len(bucket_store.sizes) == 1 else bucket_store.torch_pg[i] flat_grads_list = list(cur_flat_grads.split(len(cur_flat_grads) // sz)) - received_grad = torch.zeros_like(flat_grads_list[0]) + received_grad = torch.empty_like(flat_grads_list[0]) if self._fp8_communication: reduce_scatter_fp8( received_grad,