mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56:21 +08:00
28fb291db1
No behaviour change. `CONTRIBUTING.md` requires `ruff check --no-cache` and `ruff format --check --no-cache` to pass, and the SM90 branch predates that requirement: 72 findings, all in the files added here. Most of it is mechanical (`ruff format`), plus: - annotations for the helper signatures (`torch.Tensor` / `int` / `Generator` / return types); hooks are typed as `Callable[[nn.Module, tuple, object], None]` rather than left bare; - two `if`/`else` assignments become ternaries (SIM108), and the index read in `h20_real_requant_check.py` uses a context manager (SIM115). Also corrects the `sparse_mla_fwd` docstring, which still described the old `0 <= idx <= q_position` masking contract. Passing `is_causal=True` there is what produced the all-NaN layer 0: indices address the assembled `[swa cache | current swa | compressed]` buffer, so the row number is not a valid causal bound and the engine must pass False.
103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
"""One short prompt, greedy, few tokens — a fast semantic sanity check on SM90.
|
|
|
|
`offline_inference.py` runs five prompts for up to 128 tokens each, which is minutes of decode
|
|
per attempt. This asks a single question and stops after `--max_tokens`, so a bring-up run tells
|
|
us what we actually care about: whether the first tokens are words rather than noise.
|
|
|
|
python scripts/h20_ask.py "问题：中国的首都是哪里？\n答案：" --max_tokens 8
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env
|
|
from _launch import launch_dp
|
|
|
|
from deepinfer.sampling_params import SamplingParams
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"prompt",
|
|
nargs="*",
|
|
default=["问题：中国的首都是哪里？\n答案：", "问题：中国的首都是哪里？\n答案：北京"],
|
|
)
|
|
parser.add_argument("--max_tokens", type=int, default=8)
|
|
parser.add_argument("--ep", type=int, default=8)
|
|
parser.add_argument("--cp", type=int, default=1)
|
|
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
|
|
parser.add_argument("--temperature", type=float, default=0.0, help="0 = greedy, so the answer is reproducible")
|
|
parser.add_argument(
|
|
"--ids",
|
|
default="",
|
|
help="Comma-separated prompt token ids. Use this to feed *exactly* the ids the reference model saw "
|
|
"(the official generate.py wraps text in a chat template, so raw strings are not comparable).",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
ep_size, cp_size = args.ep, args.cp
|
|
dp_size = max(ep_size // cp_size, 1)
|
|
world_size = dp_size * cp_size
|
|
assert world_size % ep_size == 0
|
|
|
|
|
|
def main(dp_size: int, dp_rank: int) -> None:
|
|
setup_common_env(
|
|
model_name=args.model_name,
|
|
sampler_force_topk=4096,
|
|
set_override_experts=False,
|
|
enable_eplb=False,
|
|
)
|
|
os.environ.setdefault("DEEPINFER_BATCH_INVARIANT", "1")
|
|
|
|
llm = create_llm(
|
|
args.model_name,
|
|
dp_rank,
|
|
ep_size=ep_size,
|
|
cp_size=cp_size,
|
|
env={"DEEPINFER_LOGGING_INTERVAL_STEP": "10"},
|
|
data_parallel_size=dp_size,
|
|
max_num_batched_tokens=4096,
|
|
max_model_len=4096,
|
|
max_num_seqs=max(1, len(args.prompt)),
|
|
gpu_memory_utilization=0.85,
|
|
lookahead_tokens=0,
|
|
)
|
|
|
|
sampling_params = [
|
|
SamplingParams(
|
|
temperature=args.temperature,
|
|
max_tokens=args.max_tokens,
|
|
ignore_eos=True,
|
|
top_logprobs=4,
|
|
seq_seed=0,
|
|
)
|
|
for _ in args.prompt
|
|
]
|
|
|
|
if args.ids:
|
|
ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
|
|
outputs = llm.generate(prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False)
|
|
shown = [f"<ids len={len(ids)}>"]
|
|
else:
|
|
outputs = llm.generate(prompts=args.prompt, sampling_params=sampling_params, use_tqdm=False)
|
|
shown = args.prompt
|
|
for i, output in enumerate(outputs):
|
|
assert output, f"sample {i}: no generated tokens"
|
|
tokens = [x.token_id for x in output]
|
|
logps = [x.token_logprobs for x in output]
|
|
text = llm.llm_engine.tokenizer.decode(tokens)
|
|
label = shown[i] if args.ids else args.prompt[i]
|
|
print(f"\n=== DP {dp_rank} sample {i}")
|
|
print(f"=== prompt: {label!r}")
|
|
print(f"=== token ids ({len(tokens)}): {tokens}")
|
|
print(f"=== per-token logprob: {[round(x, 3) for x in logps]}")
|
|
print(f"=== continuation: {text!r}", flush=True)
|
|
if args.ids:
|
|
print(f"=== ids in / out: {len(ids)} -> {tokens}")
|
|
|
|
llm.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
launch_dp(main, dp_size)
|