docstr: Tests training for TP + FSDP integration by comparing an FSDP-only
  model with a TP + FSDP model.
docstr_tok: ["Tests","training","for","TP","+","FSDP","integration","by","comparing","an","FSDP-only","model","with","a","TP","+","FSDP","model","."]
code: def test_fsdp_tp_integration(self, tensor_parallel_size, cpu_offload):
      """
      Tests training for TP + FSDP integration by comparing an FSDP-only
      model with a TP + FSDP model.
      """
      LR = 3e-5
      torch.manual_seed(0)
      model = SimpleModel().cuda(self.rank)
      tp_fsdp_model = copy.deepcopy(model)
      sharded_param_names = SimpleModel.get_sharded_param_names()
      non_sharded_param_names = (
          SimpleModel.get_non_sharded_param_names()
      )
      (
          param_name_to_numel,
          param_name_to_sharding_info,
      ) = self._get_params_and_sharding_info(
          model,
          sharded_param_names,
          tensor_parallel_size,
      )
  
      input_seed = self.rank
      torch.manual_seed(input_seed + 1)
      inp_size = [2, 3, 5]
      inp = torch.rand(*inp_size).cuda(self.rank)
      self.assertEqual(model(inp), tp_fsdp_model(inp))  # sanity check
  
      mesh_1d = init_device_mesh("cuda", (self.world_size,))
      fsdp_model = FSDP(
          model, cpu_offload=cpu_offload, device_mesh=mesh_1d
      )
      mesh_2d = init_device_mesh(
          "cuda",
          (
              self.world_size // tensor_parallel_size,
              tensor_parallel_size,
          ),
          mesh_dim_names=["dp", "tp"],
      )
      # Shard with TP and then wrap with FSDP
      sequence_parallelize_plan = {
          "net1": ColwiseParallel(input_layouts=Shard(0)),
          "net2": RowwiseParallel(output_layouts=Shard(0)),
      }
      tp_fsdp_model = parallelize_module(
          tp_fsdp_model,
          mesh_2d["tp"],
          sequence_parallelize_plan,
      )
      tp_pg = mesh_2d["tp"].get_group(mesh_dim=0)
      assert isinstance(tp_fsdp_model.net1.weight, DTensor)
      assert isinstance(tp_fsdp_model.net2.weight, DTensor)
      tp_fsdp_model = FSDP(
          tp_fsdp_model,
          cpu_offload=cpu_offload,
          device_mesh=mesh_2d["dp"],
      )
      fsdp_pg = mesh_2d["dp"].get_group(mesh_dim=0)
  
      # Check the forward by checking output equality
      fsdp_out = fsdp_model(inp)
      tp_fsdp_out = tp_fsdp_model(inp)
      self.assertEqual(fsdp_out, tp_fsdp_out)
  
      # Check the backward by checking gradient equality
      fsdp_out.sum().backward()
      tp_fsdp_out.sum().backward()
      self._sync_tp_grads(
          tp_fsdp_model,
          tp_pg,
          param_name_to_numel,
          non_sharded_param_names,
      )
      model_grads = self._get_grads_as_flattened(
          fsdp_model,
          False,
          param_name_to_numel,
          param_name_to_sharding_info,
          None,
          self.process_group,
          None,
      )
      model_tp_grads = self._get_grads_as_flattened(
          tp_fsdp_model,
          True,
          param_name_to_numel,
          param_name_to_sharding_info,
          tp_pg,
          fsdp_pg,
          sharded_param_names,
      )
      self.assertEqual(model_grads, model_tp_grads)
  
      # Check the optimizer step by performing a second forward pass
      fsdp_optim = torch.optim.SGD(fsdp_model.parameters(), lr=LR)
      tp_fsdp_optim = torch.optim.SGD(tp_fsdp_model.parameters(), lr=LR)
      fsdp_optim.step()
      tp_fsdp_optim.step()
      torch.manual_seed(input_seed + 16)
      inp = torch.rand(*inp_size).cuda(self.rank)
      fsdp_out = fsdp_model(inp)
      tp_fsdp_out = tp_fsdp_model(inp)
      self.assertEqual(fsdp_out, tp_fsdp_out)
  
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 5 for calling the TestTPFSDPIntegration obj's underlying member method code with example usage: obj.test_fsdp_tp_integration(tensor_parallel_size, cpu_offload) without return types