File size: 2,927 Bytes
454b3e6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """基线：单问题 / 多问题 延迟拆解（tokenize vs GPU forward），以及 profiler top kernels。"""
import os, sys, time, torch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "apps"))
from common import get_agent
from laya.common import build_sequence, collate_items, QTYPES, render_options
agent = get_agent(os.environ.get("LAYA_VARIANT", "multilingual"))
state = {"subject": "Duplicate charge on invoice 4411", "body": "We were billed twice for March. Please refund the duplicate or we're moving to a competitor. " * 3}
Q = {"department": {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"billing": "invoices, refunds", "technical": "bugs, outages", "sales": "pricing"}},
"urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["not urgent", "soon", "blocking"]},
"churn": {"type": "noul", "instructions": "Does the user threaten to cancel?"}}
def qs(n): return {f"{k}{i}": v for i in range(n) for k, v in Q.items()}
def timeit(fn, iters=20):
for _ in range(3): fn()
torch.cuda.synchronize(); t = time.perf_counter()
for _ in range(iters): fn()
torch.cuda.synchronize(); return (time.perf_counter() - t) / iters * 1000
for n in (1, 3, 10, 30):
q = qs(n) if n != 1 else {"department": Q["department"]}
ms = timeit(lambda: agent.predict(state, q))
# tokenization-only cost
def tok_only():
items = []
for qid in q:
qq = agent._to_internal(q[qid]); seq, m = build_sequence(agent.tok, state, qq, agent.cfg["max_len"], agent.cfg["head_max_len"])
items.append({"ids": seq, "markers": m, "qtype": QTYPES[qq["t"]]})
return collate_items([items], agent.tok.pad_token_id)
tms = timeit(tok_only)
b = tok_only()
print(f"questions={len(q):3d} seq_len={b['input_ids'].shape[1]:4d} predict={ms:7.2f} ms tokenize={tms:6.2f} ms gpu+post={ms-tms:7.2f} ms per-q={ms/len(q):5.2f} ms")
# profile the model forward at n=10
q = qs(10); b = None
items = []
for qid in q:
qq = agent._to_internal(q[qid]); seq, m = build_sequence(agent.tok, state, qq, agent.cfg["max_len"], agent.cfg["head_max_len"])
items.append({"ids": seq, "markers": m, "qtype": QTYPES[qq["t"]]})
b = {k: v.to("cuda") for k, v in collate_items([items], agent.tok.pad_token_id).items() if torch.is_tensor(v)}
def fwd():
with torch.no_grad(), torch.autocast("cuda", dtype=agent.dtype):
return agent.model(b["input_ids"], b["attention_mask"], b["marker_pos"], b["marker_mask"], b["qtype"])
print(f"\nmodel forward only (30 q, L={b['input_ids'].shape[1]}): {timeit(fwd):.2f} ms dtype={agent.dtype}")
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for _ in range(5): fwd()
torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=25))
|