Spaces:
Sleeping
Sleeping
File size: 6,011 Bytes
28a1a01 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """Fair batch-one decode comparison: HF eager, torch.compile, and CISM.
All baselines run the SAME model checkpoint in FP32 with a KV cache, greedy
decoding, one torch thread, warmup discarded. torch.compile uses a manual
decode loop with StaticCache (its recommended fast path), not slow
use_cache=False generate(). CISM numbers come from the same machine via
cism bench. Output: JSON to stdout. Exit 0 on success.
"""
import argparse
import json
import time
import torch
def timed(fn, warmup_tokens: int, measured_tokens: int):
fn(warmup_tokens)
start = time.perf_counter()
generated = fn(measured_tokens)
elapsed = time.perf_counter() - start
return {
"tokens": generated,
"seconds": elapsed,
"tokens_per_second": generated / elapsed if elapsed > 0 else 0,
}
def run_eager(model, input_ids, eos_pad, count, warmup):
def loop(tokens):
out = model.generate(
input_ids, attention_mask=None, max_new_tokens=tokens,
min_new_tokens=tokens, do_sample=False, num_beams=1,
use_cache=True, pad_token_id=eos_pad,
)
return tokens
result = timed(loop, warmup, count)
result["method"] = "transformers generate(use_cache=True) eager"
return result
def run_manual_eager(model, input_ids, count, warmup):
"""Raw decode loop without generate()'s framework overhead."""
from transformers import StaticCache
cache = StaticCache(config=model.config, max_batch_size=1,
max_cache_len=warmup + count + len(input_ids[0]),
device="cpu", dtype=torch.float32)
def loop(tokens):
with torch.inference_mode():
cache.reset()
generated = 0
out = model(input_ids, past_key_values=cache, use_cache=True)
token = int(out.logits[0, -1].argmax())
generated += 1
one = torch.tensor([[token]])
while generated < tokens:
out = model(one, past_key_values=cache, use_cache=True)
token = int(out.logits[0, -1].argmax())
generated += 1
return generated
result = timed(loop, warmup, count)
result["method"] = "manual eager decode loop + StaticCache (no generate overhead)"
return result
def _dynamo_snapshot():
import torch._dynamo
return {name: sum(counts.values()) for name, counts in torch._dynamo.utils.counters.items()}
def _dynamo_delta(before, after):
keys = set(before) | set(after)
return {key: after.get(key, 0) - before.get(key, 0) for key in keys if after.get(key, 0) != before.get(key, 0)}
def run_compiled(model, input_ids, count, warmup, dynamic):
from transformers import StaticCache
cache = StaticCache(config=model.config, max_batch_size=1,
max_cache_len=warmup + count + len(input_ids[0]),
device="cpu", dtype=torch.float32)
compiled = torch.compile(model, dynamic=dynamic)
def loop(tokens):
with torch.inference_mode():
cache.reset()
generated = 0
out = compiled(input_ids, past_key_values=cache, use_cache=True)
token = int(out.logits[0, -1].argmax())
generated += 1
one = torch.tensor([[token]])
while generated < tokens:
out = compiled(one, past_key_values=cache, use_cache=True)
token = int(out.logits[0, -1].argmax())
generated += 1
return generated
# Warmup must absorb ALL compilation and recompilation; counters are reset
# afterwards so the measured window only reports steady-state guard hits.
loop(warmup)
loop(warmup)
before = _dynamo_snapshot()
generated = timed(loop, warmup, count)
after = _dynamo_snapshot()
generated["dynamo_counter_delta_during_measurement"] = _dynamo_delta(before, after)
generated["method"] = f"torch.compile(model, dynamic={dynamic}) + StaticCache decode loop"
return generated
def main():
parser = argparse.ArgumentParser()
parser.add_argument("model")
parser.add_argument("--revision", default=None)
parser.add_argument("--threads", type=int, default=1)
parser.add_argument("--tokens", type=int, default=64)
parser.add_argument("--warmup", type=int, default=32)
parser.add_argument("--skip-compile", action="store_true")
args = parser.parse_args()
torch.set_num_threads(args.threads)
torch.set_num_interop_threads(1)
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision)
model = AutoModelForCausalLM.from_pretrained(
args.model, revision=args.revision, dtype=torch.float32,
attn_implementation="eager",
).to("cpu").eval()
config = {"max_position_embeddings": model.config.max_position_embeddings}
ids = tokenizer("The future of artificial intelligence is", return_tensors="pt").input_ids
eos_pad = tokenizer.eos_token_id or 0
report = {
"model": args.model, "torch": torch.__version__, "threads": args.threads,
"precision": "fp32", "greedy": True, "warmup_tokens": args.warmup,
"measured_tokens": args.tokens,
}
report["eager_generate"] = run_eager(model, ids, eos_pad, args.tokens, args.warmup)
report["eager_manual_loop"] = run_manual_eager(model, ids, args.tokens, args.warmup)
if args.skip_compile:
report["compiled_dynamic_false"] = {"skipped": True}
report["compiled_dynamic_true"] = {"skipped": True}
else:
for dynamic in (False, True):
key = f"compiled_dynamic_{str(dynamic).lower()}"
try:
report[key] = run_compiled(model, ids, args.tokens, args.warmup, dynamic)
except Exception as error:
report[key] = {"failed": f"{type(error).__name__}: {error}"}
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|