Repository: EleutherAI/gpt-neox

Issue:
Upstream DeepSpeed breaks HF conversion script
The tools/convert_to_hf.py script will need to be updated / a different version may need to be created for checkpoints saved with DeepSpeed. Checkpoints are no longer saved layer-by-layer, it seems, and now all weights are in several `mp_rank_{MP_RANK}_model_states.pt` files for each Model Parallel partition.

Upstream DeepSpeed checkpoint:
```
drwxr-xr-x 2 hailey eleuther     33280 Dec 18 14:29 configs
-rw-r--r-- 1 hailey eleuther 810771646 Dec 18 14:29 mp_rank_00_model_states.pt
-rw-r--r-- 1 hailey eleuther 608006863 Dec 18 14:29 zero_pp_rank_0_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008079 Dec 18 14:29 zero_pp_rank_1_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008079 Dec 18 14:29 zero_pp_rank_2_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008143 Dec 18 14:29 zero_pp_rank_3_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008143 Dec 18 14:29 zero_pp_rank_4_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008143 Dec 18 14:29 zero_pp_rank_5_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608008079 Dec 18 14:29 zero_pp_rank_6_mp_rank_00_optim_states.pt
-rw-r--r-- 1 hailey eleuther 608006863 Dec 18 14:29 zero_pp_rank_7_mp_rank_00_optim_states.pt
```

DeeperSpeed checkpoint:
```
drwxrwxrwx 2 hailey eleuther     33280 Nov 18 04:55 configs
-rwxrwxrwx 1 hailey eleuther 206045931 Nov 18 04:55 layer_00-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_02-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_03-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_04-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_05-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_06-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_07-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_08-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_09-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_10-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_11-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_12-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_13-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_14-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_15-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_16-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_17-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_18-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_19-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_20-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_21-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_22-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_23-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_24-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 100720126 Nov 18 04:55 layer_25-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther      9127 Nov 18 04:55 layer_27-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther 206045931 Nov 18 04:55 layer_28-model_00-model_states.pt
-rwxrwxrwx 1 hailey eleuther     16291 Nov 18 04:55 mp_rank_00_model_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_0_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_10_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_11_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_12_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_13_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_14_mp_rank_00_optim_states.pt
-rwxrwxrwx 1 hailey eleuther 287605953 Nov 18 04:55 zero_pp_rank_15_mp_rank_00_optim_states.pt
...
```

Updating the script shouldn't be too hard at all though. 


Discussion:

