Upload inference/generate.py with huggingface_hub

inference/generate.py

@@ -1,4 +1,5 @@
 import os
+import re
 import json
 from argparse import ArgumentParser
 from typing import List
@@ -50,18 +51,18 @@ def hf_to_deepseek_key(hf_key: str) -> str:
     key = key.replace(".mlp.up_proj.", ".ffn.w3.")
     key = key.replace(".mlp.down_proj.", ".ffn.w2.")
 
-    # MoE
-    key = key.replace(".mlp.shared_experts.gate_proj.", ".moe.shared_experts.w1.")
-    key = key.replace(".mlp.shared_experts.up_proj.", ".moe.shared_experts.w3.")
-    key = key.replace(".mlp.shared_experts.down_proj.", ".moe.shared_experts.w2.")
-    key = key.replace(".mlp.experts.", ".moe.experts.")
-    key = key.replace(".mlp.gate.weight", ".moe.gate.weight")
+    # MoE (uses "ffn" module name in model, not "moe")
+    key = key.replace(".mlp.shared_experts.gate_proj.", ".ffn.shared_experts.w1.")
+    key = key.replace(".mlp.shared_experts.up_proj.", ".ffn.shared_experts.w3.")
+    key = key.replace(".mlp.shared_experts.down_proj.", ".ffn.shared_experts.w2.")
+    key = key.replace(".mlp.experts.", ".ffn.experts.")
+    key = key.replace(".mlp.gate.weight", ".ffn.gate.weight")
+    key = key.replace(".mlp.gate.e_score_correction_bias", ".ffn.gate.bias")
 
     # Expert weights
-    import re
-    key = re.sub(r"\.moe\.experts\.(\d+)\.gate_proj\.", r".moe.experts.\1.w1.", key)
-    key = re.sub(r"\.moe\.experts\.(\d+)\.up_proj\.", r".moe.experts.\1.w3.", key)
-    key = re.sub(r"\.moe\.experts\.(\d+)\.down_proj\.", r".moe.experts.\1.w2.", key)
+    key = re.sub(r"\.ffn\.experts\.(\d+)\.gate_proj\.", r".ffn.experts.\1.w1.", key)
+    key = re.sub(r"\.ffn\.experts\.(\d+)\.up_proj\.", r".ffn.experts.\1.w3.", key)
+    key = re.sub(r"\.ffn\.experts\.(\d+)\.down_proj\.", r".ffn.experts.\1.w2.", key)
 
     return key
 
@@ -86,7 +87,7 @@ def load_sharded_model(model, ckpt_path):
         for i, shard_file in enumerate(shard_files):
             shard_path = os.path.join(ckpt_path, shard_file)
             print(f"  [{i+1}/{len(shard_files)}] {shard_file}", end="", flush=True)
-            shard_dict = load_file(shard_path, device="cuda")
+            shard_dict = load_file(shard_path, device="cpu")
 
             # Copy matching tensors to model (with key mapping)
             matched = 0
@@ -151,11 +152,12 @@ def generate(
     prompt_lens = [len(t) for t in prompt_tokens]
     assert max(prompt_lens) <= model.max_seq_len, f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})"
     total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
-    tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long, device="cuda")
+    device = next(model.parameters()).device
+    tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long, device=device)
     for i, t in enumerate(prompt_tokens):
-        tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
+        tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long, device=device)
     prev_pos = 0
-    finished = torch.tensor([False] * len(prompt_tokens), device="cuda")
+    finished = torch.tensor([False] * len(prompt_tokens), device=device)
     prompt_mask = tokens != -1
     for cur_pos in range(min(prompt_lens), total_len):
         logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
@@ -205,14 +207,14 @@ def main(
     global print
     if rank != 0:
         print = lambda *_, **__: None
-    torch.cuda.set_device(local_rank)
     torch.set_default_dtype(torch.bfloat16)
-    torch.set_num_threads(8)
+    torch.set_num_threads(96)  # Use all CPU threads
     torch.manual_seed(33377335)
     with open(config) as f:
         args = ModelArgs(**json.load(f))
     print(args)
-    with torch.device("cuda"):
+    print("Creating model on CPU (this may take a while)...")
+    with torch.device("cpu"):
         model = Transformer(args)
     tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
     print("Loading model weights...")