code: def _unshard_param(
      ref_state_dict,
      fqn,
      model_parallel_group,
      grp_gloo,
      data_tensor,
      tp_sharded_shape,
  ):
      """
      Unshards the row or col-wise sharded parameter.
      For rowwise, this is done by reshaping into the local shape, allgathering,
      and stacking rows. For colwise, the only difference is we stack columns.
      This is done via vstack and column_stack respectively.
      """
      mp_size = dist.get_world_size(model_parallel_group)
  
      ref_shape = ref_state_dict[fqn].shape
      assert (
          ref_shape[0] == tp_sharded_shape[0]
          or ref_shape[1] == tp_sharded_shape[1]
      ), f"Expected sharded shape to match either row or col-wise, but does not: {ref_shape} {tp_sharded_shape}"
      _verify_fqn_across_ranks(fqn, grp_gloo)
      if ref_shape[0] != tp_sharded_shape[0]:
          assert ref_shape[0] == tp_sharded_shape[0] * mp_size
          # reshape the flat data_tensor into the rowwise shape
          data_tensor = data_tensor.reshape(tp_sharded_shape)
          # now, all_gather such tensors
          tensor_list = _all_gather_into_list(
              data_tensor, model_parallel_group
          )
          # stack rowwise to produce the final unsharded tensor
          data_tensor = torch.vstack(tensor_list).cpu()
          assert data_tensor.shape == ref_shape
          full_shape = data_tensor.shape
      elif (
          len(ref_shape) > 1
          and len(tp_sharded_shape) > 1
          and ref_shape[1] != tp_sharded_shape[1]
      ):
          assert ref_shape[1] == mp_size * tp_sharded_shape[1]
          # first, reshape the flat data_tensor into the colwise shape
          data_tensor = data_tensor.reshape(tp_sharded_shape)
          tensor_list = _all_gather_into_list(
              data_tensor, model_parallel_group
          )
          data_tensor = torch.column_stack(tensor_list).cpu()
          assert (
              data_tensor.shape == ref_shape
          ), f"{data_tensor.shape} vs {ref_shape}"
          full_shape = data_tensor.shape
      else:
          assert (
              ref_shape == tp_sharded_shape
          )  # not tensor parallel sharded
          full_shape = tp_sharded_shape
          logging.warning(
              f"{fqn} {ref_shape} {full_shape} - not sharded"
          )
      return data_tensor, full_shape
  
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 4 for calling a global function with example usage: _unshard_param(ref_state_dict, fqn, model_parallel_group, grp_gloo, data_tensor, tp_sharded_shape) and returns: data_tensor,  full_shape