Better shortcut for shared_param = True

apex/parallel/distributed.py

@@ -79,17 +79,19 @@ class DistributedDataParallel(Module):
     def create_hooks(self):
         #all reduce gradient hook
         def allreduce_params():
-            if(self.needs_reduction):
-                self.needs_reduction = False
-                self.needs_refresh = False
-            else:
+            if not self.needs_reduction:
                 return
-            grads = [param.grad.data for param in self.module.parameters() if param.grad is not None]
-            flat_dist_call(grads, dist.all_reduce)
+            self.needs_reduction = False
+
+            #parameter ordering refresh
+            if self.needs_refresh and not self.shared_param:
                 t_record = torch.cuda.IntTensor(self.record)
                 dist.broadcast(t_record, 0)
                 self.record = [int(entry) for entry in t_record]
+                self.needs_refresh = False
 
+            grads = [param.grad.data for param in self.module.parameters() if param.grad is not None]
+            flat_dist_call(grads, dist.all_reduce)
 
         def flush_buckets():
             if not self.needs_reduction:
@@ -184,10 +186,10 @@ class DistributedDataParallel(Module):
         #Force needs_refresh to True if there are shared params
         #this will force it to always, only call flush_buckets which is safe
         #for shared parameters in the model.
-        if self.shared_param:
-            self.param_refs = []
-            
-        self.needs_refresh = True if not self.param_refs else any(
+        if not self.param_refs or self.shared_param:
+            self.needs_refresh = True
+        else:
+            self.needs_refresh = any(
                 [param1 is not param2 for param1, param2 in zip(param_list, self.param_refs)]
             )