Change zero_grad() argument to match pytorch

deepspeed/runtime/fp16/fused_optimizer.py

@@ -130,14 +130,14 @@ def initialize_optimizer_states(self):
 
         return
 
-    def zero_grad(self, set_grads_to_None=True):
+    def zero_grad(self, set_to_none=False):
         """
         Zero FP16 parameter grads.
         """
         # For speed, set model fp16 grad to None by default
         for group in self.fp16_groups:
             for p in group:
-                if set_grads_to_None:
+                if set_to_none:
                     p.grad = None
                 else:
                     if p.grad is not None:

deepspeed/runtime/fp16/unfused_optimizer.py

@@ -110,15 +110,15 @@ def __init__(self,
 
         self.initialize_optimizer_states()
 
-    def zero_grad(self, set_grads_to_None=True):
+    def zero_grad(self, set_to_none=False):
         """
         Zero FP16 parameter grads.
         """
         # FP32 grad should never exist outside of the step function
         # For speed, set model fp16 grad to None by default
         for group in self.fp16_groups:
             for p in group:
-                if set_grads_to_None:
+                if set_to_none:
                     p.grad = None
                 else:
                     if p.grad is not None:

deepspeed/runtime/zero/stage3.py

@@ -1509,7 +1509,7 @@ def get_partition_info(self, tensor_list, partition_size, partition_id):
         return params_in_partition, params_not_in_partition, first_offset
 
     @instrument_w_nvtx
-    def zero_grad(self, set_grads_to_None=True):
+    def zero_grad(self, set_to_none=False):
         """
         Zero FP16 parameter grads.
         """
@@ -1519,7 +1519,7 @@ def zero_grad(self, set_grads_to_None=True):
         # For speed, set model fp16 grad to None by default
         for group in self.fp16_groups:
             for p in group:
-                if set_grads_to_None:
+                if set_to_none:
                     if p.grad is not None and p.grad.is_cuda:
                         p.grad.record_stream(torch.cuda.current_stream())
                     p.grad = None

deepspeed/runtime/zero/stage_1_and_2.py

@@ -1526,15 +1526,15 @@ def get_partition_info(self, tensor_list, partition_size, partition_id):
 
         return params_in_partition, params_not_in_partition, first_offset
 
-    def zero_grad(self, set_grads_to_None=True):
+    def zero_grad(self, set_to_none=False):
         """
         Zero FP16 parameter grads.
         """
         # FP32 grad should never exist.
         # For speed, set model fp16 grad to None by default
         for group in self.bit16_groups:
             for p in group:
-                if set_grads_to_None:
+                if set_to_none:
                     p.grad = None  # epilogue and in step
                 else:
                     if p.grad is not None: