From 5ca2627743162a5e8efa934782150f40a38f31ed Mon Sep 17 00:00:00 2001 From: Reza Yazdani <44502768+RezaYazdaniAminabadi@users.noreply.github.com> Date: Tue, 15 Feb 2022 15:58:12 -0800 Subject: [PATCH] Fix CPU-Offload: Send groups of parameter lists as the FP16 parameters (#1774) Co-authored-by: Olatunji Ruwase --- deepspeed/runtime/zero/stage_1_and_2.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index 81ca62af..59020684 100755 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1721,10 +1721,9 @@ class DeepSpeedZeroOptimizer(object): if self.deepspeed_adam_offload: from deepspeed.ops.adam import DeepSpeedCPUAdam if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half: - bit16_param_groups = [ + bit16_param_groups = [[ bit16_partitions[partition_id] - for bit16_partitions in self.parallel_partitioned_bit16_groups - ] + ] for bit16_partitions in self.parallel_partitioned_bit16_groups] self.optimizer.step(fp16_param_groups=bit16_param_groups) else: self.optimizer.step() -- GitLab