code: def test_baichuan_state_dict(model_name):
      config = baichuan_config_to_gpt2_config(
          AutoConfig.from_pretrained(model_name, trust_remote_code=True)
      )
      pretrained_state_dict = remap_state_dict_hf_baichuan(
          state_dict_from_pretrained(model_name), config
      )
      model = GPTLMHeadModel(
          config, device="meta"
      )  # Without device='meta' init is very slow
      state_dict = model.state_dict()
      assert len(state_dict.keys()) == len(pretrained_state_dict.keys())
      assert state_dict.keys() == pretrained_state_dict.keys()
      for k in state_dict.keys():
          assert state_dict[k].shape == pretrained_state_dict[k].shape
  
file_path: flash-attention/tests/models/test_baichuan.py
imports: import os
  import time
  from pathlib import Path
  import torch
  import pytest
  from einops import rearrange
  from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
  from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp, shard_state_dict_tp
  from flash_attn.models.baichuan import remap_state_dict_hf_baichuan, baichuan_config_to_gpt2_config
  from flash_attn.utils.distributed import all_gather_raw
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from flash_attn.utils.generation import update_graph_cache
label_desc: Use image node_id 1 for calling a global function with example usage: test_baichuan_state_dict(model_name) without return types