Skip to content

Commit 4a8bdd5

Browse files
lshpkulshpku
committed
Add EB4.5 PT-Paddle bidirectional convertion tools
1 parent 18a3a1d commit 4a8bdd5

2 files changed

Lines changed: 214 additions & 0 deletions

File tree

‎examples/pre-training/tools/cvt_eb45_paddle2pt.py‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
import json
17+
import paddle
18+
from safetensors import safe_open
19+
from safetensors.paddle import save_file
20+
from collections import defaultdict
21+
22+
SRC_PATH = "./ERNIE-4.5-300B-A47B-Base-Paddle"
23+
DST_PATH = "./ERNIE-4.5-300B-A47B-Base-PT-out"
24+
25+
os.makedirs(DST_PATH, exist_ok=True)
26+
27+
with open(os.path.join(SRC_PATH, "model.safetensors.index.json")) as f:
28+
src_map = json.load(f)["weight_map"]
29+
with open(os.path.join(SRC_PATH, "config.json")) as f:
30+
config = json.load(f)
31+
32+
assert (
33+
config["hidden_size"] % config["num_attention_heads"] == 0
34+
), "head_dim not divisible"
35+
head_dim = config["hidden_size"] // config["num_attention_heads"]
36+
q_size = head_dim * config["num_attention_heads"]
37+
kv_size = head_dim * config["num_key_value_heads"]
38+
39+
total_size = 0
40+
dst_map = {}
41+
42+
src_rev_map = defaultdict(set)
43+
for k, v in src_map.items():
44+
src_rev_map[v].add(k)
45+
46+
for src_file, src_keys in src_rev_map.items():
47+
print("reading:", src_file, "size:", len(src_keys))
48+
dst_weight = {}
49+
50+
for src_key in sorted(src_keys):
51+
with safe_open(
52+
os.path.join(SRC_PATH, src_file), framework="paddle", device="cpu"
53+
) as f:
54+
tensor = f.get_tensor(src_key)
55+
56+
base_key = src_key
57+
if base_key.startswith("ernie."):
58+
base_key = "model." + base_key[6:]
59+
60+
if ".up_gate_proj." in src_key:
61+
# split gate_proj / up_proj (equal halves)
62+
half = tensor.shape[-1] // 2
63+
gate_tensor = tensor[:, :half]
64+
up_tensor = tensor[:, half:]
65+
66+
# transpose back
67+
gate_tensor = gate_tensor.T.contiguous()
68+
up_tensor = up_tensor.T.contiguous()
69+
70+
gate_key = base_key.replace(".up_gate_proj.", ".gate_proj.")
71+
up_key = base_key.replace(".up_gate_proj.", ".up_proj.")
72+
73+
dst_weight[gate_key] = gate_tensor
74+
dst_weight[up_key] = up_tensor
75+
dst_map[gate_key] = src_file
76+
dst_map[up_key] = src_file
77+
78+
elif ".qkv_proj." in src_key:
79+
# split q / k / v (unequal: q_size, kv_size, kv_size)
80+
q_tensor, k_tensor, v_tensor = paddle.split(
81+
tensor, [q_size, kv_size, kv_size], axis=-1
82+
)
83+
84+
# transpose back
85+
q_tensor = q_tensor.T.contiguous()
86+
k_tensor = k_tensor.T.contiguous()
87+
v_tensor = v_tensor.T.contiguous()
88+
89+
q_key = base_key.replace(".qkv_proj.", ".q_proj.")
90+
k_key = base_key.replace(".qkv_proj.", ".k_proj.")
91+
v_key = base_key.replace(".qkv_proj.", ".v_proj.")
92+
93+
dst_weight[q_key] = q_tensor
94+
dst_weight[k_key] = k_tensor
95+
dst_weight[v_key] = v_tensor
96+
dst_map[q_key] = src_file
97+
dst_map[k_key] = src_file
98+
dst_map[v_key] = src_file
99+
100+
else:
101+
# no merge, just possibly transpose
102+
if "_proj." in src_key or ".gate." in src_key or "lm_head" in src_key:
103+
tensor = tensor.T.contiguous()
104+
105+
dst_weight[base_key] = tensor
106+
dst_map[base_key] = src_file
107+
108+
print(end=".", flush=True)
109+
110+
save_file(dst_weight, os.path.join(DST_PATH, src_file))
111+
print()
112+
113+
with open(os.path.join(DST_PATH, "model.safetensors.index.json"), "w") as f:
114+
data = {
115+
"metadata": {
116+
"total_size": total_size,
117+
},
118+
"weight_map": dst_map,
119+
}
120+
json.dump(data, f, ensure_ascii=False, indent=2)
121+
print("done")

