support llama3.2

flexflow · Sep 29, 2024 · 6bc1eab · 6bc1eab
1 parent 7848871
commit 6bc1eab
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/flexflow/serve/models/llama.py b/python/flexflow/serve/models/llama.py
@@ -281,3 +281,7 @@ def convert_hf_model(model, dst_folder):
         for name, params in model.named_parameters():
             name = FlexFlowLLAMA.convert_hf_weight_name(name)
             params.detach().cpu().numpy().tofile(f"{dst_folder}/{name}")
+        # LM head weight
+        model.lm_head.weight.detach().cpu().numpy().tofile(
+            os.path.join(dst_folder, "lm_head.weight")
+        )