format

deepspeedai · Apr 14, 2024 · 8edeea6 · 8edeea6
1 parent 5da0da8
commit 8edeea6
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -392,17 +392,14 @@ def __init__(self,
             if not fp16_master_weights_and_gradients:
                 weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                     self.device).clone().float().detach()
-                if self.cpu_offload:
-                    self.single_partition_of_fp32_groups.append(get_accelerator().pin_memory(weights_partition))
-                else:
-                    self.single_partition_of_fp32_groups.append(weights_partition)
             else:
-                weights_partition_half = self.parallel_partitioned_bit16_groups[i][partition_id].to(
+                weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to(
                     self.device).clone().half().detach()
-                if self.cpu_offload:
-                    self.single_partition_of_fp32_groups.append(get_accelerator().pin_memory(weights_partition_half))
-                else:
-                    self.single_partition_of_fp32_groups.append(weights_partition_half)
+
+            if self.cpu_offload:
+                weights_partition = get_accelerator().pin_memory(weights_partition)
+
+            self.single_partition_of_fp32_groups.append(weights_partition)
 
             # Set local optimizer to have flat params of its own partition.
             # After this, the local optimizer will only contain its own partition of params.