Change inplace to no_copy

apex/contrib/optimizers/distributed_fused_adam.py

@@ -154,11 +154,11 @@ class DistributedFusedAdam(torch.optim.Optimizer):
         self._works = []
 
         import inspect
-        if if 'inplace' in inspect.getfullargspec(torch.distributed.reduce_scatter).args:
-            self._pg_supports_inplace = True
+        if 'no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args:
+            self._pg_supports_no_copy = True
         else:
-            self._pg_supports_inplace = False
-            print("WARNING! torch.distributed.reduce_scatter does not support inplace op.")
+            self._pg_supports_no_copy = False
+            print("WARNING! torch.distributed.reduce_scatter does not support no_copy op.")
 
 
     def set_last_step(self, last_step):
@@ -188,8 +188,8 @@ class DistributedFusedAdam(torch.optim.Optimizer):
         end = start + self._block_size
         grad_block = flat_grads[start:end]
         grad_shards = [grad_block[i*self._shard_size:(i+1)*self._shard_size] for i in range(self._group_size)]
-        if self._pg_supports_inplace:
-            work = torch.distributed.reduce_scatter(grad_shards[self._rank_in_group],grad_shards,group=self._rs_pg[block_id%len(self._rs_pg)],async_op=True,inplace=True)
+        if self._pg_supports_no_copy:
+            work = torch.distributed.reduce_scatter(grad_shards[self._rank_in_group],grad_shards,group=self._rs_pg[block_id%len(self._rs_pg)],async_op=True,no_copy=True)
         else:
             work = torch.distributed.reduce_scatter(grad_shards[self._rank_in_group],grad_shards,group=self._rs_pg[block_id%len(self._rs_pg)],async_op=True)
         if self._num_groups > 1:
@@ -210,8 +210,8 @@ class DistributedFusedAdam(torch.optim.Optimizer):
         shard_end = shard_start + self._shard_size
         block_id = start // self._block_size
         self._partial_step_single_shard(block_id)
-        if self._pg_supports_inplace:
-            work = torch.distributed.all_gather(new_params_shards,new_params_shards[self._rank_in_group],group=self._ag_pg[block_id%len(self._ag_pg)],async_op=True,inplace=True)
+        if self._pg_supports_no_copy:
+            work = torch.distributed.all_gather(new_params_shards,new_params_shards[self._rank_in_group],group=self._ag_pg[block_id%len(self._ag_pg)],async_op=True,no_copy=True)
         else:
             work = torch.distributed.all_gather(new_params_shards,new_params_shards[self._rank_in_group],group=self._ag_pg[block_id%len(self._ag_pg)],async_op=True)
         return work