code: def test_tp_convnext_train(rank, world_size):
      os.environ["MASTER_ADDR"] = "localhost"
      os.environ["MASTER_PORT"] = "12355"
  
      dist.init_process_group("nccl", rank=rank, world_size=world_size)
  
      in_shape = [7, 3, 512, 1024]
      output_shape = [7, 1000]
      device = torch.device("cuda", rank)
      torch.cuda.set_device(device)
      torch.cuda.set_per_process_memory_fraction(1.0, device)
      mesh = DeviceMesh("cuda", torch.arange(world_size))
  
      torch.manual_seed(12)
      model = ConvNeXt(
          depths=[3, 3, 27, 3],
          dims=[256, 512, 1024, 2048],
          drop_path_rate=0.0,
          num_classes=1000,
      ).to(device)
      model = distribute_module(
          model, mesh, _conv_fn, input_fn=None, output_fn=None
      )
  
      criterion = torch.nn.CrossEntropyLoss()
      optimizer = torch.optim.Adam(
          model.parameters(), lr=1e-4, amsgrad=False
      )
  
      x = torch.randn(*in_shape).to(device).requires_grad_()
      y_target = (
          torch.empty(output_shape[0], dtype=torch.long)
          .random_(output_shape[1])
          .to(device)
      )
      x = distribute_tensor(x, mesh, [Shard(3)])
      y_target = distribute_tensor(y_target, mesh, [Replicate()])
  
      # warm up
      y = model(x)
      loss = criterion(y, y_target)
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      torch.cuda.synchronize(device)
  
      forward_time = 0.0
      backward_time = 0.0
      start = time.time()
      for i in range(ITER_TIME):
          t1 = time.time()
          y = model(x)
          torch.cuda.synchronize(device)
          t2 = time.time()
  
          loss = criterion(y, y_target)
          optimizer.zero_grad()
  
          t3 = time.time()
          loss.backward()
          torch.cuda.synchronize(device)
          t4 = time.time()
  
          optimizer.step()
  
          forward_time += t2 - t1
          backward_time += t4 - t3
      torch.cuda.synchronize(device)
      end = time.time()
      max_reserved = torch.cuda.max_memory_reserved(device)
      max_allocated = torch.cuda.max_memory_allocated(device)
      print(
          f"rank {rank}, {ITER_TIME} iterations, average latency {(end - start)/ITER_TIME*1000:10.2f} ms"
      )
      print(
          f"rank {rank}, forward {forward_time/ITER_TIME*1000:10.2f} ms, backward {backward_time/ITER_TIME*1000:10.2f} ms"
      )
      print(
          f"rank {rank}, max reserved {max_reserved/1024/1024/1024:8.2f} GiB, max allocated {max_allocated/1024/1024/1024:8.2f} GiB"
      )
      dist.destroy_process_group()
  
imports: import os
  import time
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.distributed._tensor import DeviceMesh, distribute_module, distribute_tensor, Replicate, Shard
label_desc: Use image node_id 3 for calling a global function with example usage: test_tp_convnext_train(rank, world_size) without return types