code: def _test_mlp_training_e2e(
      self, is_seq_parallel=False, recompute_activation=False
  ):
      inp_size = [8, 10]
      # Ensure all tp ranks have same input.
      rng_seed = self.rank if is_seq_parallel else 0
      torch.manual_seed(rng_seed)
      inp = torch.rand(*inp_size, device=self.device_type)
      model = MLPModule(self.device_type)
      model_tp = deepcopy(model)
  
      # Ensure model are initialized the same way.
      self._check_module(model, model_tp)
  
      # Shard module and initialize optimizer.
      LR = 0.25
      device_mesh = DeviceMesh(
          self.device_type,
          torch.arange(0, NUM_DEVICES),
      )
      parallelize_plan = {
          "net1": ColwiseParallel(input_layouts=Shard(0))
          if is_seq_parallel
          else ColwiseParallel(),
          "net2": RowwiseParallel(output_layouts=Shard(0))
          if is_seq_parallel
          else RowwiseParallel(),
      }
      model_tp = parallelize_module(
          model_tp, device_mesh, parallelize_plan
      )
      if recompute_activation:
          model_tp = input_reshard(
              checkpoint_wrapper(
                  model_tp, checkpoint_impl=CheckpointImpl.NO_REENTRANT
              ),
              device_mesh,
              None if is_seq_parallel else 0,
          )
      optim = torch.optim.SGD(model.parameters(), lr=LR)
      optim_tp = torch.optim.SGD(model_tp.parameters(), lr=LR)
  
      output = model(inp)
      output_tp = model_tp(inp)
      self.assertEqual(output, output_tp)
  
      output.sum().backward()
      output_tp.sum().backward()
  
      if is_seq_parallel:
          # Sum gradients from different ranks, since input
          # are different across ranks for sequence parallel.
          dist.all_reduce(model.net1.weight.grad)
          dist.all_reduce(model.net1.bias.grad)
          dist.all_reduce(model.net2.weight.grad)
          dist.all_reduce(model.net2.bias.grad)
  
      # Ensure gradients are same.
      self._check_module(model, model_tp, check_grad=True)
  
      optim.step()
      optim_tp.step()
  
      # Ensure model weights are still same after update.
      # Due to the trick we use for Partial aggregation, we only check the weight when local_rank = 0.
      self._check_module(model, model_tp)
  
      inp = torch.rand(*inp_size, device=self.device_type)
      output = model(inp)
      output_tp = model_tp(inp)
      self.assertEqual(output, output_tp)
  
file_path: pytorch/test/distributed/tensor/parallel/test_tp_examples.py
imports: from copy import deepcopy
  import torch
  import torch.distributed
  from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard
  import torch.distributed._functional_collectives
  from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper, CheckpointImpl
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, PrepareModuleInput, RowwiseParallel
  from torch.distributed.tensor.parallel.input_reshard import input_reshard
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, MLPModule, ModelArgs, NUM_DEVICES, skip_unless_torch_gpu, Transformer, with_comms
label_desc: Use image node_id 2 for calling the DistTensorParallelExampleTest obj's underlying member method code with example usage: obj._test_mlp_training_e2e(is_seq_parallel, recompute_activation) without return types