‎examples/pre-training/tools/cvt_eb45_pt2paddle.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
import json
17+
import paddle
18+
from safetensors import safe_open
19+
from safetensors.paddle import save_file
20+
from collections import defaultdict
21+
22+
src_path = "./ERNIE-4.5-300B-A47B-Base-PT"
23+
dst_path = "./ERNIE-4.5-300B-A47B-Base-Paddle-out"
24+
25+
os.makedirs(dst_path, exist_ok=True)
26+
27+
with open(os.path.join(src_path, "model.safetensors.index.json")) as f:
28+
src_map = json.load(f)["weight_map"]
29+
30+
total_size = 0
31+
dst_map = {}
32+
33+
src_rev_map = defaultdict(set)
34+
for k, v in src_map.items():
35+
src_rev_map[v].add(k)
36+
37+
for src_file, src_keys in src_rev_map.items():
38+
print("reading:", src_file, "size:", len(src_keys))
39+
dst_weight = {}
40+
41+
for src_key in src_keys:
42+
if ".k_proj." in src_key or ".v_proj." in src_key or ".up_proj." in src_key:
43+
continue
44+
45+
dst_key = src_key
46+
if dst_key.startswith("model."):
47+
dst_key = "ernie." + dst_key[6:]
48+
49+
key_decomp = [src_key]
50+
if ".gate_proj." in src_key:
51+
dst_key = dst_key.replace(".gate_proj.", ".up_gate_proj.")
52+
key_decomp = [
53+
src_key,
54+
src_key.replace(".gate_proj.", ".up_proj."),
55+
]
56+
elif ".q_proj." in src_key:
57+
dst_key = dst_key.replace(".q_proj.", ".qkv_proj.")
58+
key_decomp = [
59+
src_key,
60+
src_key.replace(".q_proj.", ".k_proj."),
61+
src_key.replace(".q_proj.", ".v_proj."),
62+
]
63+
64+
weight_decomp = []
65+
for key in key_decomp:
66+
with safe_open(
67+
os.path.join(src_path, src_file), framework="paddle", device="cpu"
68+
) as f:
69+
tensor = f.get_tensor(key)
70+
if "_proj." in key or ".gate." in key or "lm_head" in key:
71+
tensor = tensor.T.contiguous()
72+
weight_decomp.append(tensor)
73+
74+
dst_weight[dst_key] = (
75+
weight_decomp[0]
76+
if len(weight_decomp) == 1
77+
else paddle.concat(weight_decomp, axis=-1)
78+
)
79+
dst_map[dst_key] = src_file
80+
print(end=".", flush=True)
81+
82+
save_file(dst_weight, os.path.join(dst_path, src_file))
83+
print()
84+
85+
with open(os.path.join(dst_path, "model.safetensors.index.json"), "w") as f:
86+
data = {
87+
"metadata": {
88+
"total_size": total_size,
89+
},
90+
"weight_map": dst_map,
91+
}
92+
json.dump(data, f, ensure_ascii=False, indent=2)
93+
print("done")

0 commit comments

Comments
 (0)