diff --git a/benchmark/benchmark_batch/benchmark_tokenizer.py b/benchmark/benchmark_batch/benchmark_tokenizer.py index 88a5820b6..4439f3a05 100644 --- a/benchmark/benchmark_batch/benchmark_tokenizer.py +++ b/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"Tokens per prompt: {NUM_TOKENS}") - print(f"Number of runs per batch size: {NUM_RUNS}") + print(f"Tokenizer: {args.tokenizer}") + print(f"Functions: {', '.join(args.function)}") + 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(TOKENIZER_DIR) + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + max_batch_size = max(args.batch_sizes) - # The largest batch size determines how many prompts we need - max_batch_size = max(BATCH_SIZES) - all_prompts = generate_random_prompts(max_batch_size, NUM_TOKENS, tokenizer) + token_ids = generate_random_token_ids(max_batch_size, args.num_tokens, tokenizer) - results = [] - print("\nRunning benchmark...") + if "encode" in args.function: + 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: - print(f"\nBenchmarking batch size: {batch_size}") - result = benchmark_sequential_vs_batch(all_prompts, batch_size, tokenizer) - results.append(result) + if "decode" in args.function: + # mimic DetokenizerManager's usual case + decode_kwargs = dict( + 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):") - for i, run_time in enumerate(result["batch_runs"]): - print(f" Run {i+1}: {run_time:.2f} ms") - print(f" Average: {result['avg_batch_ms']:.2f} ms") +def run_benchmark( + name, data, sequential_fn, batch_fn, batch_sizes, num_runs, skip_batch +): + 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: - print( - 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 Size", "Sequential (ms)"] + if not skip_batch: + 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__":