Tiny add detokenization benchmarks (#16400)

benchmark/benchmark_batch/benchmark_tokenizer.py

@@ -1,124 +1,194 @@
+import argparse
 import random
 import time
 from statistics import mean
 
 from transformers import AutoTokenizer
 
-# CONFIG
-TOKENIZER_DIR = (
-    "/shared/public/sharing/fait360brew/training/models/meta-llama/Llama-3.2-3B"
-)
-NUM_TOKENS = 20000  # Each prompt should contain this many tokens
-BATCH_SIZES = [1, 2, 4, 8]  # Test different batch sizes
-NUM_RUNS = 5  # Number of runs for each batch size to get reliable measurements
-
-
-def generate_random_prompts(num_prompts, num_tokens, tokenizer):
-    """Generate random prompts with specified token count."""
-    vocab_size = tokenizer.vocab_size
-    all_prompts = []
-
-    print(f"Generating {num_prompts} random prompts with {num_tokens} tokens each...")
-    for i in range(num_prompts):
-        # Generate random token IDs - this directly gives us the exact token count
-        random_token_ids = [
-            random.randint(0, vocab_size - 1) for _ in range(num_tokens)
-        ]
-        random_text = tokenizer.decode(
-            random_token_ids, clean_up_tokenization_spaces=True
-        )
-
-        prompt = f"Prompt {i}: {random_text}"
-        tokens = tokenizer.encode(prompt)
-        print(f"  Prompt {i}: {len(tokens)} tokens")
-        all_prompts.append(prompt)
-
-    return all_prompts
-
-
-def benchmark_sequential_vs_batch(prompts, batch_size, tokenizer):
-    """Compare sequential vs batch tokenization for a given batch size."""
-
-    # Sequential tokenization using encode()
-    sequential_times = []
-    for run in range(NUM_RUNS):
-        batch_prompts = prompts[:batch_size]  # Use same prompts for fair comparison
-
-        start_time = time.perf_counter()
-        for prompt in batch_prompts:
-            tokens = tokenizer.encode(prompt)
-        sequential_time = (time.perf_counter() - start_time) * 1000
-        sequential_times.append(sequential_time)
-
-    # Batch tokenization using tokenizer()
-    batch_times = []
-    for run in range(NUM_RUNS):
-        batch_prompts = prompts[:batch_size]  # Use same prompts for fair comparison
-
-        start_time = time.perf_counter()
-        tokens = tokenizer(batch_prompts)
-        batch_time = (time.perf_counter() - start_time) * 1000
-        batch_times.append(batch_time)
-
-    return {
-        "batch_size": batch_size,
-        "avg_sequential_ms": mean(sequential_times),
-        "avg_batch_ms": mean(batch_times),
-        "speedup_factor": (
-            mean(sequential_times) / mean(batch_times) if mean(batch_times) > 0 else 0
-        ),
-        "sequential_runs": sequential_times,
-        "batch_runs": batch_times,
-    }
-
 
 def main():
+    args = parse_args()
+
     print("Tokenizer Benchmark: Sequential vs Batch Processing")
     print("-" * 60)
-    print(f"Tokenizer: {TOKENIZER_DIR}")
+    print(f"Tokenizer: {args.tokenizer}")
-    print(f"Tokens per prompt: {NUM_TOKENS}")
+    print(f"Functions: {', '.join(args.function)}")
-    print(f"Number of runs per batch size: {NUM_RUNS}")
+    print(f"Tokens per prompt: {args.num_tokens}")
+    print(f"Number of runs per batch size: {args.num_runs}")
+    print(f"Skip batch: {args.no_batch}")
     print("-" * 60)
 
-    # Load tokenizer once for all operations
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
-    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
+    max_batch_size = max(args.batch_sizes)
 
-    # The largest batch size determines how many prompts we need
+    token_ids = generate_random_token_ids(max_batch_size, args.num_tokens, tokenizer)
-    max_batch_size = max(BATCH_SIZES)
-    all_prompts = generate_random_prompts(max_batch_size, NUM_TOKENS, tokenizer)
 
-    results = []
+    if "encode" in args.function:
-    print("\nRunning benchmark...")
+        prompts = [
+            tokenizer.decode(ids, clean_up_tokenization_spaces=True)
+            for ids in token_ids
+        ]
+        run_benchmark(
+            name="encode",
+            data=prompts,
+            sequential_fn=lambda batch: [tokenizer.encode(p) for p in batch],
+            batch_fn=lambda batch: tokenizer(batch),
+            batch_sizes=args.batch_sizes,
+            num_runs=args.num_runs,
+            skip_batch=args.no_batch,
+        )
 
-    for batch_size in BATCH_SIZES:
+    if "decode" in args.function:
-        print(f"\nBenchmarking batch size: {batch_size}")
+        # mimic DetokenizerManager's usual case
-        result = benchmark_sequential_vs_batch(all_prompts, batch_size, tokenizer)
+        decode_kwargs = dict(
-        results.append(result)
+            skip_special_tokens=True,
+            spaces_between_special_tokens=True,
+        )
+        run_benchmark(
+            name="decode",
+            data=token_ids,
+            sequential_fn=lambda batch: [
+                tokenizer.decode(ids, **decode_kwargs) for ids in batch
+            ],
+            batch_fn=lambda batch: tokenizer.batch_decode(batch, **decode_kwargs),
+            batch_sizes=args.batch_sizes,
+            num_runs=args.num_runs,
+            skip_batch=args.no_batch,
+        )
 
-        print(f"  Sequential tokenization (encode):")
-        for i, run_time in enumerate(result["sequential_runs"]):
-            print(f"    Run {i+1}: {run_time:.2f} ms")
-        print(f"    Average: {result['avg_sequential_ms']:.2f} ms")
 
-        print(f"  Batch tokenization (tokenizer):")
+def run_benchmark(
-        for i, run_time in enumerate(result["batch_runs"]):
+    name, data, sequential_fn, batch_fn, batch_sizes, num_runs, skip_batch
-            print(f"    Run {i+1}: {run_time:.2f} ms")
+):
-        print(f"    Average: {result['avg_batch_ms']:.2f} ms")
+    print("\n" + "=" * 60)
+    print(f"{name.upper()} BENCHMARK")
+    print("=" * 60)
 
-        print(f"  Speedup factor: {result['speedup_factor']:.2f}x")
+    results = [
+        benchmark(data, bs, sequential_fn, batch_fn, num_runs, skip_batch)
+        for bs in batch_sizes
+    ]
+    print_results(results, name, skip_batch)
+
+
+def benchmark(data, batch_size, sequential_fn, batch_fn, num_runs, skip_batch):
+    batch_data = data[:batch_size]
+    sequential_times = measure_times(lambda: sequential_fn(batch_data), num_runs)
+    avg_seq = mean(sequential_times)
+
+    out = {
+        "batch_size": batch_size,
+        "avg_sequential_ms": avg_seq,
+        "sequential_runs": sequential_times,
+    }
+
+    if not skip_batch:
+        batch_times = measure_times(lambda: batch_fn(batch_data), num_runs)
+        avg_batch = mean(batch_times)
+        out |= {
+            "avg_batch_ms": avg_batch,
+            "speedup_factor": avg_seq / avg_batch if avg_batch > 0 else 0,
+            "batch_runs": batch_times,
+        }
+
+    return out
+
+
+def print_results(results, func_name, skip_batch):
+    for r in results:
+        print(f"\nBatch size: {r['batch_size']}")
+        print_runs(
+            f"Sequential {func_name}", r["sequential_runs"], r["avg_sequential_ms"]
+        )
+        if not skip_batch:
+            print_runs(f"Batch {func_name}", r["batch_runs"], r["avg_batch_ms"])
+            print(f"  Speedup factor: {r['speedup_factor']:.2f}x")
 
     print("\n" + "=" * 60)
-    print("SUMMARY OF RESULTS")
+    print(f"SUMMARY: {func_name.upper()}")
     print("=" * 60)
-    print(
-        f"{'Batch Size':<10} {'Sequential (ms)':<18} {'Batch (ms)':<18} {'Speedup':<10}"
-    )
-    print("-" * 60)
 
-    for result in results:
+    headers = ["Batch Size", "Sequential (ms)"]
-        print(
+    if not skip_batch:
-            f"{result['batch_size']:<10} {result['avg_sequential_ms']:.2f} ms{' ' * 8} {result['avg_batch_ms']:.2f} ms{' ' * 8} {result['speedup_factor']:.2f}x"
+        headers += ["Batch (ms)", "Speedup"]
-        )
+    print("".join(f"{h:<18}" for h in headers))
+    print("-" * (18 * len(headers)))
+
+    for r in results:
+        row = [f"{r['batch_size']}", f"{r['avg_sequential_ms']:.2f} ms"]
+        if not skip_batch:
+            row += [f"{r['avg_batch_ms']:.2f} ms", f"{r['speedup_factor']:.2f}x"]
+        print("".join(f"{v:<18}" for v in row))
+
+
+def print_runs(label, runs, avg):
+    print(f"  {label}:")
+    for i, t in enumerate(runs):
+        print(f"    Run {i+1}: {t:.2f} ms")
+    print(f"    Average: {avg:.2f} ms")
+
+
+def measure_times(fn, num_runs):
+    times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        fn()
+        times.append((time.perf_counter() - start) * 1000)
+    return times
+
+
+def generate_random_token_ids(num_prompts, num_tokens, tokenizer):
+    vocab_size = tokenizer.vocab_size
+    print(f"Generating {num_prompts} random sequences with {num_tokens} tokens each...")
+    return [
+        [random.randint(0, vocab_size - 1) for _ in range(num_tokens)]
+        for _ in range(num_prompts)
+    ]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Tokenizer Benchmark: Sequential vs Batch Processing"
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=str,
+        required=True,
+        help="Tokenizer name or path (e.g. nvidia/Kimi-K2-Thinking-NVFP4)",
+    )
+    parser.add_argument(
+        "--function",
+        type=str,
+        nargs="+",
+        choices=["encode", "decode"],
+        default=["encode", "decode"],
+        help="Functions to benchmark (default: encode decode)",
+    )
+    parser.add_argument(
+        "--num-tokens",
+        type=int,
+        default=20000,
+        help="Number of tokens per prompt (default: 20000)",
+    )
+    parser.add_argument(
+        "--batch-sizes",
+        type=int,
+        nargs="+",
+        default=[1, 2, 4, 8],
+        help="Batch sizes to test (default: 1 2 4 8)",
+    )
+    parser.add_argument(
+        "--no-batch",
+        action="store_true",
+        help="Skip batch benchmark, only run sequential",
+    )
+    parser.add_argument(
+        "--num-runs",
+        type=int,
+        default=5,
+        help="Number of runs per batch size (default: 5)",
+    )
+    return parser.parse_args()
 
 
 if __name__ == "__main__":