code: def make_sharded_tensors_for_checkpoint(
      state_dict: StateDict,
      state_dict_prefix: str,
      sharded_key_prefix: Optional[str] = None,
      tensor_parallel_layers_axis_map: Optional[Dict[str, int]] = None,
      sharded_offsets: Iterable[Tuple[int, int, int]] = (),
      extra_state_suffix: str = "_extra_state",
  ):
      """Wraps tensors from transformer layers with ShardedTensor or ShardedObject.
  
      For a given `state_dict`, wraps:
      - all _extra_states with ShardedObject
      - all tensors specified in tensor_parallel_layers_axis_map with TP and DP sharded ShardedTensor
      - other values with DP sharded ShardedTensor
  
      Args:
          state_dict (StateDict): state_dict to convert
          state_dict_prefix (str): prefix appended to keys in final state dict
          sharded_key_prefix (str, optional): prefix appended to ShardedTensor keys
          tensor_parallel_layers_axis_map (Dict[str, int], optional): dict mapping layer
              names to the axis for TP sharding
          sharded_offsets (Iterable[Tuple[int, int, int]], optional): sharding already
              applied (e.g. PP related), passed along to ShardedTensor
          extra_state_suffix (str, default = '_extra_state'): layers with this
              suffix will be wrapped with ShardedObject instead of ShardedTensor.
  
      """
      if sharded_key_prefix is None:
          sharded_key_prefix = state_dict_prefix
  
      if tensor_parallel_layers_axis_map is None:
          tensor_parallel_layers_axis_map = {}
  
      sharded_state_dict = {}
      for layer_name in state_dict.keys():
          tensor = state_dict[layer_name]
          layer_key = f"{state_dict_prefix}{layer_name}"
          sharded_key = f"{sharded_key_prefix}{layer_name}"
  
          if layer_name.endswith(extra_state_suffix):
              sharded_state_dict[
                  layer_key
              ] = make_sharded_object_for_checkpoint(
                  tensor, sharded_key, sharded_offsets
              )
  
          elif layer_name in tensor_parallel_layers_axis_map:
              tp_axis = tensor_parallel_layers_axis_map[layer_name]
              sharded_state_dict[
                  layer_key
              ] = make_tp_sharded_tensor_for_checkpoint(
                  tensor,
                  sharded_key,
                  tp_axis,
                  prepend_offsets=sharded_offsets,
              )
  
          else:
              sharded_state_dict[
                  layer_key
              ] = make_sharded_tensor_for_checkpoint(
                  tensor,
                  sharded_key,
                  prepend_offsets=sharded_offsets,
              )
  
      return sharded_state_dict
  
file_path: megatron-lm/megatron/core/transformer/utils.py
imports: from operator import itemgetter
  from typing import Any, Dict, Iterable, Optional, Tuple, Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, StateDict
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 6 for calling a global function with example usage: make_sharded_tensors_for_checkpoint(state_dict, state_dict_prefix, sharded_key_prefix, tensor_parallel_layers_axis_map, sharded_offsets, extra_state_suffix) and returns: sharded_state_dict