update

janelu9 · Jan 26, 2025 · 015b155 · 015b155
1 parent 2cc6c2a
commit 015b155
Show file tree

Hide file tree

Showing 6 changed files with 4 additions and 4 deletions.
diff --git a/jllm/cat2hf.py b/jllm/cat2hf.py
@@ -36,10 +36,10 @@ def func(meta_data):
         model_file =f"model-{pipe_rank:05d}-of-"+f"{num_stages:05d}.safetensors"
 
         for k in tqdm.tqdm(keys):
-            if  "o_proj" in k or "down_proj" in k in k:
+            if  "o_proj" in k or "down_proj" in k or 'attn.proj.weight' in k or 'mlp.fc2.weight' in k or 'mlp.2.weight' in k :
                 state_dict[k] = torch.cat([p[1].pop(k) for p in pts],1)
-            elif "lm_head" in k or "gate_proj" in k or "up_proj" in k or "embed_tokens" in k\
-            or "q_proj" in k or "k_proj" in k or "v_proj" in k:
+            elif "lm_head" in k or "gate_proj" in k or "up_proj" in k or "embed_tokens" in k or "q_proj" in k \
+            or "k_proj" in k or "v_proj" in k or 'attn.qkv' in k or 'mlp.fc1' in k or 'mlp.0' in k:
                 state_dict[k] = torch.cat([p[1].pop(k) for p in pts])
             else:
                 state_dict[k] = pts[0][1].pop(k)

diff --git a/jllm/data/utils.cpython-311-x86_64-linux-gnu.so b/jllm/data/utils.cpython-311-x86_64-linux-gnu.so
diff --git a/jllm/data/utils.cpython-39-aarch64-linux-gnu.so b/jllm/data/utils.cpython-39-aarch64-linux-gnu.so
diff --git a/jllm/model/qwen2_vl/parallel_qwen2_vl.cpython-311-x86_64-linux-gnu.so b/jllm/model/qwen2_vl/parallel_qwen2_vl.cpython-311-x86_64-linux-gnu.so
diff --git a/jllm/model/qwen2_vl/parallel_qwen2_vl.cpython-39-aarch64-linux-gnu.so b/jllm/model/qwen2_vl/parallel_qwen2_vl.cpython-39-aarch64-linux-gnu.so
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 import io
 
 project_name = "jllm"  
-version = "4.0.2" 
+version = "4.0.3" 
 
 setuptools.setup(
     name=project_name,