Repository: pytorch/PiPPy

Issue:
Implement shared parameter synchronization as part of runtime
Right now we're just hacking it up and doing it manually in the test:

https://github.com/jamesr66a/PiPPy/blob/5f4c6cd4676d6135dec9ee86341286416afd296f/test/local_test_forward_backward.py#L118

We could possibly implement a SyncTiedParameters phase in WorkItem, similar to DeepSpeed, to automatically do the necessary allreduce to synchronize parameters shared across stages


Discussion:

Resolving patch:
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -340,6 +340,14 @@ def extract_rref_args(arg):
         return future
 
 
+def get_grad_from_executor(executor, qualname):
+    return executor.local_value().mod.get_parameter(qualname).grad
+
+def set_grad_in_executor(executor, qualname, value):
+    param = executor.local_value().mod.get_parameter(qualname)
+    param.grad = value
+
+
 class PipelineDriverBase:
     def __init__(self, pipe : Pipe, args_chunk_spec, kwargs_chunk_spec, output_chunk_spec, world_size : int,
                  all_ranks : List[int] = None):
@@ -415,6 +423,25 @@ def run(self, *args, chunks : int, batch_dims : Optional[List[Optional[int]]] =
         raise NotImplementedError('PipelineDriverBase is an abstract base class, please use a concrete '
                                   'implementation class.')
 
+
+    def _sync_replicated_params(self):
+        logging.info(f'[root] Synchronizing gradients for {len(self.pipe.replicated_params)} sets of replicated parameters')
+        for param_set in self.pipe.replicated_params:
+            grad_values = []
+            for module_name, param_qualname in param_set.items():
+                assert module_name in self.remote_stage_executor_rrefs
+                rank, module_rref = self.remote_stage_executor_rrefs[module_name]
+                grad_value = rpc.rpc_sync(rank, get_grad_from_executor, (module_rref, param_qualname))
+                grad_values.append(grad_value)
+
+            synced_value = torch.sum(torch.stack(grad_values), dim=0)
+
+            for module_name, param_qualname in param_set.items():
+                assert module_name in self.remote_stage_executor_rrefs
+                rank, module_rref = self.remote_stage_executor_rrefs[module_name]
+                rpc.rpc_sync(rank, set_grad_in_executor, (module_rref, param_qualname, synced_value))
+
+
 class RemoteInterpreter(torch.fx.Interpreter):
     def __init__(self, remote_stage_executor_rrefs, module, cur_microbatch : int,   
                  args, kwargs, garbage_collect_values=True):
@@ -551,7 +578,8 @@ def node_is_loss_or_output(n):
         if last_nodes[0].op == 'output':
             logging.info('[root] Program does not have loss/backward, returning outputs directly')
             # Forward-only; return output values
-            return self._retrieve_output_values(microbatch_interpreters, last_nodes, _debug_mask_minibatches)
+            local_results = self._retrieve_output_values(microbatch_interpreters, last_nodes)
+            return merge_chunks(local_results, self.output_chunk_spec, _debug_mask_minibatches)
 
         logging.info('[root] Executing loss + backward stages')
         last_nodes = []
@@ -560,15 +588,21 @@ def node_is_loss_or_output(n):
 
         assert all(n == last_nodes[0] for n in last_nodes)
         assert last_nodes[0].op == 'output'
-        return self._retrieve_output_values(microbatch_interpreters, last_nodes, _debug_mask_minibatches)
 
-    def _retrieve_output_values(self, microbatch_interpreters, last_nodes, _debug_mask_minibatches):
-        logging.info(f'[root] Combining output values from {len(microbatch_interpreters)} chunks')
+        local_results = self._retrieve_output_values(microbatch_interpreters, last_nodes)
+
+        # Shared parameter sync
+        # At this point, all of the gradient jobs should have been run
+        # (by way of the synchronization dependency earlier)
+        self._sync_replicated_params()
+
+        return merge_chunks(local_results, self.output_chunk_spec, _debug_mask_minibatches)
+
+    def _retrieve_output_values(self, microbatch_interpreters, last_nodes):
+        logging.info(f'[root] Retrieving output values from {len(microbatch_interpreters)} chunks')
         output_vals = []
         for interp, last_node in zip(microbatch_interpreters, last_nodes):
             interp.run_until(lambda n : False)
             output_vals.append(interp.env[last_node])
 
-        local_results = torch.fx.node.map_aggregate(output_vals, to_here)
-
-        return merge_chunks(local_results, self.output_chunk_spec, _debug_mask_minibatches)
+        return torch.fx.node.map_aggregate(output_vals, to_here)


Test patch:
diff --git a/test/local_test_forward_backward.py b/test/local_test_forward_backward.py
--- a/test/local_test_forward_backward.py
+++ b/test/local_test_forward_backward.py
@@ -117,27 +117,6 @@ def forward(self, x):
 
     all_grad_qualnames = {k: None for k, v in ec_pipe.named_parameters()}
 
-    replicated_params_qualnames = {}
-
-    # Shared parameter sync. TODO: move this to actual runtime
-    for param_set in ec_pipe.replicated_params:
-        grad_values = []
-        for module_name, param_qualname in param_set.items():
-            assert module_name in pipe_driver.remote_stage_executor_rrefs
-            rank, module_rref = pipe_driver.remote_stage_executor_rrefs[module_name]
-            grad_value = rpc.rpc_sync(rank, get_grad_from_executor, (module_rref, param_qualname))
-            grad_values.append(grad_value)
-            all_grad_qualnames.setdefault(f'split_gm.{module_name}.{param_qualname}')
-
-        synced_value = torch.sum(torch.stack(grad_values), dim=0)
-
-        for module_name, param_qualname in param_set.items():
-            assert module_name in pipe_driver.remote_stage_executor_rrefs
-            rank, module_rref = pipe_driver.remote_stage_executor_rrefs[module_name]
-            rpc.rpc_sync(rank, set_grad_in_executor, (module_rref, param_qualname, synced_value))
-
-            replicated_params_qualnames.setdefault(f'split_gm.{module_name}.{param_qualname}')
-
     pipe_grads = {}
 
     for name in all_grad_qualnames: