Upload 3 files

convert_llama_megatron_hf.py

@@ -0,0 +1,382 @@
+import argparse
+import os
+from collections import OrderedDict
+
+import torch
+from transformers import LlamaConfig, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+import accelerate
+
+transformer_layer_name_list = {
+    "input_norm": [
+        "input_norm.weight",
+        "self_attention.norm_qkv.layer_norm_weight",
+    ],
+    "query_key_value": [
+        "self_attention.query_key_value.weight",
+        "self_attention.norm_qkv.weight",
+    ],
+    "query": ["self_attention.query.weight"],
+    "key_value": ["self_attention.key_value.weight"],
+    "o_proj": ["self_attention.dense.weight", "self_attention.proj.weight"],
+    "mlp_gate_up": ["mlp.dense_h_to_4h.weight", "norm_mlp.fc1_weight"],
+    "mlp_down": ["mlp.dense_4h_to_h.weight", "norm_mlp.fc2_weight"],
+    "post_attention_norm": [
+        "post_attention_norm.weight",
+        "norm_mlp.layer_norm_weight",
+    ],
+}
+
+
+def recursive_print(name, val, spaces=0):
+    # Format the message.
+    if name is None:
+        msg = None
+    else:
+        fmt = "." * max(0, spaces - 2) + "# {:" + str(50 - spaces) + "s}"
+        msg = fmt.format(name)
+
+    # Print and recurse (if needed).
+    if isinstance(val, dict):
+        if msg is not None:
+            print(msg)
+        for k in val.keys():
+            recursive_print(k, val[k], spaces + 2)
+    elif isinstance(val, torch.Tensor):
+        print(msg, ":", val.size())
+    else:
+        print(msg, ":", val)
+
+
+def get(dicts, key):
+    return [dict[key] for dict in dicts]
+
+
+def check_get(dicts, prefix, key_list):
+    return [
+        dict[prefix + key] for dict in dicts for key in key_list if prefix + key in dict
+    ]
+
+
+def check_assign(encoder, this_layer_index, this_encoder, layer_index, key_list):
+    for key in key_list:
+        full_key = f"layers.{layer_index}." + key
+        if full_key in this_encoder:
+            encoder[f"layers.{this_layer_index}." + key] = this_encoder[full_key]
+            break
+    return encoder
+
+
+def merge_col(tensors):
+    return torch.cat(
+        [
+            tensor["weight"] if type(tensor) is OrderedDict else tensor
+            for tensor in tensors
+        ],
+        dim=0,
+    )
+
+
+def merge_row(tensors):
+    return torch.cat(
+        [
+            tensor["weight"] if type(tensor) is OrderedDict else tensor
+            for tensor in tensors
+        ],
+        dim=1,
+    )
+
+
+def convert_megatron_checkpoint(hf_model, state_dicts, model_config: LlamaConfig):
+    # The model.
+    models = get(state_dicts, "model")
+
+    # The language model.
+    lms = get(models, "language_model")
+
+    # The embeddings.
+    embeddings = get(lms, "embedding")
+
+    # The word embeddings.
+    word_embeddings = get(embeddings, "word_embeddings")
+
+    # Truncate the embedding table to vocab_size rows.
+    merged_padded_word_embeddings = merge_col(word_embeddings)
+    merged_word_embeddings = merged_padded_word_embeddings[: model_config.vocab_size, :]
+    hf_model.model.embed_tokens.load_state_dict(
+        {"weight": merged_word_embeddings}, strict=True
+    )
+
+    # The transformer.
+    transformers = get(lms, "encoder")
+
+    for i in range(model_config.num_hidden_layers):
+        print("Converting layer", i)
+        prefix = f"layers.{i}."
+        layer: LlamaDecoderLayer = hf_model.model.layers[i]
+
+        layer.input_layernorm.load_state_dict(
+            {
+                "weight": check_get(
+                    transformers, prefix, transformer_layer_name_list["input_norm"]
+                )[0]
+            },
+            strict=True,
+        )
+
+        hidden_size = model_config.hidden_size
+        inter_size = model_config.intermediate_size
+        num_heads = model_config.num_attention_heads
+        kv_heads = model_config.num_key_value_heads
+        kv_hidden_size = hidden_size // num_heads * kv_heads
+        if num_heads == kv_heads:
+            qkv = merge_col(
+                check_get(
+                    transformers, prefix, transformer_layer_name_list["query_key_value"]
+                )
+            )
+            qkv = qkv.view(num_heads, 3, hidden_size // num_heads, hidden_size)
+            q, k, v = torch.chunk(qkv, 3, dim=1)
+            q, k, v = (
+                q.reshape(hidden_size, hidden_size),
+                k.reshape(hidden_size, hidden_size),
+                v.reshape(hidden_size, hidden_size),
+            )
+        else: 
+            qkv = merge_col(
+                check_get(
+                    transformers, prefix, transformer_layer_name_list["query_key_value"]
+                )
+            )
+ 
+            num_queries_per_key_value = num_heads // kv_heads
+            qkv = qkv.view(
+                kv_heads,
+                num_queries_per_key_value + 2,
+                hidden_size // num_heads,
+                hidden_size,
+            )
+            q, k, v = torch.split(qkv, [num_queries_per_key_value, 1, 1], dim=1)
+                 
+            
+            q, k, v = (
+                q.reshape(hidden_size, hidden_size),
+                k.reshape(kv_hidden_size, hidden_size),
+                v.reshape(kv_hidden_size, hidden_size),
+            )
+
+        layer.self_attn.q_proj.load_state_dict({"weight": q}, strict=True)
+        layer.self_attn.k_proj.load_state_dict({"weight": k}, strict=True)
+        layer.self_attn.v_proj.load_state_dict({"weight": v}, strict=True)
+
+        layer.self_attn.o_proj.load_state_dict(
+            {
+                "weight": merge_row(
+                    check_get(
+                        transformers, prefix, transformer_layer_name_list["o_proj"]
+                    )
+                )
+            },
+            strict=True,
+        )
+
+        gate, up = (
+            merge_col(
+                check_get(
+                    transformers, prefix, transformer_layer_name_list["mlp_gate_up"]
+                )
+            )
+            .view(len(state_dicts), 2, -1, hidden_size)
+            .chunk(2, dim=1)
+        )
+        gate, up = gate.reshape(inter_size, hidden_size), up.reshape(
+            inter_size, hidden_size
+        )
+        layer.mlp.gate_proj.load_state_dict({"weight": gate}, strict=True)
+        layer.mlp.up_proj.load_state_dict({"weight": up}, strict=True)
+        layer.mlp.down_proj.load_state_dict(
+            {
+                "weight": merge_row(
+                    check_get(
+                        transformers, prefix, transformer_layer_name_list["mlp_down"]
+                    )
+                )
+            },
+            strict=True,
+        )
+
+        layer.post_attention_layernorm.load_state_dict(
+            {
+                "weight": check_get(
+                    transformers,
+                    prefix,
+                    transformer_layer_name_list["post_attention_norm"],
+                )[0]
+            },
+            strict=True,
+        )
+
+    # The final norm.
+    hf_model.model.norm.load_state_dict(
+        {"weight": transformers[0]["final_norm.weight"]}, strict=True
+    )
+
+    # For LM head, transformers' wants the matrix to weight embeddings.
+    output_layers = get(lms, "output_layer")
+    merged_padded_output_layers = merge_col(output_layers)
+    merged_output_layers = merged_padded_output_layers[: model_config.vocab_size, :]
+    hf_model.lm_head.load_state_dict({"weight": merged_output_layers}, strict=True)
+
+
+def check_padded_vocab_size(train_args, orig_vocab_size):
+    """Pad vocab size so it is divisible by model parallel size and
+    still having GPU friendly size."""
+
+    after = orig_vocab_size
+    multiple = (
+        train_args.make_vocab_size_divisible_by * train_args.tensor_model_parallel_size
+    )
+    while (after % multiple) != 0:
+        after += 1
+    assert (
+        train_args.padded_vocab_size == after
+    ), "Mismatched vocab size and padded vocab size."
+
+
+def get_train_args(state_dict):
+    args = state_dict.get("args", None)
+    assert args is not None
+    return args
+
+
+def get_model_config(train_args, vocab_size):
+    config = LlamaConfig()
+    check_padded_vocab_size(train_args, vocab_size)
+    config.vocab_size = vocab_size
+    # config.vocab_size = train_args.padded_vocab_size
+    config.max_position_embeddings = train_args.max_position_embeddings
+    config.hidden_size = train_args.hidden_size
+    config.num_hidden_layers = train_args.num_layers
+    config.num_attention_heads = train_args.num_attention_heads
+    config.num_key_value_heads = train_args.num_query_groups
+    config.intermediate_size = train_args.ffn_hidden_size
+    if hasattr(train_args, "rope_base"):
+        config.rope_theta = train_args.rope_base
+    config.pad_token_id = 0
+    config.torch_dtype  = train_args.params_dtype
+    return config
+
+
+def load_state_dicts(input_dir):
+    state_dicts = [
+        torch.load(os.path.join(f.path, "model_optim_rng.pt"), map_location="cpu")
+        for f in os.scandir(input_dir)
+        if f.is_dir()
+    ]
+    args = get_train_args(state_dicts[0])
+    if args.transformer_pipeline_model_parallel_size == 1:
+        return state_dicts, args
+
+    state_dicts = []
+    tp_size = args.tensor_model_parallel_size
+    pp_size = args.transformer_pipeline_model_parallel_size
+    num_layers_per_pile = args.num_layers // pp_size
+    for tp_index in range(tp_size):
+        model_file = f"{input_dir}/mp_rank_{tp_index:02d}_000/model_optim_rng.pt"
+        print(f"loading {model_file}")
+        state_dict = torch.load(
+            model_file,
+            map_location="cpu",
+        )
+        lm = state_dict["model"]["language_model"]
+        encoder = lm["encoder"]
+        for pp_index in range(1, pp_size):
+            model_file = f"{input_dir}/mp_rank_{tp_index:02d}_{pp_index:03d}/model_optim_rng.pt"
+            this_state_dict = torch.load(
+                model_file,
+                map_location="cpu",
+            )
+            print(f"loading {model_file}")
+            this_lm = this_state_dict["model"]["language_model"]
+            this_encoder = this_lm["encoder"]
+
+            if pp_index == pp_size - 1:
+                lm["output_layer"] = this_lm["output_layer"]
+                encoder["final_norm.weight"] = this_encoder[
+                    "final_norm.weight"
+                ]
+
+            for layer_index in range(num_layers_per_pile):
+                this_layer_index = layer_index + num_layers_per_pile * pp_index
+                if args.num_attention_heads == args.num_query_groups:
+                    encoder = check_assign(
+                        encoder,
+                        this_layer_index,
+                        this_encoder,
+                        layer_index,
+                        key_list=transformer_layer_name_list["query_key_value"],
+                    )
+                else:
+                    for key in ["query", "key_value", "query_key_value"]:
+                        encoder = check_assign(
+                            encoder,
+                            this_layer_index,
+                            this_encoder,
+                            layer_index,
+                            key_list=transformer_layer_name_list[key],
+                        )
+                for key in transformer_layer_name_list.keys():
+                    if key not in ("query_key_value", "query", "key_value"):
+                        encoder = check_assign(
+                            encoder,
+                            this_layer_index,
+                            this_encoder,
+                            layer_index,
+                            key_list=transformer_layer_name_list[key],
+                        )
+        state_dicts.append(state_dict)
+
+    return state_dicts, args
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        help="Path to the megatron checkpoint dir",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        help="Path to the huggingface checkpoint dir",
+    )
+    parser.add_argument(
+        "--vocab-size",
+        type=int,
+        default=64000,
+        help="unpadded tokenizer vocab size",
+    )
+    args = parser.parse_args()
+
+    print("Load megatron checkpoint")
+    state_dicts, train_args = load_state_dicts(args.input_dir)
+
+    model_config = get_model_config(train_args, args.vocab_size)
+    print(f"Model config: {model_config}", flush=True)
+    
+    
+    print("Create hf model", flush=True)
+    # with accelerate.init_empty_weights():
+    hf_model = LlamaForCausalLM(model_config)
+    hf_model = hf_model.to(torch.bfloat16)
+    
+    print("convert megatron to hf", flush=True)
+    convert_megatron_checkpoint(hf_model, state_dicts, model_config)
+
+    print("save hf model", flush=True)
+    hf_model.save_pretrained(args.output_dir, safe_serialization=False)
+
+
+if __name__ == "__main__":
+    main()