code: def build_distributed_state_dict_from_consolidated(
      model: nn.Module,
      consolidated_state_dict: Dict[str, Tensor],
      model_parallel_world_size: int,
      offload_to_cpu: bool = False,
      use_dtensor: bool = False,
  ) -> Dict[str, Union[Tensor, DTensor, ShardedTensor]]:
      """
      Main API that takes a model (with no parallelism applied) and a fairscale checkpoint
      and builds a PT-D compliant distributed state dict. Note that this expects a consolidated
      checkpoint.
  
      Args:
          model (torch.nn.Module): module with no parallelism applied (i.e. result of `build_model` with parallel_impl=ParallelImpl.NONE)
          fs_state_dict (Dict[str, Any]): Fairscale consolidated
          offload_to_cpu (bool): Whether to offload the resulting state_dict to CPU (default: False)
          use_dtensor (bool): Whether to use PyTorch Distributed Tensor instead of ShardedTensor (default: False)
              (this will eventually default to True)
          model_parallel_world_size: Model parallel world size that was used to create the consolidated checkpoint.
              This can be obtained by checking the number of consolidated0x.pth files in the checkpoint directory.
  
      Example usage::
          ```
  
          MODEL_PARALLEL_SIZE = 8
          ckpt_path = get_consolidated_ckpt_path(
              ckpt_dir=PTH_65b, mp_rank=local_rank, mp_size=MODEL_PARALLEL_SIZE
          )
          state_dict = torch.load(ckpt_path)
          # Build a local LLaMA with no parallelism
          model = build_model(...)
          sharded_state_dict = build_distributed_state_dict_from_consolidated(
              model, state_dict, model_parallel_world_size=MODEL_PARALLEL_SIZE,
          )
          # Wrap model with PT-native APIs + load
          model = FSDP(model)
          FSDP.set_state_dict_type(StateDictType.SHARDED_STATE_DICT)
          model.load_state_dict(sharded_state_dict)
          ```
  
      Note: Please make sure to pass an unsharded model as the model arg! Otherwise, things will not
      work.
  
      This distributed state dict is a mapping of FQN: ShardedTensor/DTensor. It will be replaced with
      DTensor once DTensor 2D checkpoint format is fully rolled out.
  
      Note: This has only been tested for loading state_dict into PT-D FSDP sharded_state_dict for now.
      """
      torch._C._log_api_usage_once("build_distributed_state_dict")
      dist_state_dict = {}
      ref_state_dict = model.state_dict()
      grp_gloo = dist.new_group(backend="gloo")
      # TODO: this should be the FSDP device mesh
      mesh = (
          DeviceMesh(
              device_type="cuda",
              mesh=list(range(dist.get_world_size())),
          )
          if use_dtensor
          else None
      )
      input_dtypes = {v.dtype for v in consolidated_state_dict.values()}
      logging.warning(f"input_dtypes {input_dtypes}")
      model_parallel_group, _ = dist.new_subgroups(
          group_size=model_parallel_world_size
      )
      for fqn, tensor in consolidated_state_dict.items():
          # Hack for buffer
          if "rope.freqs" in fqn:
              dist_state_dict[fqn] = tensor.clone()
              continue
          if _is_tp_sharded(fqn):
              tensor, _ = _unshard_param(
                  ref_state_dict,
                  fqn,
                  model_parallel_group,
                  grp_gloo,
                  tensor,
                  tensor.shape,
              )
          if use_dtensor:
              assert mesh is not None
              tensor = _ext_chunk_dtensor(
                  tensor=tensor.contiguous(),
                  rank=dist.get_rank(),
                  device_mesh=mesh,
              )
  
          else:
              tensor = _ext_chunk_tensor(
                  tensor=tensor.contiguous(),
                  rank=dist.get_rank(),
                  world_size=dist.get_world_size(),
                  num_devices_per_node=torch.cuda.device_count(),  # TODO: this is not accurate if user set CUDA_VISIBLE_DEVICES
                  pg=dist.distributed_c10d._get_default_group(),  # TODO: this should be the FSDP process group
              )
  
          dist_state_dict[fqn] = tensor
      dtypes = {v.dtype for v in dist_state_dict.values()}
      logging.warning(f"Made dist_state_dict with dtypes {dtypes}")
      return dist_state_dict
  
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 5 for calling a global function with example usage: build_distributed_state_dict_from_consolidated(model, consolidated_state_dict, model_parallel_world_size, offload_to_cpu, use_dtensor) and returns: dist_state_dict