code: def _all_gather_into_list(data_tensor, model_parallel_group):
      tensor_list = [
          torch.zeros_like(data_tensor).cuda()
          for _ in range(dist.get_world_size(model_parallel_group))
      ]
      dist.all_gather(
          tensor_list, data_tensor.cuda(), group=model_parallel_group
      )
      return tensor_list
  
file_path: serve/examples/large_models/tp_llama/checkpoint_converter.py
imports: import logging
  from dataclasses import dataclass
  from typing import Dict, List, Union
  import torch
  import torch.distributed
  from torch import nn, Tensor
  from torch.distributed._shard.sharded_tensor.api import ShardedTensor
  from torch.distributed._tensor import DeviceMesh, DTensor
  from torch.distributed.fsdp._fsdp_extensions import _ext_chunk_dtensor, _ext_chunk_tensor
label_desc: Use image node_id 2 for calling a global function with example usage: _all_gather_into_list(data_tensor, model_parallel_group) and returns: tensor_list