Super tiny add gsp-fast-prepare (#14992)

This commit is contained in:
fzyzcjy
2025-12-13 09:45:21 +08:00
committed by GitHub
parent 6d2b3324ef
commit df111bc0fe

View File

@@ -1661,7 +1661,11 @@ def sample_generated_shared_prefix_requests(
flat_index = group_idx * prompts_per_group + prompt_idx
question = questions[flat_index]
full_prompt = f"{system_prompt}\n\n{question}"
prompt_len = len(tokenizer.encode(full_prompt))
prompt_len = (
1
if getattr(args, "gsp_fast_prepare", False)
else len(tokenizer.encode(full_prompt))
)
input_requests.append(
DatasetRow(
@@ -1681,14 +1685,15 @@ def sample_generated_shared_prefix_requests(
print(f"Number of groups: {num_groups}")
print(f"Prompts per group: {prompts_per_group}")
print(f"Total prompts: {len(input_requests)}")
print(f"Total input tokens: {total_input_tokens}")
print(f"Total output tokens: {total_output_tokens}")
print(
f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
)
print(
f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n"
)
if not getattr(args, "gsp_fast_prepare", False):
print(f"Total input tokens: {total_input_tokens}")
print(f"Total output tokens: {total_output_tokens}")
print(
f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
)
print(
f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n"
)
# Save to cache
cache_path.parent.mkdir(parents=True, exist_ok=True)
@@ -2915,6 +2920,11 @@ if __name__ == "__main__":
default=1.0,
help="Range of sampled ratio of input/output length, used only for gsp dataset.",
)
group.add_argument(
"--gsp-fast-prepare",
action="store_true",
help="Speedup preparing by removing statistics computation, which will make some output statistics inaccurate but suitable for pressure tests.",
)
mooncake_group = parser.add_argument_group("mooncake dataset arguments")
mooncake_group.add_argument(
"--mooncake-slowdown-factor",