Modernize RPC example

Author: msaroufim

distributed/rpc/pipeline/main.py

@@ -9,7 +9,6 @@
 import torch.distributed.rpc as rpc
 import torch.multiprocessing as mp
 import torch.optim as optim
-from torch.distributed.optim import DistributedOptimizer
 from torch.distributed.rpc import RRef
 
 from torchvision.models.resnet import Bottleneck
@@ -185,15 +184,35 @@ def parameter_rrefs(self):
 image_h = 128
 
 
+def create_optimizer_for_remote_params(worker_name, param_rrefs, lr=0.05):
+    """Create optimizer on remote worker for given parameters"""
+    params = [p.to_here() for p in param_rrefs]
+    opt = optim.SGD(params, lr=lr)
+    # Use torch.compile to optimize the optimizer step
+    opt.step = torch.compile(opt.step)
+    return opt
+
+
 def run_master(split_size):
 
     # put the two model parts on worker1 and worker2 respectively
     model = DistResNet50(split_size, ["worker1", "worker2"])
     loss_fn = nn.MSELoss()
-    opt = DistributedOptimizer(
-        optim.SGD,
-        model.parameter_rrefs(),
-        lr=0.05,
+    
+    # Get parameter RRefs for each model shard
+    p1_param_rrefs = model.p1_rref.remote().parameter_rrefs().to_here()
+    p2_param_rrefs = model.p2_rref.remote().parameter_rrefs().to_here()
+    
+    # Create optimizers on remote workers
+    opt1_rref = rpc.remote(
+        "worker1",
+        create_optimizer_for_remote_params,
+        args=("worker1", p1_param_rrefs)
+    )
+    opt2_rref = rpc.remote(
+        "worker2",
+        create_optimizer_for_remote_params,
+        args=("worker2", p2_param_rrefs)
     )
 
     one_hot_indices = torch.LongTensor(batch_size) \
@@ -213,7 +232,14 @@ def run_master(split_size):
         with dist_autograd.context() as context_id:
             outputs = model(inputs)
             dist_autograd.backward(context_id, [loss_fn(outputs, labels)])
-            opt.step(context_id)
+            
+            # Step both optimizers
+            opt1_rref.rpc_sync().step()
+            opt2_rref.rpc_sync().step()
+            
+            # Zero gradients
+            opt1_rref.rpc_sync().zero_grad()
+            opt2_rref.rpc_sync().zero_grad()
 
 
 def run_worker(rank, world_size, num_split):

distributed/rpc/pipeline/requirements.txt

@@ -1,2 +1,2 @@
-torch==1.9.0
-torchvision==0.7.0
+torch
+torchvision