code: def test_inv_remap_state_dict_hf_llama(model_name):
      config = llama_config_to_gpt2_config(
          AutoConfig.from_pretrained(model_name, trust_remote_code=True)
      )
      state_dict = state_dict_from_pretrained(model_name)
      # inv_remap_state_dict_hf_llama should be the inverse of remap_state_dict_hf_llama
      state_dict = {
          key: val
          for key, val in state_dict.items()
          if "rotary_emb.inv_freq" not in key
      }
      pretrained_state_dict = remap_state_dict_hf_llama(
          state_dict, config
      )
      state_dict_recover = inv_remap_state_dict_hf_llama(
          pretrained_state_dict, config
      )
      assert set(state_dict_recover.keys()) == set(state_dict.keys())
      for key in state_dict_recover.keys():
          torch.testing.assert_close(
              state_dict_recover[key], state_dict[key]
          )
  
imports: import os
  import time
  from pathlib import Path
  import shutil
  import pytest
  import torch
  from einops import rearrange
  from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp, shard_state_dict_tp
  from flash_attn.models.llama import config_from_checkpoint, inv_remap_state_dict_hf_llama, llama_config_to_gpt2_config, remap_state_dict_hf_llama, remap_state_dict_meta_llama, state_dicts_from_checkpoint
  from flash_attn.utils.distributed import all_gather_raw
  from flash_attn.utils.generation import update_graph_cache
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from transformers import LlamaConfig, LlamaTokenizer
  from transformers.models.llama.modeling_llama import LlamaForCausalLM
  from transformers import AutoConfig
label_desc: Use image node_id 3 for calling a global function with example usage: test_inv_remap_state_dict_hf_llama(model_name) without return types