mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
"""Batch-invariant test in pytest form.
|
|
|
|
Run directly:
|
|
python -m pytest tests/models/test_batch_invariant.py -s
|
|
|
|
Verifies that with DEEPINFER_BATCH_INVARIANT=1, running generate on the same prompt batched and singly
|
|
gives identical results. Models / parallel sizes are parameterized by CASES.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import random
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
|
|
import pytest
|
|
|
|
from examples._config import DEFAULT_MODEL_NAME, create_llm
|
|
from tests.models._runner import run_model_test
|
|
|
|
# ===== Parameter tables =====
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Case:
|
|
model_name: str
|
|
ep_size: int
|
|
use_nextn_decode: int = 0 # the backbone only by default; cases that need NextN coverage enable it explicitly
|
|
|
|
|
|
CASES = [
|
|
Case(DEFAULT_MODEL_NAME, ep_size=8, use_nextn_decode=5),
|
|
]
|
|
NUM_SEQS = 64 # number of sequences generated per case
|
|
|
|
# ===== Test logic =====
|
|
|
|
|
|
def _get_result(output: list) -> tuple[list, list, list, list]:
|
|
"""Extract (token_ids, logprobs, top_token_ids, top_logprobs) from the RequestOutput list."""
|
|
token_ids = [x.token_id for x in output]
|
|
logprobs = [x.token_logprobs for x in output]
|
|
top_token_ids = [x.top_token_ids for x in output]
|
|
top_logprobs = [x.top_logprobs for x in output]
|
|
return token_ids, logprobs, top_token_ids, top_logprobs
|
|
|
|
|
|
# ===== Subprocess entry =====
|
|
|
|
|
|
def _run_batch_invariant(case: Case, dp_size: int, dp_rank: int) -> None:
|
|
"""Verify batch invariance on one DP rank (called inside a subprocess by run_dp)."""
|
|
llm = create_llm(
|
|
case.model_name,
|
|
dp_rank,
|
|
ep_size=case.ep_size,
|
|
env={
|
|
"DEEPINFER_LOGGING_INTERVAL_STEP": "1000",
|
|
"DEEPINFER_USE_NEXTN_DECODE": str(case.use_nextn_decode),
|
|
"DEEPINFER_BATCH_INVARIANT": "1",
|
|
},
|
|
max_num_batched_tokens=16384,
|
|
max_model_len=16384,
|
|
max_num_seqs=NUM_SEQS,
|
|
gpu_memory_utilization=0.9,
|
|
lookahead_tokens=min(1, case.use_nextn_decode),
|
|
)
|
|
|
|
from deepinfer.sampling_params import SamplingParams
|
|
|
|
random.seed(dp_rank)
|
|
prompt_lens = [max(1, int(random.normalvariate(2500, 1000))) for _ in range(NUM_SEQS)]
|
|
print(f"{dp_rank=}, {prompt_lens=}")
|
|
# Engram input token ids must be within the original→compressed token mapping table.
|
|
vocab_size = llm.llm_engine.model_config.get_vocab_size()
|
|
if llm.llm_engine.model_config.use_compressed_tokenizer:
|
|
from deepinfer.model_executor.models.v4_modules.engram import CompressedTokenizer
|
|
|
|
vocab_size = len(CompressedTokenizer.lookup_table)
|
|
|
|
# avoid special tokens
|
|
special_tokens_dict = llm.llm_engine.tokenizer.added_tokens_decoder
|
|
non_special_tokens = [token for token in range(vocab_size) if token not in special_tokens_dict]
|
|
prompt_token_ids = [[random.choice(non_special_tokens) for _ in range(length)] for length in prompt_lens]
|
|
|
|
sampling_params = SamplingParams(
|
|
temperature=1.0,
|
|
top_p=0.7,
|
|
max_tokens=50,
|
|
ignore_eos=True,
|
|
top_logprobs=5,
|
|
prompt_logprobs=False,
|
|
)
|
|
|
|
outputs = llm.generate(None, sampling_params, prompt_token_ids, use_tqdm=False)
|
|
llm.llm_engine.log_stats = False
|
|
|
|
for token_ids, output in zip(prompt_token_ids, outputs, strict=True):
|
|
sp = copy.copy(sampling_params)
|
|
sp.golden_token_ids = [x.token_id for x in output]
|
|
output_single = llm.generate(None, sp, [token_ids], use_tqdm=False)[0]
|
|
result_batched = _get_result(output)
|
|
result_single = _get_result(output_single)
|
|
assert result_batched == result_single, f"{dp_rank=}, {result_batched=}, {result_single=}"
|
|
|
|
llm.close()
|
|
|
|
|
|
# ===== pytest entry =====
|
|
|
|
|
|
@pytest.mark.parametrize("case", CASES, ids=lambda c: c.model_name)
|
|
def test_batch_invariant(case: Case) -> None:
|
|
"""pytest entry: verify batch invariance across all DP ranks."""
|
|
run_model_test(partial(_run_batch_invariant, case), case.ep_size)
|