From 43dac83effb2ab3b19cae21c4cf0c267b4cb3359 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Tue, 28 Jun 2022 02:17:54 +0500 Subject: [PATCH] correct partition_id in fp32 param -> fp16 param for MoE+z2 --- deepspeed/runtime/zero/stage_1_and_2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index 7b683d12eeb4..c6a2e26aac8f 100755 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1762,7 +1762,8 @@ def step(self, closure=None): for group in self.single_partition_of_fp32_groups: group.grad = None # in step - for bit16_partitions, fp32_partition in zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups): + for group_id, (bit16_partitions, fp32_partition) in enumerate(zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) bit16_partitions[partition_id].data.copy_(fp32_partition.data) self.stop_timers([OPTIMIZER_STEP])