Add parameter broadcasting to PJRT examples. (#3836)

* Add parameter broadcasting to PJRT examples.

* Naming fix in tests

Author: darisoy

test/pjrt/test_parameter_broadcast_tpu.py

@@ -0,0 +1,36 @@
+from numpy.testing import assert_array_equal, assert_raises
+import torch
+import torch.nn as nn
+import torch_xla.core.xla_model as xm
+from torch_xla.experimental import pjrt
+from absl.testing import absltest, parameterized
+
+
+def broadcast(sync):
+  torch.manual_seed(xm.get_ordinal())
+  device = xm.xla_device()
+  model = nn.Linear(5, 5).to(device)
+  if sync:
+    pjrt.broadcast_master_param(model)
+  return next(model.parameters()).detach().cpu().numpy()
+
+
+class TestBroadcastParametersPjrt(parameterized.TestCase):
+
+  @parameterized.named_parameters(('synchronized_parameters', True),
+                                  ('unsynchronized_parameters', False))
+  def test_broadcast_parameter_sync(self, sync):
+    torch.set_default_tensor_type('torch.FloatTensor')
+    results = pjrt.run_multiprocess(broadcast, sync)
+    master_params = results[0][0]
+    for process_key in results:
+      worker_params = results[process_key][0]
+      if sync:
+        assert_array_equal(master_params, worker_params)
+      elif process_key != 0:
+        assert_raises(AssertionError, assert_array_equal, master_params,
+                      worker_params)
+
+
+if __name__ == '__main__':
+  absltest.main()

test/pjrt/test_train_pjrt_imagenet.py

@@ -114,7 +114,7 @@ def _train_update(device, step, loss, tracker, epoch, writer):
       summary_writer=writer)
 
 
-def train_imagenet(state_dict):
+def train_imagenet():
   print('==> Preparing data..')
   img_dim = get_model_property('img_dim')
   if FLAGS.fake_data:
@@ -182,8 +182,8 @@ def train_imagenet(state_dict):
 
   device = xm.xla_device()
   model = get_model_property('model_fn')()
-  model.load_state_dict(state_dict)
   model = model.to(device)
+  pjrt.broadcast_master_param(model)
   writer = None
   if xm.is_master_ordinal():
     writer = test_utils.get_summary_writer(FLAGS.logdir)
@@ -262,10 +262,8 @@ def test_loop_fn(loader, epoch):
 
 if __name__ == '__main__':
   torch.set_default_tensor_type('torch.FloatTensor')
-  torch.manual_seed(42)
-  model = get_model_property('model_fn')()
 
-  results = pjrt.run_multiprocess(train_imagenet, model.state_dict())
+  results = pjrt.run_multiprocess(train_imagenet)
   print('Replica max_accuracy:', pprint.pformat(results))
   accuracy = np.mean([
       np.mean(list(thread_results.values()))

test/pjrt/test_train_pjrt_mnist.py

@@ -59,7 +59,7 @@ def _train_update(device, x, loss, tracker, writer):
       summary_writer=writer)
 
 
-def train_mnist(flags, state_dict):
+def train_mnist(flags):
   if flags.fake_data:
     train_loader = xu.SampleGenerator(
         data=(torch.zeros(flags.batch_size, 1, 28,
@@ -112,8 +112,8 @@ def train_mnist(flags, state_dict):
 
   device = xm.xla_device()
   model = MNIST()
-  model.load_state_dict(state_dict)
   model = model.to(device)
+  pjrt.broadcast_master_param(model)
   writer = None
   if xm.is_master_ordinal():
     writer = test_utils.get_summary_writer(flags.logdir)
@@ -177,10 +177,8 @@ def test_loop_fn(loader):
 
 if __name__ == '__main__':
   torch.set_default_tensor_type('torch.FloatTensor')
-  torch.manual_seed(1)
-  model = MNIST()
 
-  results = pjrt.run_multiprocess(train_mnist, FLAGS, model.state_dict())
+  results = pjrt.run_multiprocess(train_mnist, FLAGS)
   print('Replica max_accuracy:', pprint.pformat(results))
   accuracy = np.mean([
       np.mean(list(thread_results.values()))

torch_xla/experimental/pjrt.py

@@ -2,9 +2,11 @@
 import functools
 import os
 import threading
+from itertools import chain
 from typing import Callable, Dict, Optional, TypeVar
 
 import torch
+import torch.nn as nn
 import torch_xla
 import torch_xla.core.xla_env_vars as xenv
 import torch_xla.core.xla_model as xm
@@ -214,3 +216,20 @@ def run_multiprocess(fn: Callable[..., R], *args,
     }
 
   return results
+
+
+def broadcast_master_param(model: nn.Module) -> None:
+  """
+  Broadcast the model parameters from master process to other processes
+  """
+  parameters_and_buffers = []
+  is_master = xm.is_master_ordinal(local=False)
+  for p in chain(model.parameters(), model.buffers()):
+    # Set all params in non-master devices to zero so that all_reduce is
+    # equivalent to broadcasting parameters from master to other devices.
+    scale = torch.tensor(1 if is_master else 0, dtype=p.data.dtype)
+    scale = scale.to(p.data.device)
+    p.data.mul_(scale)
+    parameters_and_buffers.append(p.data)
+  xm.all_reduce(xm.REDUCE_SUM, parameters_and_buffers)
+  xm.mark_step()