code: def remap_state_dict_hf_bigcode(state_dict, config: PretrainedConfig):
      """
      Map the state_dict of a Huggingface BigCode model to be flash_attn compatible.
      """
  
      # Word embedding and position embedding
      def key_mapping_pos_emb(key):
          return re.sub(
              r"^transformer.wpe.",
              "transformer.embeddings.position_embeddings.",
              key,
          )
  
      state_dict = OrderedDict(
          (key_mapping_pos_emb(k), v) for k, v in state_dict.items()
      )
      word_embeddings = state_dict.pop("transformer.wte.weight")
      # It's possible that vocab_size is padded to be a multiple of 8, for example.
      pad_vocab_size_multiple = getattr(
          config, "pad_vocab_size_multiple", 1
      )
      vocab_size = (
          math.ceil(config.vocab_size / pad_vocab_size_multiple)
          * pad_vocab_size_multiple
      )
      state_dict[
          "transformer.embeddings.word_embeddings.weight"
      ] = F.pad(
          word_embeddings,
          (0, 0, 0, vocab_size - word_embeddings.shape[0]),
      )
      state_dict["lm_head.weight"] = state_dict[
          "transformer.embeddings.word_embeddings.weight"
      ]
  
      # LayerNorm
      def key_mapping_ln(key):
          key = re.sub(
              r"^transformer.ln_f.(weight|bias)",
              r"transformer.ln_f.\1",
              key,
          )
          key = re.sub(
              r"^transformer.h.(\d+).ln_(1|2).(weight|bias)",
              r"transformer.layers.\1.norm\2.\3",
              key,
          )
          return key
  
      state_dict = OrderedDict(
          (key_mapping_ln(k), v) for k, v in state_dict.items()
      )
  
      def key_mapping_mlp(key):
          key = re.sub(
              r"^transformer.h.(\d+).mlp.c_fc.weight",
              r"transformer.layers.\1.mlp.fc1.weight",
              key,
          )
          key = re.sub(
              r"^transformer.h.(\d+).mlp.c_proj.weight",
              r"transformer.layers.\1.mlp.fc2.weight",
              key,
          )
          key = re.sub(
              r"^transformer.h.(\d+).mlp.c_fc.bias",
              r"transformer.layers.\1.mlp.fc1.bias",
              key,
          )
          key = re.sub(
              r"^transformer.h.(\d+).mlp.c_proj.bias",
              r"transformer.layers.\1.mlp.fc2.bias",
              key,
          )
          return key
  
      state_dict = OrderedDict(
          (key_mapping_mlp(k), v) for k, v in state_dict.items()
      )
  
      # TODO: add support for multi-head attention
      assert (
          config.multi_query
      ), "Only multi-query attention is supported"
  
      # Attention
      for d in range(config.num_hidden_layers):
          embed_dim = config.n_embd
          head_dim = embed_dim // config.n_head
  
          c_attn_weight = state_dict.pop(
              f"transformer.h.{d}.attn.c_attn.weight"
          )
          # with multi-query attention, the weights have shape (embed_dim, embed_dim + head_dim + head_dim)
          # see https://github.com/huggingface/transformers/blob/95b374952dc27d8511541d6f5a4e22c9ec11fb24/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py#L112
          # see also https://github.com/ggerganov/ggml/blob/dd1d575956e54c5bdc07632f25506b3b1884dbd2/examples/starcoder/convert-hf-to-ggml.py#L183
          # ((n_head + 2) * head_dim, embed_dim) -> (3 * n_heads * head_dim, hidden_dim)
          q, k, v = torch.split(
              c_attn_weight, [embed_dim, head_dim, head_dim], dim=0
          )
          # duplicate k, v along the first axis (head_dim, hidden_dim) -> (n_heads * head_dim, hidden_dim)
          k = torch.tile(k, (config.n_head, 1))
          v = torch.tile(v, (config.n_head, 1))
          state_dict[
              f"transformer.layers.{d}.mixer.Wqkv.weight"
          ] = torch.cat((q, k, v), dim=0)
  
          # same deal with the bias
          c_attn_bias = state_dict.pop(
              f"transformer.h.{d}.attn.c_attn.bias"
          )
          # ((n_head + 2) * head_dim, embed_dim) -> (3 * n_heads * head_dim, hidden_dim)
          q, k, v = torch.split(
              c_attn_bias, [embed_dim, head_dim, head_dim], dim=0
          )
          # duplicate k, v along the first axis (head_dim, hidden_dim) -> (n_heads * head_dim, hidden_dim)
          k = torch.tile(k, (config.n_head,))
          v = torch.tile(v, (config.n_head,))
          state_dict[
              f"transformer.layers.{d}.mixer.Wqkv.bias"
          ] = torch.cat((q, k, v), dim=0)
  
      def key_mapping_attn(key):
          key = re.sub(
              r"^transformer.h.(\d+).attn.c_proj.weight",
              r"transformer.layers.\1.mixer.out_proj.weight",
              key,
          )
          key = re.sub(
              r"^transformer.h.(\d+).attn.c_proj.bias",
              r"transformer.layers.\1.mixer.out_proj.bias",
              key,
          )
          return key
  
      state_dict = OrderedDict(
          (key_mapping_attn(k), v) for k, v in state_dict.items()
      )
  
      return state_dict
  
file_path: flash-attention/flash_attn/models/bigcode.py
imports: import math
  import re
  from collections import OrderedDict
  import torch
  import torch.nn.functional
  from transformers import GPT2Config, GPTBigCodeConfig, PretrainedConfig
label_desc: Use image node_id 1 for calling a global function with example usage: remap_state_dict_hf_bigcode(state_dict, config) and returns: state_dict, re, key, key, key