code: def get_param_id_to_sharded_param_map(
      model_sharded_state_dict: ShardedStateDict,
      optim_params_iter: Iterable[torch.nn.Parameter],
  ) -> Dict[int, Union[ShardedTensor, ShardedTensorFactory]]:
      (
          model_sharded_state_dict,
          _,
      ) = extract_sharded_tensors_and_factories(
          model_sharded_state_dict
      )
      id_to_sharded_param_map = {}
      param_to_id_map = get_optim_param_to_id_map(optim_params_iter)
      for ten in nested_values(model_sharded_state_dict):
          if id(ten.data) in param_to_id_map:
              id_to_sharded_param_map[
                  param_to_id_map[id(ten.data)]
              ] = ten
          else:
              logger.debug(f"{ten} is not tracked by the optimizer")
  
      if not id_to_sharded_param_map:
          logger.warning(
              "Sharded parameters mapping is empty. It means tensors in model state dict"
              " do not correspond to tensors in optimizer parameters map."
              " Make sure to call state_dict with `keep_vars=True`."
          )
      return id_to_sharded_param_map
  
file_path: megatron-lm/megatron/core/dist_checkpointing/optimizer.py
imports: import logging
  from copy import deepcopy
  from dataclasses import replace
  from itertools import chain
  from typing import Dict, Iterable, List, Tuple, Union
  import torch
  from .dict_utils import nested_values
  from .mapping import LocalNonpersitentObject, ShardedStateDict, ShardedTensor, ShardedTensorFactory, StateDict
  from .utils import extract_sharded_tensors, extract_sharded_tensors_and_factories
label_desc: Use image node_id 2 for calling a global function with example usage: get_param_id_to_sharded_param_map(model_sharded_state_dict, optim_params_iter) and returns: id_to_sharded_param_map