Repository: EleutherAI/gpt-neox

Issue:
1.0 HF conversion script fails on Python 3.8
**Describe the bug**

While the official recommended Python version for NeoX is 3.8, **convert_v1.0_to_hf.py** only works with Python 3.9 or later due to minor incompatibility (tested on Python 3.8.10).

>TypeError: 'type' object is not subscriptable

**Suggested change**

Modify `list[torch.Tensor]` at the line 42 to `List[torch.Tensor]` (capitalized L)
Add import `from typing import List`


Discussion:

Resolving patch:
diff --git a/tools/convert_sequential_to_hf.py b/tools/convert_sequential_to_hf.py
--- a/tools/convert_sequential_to_hf.py
+++ b/tools/convert_sequential_to_hf.py
@@ -58,7 +58,7 @@ def load_partitions(input_checkpoint_path, mp_partitions) -> List[torch.Tensor]:
 
 
 def get_state(
-    state_dicts: list[torch.Tensor],
+    state_dicts: List[torch.Tensor],
     key: str,
     layer_idx: int,
 ) -> torch.Tensor:
@@ -157,7 +157,27 @@ def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
 
     hf_model = GPTNeoXForCausalLM(
         hf_config
-    ).half()  # nice-to-have: lazy init weights somehow?
+    )
+
+    # save model in FP16 if Deepspeed fp16 was used in config, else 32 bit
+    fp16 = get_key(loaded_config, "fp16")
+    # save model in fp16/bf16 if Deepspeed fp16 or bf16 mixed precision was used in config, else 32 bit weights
+    fp16 = get_key(loaded_config, "fp16")
+    if fp16:
+        try:
+            # current behavior is to pass "fp16": {"enabled": true}, when using upstream Deepspeed
+            if fp16["enabled"]:
+                hf_model.half()
+                print("Saving weights in fp16 precision...")
+        except:
+            try:
+                # attempt to access bf16 dict in yaml file, if fp16 not enabled
+                bf16 = get_key(loaded_config, "bf16")
+                if bf16:
+                    hf_model.to(dtype=torch.bfloat16)
+                    print("Saving weights in bf16 precision...")
+            except:
+                print("Model not trained in fp16 / bf16 mixed precision, saving weights in fp32...")  
 
     mp_partitions = get_key(loaded_config, "model-parallel-size")
 
diff --git a/tools/convert_v1.0_to_hf.py b/tools/convert_v1.0_to_hf.py
--- a/tools/convert_v1.0_to_hf.py
+++ b/tools/convert_v1.0_to_hf.py
@@ -18,6 +18,7 @@
 import yaml
 import argparse
 from tqdm import tqdm
+from typing import List
 
 import torch
 from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
@@ -41,7 +42,7 @@
 
 def load_partitions(
     input_checkpoint_path, mp_partitions, layer_idx
-) -> list[torch.Tensor]:
+) -> List[torch.Tensor]:
     """Returns a list containing all weights in a given layer from a model (across MP partitions)"""
 
     loaded_tp_ranks = [
@@ -146,12 +147,21 @@ def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
 
     hf_model = GPTNeoXForCausalLM(hf_config)
 
-    # save model in FP16 if Deepspeed fp16 was used in config, else 32 bit
+    # save model in fp16/bf16 if Deepspeed fp16 or bf16 mixed precision was used in config, else 32 bit weights
     fp16 = get_key(loaded_config, "fp16")
     if fp16:
-        if fp16["fp16"]:
-            hf_model.half()
-
+        try:
+            # this conditional is quite messy because there were a number of ways to specify bf16 or fp16 training
+            # in DeeperSpeed v1.0 .
+            if (fp16.get("fp16", None) or fp16["enabled"]) and not (fp16.get("type", None) == "bfloat16"):
+                hf_model.half()
+                print("Saving weights in fp16 precision...")
+            elif fp16.get("type", None) == "bfloat16":
+                hf_model.to(dtype=torch.bfloat16)
+                print("Saving weights in bf16 precision...")
+        except:
+            print("Model not trained in fp16 / bf16 mixed precision, saving weights in fp32...")
+    
     mp_partitions = get_key(loaded_config, "model-parallel-size")
 
     ### Embedding layer ###


Test patch: