From 7606795a2c49c28c275441edc5a81ead2325e4a5 Mon Sep 17 00:00:00 2001 From: Nadav Elyahu <88962733+nelyahu@users.noreply.github.com> Date: Tue, 17 Dec 2024 09:51:46 +0200 Subject: [PATCH] Update stage_1_and_2.py fix index value to 'i' instead of 'group_idx' --- deepspeed/runtime/zero/stage_1_and_2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index c4163d6a850f..ecb2a527f870 100755 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1069,7 +1069,7 @@ def average_tensor(self, tensor): process_group = self.dp_process_group # count = 0 for i, param_idx_in_group, param_id in self.params_in_ipg_bucket: - param = self.bit16_groups[group_idx][param_idx_in_group] + param = self.bit16_groups[i][param_idx_in_group] process_group = self.dp_process_group