A benchmark runner script with pruna applying optimization algorithms (HQQ (4bit) + torch.compile ) for Llama-3.2-1B-Instruct

hqq_torch_compile_pruna.py

import os
import gc
import time
import json
import platform

import torch
import numpy as np
import pandas as pd

from tqdm.auto import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
from pruna import smash, SmashConfig

# ============================================================
# Configuration
# ============================================================

DEVICE = "cuda:0"

MODEL_ID = "meta-llama/Llama-3.2-1B-Instruct"

# ------------------------------------------------------------
# Prompt Length Sweep
# ------------------------------------------------------------

PROMPT_LENGTHS = [
    32,
    128,
]

# ------------------------------------------------------------
# Decode Length Sweep
# ------------------------------------------------------------

GENERATION_LENGTHS = [
    64,
    128,
    256,
    512,
    1024,
    2048,
    4096,
]

NUM_RUNS = 1
WARMUP_RUNS = 2

OUTPUT_DIR = "benchmark_results_pruna"

os.makedirs(OUTPUT_DIR, exist_ok=True)

OUTPUT_CSV = os.path.join(
    OUTPUT_DIR,
    "raw_benchmark_results.csv",
)

OUTPUT_METADATA = os.path.join(
    OUTPUT_DIR,
    "experiment_metadata.json",
)

# ============================================================
# CUDA Helpers
# ============================================================

def reset_cuda():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()


def synchronize():
    torch.cuda.synchronize()


# ============================================================
# Prompt Generation
# ============================================================

BASE_PROMPT = (
    "Paris is one of the most historically significant cities in Europe. "
)

def build_prompt(target_tokens: int) -> str:
    """
    Repeats base prompt until approximate token count is reached.
    """

    prompt = BASE_PROMPT

    while True:
        token_count = len(
            tokenizer(
                prompt,
                return_tensors="pt",
            )["input_ids"][0]
        )

        if token_count >= target_tokens:
            break

        prompt += BASE_PROMPT

    return prompt


# ============================================================
# Environment Metadata
# ============================================================

def collect_metadata():

    metadata = {
        "model_id": MODEL_ID,
        "device": DEVICE,
        "gpu": torch.cuda.get_device_name(0),
        "torch_version": torch.__version__,
        "cuda_version": str(torch.version.cuda),
        "python_version": platform.python_version(),

        "prompt_lengths": PROMPT_LENGTHS,
        "generation_lengths": GENERATION_LENGTHS,

        "num_runs": NUM_RUNS,
        "warmup_runs": WARMUP_RUNS,

        "dtype": "bfloat16",

        "quantization": {
            "method": "HQQ",
            "weight_bits": 4,
            "group_size": 64,
            "torchao_kernels": True,
        },

        "torch_compile": {
            "backend": "inductor",
            "mode": "reduce-overhead",
            "dynamic": True,
            "fullgraph": False,
        },
    }

    return metadata


# ============================================================
# Load Tokenizer + Base Model
# ============================================================

print("Loading tokenizer and base model...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
)

base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map=DEVICE,
)

# ============================================================
# Pruna Smash Config
# ============================================================

print("\nSmashing model with HQQ + torch.compile...")

smash_config = SmashConfig()

smash_config.add_tokenizer(tokenizer)

# ------------------------------------------------------------
# HQQ
# ------------------------------------------------------------

smash_config.add("hqq")

smash_config.add({
    "hqq_weight_bits": 4,
})

smash_config.add({
    "hqq_group_size": 64,
})

smash_config.add({
    "hqq_compute_dtype": "torch.bfloat16",
})

smash_config.add({
    "hqq_use_torchao_kernels": True,
})

# ------------------------------------------------------------
# torch.compile
# ------------------------------------------------------------

smash_config.add("torch_compile")

smash_config.add({
    "torch_compile_mode": "reduce-overhead",
})

smash_config.add({
    "torch_compile_backend": "inductor",
})

smash_config.add({
    "torch_compile_dynamic": True,
})

smash_config.add({
    "torch_compile_fullgraph": False,
})

model = smash(
    model=base_model,
    smash_config=smash_config,
)

model.eval()

# ============================================================
# Warmup
# ============================================================

print("\nRunning warmup...")

warmup_inputs = tokenizer(
    "Warmup prompt for CUDA graph stabilization.",
    return_tensors="pt",
).to(DEVICE)

for _ in range(WARMUP_RUNS):

    with torch.inference_mode():

        _ = model.generate(
            **warmup_inputs,
            max_new_tokens=32,
            do_sample=False,
            use_cache=True,
        )

synchronize()

print("Warmup complete.")

# ============================================================
# Benchmark
# ============================================================

results = []

print("\nStarting benchmark...")

