fix(xla): convert group-local to global ranks in broadcast (#9657)

Related AWS Neuron ticket:
https://t.corp.amazon.com/V1941917988/overview

broadcast was passing group-local ranks directly to
xm.collective_broadcast() which expects global ranks, causing data
curroption in single-member process groups

TEST:

```
import os
import torch
import torch.distributed as dist
import torch_xla as xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.runtime as xr

def main():
    dist.init_process_group(backend="xla")

    rank = dist.get_rank()
    world_size = dist.get_world_size()

    tp = dist.new_group(ranks=[rank])
    tp_rank = dist.get_rank(group=tp)
    tp_size = dist.get_world_size(group=tp)

    print(
        f">>>> pid={os.getpid()}, rank={rank}\n"
        f">>> world_size={world_size}\n"
        f">>> tp_rank={tp_rank}, tp_size={tp_size}, tp_members={dist.get_process_group_ranks(tp)}"
    )

    do_train, do_valid, do_test = 0.1, 0.2, 0.3
    # breakpoint()
    flags = torch.tensor([do_train, do_valid, do_test], dtype=torch.float32, device='xla')
    # breakpoint()
    dist.broadcast(flags, rank, group=tp)

    print(f">>>> pid={os.getpid()}, rank={rank}\n"
          f">>> do_train={flags[0].item()}, do_valid={flags[1].item()}, do_test={flags[2].item()}\n"
          f">>> global_ordinal={xr.global_ordinal()}")

if __name__ == "__main__":
    main()
```

Results after this fix:
```
torchrun --nproc-per-node=2 --nnodes=1 ./bug.py
W0926 18:50:41.903000 1081605 torch/distributed/run.py:766] 
W0926 18:50:41.903000 1081605 torch/distributed/run.py:766] *****************************************
W0926 18:50:41.903000 1081605 torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W0926 18:50:41.903000 1081605 torch/distributed/run.py:766] *****************************************
>>>> pid=1081679, rank=0
>>> world_size=2
>>> tp_rank=0, tp_size=1, tp_members=[0]
>>>> pid=1081680, rank=1
>>> world_size=2
>>> tp_rank=0, tp_size=1, tp_members=[1]
.
.
.
2.19.8089.0+8ab9f450/MODULE_10344927339446294134+e30acd3a/model.neff
>>>> pid=1081680, rank=1
>>> do_train=0.10000000149011612, do_valid=0.20000000298023224, do_test=0.30000001192092896
>>> global_ordinal=1
>>>> pid=1081679, rank=0
>>> do_train=0.10000000149011612, do_valid=0.20000000298023224, do_test=0.30000001192092896
```

Now both ranks have the correct values. Previously Rank1 was all zeros.

Author: Hoomaaan

test/test_torch_distributed_xla_backend.py

@@ -44,6 +44,18 @@ def patch_world(rank, size):
     yield
 
 
+@contextlib.contextmanager
+def patch_world_with_xla_runtime(rank, size):
+  assert isinstance(dist.group.WORLD,
+                    torch_xla.distributed.xla_backend.ProcessGroupXla)
+
+  with mock.patch.object(dist.group.WORLD, 'rank', return_value=rank), \
+       mock.patch.object(dist.group.WORLD, 'size', return_value=size), \
+       mock.patch.object(xr, 'global_ordinal', return_value=rank), \
+       mock.patch.object(xr, 'world_size', return_value=size):
+    yield
+
+
 class XlaBackendTest(parameterized.TestCase):
 
   @classmethod
@@ -328,6 +340,81 @@ def test_unimplemented_op(self, op):
     with self.assertRaises(NotImplementedError):
       getattr(pg_xla, op)(tensor)
 
+  @patch_world_with_xla_runtime(rank=0, size=2)
+  def test_broadcast_single_rank_group_rank0(self):
+    """Test broadcast in single-member process group for rank 0"""
+    device = torch_xla.device()
+
+    with new_group_barrier_disabled():
+      tp = dist.new_group(ranks=[0])
+
+    # Create flags tensor with initial values (simulating rank 0's values)
+    flags = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, device=device)
+
+    # Broadcast within the single-member group (should be a no-op but shouldn't crash)
+    dist.broadcast(flags, src=0, group=tp)
+
+    # Values should remain unchanged since it's a single-member group
+    self.assertAlmostEqual(flags[0].item(), 0.1, places=5)
+    self.assertAlmostEqual(flags[1].item(), 0.2, places=5)
+    self.assertAlmostEqual(flags[2].item(), 0.3, places=5)
+
+    # Verify the process group properties
+    self.assertEqual(dist.get_rank(group=tp), 0)
+    self.assertEqual(dist.get_world_size(group=tp), 1)
+
+  @patch_world_with_xla_runtime(rank=1, size=2)
+  def test_broadcast_single_rank_group_rank1(self):
+    """Test broadcast in single-member process group for rank 1"""
+    device = torch_xla.device()
+
+    with new_group_barrier_disabled():
+      tp = dist.new_group(ranks=[1])
+
+    # Create flags tensor with initial values (simulating rank 1's values)
+    flags = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, device=device)
+
+    # Broadcast within the single-member group (should be a no-op but shouldn't crash)
+    dist.broadcast(flags, src=1, group=tp)
+
+    # Values should remain unchanged since it's a single-member group
+    self.assertAlmostEqual(flags[0].item(), 0.1, places=5)
+    self.assertAlmostEqual(flags[1].item(), 0.2, places=5)
+    self.assertAlmostEqual(flags[2].item(), 0.3, places=5)
+
+    # Verify the process group properties
+    self.assertEqual(dist.get_rank(group=tp),
+                     0)  # Local rank in single-member group is 0
+    self.assertEqual(dist.get_world_size(group=tp), 1)
+
+  @patch_world_with_xla_runtime(rank=0, size=2)
+  def test_broadcast_global_rank_conversion_single_member(self):
+    """Test that global rank conversion works correctly for single-member groups"""
+    device = torch_xla.device()
+
+    # Create single-member group for rank 0
+    with new_group_barrier_disabled():
+      tp = dist.new_group(ranks=[0])
+
+    flags = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, device=device)
+
+    # Get the ProcessGroupXla instance to test directly
+    self.assertIsInstance(tp, torch_xla.distributed.xla_backend.ProcessGroupXla)
+
+    # Test broadcast options - local rank 0 should map to global rank 0
+    opts = dist.BroadcastOptions()
+    opts.rootRank = 0
+    opts.rootTensor = 0
+
+    # This should work without variable name errors
+    work = tp.broadcast([flags], opts)
+    self.assertIsNotNone(work)
+
+    # Values should be preserved
+    self.assertAlmostEqual(flags[0].item(), 0.1, places=5)
+    self.assertAlmostEqual(flags[1].item(), 0.2, places=5)
+    self.assertAlmostEqual(flags[2].item(), 0.3, places=5)
+
 
 if __name__ == '__main__':
   if xr.device_type() != 'CPU':

torch_xla/distributed/xla_backend.py

@@ -131,9 +131,14 @@ def allgather_coalesced(self, output_tensors_list, input_tensors, opts=None):
   # Call site:
   # https://github.com/pytorch/pytorch/blob/release/1.10/torch/distributed/distributed_c10d.py#L1129
   def broadcast(self, tensors, opts):
+    import torch.distributed as dist
+
     root_tensor = tensors[opts.rootTensor]
+    # Convert group local rank to global rank for xla collectives
+    group_source = opts.rootRank
+    global_src = dist.get_global_rank(self, group_source)
     xm.collective_broadcast([root_tensor],
-                            opts.rootRank,
+                            global_src,
                             groups=self._mesh,
                             pin_layout=False)