[Distributed] Make xm.all_gather a single graph in Dynamo

test/test_mp_all_gather.py

@@ -6,6 +6,10 @@
 import torch_xla.distributed.xla_multiprocessing as xmp
 
 
+def all_gather(tensor, dim):
+ return xm.all_gather(tensor, dim=dim)
+
+
 def _mp_fn(index):
  device = xm.xla_device()
  world_size = xm.xrt_world_size()
@@ -14,6 +18,18 @@ def _mp_fn(index):
  ordinal_tensor = torch.tensor([index], dtype=torch.float).to(device)
  result = xm.all_gather(ordinal_tensor, dim=0)
 
+ cpu_result = result.cpu()
+ expected = torch.arange(0, world_size, dtype=torch.float)
+ if not cpu_result.allclose(expected):
+ print('xm.all_gather() produced wrong reductions', file=sys.stderr)
+ print(f'[{index}] {cpu_result}', file=sys.stderr)
+ sys.exit(1)
+
+ compiled_all_gather = torch.compile(
+ all_gather, backend='torchxla_trace_once', fullgraph=True)
+ ordinal_tensor = torch.tensor([index], dtype=torch.float).to(device)
+ result = compiled_all_gather(ordinal_tensor, dim=0)
+
  cpu_result = result.cpu()
  expected = torch.arange(0, world_size, dtype=torch.float)
  if not cpu_result.allclose(expected):

torch_xla/core/xla_model.py

@@ -9,6 +9,7 @@
 import torch.nn.functional as F
 import torch_xla
 from torch_xla.experimental import pjrt
+from torch_xla.experimental import tpu
 import torch_xla.core.xla_env_vars as xenv
 import torch_xla.debug.metrics_saver as ms
 import torch_xla.utils.utils as xu
@@ -26,6 +27,27 @@
 _DEVICE_CONTEXTS = dict()
 _DEVICE_CONTEXTS_LOCK = threading.Lock()
 
+# Note [Dynamo WORLD_SIEZ and ORDINAL]
+# Belows are workaround to cache the ordinal and world_size such that
+# Dynamo won't do graph breaks when xm.xrt_world_size() and xm.get_ordinal() are called.
+_WORLD_SIZE = None
+_ORDINAL = None
+
+
+def _init_world_size_ordinal():
+ global _WORLD_SIZE, _ORDINAL
+
+ if not pjrt.using_pjrt():
+ return
+
+ # We don't support V3-8. See Note [V3-8 Threading]
+ if pjrt.device_type() == 'TPU' and tpu.version() < 4:
+ return
+
+ if _WORLD_SIZE is None:
+ _WORLD_SIZE = xrt_world_size()
+ _ORDINAL = get_ordinal()
+
 
 class DeviceContext(object):
 
@@ -90,6 +112,10 @@ def xrt_world_size(defval=1):
  Returns:
  The number of devices which is taking part of the replication.
  """
+ global _WORLD_SIZE
+ if _WORLD_SIZE is not None:
+ return _WORLD_SIZE
+
  if pjrt.using_pjrt():
  return pjrt.world_size()
 
@@ -109,6 +135,10 @@ def get_ordinal(defval=0):
  Returns:
  The replication ordinal of the current thread.
  """
+ global _ORDINAL
+ if _ORDINAL is not None:
+ return _ORDINAL
+
  if pjrt.using_pjrt():
  return pjrt.global_ordinal()
 
@@ -533,8 +563,7 @@ def all_gather(value, dim=0, groups=None, output=None, pin_layout=True):
  A tensor which has, in the ``dim`` dimension, all the values from the
  participating replicas.
  """
- if pin_layout and xla_device_hw(
- value.device) in ('TPU', 'GPU', 'XPU') and output == None:
+ if pin_layout and output == None:
  # There is not an easy way to pin the all_gather layout on TPU and GPU, use
  # all_reduce based all_gather for this purpose.
  return _all_gather_using_all_reduce(

torch_xla/csrc/cross_replica_reduces.cpp

@@ -18,6 +18,7 @@
 namespace torch_xla {
 namespace {
 
+// Note [V3-8 Threading]
 // For V3-8 + PJRT, we have 4 processes and each process has 2 threads to manage
 // the 8 cores. Therefore, we need different tokens for different threads.
 std::unordered_map<int64_t, std::shared_ptr<torch::lazy::Value>>

torch_xla/experimental/pjrt.py

@@ -235,6 +235,9 @@ def _run_thread_per_device(
  def _thread_fn(device: torch.device):
  torch_xla._XLAC._xla_set_default_device(device)
 
+ # See Note Note [Dynamo WORLD_SIEZ and ORDINAL].
+ xm._init_world_size_ordinal()
+
  return fn()
 
  with concurrent.futures.ThreadPoolExecutor(