for prompt_target_length in PROMPT_LENGTHS:

    print(
        f"\nPreparing prompt with target length "
        f"{prompt_target_length} tokens..."
    )

    prompt = build_prompt(prompt_target_length)

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    ).to(DEVICE)

    actual_prompt_length = inputs["input_ids"].shape[1]

    print(
        f"Actual prompt length: "
        f"{actual_prompt_length} tokens"
    )

    for generation_length in tqdm(
        GENERATION_LENGTHS,
        desc=f"Prompt {actual_prompt_length}",
    ):

        print(
            f"\nPrompt Length: {actual_prompt_length} | "
            f"Generation Length: {generation_length}"
        )

        for run_idx in tqdm(
            range(NUM_RUNS),
            desc="Runs",
            leave=False,
        ):

            try:

                # ============================================
                # Reset CUDA
                # ============================================

                reset_cuda()
                synchronize()

                # ============================================
                # PREFILL PHASE
                # ============================================

                prefill_start = time.perf_counter()

                with torch.inference_mode():

                    prefill_outputs = model(
                        **inputs,
                        use_cache=True,
                    )

                synchronize()

                prefill_end = time.perf_counter()

                prefill_latency_s = (
                    prefill_end - prefill_start
                )

                prefill_peak_memory_gb = (
                    torch.cuda.max_memory_allocated()
                    / 1024**3
                )

                past_key_values = (
                    prefill_outputs.past_key_values
                )

                # ============================================
                # DECODE PHASE
                # ============================================

                decode_input_ids = (
                    inputs["input_ids"][:, -1:]
                )

                generated_token_count = 0

                decode_start = time.perf_counter()

                with torch.inference_mode():

                    for _ in range(generation_length):

                        decode_outputs = model(
                            input_ids=decode_input_ids,
                            past_key_values=past_key_values,
                            use_cache=True,
                        )

                        next_token = torch.argmax(
                            decode_outputs.logits[:, -1, :],
                            dim=-1,
                            keepdim=True,
                        )

                        decode_input_ids = next_token

                        past_key_values = (
                            decode_outputs.past_key_values
                        )

                        generated_token_count += 1

                synchronize()

                decode_end = time.perf_counter()

                decode_time_s = (
                    decode_end - decode_start
                )

                # ============================================
                # Metrics
                # ============================================

                decode_tokens_per_sec = (
                    generated_token_count
                    / decode_time_s
                )

                avg_decode_latency_per_token_ms = (
                    (decode_time_s / generated_token_count)
                    * 1000
                )

                peak_memory_gb = (
                    torch.cuda.max_memory_allocated()
                    / 1024**3
                )

                # Relative decode memory growth
                decode_memory_growth_gb = (
                    peak_memory_gb
                    - prefill_peak_memory_gb
                )

                memory_per_generated_token_mb = (
                    (decode_memory_growth_gb * 1024)
                    / generated_token_count
                )

                # ============================================
                # Save Result
                # ============================================

                result = {

                    # Run Metadata
                    "run": run_idx + 1,

                    # Model
                    "model_id": MODEL_ID,

                    # Prompt
                    "prompt_target_length": (
                        prompt_target_length
                    ),

                    "actual_prompt_length": (
                        actual_prompt_length
                    ),

                    # Decode
                    "generation_length": (
                        generation_length
                    ),

                    # Prefill Metrics
                    "prefill_latency_s": (
                        prefill_latency_s
                    ),

                    "prefill_peak_memory_gb": (
                        prefill_peak_memory_gb
                    ),

                    # Decode Metrics
                    "decode_time_s": (
                        decode_time_s
                    ),

                    "decode_tokens_per_sec": (
                        decode_tokens_per_sec
                    ),

                    "avg_decode_latency_per_token_ms": (
                        avg_decode_latency_per_token_ms
                    ),

                    # Memory Metrics
                    "peak_memory_gb": (
                        peak_memory_gb
                    ),

                    "decode_memory_growth_gb": (
                        decode_memory_growth_gb
                    ),

                    "memory_per_generated_token_mb": (
                        memory_per_generated_token_mb
                    ),

                    # Config
                    "compile_backend": "inductor",

                    "compile_mode": (
                        "reduce-overhead"
                    ),

                    "quantization": "HQQ_4bit",

                    "dtype": "bfloat16",
                }

                results.append(result)

                print(
                    f"Run {run_idx + 1} | "
                    f"Prefill: "
                    f"{prefill_latency_s:.4f}s | "
                    f"Decode TPS: "
                    f"{decode_tokens_per_sec:.2f} tok/s | "
                    f"Decode Latency/token: "
                    f"{avg_decode_latency_per_token_ms:.2f} ms | "
                    f"Decode Memory Growth: "
                    f"{decode_memory_growth_gb:.4f} GB"
                )

            except torch.cuda.OutOfMemoryError:

                print(
                    f"OOM | "
                    f"Prompt={actual_prompt_length} | "
                    f"Generation={generation_length}"
                )

                reset_cuda()

                continue

# ============================================================
# Save Results
# ============================================================

results_df = pd.DataFrame(results)

results_df.to_csv(
    OUTPUT_CSV,
    index=False,
)

metadata = collect_metadata()

with open(OUTPUT_METADATA, "w") as f:
    json.dump(
        metadata,
        f,
        indent=4,
    )

# ============================================================
# Summary
# ============================================================

print("\nBenchmark complete.")

print(
    f"\nSaved raw benchmark results to:\n"
    f"{OUTPUT_CSV}"
)

print(
    f"\nSaved experiment metadata to:\n"
    f"{OUTPUT_METADATA}"
)

print("\nResults Preview:\n")

print(results_df.head())