Resolving patch:
diff --git a/tools/convert_sequential_to_hf.py b/tools/convert_sequential_to_hf.py
new file mode 100644
--- /dev/null
+++ b/tools/convert_sequential_to_hf.py
@@ -0,0 +1,353 @@
+# Copyright (c) 2023, EleutherAI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+
+import yaml
+import argparse
+from tqdm import tqdm
+
+import torch
+from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
+
+from typing import List
+
+sys.path.append(
+    os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
+)
+from megatron.tokenizer import build_tokenizer
+
+
+"""
+A script for converting saved NeoX Checkpoints to Huggingface (HF) compatible GPT-NeoX type models.
+
+Note that this script does not support all NeoX features.
+Please investigate carefully whether your model is compatible with all architectures supported by the GPTNeoXForCausalLM class in HF.
+
+(e.g. position embeddings such as AliBi may not be supported by Huggingface's GPT-NeoX architecture.
+"""
+
+
+def load_partitions(input_checkpoint_path, mp_partitions) -> List[torch.Tensor]:
+    """Returns a list containing all states from a model (across MP partitions)"""
+
+    loaded_tp_ranks = [
+        torch.load(
+            os.path.join(
+                input_checkpoint_path,
+                f"mp_rank_{i:02}_model_states.pt",
+            ),
+            map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        )
+        for i in range(mp_partitions)
+    ]
+
+    return loaded_tp_ranks
+
+
+def get_state(
+    state_dicts: list[torch.Tensor],
+    key: str,
+    layer_idx: int,
+) -> torch.Tensor:
+    """Accesses all MP partitions of a given layer/weight's state."""
+    # main DeepSpeed saves each MP partition
+    key = f"sequential.{layer_idx}.{key}"
+
+    return [state_dict["module"][key] for state_dict in state_dicts]
+
+
+def get_key(loaded_config, key, default=None):
+    """
+    Search for a given key in a NeoX yaml. normalizes underscores -> hyphens
+    """
+    key = key.replace("_", "-")
+    try:
+        return loaded_config[key]
+    except KeyError:
+        key = key.replace("-", "_")
+        try:
+            return loaded_config[key]
+        except KeyError:
+            return default
+
+
+def create_config(neox_config):
+    """take in a loaded yaml from NeoX and assign relevant values to HF config.
+    Returns: GPTNeoXConfig() object
+    """
+
+    class TokenizerArgs:
+        # kinda hacky.
+        # this is to get something with the same interface as is used in build_tokenizer()
+        # without diving into loading a neox_args object or using argparse etc.
+        def __init__(self, neox_config):
+            self.make_vocab_size_divisible_by = get_key(
+                neox_config, "make-vocab-size-divisible-by", default=128
+            )
+            self.model_parallel_size = get_key(neox_config, "model-parallel-size")
+            self.vocab_file = get_key(neox_config, "vocab-file")
+            self.merge_file = get_key(neox_config, "merge-file")
+            self.tokenizer_type = get_key(neox_config, "tokenizer-type")
+
+            self.rank = 0
+
+    args = TokenizerArgs(neox_config)
+    tokenizer = build_tokenizer(args)
+    try:  # GPT2TokenizerFast raises NotImplementedError
+        pad_token = tokenizer.pad
+    except:
+        pad_token = (
+            1  # pad defaulting to 1. follows convention from GPT-NeoX-20b tokenizer
+        )
+
+    # TODO: change the default value here based on discussion regarding `gpt_j_tied` config parameter's default
+    use_tied_lns = get_key(neox_config, "gpt-j-tied", False)
+
+    if use_tied_lns:
+        raise NotImplementedError(
+            """ERROR: Huggingface Transformers does not yet support a single shared layernorm
+                per transformer block for GPT-NeoX models trained  w/ GPT-J parallel residuals.
+                See https://github.com/EleutherAI/gpt-neox/pull/481 for further details."""
+        )
+
+    # set all config values.
+    hf_config = GPTNeoXConfig(
+        vocab_size=args.padded_vocab_size,
+        hidden_size=get_key(neox_config, "hidden-size"),
+        num_hidden_layers=get_key(neox_config, "num-layers"),
+        num_attention_heads=get_key(neox_config, "num-attention-heads"),
+        intermediate_size=(get_key(neox_config, "hidden-size") * 4),
+        hidden_act=get_key(neox_config, "activation", default="gelu"),
+        rotary_pct=get_key(neox_config, "rotary-pct", default=1.0),
+        rotary_emb_base=get_key(neox_config, "rotary-emb-base", default=10000),
+        max_position_embeddings=get_key(neox_config, "max-position-embeddings"),
+        initializer_range=get_key(neox_config, "init-method-std", 0.02),
+        layer_norm_eps=get_key(neox_config, "layernorm-epsilon", 1e-5),
+        use_cache=True,
+        bos_token_id=tokenizer.eod,
+        eos_token_id=tokenizer.eod,
+        tie_word_embeddings=(not get_key(neox_config, "no-weight-tying", False)),
+        use_parallel_residual=get_key(neox_config, "gpt-j-residual", False),
+    )
+    return hf_config
+
+
+def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
+    """convert a NeoX checkpoint to a HF model format.
+    should perform model-parallel merging correctly
+    but only supports features allowed by HF GPT-NeoX implementation (e.g. rotary embeddings)
+    """
+
+    hf_config = GPTNeoXConfig()
+
+    hf_config = create_config(loaded_config)
+
+    hf_model = GPTNeoXForCausalLM(
+        hf_config
+    ).half()  # nice-to-have: lazy init weights somehow?
+
+    mp_partitions = get_key(loaded_config, "model-parallel-size")
+
+    # DeepSpeed main saves all model states from an MP rank in one file. load the MP ranks only once and index into them with get_state()
+    loaded_tp_ranks = load_partitions(input_checkpoint_path, mp_partitions)
+
+    ### Embedding layer ###
+    # Embedding is layer idx 0
+    hf_model.gpt_neox.embed_in.load_state_dict(
+        {
+            "weight": torch.cat(
+                get_state(loaded_tp_ranks, "word_embeddings.weight", 0), dim=0
+            )
+        }
+    )
+    assert (
+        hf_config.vocab_size == hf_model.gpt_neox.embed_in.weight.shape[0]
+    ), f"ERROR: calculated vocab size {hf_config.vocab_size} != embed param size {hf_model.gpt_neox.embed_in.shape[0]}"
+    ### End Embedding Layer ###
+
+    for layer_i in tqdm(range(get_key(loaded_config, "num-layers"))):
+
+        # get layer from hf model
+        hf_layer = hf_model.gpt_neox.layers[layer_i]
+
+        # + 2 bc of embed layer and a dummy _pre_transformer_block
+        state_dict = {}
+        for key in [
+            "attention.dense.weight",
+            "mlp.dense_4h_to_h.weight",
+        ]:
+            state_dict[key] = torch.cat(
+                get_state(loaded_tp_ranks, key, layer_i + 2), dim=1
+            )
+
+        # average layernorm stats over mp ranks
+        for key in [
+            "input_layernorm.weight",
+            "input_layernorm.bias",
+            "post_attention_layernorm.weight",
+            "post_attention_layernorm.bias",
+        ]:
+            state_dict[key] = sum(get_state(loaded_tp_ranks, key, layer_i + 2)) / len(
+                loaded_tp_ranks
+            )
+
+        # LinearWithTPMerge
+        for key in [
+            "mlp.dense_h_to_4h.weight",
+            "mlp.dense_h_to_4h.bias",
+            "attention.query_key_value.weight",
+            "attention.query_key_value.bias",
+        ]:
+            state_dict[key] = torch.cat(
+                get_state(loaded_tp_ranks, key, layer_i + 2), dim=0
+            )
+
+        # LinearWithTPSplitBias
+        for key in [
+            "mlp.dense_4h_to_h.bias",
+            "attention.dense.bias",
+        ]:
+            state_dict[key] = sum(get_state(loaded_tp_ranks, key, layer_i + 2))
+
+        # Just take one
+        state_dict["attention.rotary_emb.inv_freq"] = get_state(
+            loaded_tp_ranks, "attention.rotary_emb.inv_freq", layer_i + 2
+        )[0]
+
+        state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
+        state_dict["attention.masked_bias"] = hf_layer.state_dict()[
+            "attention.masked_bias"
+        ]
+
+        # load state_dict into layer
+        hf_layer.load_state_dict(state_dict)
+
+    # Load final layer norm
+    hf_model.gpt_neox.final_layer_norm.load_state_dict(
+        {
+            "weight": (
+                sum(
+                    get_state(
+                        loaded_tp_ranks,
+                        "norm.weight",
+                        get_key(loaded_config, "num-layers") + 3,
+                    )
+                )
+            )
+            / len(loaded_tp_ranks),
+            "bias": (
+                sum(
+                    get_state(
+                        loaded_tp_ranks,
+                        "norm.bias",
+                        get_key(loaded_config, "num-layers") + 3,
+                    )
+                )
+            )
+            / len(loaded_tp_ranks),
+        }
+    )
+    # output embedding / LM head
+    hf_model.embed_out.load_state_dict(
+        {
+            "weight": torch.cat(
+                get_state(
+                    loaded_tp_ranks,
+                    "final_linear.weight",
+                    get_key(loaded_config, "num-layers") + 4,
+                ),
+                dim=0,
+            ),
+        }
+    )
+
+    del loaded_tp_ranks
+
+    return hf_model
+
+
+if __name__ == "__main__":
+
+    # before running script:
+    # `pip install --upgrade transformers`
+    # `huggingface-cli login`
+    #
+    from huggingface_hub import create_repo, HfApi
+
+    parser = argparse.ArgumentParser(
+        description="Merge MP partitions and convert to HF Model."
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        help="Path to NeoX checkpoint, e.g. /path/to/model/global_step143000",
+    )
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        help="Path to config file for the input NeoX checkpoint.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        help="Output dir, where to save the HF Model, tokenizer, and configs",
+    )
+    parser.add_argument(
+        "--upload",
+        action="store_true",
+        help="Set to true in order to upload to the HF Hub directly.",
+    )
+    args = parser.parse_args()
+
+    with open(args.config_file) as f:
+        loaded_config = yaml.full_load(f)
+
+    hf_model = convert(args.input_dir, loaded_config, args.output_dir)
+
+    hf_model.save_pretrained(args.output_dir)
+
+    # save tokenizer to directory as well, for easy loading of model as a HF model
+    tokenizer_type = get_key(loaded_config, "tokenizer-type")
+
+    if tokenizer_type == "HFTokenizer":
+        print(f"saving tokenizer from file {get_key(loaded_config, 'vocab-file')}")
+        from transformers import PreTrainedTokenizerFast
+
+        tokenizer = PreTrainedTokenizerFast(
+            tokenizer_file=get_key(loaded_config, "vocab-file")
+        )
+        print("loaded tokenizer: ", tokenizer)
+        tokenizer.save_pretrained(args.output_dir)
+        print("tokenizer saved!")
+
+        print(
+            tokenizer.decode(
+                hf_model.generate(
+                    tokenizer.encode("Hello, I am testing ", return_tensors="pt")
+                )[0]
+            )
+        )
+
+    if args.upload:
+        repo_name = input("Provide a repository name for the HF Hub: ")
+        create_repo(repo_name, repo_type="model", private=False, use_auth_token=True)
+
+        api = HfApi()
+        api.upload_folder(
+            folder_path=args.output_dir,
+            repo_id=repo_name,
+            repo_type="model",
+        )
diff --git a/tools/convert_to_hf.py b/tools/convert_v1.0_to_hf.py
similarity index 97%
rename from tools/convert_to_hf.py
rename to tools/convert_v1.0_to_hf.py
--- a/tools/convert_to_hf.py
+++ b/tools/convert_v1.0_to_hf.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021, EleutherAI
+# Copyright (c) 2023, EleutherAI
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -50,7 +50,7 @@ def load_partitions(
                 input_checkpoint_path,
                 f"layer_{layer_idx:02}-model_{i:02}-model_states.pt",
             ),
-            map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
         )
         for i in range(mp_partitions)
     ]
@@ -144,12 +144,10 @@ def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
 
     hf_config = create_config(loaded_config)
 
-    hf_model = GPTNeoXForCausalLM(
-        hf_config
-    )
-    
+    hf_model = GPTNeoXForCausalLM(hf_config)
+
     # save model in FP16 if Deepspeed fp16 was used in config, else 32 bit
-    fp16 =  get_key(loaded_config, "fp16")
+    fp16 = get_key(loaded_config, "fp16")
     if fp16:
         if fp16["fp16"]:
             hf_model.half()


Test patch: