docstr: Generates TP and FSDP subprocess groups. ``tensor_parallel_size`` gives
  the TP process group size.
  
  For example, if the global world size is 8 and the tensor parallel size
  is 2, then this creates:
  - 4 TP subprocess groups: [0, 1], [2, 3], [4, 5], [6, 7]
  - 2 FSDP subprocess groups: [0, 2, 4, 6], [1, 3, 5, 7]
docstr_tok: ["Generates","TP","and","FSDP","subprocess","groups",".","``","tensor_parallel_size","``","gives","the","TP","process","group","size",".","For","example",",","if","the","global","world","size","is","8","and","the","tensor","parallel","size","is","2",",","then","this","creates",":","-","4","TP","subprocess","groups",":","[","0",",","1","]",",","[","2",",","3","]",",","[","4",",","5","]",",","[","6",",","7","]","-","2","FSDP","subprocess","groups",":","[","0",",","2",",","4",",","6","]",",","[","1",",","3",",","5",",","7","]"]
code: def _get_sub_pgs(self, tensor_parallel_size: int):
      """
      Generates TP and FSDP subprocess groups. ``tensor_parallel_size`` gives
      the TP process group size.
  
      For example, if the global world size is 8 and the tensor parallel size
      is 2, then this creates:
      - 4 TP subprocess groups: [0, 1], [2, 3], [4, 5], [6, 7]
      - 2 FSDP subprocess groups: [0, 2, 4, 6], [1, 3, 5, 7]
      """
      # 2-D mesh is [dp, tp]
      twod_mesh = DeviceMesh(
          device_type="cuda",
          mesh=torch.arange(0, self.world_size).view(
              -1, tensor_parallel_size
          ),
      )
  
      fsdp_pg = twod_mesh.get_group(mesh_dim=0)
      tp_pg = twod_mesh.get_group(mesh_dim=1)
      return twod_mesh, fsdp_pg, tp_pg
  
imports: import copy
  import sys
  from collections import OrderedDict
  from typing import Dict, List, Optional, Tuple
  import torch
  from torch import distributed
  from torch.distributed._tensor import DeviceMesh, distribute_module, DTensor, init_device_mesh, Replicate, Shard
  from torch.distributed._tensor.debug import CommDebugMode
  from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN
  from torch.testing._internal.distributed._tensor.common_dtensor import MLPModule
label_desc: Use image node_id 2 for calling the TestTPFSDPIntegration obj's underlying member method code with example usage: obj._get_sub_pgs(tensor_parallel_size) and returns: twod_mesh,  fsdp_pg,  tp_pg