Retire the Qwen AR pipeline
This commit is contained in:
1 parent
5773942c8f
commit
1de78fa3ef
22 files changed
+6429
-667
No files matched your search
@@ -1627,7 +1627,7 @@ def ar_driver_fixtures():
|
||||
"""Unchanged generate_ar with a tiny GPU transition-table model, no weights.
|
||||
|
||||
Isolates orchestration (not model or throughput parity); actual prefill,
|
||||
samplers, pipeline and final-state capture are all MTPLX implementations.
|
||||
samplers and final-state capture are all MTPLX implementations.
|
||||
"""
|
||||
from mtplx.generation import generate_ar
|
||||
from mtplx.runtime import MTPLXRuntime
|
||||
@@ -1641,11 +1641,8 @@ def ar_driver_fixtures():
|
||||
table = mx.array(values.reshape(vocab, vocab))
|
||||
class Model:
|
||||
def __init__(self):
|
||||
self.calls, self.modes = [], []
|
||||
self.calls = []
|
||||
def make_cache(self): return []
|
||||
def set_ar_pipeline_mode(self, enabled):
|
||||
self.modes.append(bool(enabled))
|
||||
return True
|
||||
def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None):
|
||||
self.calls.append((ids, return_hidden))
|
||||
out = table[ids]
|
||||
@@ -1653,12 +1650,11 @@ def ar_driver_fixtures():
|
||||
tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: " ".join(map(str, ids)))
|
||||
for pipeline in (False, True):
|
||||
for asynchronous in (False, True):
|
||||
for temperature, penalty in ((0., 0.), (.6, 0.), (.6, .4)):
|
||||
for temperature, top_p, top_k, penalty in ((0., .95, 20, 0.), (.6, .95, 20, 0.), (.6, .95, 20, .4), (1., .95, 20, 0.), (1., .5, 2, 0.), (1., 0., 2, 0.), (1., 1., 2, 0.)):
|
||||
for limit in (0, 1, 2, 3, 12):
|
||||
for capture in (False, True):
|
||||
# 13/35 are the first/second sampled pipeline tokens:
|
||||
# exercise immediate EOS and EOS with an in-flight next
|
||||
# graph, not just limits and a non-triggering stop set.
|
||||
# Retain the historical stop IDs and explicitly set the
|
||||
# retired flag to prove current MTPLX ignores it.
|
||||
for stop_ids in (set(), {0, 5, 36}, {13}, {35}):
|
||||
os.environ["MTPLX_AR_PIPELINE"] = "1" if pipeline else "0"
|
||||
os.environ["MTPLX_ASYNC_AR"] = "1" if asynchronous else "0"
|
||||
@@ -1666,15 +1662,15 @@ def ar_driver_fixtures():
|
||||
rt = MTPLXRuntime(model, tokenizer, Path("."), False, None)
|
||||
emitted = []
|
||||
out = generate_ar(rt, [1, 2], max_tokens=limit,
|
||||
sampler=SamplerConfig(temperature=temperature, top_p=.95, top_k=20, presence_penalty=penalty),
|
||||
sampler=SamplerConfig(temperature=temperature, top_p=top_p, top_k=top_k, presence_penalty=penalty),
|
||||
seed=4294967297, stop_token_ids=stop_ids, token_callback=lambda ids: emitted.extend(ids), capture_final_state=capture)
|
||||
# Decode calls only: real prefill consumed [1] then [2].
|
||||
calls = [{"ids": np.array(ids).reshape(-1).tolist(), "hidden": hidden} for ids, hidden in model.calls[2:]]
|
||||
final = None if out.final_state is None else dict(shape=out.final_state.final_logits.shape,
|
||||
sha256=hashlib.sha256(np.array(out.final_state.final_logits).tobytes()).hexdigest())
|
||||
print(json.dumps(dict(pipeline=pipeline, asynchronous=asynchronous, temperature=temperature,
|
||||
print(json.dumps(dict(legacy_pipeline=pipeline, asynchronous=asynchronous, temperature=temperature, top_p=top_p, top_k=top_k,
|
||||
penalty=penalty, limit=limit, capture=capture, stops=sorted(stop_ids), tokens=out.tokens,
|
||||
emitted=emitted, calls=calls, modes=model.modes, final=final,
|
||||
emitted=emitted, calls=calls, final=final,
|
||||
verify_calls=out.stats.verify_calls, finish=out.finish_reason)), flush=True)
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user