apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-10-03 15:56:49 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

polish code with new pre-commit rule (#2923)

Browse Source
This commit is contained in:
Zero Rains
2025-07-19 23:19:27 +08:00
committed by GitHub
parent b8676d71a8
commit 25698d56d1
424 changed files with 14307 additions and 13518 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
scripts/vit_model_split.py
View File

@@ -14,12 +14,11 @@
# limitations under the License.
"""
import paddle
import paddle.distributed as dist
from paddle.distributed import fleet
import argparse
import os
import paddle
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_path",
@@ -47,12 +46,19 @@ for i in range(args.model_degree):
static_dict = {}
for k, v in input_model_state_dict.items():
if "qkv.weight" in k:
static_dict[k] = input_model_state_dict[k].reshape(
[hidden_size, 3, kv_num_heads, head_dim]
).split(args.model_degree, axis=-2)[i].reshape([hidden_size, -1])
static_dict[k] = (
input_model_state_dict[k]
.reshape([hidden_size, 3, kv_num_heads, head_dim])
.split(args.model_degree, axis=-2)[i]
.reshape([hidden_size, -1])
)
elif "qkv.bias" in k:
static_dict[k] = input_model_state_dict[k].reshape(
[3, kv_num_heads, head_dim]).split(args.model_degree, axis=-2)[i].reshape([-1])
static_dict[k] = (
input_model_state_dict[k]
.reshape([3, kv_num_heads, head_dim])
.split(args.model_degree, axis=-2)[i]
.reshape([-1])
)
elif "attn.proj.weight" in k:
static_dict[k] = input_model_state_dict[k].split(args.model_degree, axis=-2)[i]
elif "fc1.weight" in k:
@@ -64,4 +70,7 @@ for i in range(args.model_degree):
else:
static_dict[k] = v
paddle.save(static_dict, os.path.join(args.model_path, f"model_state_tp0{i}.pdparams"))
paddle.save(
static_dict,
os.path.join(args.model_path, f"model_state_tp0{i}.pdparams"),
)