code: def test_grid_pipe_data(self):
      topo = Topo(axes=["pipe", "data"], dims=[2, 2])
      grid = Grid(topology=topo)
  
      assert grid._is_grid_valid()
  
      rank = dist.get_rank()
  
      assert grid.is_first_stage == (grid.get_stage_id() == 0)
      assert grid.is_last_stage == (
          grid.get_stage_id() == grid.get_pipe_parallel_world_size() - 1
      )
  
      # Test collectives along the pipeline parallel process groups
      rank_tensor = torch.LongTensor(data=[rank]).cuda()
      dist.all_reduce(rank_tensor, group=grid.get_pipe_parallel_group())
      pipe_group = grid.pp_group
      assert torch.all(rank_tensor == sum(pipe_group))
  
      # Test collectives along the data parallel process groups
      rank_tensor = torch.LongTensor(data=[rank]).cuda()
      dist.all_reduce(rank_tensor, group=grid.get_data_parallel_group())
      data_group = grid.dp_group
      assert torch.all(rank_tensor == sum(data_group))
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/runtime/pipe/test_topology.py
imports: import pytest
  import torch
  import deepspeed.comm
  from deepspeed.runtime.pipe.topology import PipelineParallelGrid
  from deepspeed.runtime.pipe.topology import ProcessTopology
  from deepspeed.runtime.pipe.topology import _prime_factors
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestDistributedTopology obj's underlying member method code with example usage: obj.test_grid_pipe_data() without return types