Repository: EleutherAI/gpt-neox

Issue:
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! 
**Describe the bug**
A clear and concise description of what the bug is.
>**1-3B model trained with  8 gpus and 1 node, when i try to convert it to hf model：**
Traceback (most recent call last):
  File "./tools/convert_to_hf.py", line 292, in <module>
    hf_model = convert(args.input_dir, loaded_config, args.output_dir)
  File "./tools/convert_to_hf.py", line 156, in convert
    "weight": torch.cat(
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_cat)


**To Reproduce**
Steps to reproduce the behavior:
>python ./tools/convert_to_hf.py --input_dir checkpoints/global_step30000 --config_file configs/1-3B.yml --output_dir ../huggingface

**Expected behavior**
A clear and concise description of what you expected to happen.
> The expected result is  I could convert ckpt to huggingface model properly

**Proposed solution**
If you have an idea for how we can fix this problem, describe it here.
> I am not sure if the model  is incorrectly configed, or I need to modify the convert code.

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information):**
 >-  CUDA 11.3
 >- torch 1.13.0.dev20220727+cu113
 >- python 3.8.16

**Additional context**
Add any other context about the problem here.
>In training stage,  the`model-parallel-size=8` in yml conf:
   "pipe-parallel-size": 1,
   **"model-parallel-size": 8**


Discussion:
Your issue has boilerplate that hasn’t been filled out. Please complete the template.
@haileyschoelkopf it seems like there’s an issue with your HF conversation code?
Hi @zscwind ! I haven't encountered the exact error you're reporting, but adding

`map_location=torch.device('cuda:0')` to line 53 as a keyword arg to torch.load should fix this and load all your model weights to device 0. Alternatively, `map_location=torch.device('cpu')` should also have the same effect.

I'll push this fix later today :)
> Hi @zscwind ! I haven't encountered the exact error you're reporting, but adding
> 
> `map_location=torch.device('cuda:0')` to line 53 as a keyword arg to torch.load should fix this and load all your model weights to device 0. Alternatively, `map_location=torch.device('cpu')` should also have the same effect.
> 
> I'll push this fix later today :)

Yeah，I have tried your method. It works, thank you!


Resolving patch:
diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py
--- a/tools/convert_to_hf.py
+++ b/tools/convert_to_hf.py
@@ -49,7 +49,8 @@ def load_partitions(
             os.path.join(
                 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')
         )
         for i in range(mp_partitions)
     ]
@@ -145,7 +146,13 @@ 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")
+    if fp16:
+        if fp16["fp16"]:
+            hf_model.half()
 
     mp_partitions = get_key(loaded_config, "model-parallel-size")
 
diff --git a/tools/upload.py b/tools/upload.py
--- a/tools/upload.py
+++ b/tools/upload.py
@@ -13,26 +13,34 @@
 # limitations under the License.
 
 import os
+import sys
 
 from huggingface_hub import HfApi, create_repo
 
-converted_ckpt = input("Where is the checkpoint folder (HF format) you want to use? ")
-repo_name = input("Provide a repository name for the HF Hub: ")
-branch_name = input("Provide a repo branch for the HF Hub (choose main as default): ")
-create_repo(repo_name, repo_type="model", private=False)
+converted_ckpt = sys.argv[1]
+repo_name = sys.argv[2]
+branch_name = sys.argv[3]
+try:
+    create_repo(repo_name, repo_type="model", private=False)
+except:
+    print("repo {repo_name} already exists!")
+    pass
 
 files = os.listdir(converted_ckpt)
 
 api = HfApi()
 if branch_name != "main":
-    api.create_branch(
-        repo_id=repo_name,
-        repo_type="model",
-        branch=branch_name,
-    )
+    try:
+        api.create_branch(
+            repo_id=repo_name,
+            repo_type="model",
+            branch=branch_name,
+        )
+    except:
+        print(f"branch {branch_name} already exists, try again...")
 print(f"to upload: {files}")
 for file in files:
-    print(f"Uploading {file}...")
+    print(f"Uploading {file} to branch {branch_name}...")
     api.upload_file(
         path_or_fileobj=os.path.join(converted_ckpt, file),
         path_in_repo=file,


Test patch: