mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-08 10:00:29 +08:00
【Hackathon 9th No.61、65】add test_draft_model_update (#3940)
* add draft_model_update test * fix * fix * fix * fix * fix
This commit is contained in:
co63oc
248
tests/operators/test_draft_model_update.py
Normal file
248
tests/operators/test_draft_model_update.py
Normal file
@@ -0,0 +1,248 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License")
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import paddle
|
||||
|
||||
from fastdeploy.model_executor.ops.gpu import draft_model_update
|
||||
|
||||
|
||||
def is_in_end(id, end_ids, length):
|
||||
flag = False
|
||||
for i in range(length):
|
||||
if id == end_ids[i]:
|
||||
return True
|
||||
return flag
|
||||
|
||||
|
||||
# recalculate data offset, offset_new is starting from index 0
|
||||
def get_inter_next_tokens_start_offset(inter_next_tokens, max_seq_len, start_id, offset):
|
||||
offset_new = start_id + offset
|
||||
return inter_next_tokens[int(offset_new / max_seq_len)][int(offset_new % max_seq_len)]
|
||||
|
||||
|
||||
def draft_model_update_kernel(
|
||||
inter_next_tokens,
|
||||
draft_tokens,
|
||||
pre_ids,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
step_idx,
|
||||
output_cum_offsets,
|
||||
stop_flags,
|
||||
not_need_stop,
|
||||
max_dec_len,
|
||||
end_ids,
|
||||
base_model_draft_tokens,
|
||||
bsz,
|
||||
max_draft_token,
|
||||
pre_id_length,
|
||||
max_base_model_draft_token,
|
||||
end_ids_len,
|
||||
max_seq_len,
|
||||
substep,
|
||||
prefill_one_step_stop,
|
||||
):
|
||||
stop_sum = 0
|
||||
for tid in range(bsz):
|
||||
stop_flag_now_int = 0
|
||||
draft_token_now = draft_tokens[tid]
|
||||
pre_ids_now = pre_ids[tid]
|
||||
base_model_draft_tokens_now = base_model_draft_tokens[tid]
|
||||
next_tokens_start_id = tid * max_seq_len - output_cum_offsets[tid]
|
||||
# next_tokens_start =
|
||||
seq_len_this_time = seq_lens_this_time[tid]
|
||||
seq_len_encoder = seq_lens_encoder[tid]
|
||||
seq_len_decoder = seq_lens_decoder[tid]
|
||||
|
||||
# 1. update step_idx && seq_lens_dec
|
||||
if not stop_flags[tid]: # seq_lens_decoder > 0 or seq_lens_encoder > 0
|
||||
token_this_time = -1
|
||||
# decoder step
|
||||
if seq_len_decoder > 0 and seq_len_encoder <= 0:
|
||||
seq_lens_decoder[tid] += seq_len_this_time
|
||||
token_this_time = get_inter_next_tokens_start_offset(
|
||||
inter_next_tokens, max_seq_len, next_tokens_start_id, seq_len_this_time - 1
|
||||
)
|
||||
draft_token_now[0] = token_this_time
|
||||
base_model_draft_tokens_now[substep + 1] = token_this_time
|
||||
step_idx[tid] += seq_len_this_time
|
||||
pre_ids_now[step_idx[tid]] = token_this_time
|
||||
else:
|
||||
token_this_time = get_inter_next_tokens_start_offset(
|
||||
inter_next_tokens, max_seq_len, next_tokens_start_id, 0
|
||||
)
|
||||
|
||||
# seq_lens_decoder[tid] = seq_lens_encoder[tid]
|
||||
seq_lens_decoder[tid] = seq_len_encoder + seq_len_decoder
|
||||
seq_lens_encoder[tid] = 0
|
||||
pre_ids_now[1] = token_this_time
|
||||
step_idx[tid] += 1
|
||||
draft_token_now[0] = token_this_time
|
||||
base_model_draft_tokens_now[substep + 1] = token_this_time
|
||||
|
||||
# multi_end
|
||||
if is_in_end(token_this_time, end_ids, end_ids_len) or prefill_one_step_stop:
|
||||
stop_flags[tid] = True
|
||||
stop_flag_now_int = 1
|
||||
# max_dec_len
|
||||
elif step_idx[tid] >= max_dec_len[tid]:
|
||||
stop_flags[tid] = True
|
||||
draft_token_now[seq_len_this_time - 1] = end_ids[0]
|
||||
base_model_draft_tokens_now[substep + 1] = end_ids[0]
|
||||
stop_flag_now_int = 1
|
||||
else:
|
||||
draft_token_now[0] = -1
|
||||
base_model_draft_tokens_now[substep + 1] = -1
|
||||
stop_flag_now_int = 1
|
||||
|
||||
# 2. set end
|
||||
if not stop_flags[tid]:
|
||||
seq_lens_this_time[tid] = 1
|
||||
else:
|
||||
seq_lens_this_time[tid] = 0
|
||||
seq_lens_encoder[tid] = 0
|
||||
|
||||
stop_sum = stop_sum + stop_flag_now_int
|
||||
not_need_stop[0] = stop_sum < bsz
|
||||
|
||||
|
||||
def draft_model_update_ref(
|
||||
inter_next_tokens,
|
||||
draft_tokens,
|
||||
pre_ids,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
step_idx,
|
||||
output_cum_offsets,
|
||||
stop_flags,
|
||||
not_need_stop,
|
||||
max_dec_len,
|
||||
end_ids,
|
||||
base_model_draft_tokens,
|
||||
max_seq_len,
|
||||
substep,
|
||||
):
|
||||
seq_lens_this_time_shape = seq_lens_this_time.shape
|
||||
real_bsz = seq_lens_this_time_shape[0]
|
||||
end_ids_len = end_ids.shape[0]
|
||||
max_draft_token = draft_tokens.shape[1]
|
||||
pre_id_length = pre_ids.shape[1]
|
||||
max_base_model_draft_token = base_model_draft_tokens.shape[1]
|
||||
|
||||
prefill_one_step_stop = False
|
||||
import os
|
||||
|
||||
env = os.getenv("PREFILL_NODE_ONE_STEP_STOP")
|
||||
if env == "1":
|
||||
prefill_one_step_stop = True
|
||||
|
||||
draft_model_update_kernel(
|
||||
inter_next_tokens,
|
||||
draft_tokens,
|
||||
pre_ids,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
step_idx,
|
||||
output_cum_offsets,
|
||||
stop_flags,
|
||||
not_need_stop,
|
||||
max_dec_len,
|
||||
end_ids,
|
||||
base_model_draft_tokens,
|
||||
real_bsz,
|
||||
max_draft_token,
|
||||
pre_id_length,
|
||||
max_base_model_draft_token,
|
||||
end_ids_len,
|
||||
max_seq_len,
|
||||
substep,
|
||||
prefill_one_step_stop,
|
||||
)
|
||||
|
||||
|
||||
class TestDraftModelUpdate(unittest.TestCase):
|
||||
def test_draft_model_update(self):
|
||||
self._run_paddle_test()
|
||||
|
||||
def _run_paddle_test(self):
|
||||
np.random.seed(42)
|
||||
paddle.seed(42)
|
||||
|
||||
max_bsz = 128
|
||||
max_draft_token = 3
|
||||
pre_id_length = 3
|
||||
max_seq_len = 100
|
||||
max_base_model_draft_token = 4
|
||||
substep = 2
|
||||
|
||||
inter_next_tokens = paddle.randint(1, 100, shape=(max_bsz, max_seq_len), dtype="int64")
|
||||
draft_tokens = paddle.randint(1, 100, shape=(max_bsz, max_draft_token), dtype="int64")
|
||||
pre_ids = paddle.randint(1, 100, shape=(max_bsz, pre_id_length), dtype="int64")
|
||||
seq_lens_this_time = paddle.randint(1, 2, shape=(max_bsz,), dtype="int32")
|
||||
seq_lens_encoder = paddle.randint(1, 10, shape=(max_bsz,), dtype="int32")
|
||||
seq_lens_decoder = paddle.randint(1, 10, shape=(max_bsz,), dtype="int32")
|
||||
step_idx = paddle.randint(1, 10, shape=(max_bsz,), dtype="int64")
|
||||
output_cum_offsets = paddle.randint(0, 2, shape=(max_bsz,), dtype="int32")
|
||||
output_cum_offsets[0] = 0
|
||||
stop_flags = paddle.zeros([max_bsz], dtype="bool")
|
||||
not_need_stop = paddle.zeros([1], dtype="bool").to(device=paddle.CPUPlace())
|
||||
max_dec_len = paddle.randint(100, 102, shape=(max_bsz,), dtype="int64")
|
||||
end_ids = paddle.to_tensor([2], dtype="int64")
|
||||
base_model_draft_tokens = paddle.randint(1, 10, shape=(max_bsz, max_base_model_draft_token), dtype="int64")
|
||||
|
||||
inputs = (
|
||||
inter_next_tokens,
|
||||
draft_tokens,
|
||||
pre_ids,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
step_idx,
|
||||
output_cum_offsets,
|
||||
stop_flags,
|
||||
not_need_stop,
|
||||
max_dec_len,
|
||||
end_ids,
|
||||
base_model_draft_tokens,
|
||||
max_seq_len,
|
||||
substep,
|
||||
)
|
||||
# inplace modify, need to clone inputs
|
||||
inputs_clone = [x.clone() if isinstance(x, paddle.Tensor) else x for x in inputs]
|
||||
draft_model_update(*inputs)
|
||||
draft_model_update_ref(*inputs_clone)
|
||||
idx_list = (
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
5,
|
||||
6,
|
||||
8,
|
||||
9,
|
||||
12,
|
||||
)
|
||||
for i in idx_list:
|
||||
np.testing.assert_allclose(inputs[i].numpy(), inputs_clone[i].numpy())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user