Fix ZeRO-3 UnboundLocalError (#968)

* Fix UnboundLocalError * Get full partition size
microsoft · Apr 16, 2021 · 2805c39 · 2805c39
1 parent c83e49f
commit 2805c39
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
@@ -808,17 +808,12 @@ def _partition_gradient(self, param, partition_buffer=None, accumulate=False):
         start = partition_size * rank
         end = start + partition_size
 
-        dest_tensor = partition_buffer.view(-1).narrow(0, 0, partition_size)
+        dest_tensor_full_buffer = partition_buffer.view(-1).narrow(0, 0, partition_size)
 
         #print("before partition gradients")
         if start < param.ds_numel:
             elements = min(param.ds_numel - start, partition_size)
 
-            dest_tensor_full_buffer = partition_buffer.view(-1).narrow(
-                0,
-                0,
-                partition_size)
-
             dest_tensor = dest_tensor_full_buffer.narrow(0, 0, elements)
             src_tensor = param.grad.view(-1).narrow(0, start, elements)