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

hqq_torch_compile_base.py

import os
import gc
import time
import json
import platform

import torch
import pandas as pd

from tqdm.auto import tqdm
from copy import deepcopy


from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)

from transformers.cache_utils import DynamicCache

from hqq.core.quantize import (
    BaseQuantizeConfig,
)

from hqq.models.hf.base import (
    AutoHQQHFModel,
)

from hqq.utils.patching import (
    prepare_for_inference,
)

from hqq.core.quantize import (
    HQQBackend,
)

from hqq.core.quantize import (
    HQQLinear,
)

# ============================================================
# Configuration
# ============================================================
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True

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_hqq"

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:

    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


# ============================================================
# Helper: convert raw tuple KV cache -> DynamicCache and clone
# ============================================================

def to_dynamic_cache(past_key_values):
    """
    Ensure past_key_values is a DynamicCache with cloned tensors.

    Uses deepcopy for DynamicCache objects so internal tensor buffers
    are fully independent regardless of the internal attribute layout
    (which varies across transformers versions).

    For legacy tuple-of-tuples, builds a fresh DynamicCache by calling
    update() for each layer, which is the stable public API.
    """
    if isinstance(past_key_values, DynamicCache):
        # deepcopy handles any internal layout — key_cache, _cache,
        # or whatever the installed version uses.
        return deepcopy(past_key_values)

    # Legacy tuple-of-tuples: (key, value) per layer.
    # Use the public .update() API to stay version-agnostic.
    new_cache = DynamicCache()
    for layer_idx, (layer_keys, layer_values) in enumerate(past_key_values):
        new_cache.update(
            layer_keys.clone(),
            layer_values.clone(),
            layer_idx,
        )
    return new_cache


# ============================================================
# 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,
        },

        "hqq_backend": (
            "PYTORCH_COMPILE"
        ),

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

    return metadata


# ============================================================
# Load Tokenizer
# ============================================================

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
)

# ============================================================
# Load Base Model
# ============================================================

print(
    "\nLoading base model for HQQ quantization..."
)

compute_dtype = torch.bfloat16

base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=compute_dtype,
)

# ============================================================
# HQQ Quantization
# ============================================================

print(
    "\nQuantizing with HQQ "
    "(4-bit, group_size=64)..."
)

quant_config = BaseQuantizeConfig(
    nbits=4,
    group_size=64,
)

AutoHQQHFModel.quantize_model(
    base_model,
    quant_config=quant_config,
    compute_dtype=compute_dtype,
    device=DEVICE,
)

# ============================================================
# HQQ Backend
# ============================================================

print(
    "\nSetting HQQ backend "
    "to PYTORCH_COMPILE..."
)

HQQLinear.set_backend(
    HQQBackend.PYTORCH_COMPILE
)

# ============================================================
# Prepare Inference
# ============================================================

print(
    "\nPreparing inference backend..."
)

prepare_for_inference(
    base_model,
)

# ============================================================
# torch.compile
# ============================================================

print(
    "\nApplying torch.compile "
    "(reduce-overhead)..."
)

model = torch.compile(
    base_model,
    backend="inductor",
    mode="reduce-overhead",
    dynamic=True,
    fullgraph=False,
)

model.eval()

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

print("\nRunning warmup...")

warmup_inputs = tokenizer(
    "Warmup prompt for graph capture.",
    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 "
        f"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: "
            f"{actual_prompt_length} | "
            f"Generation Length: "
            f"{generation_length}"
        )

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

            try:

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

                reset_cuda()

                synchronize()

                # ====================================
                # PREFILL
                # ====================================

                prefill_start = (
                    time.perf_counter()
                )

                with torch.inference_mode():

                    torch.compiler.cudagraph_mark_step_begin()

                    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
                )

                # Convert to DynamicCache and clone tensors
                # so the first decode graph replay cannot
                # overwrite the prefill KV tensors.
                past_key_values = to_dynamic_cache(
                    prefill_outputs.past_key_values
                )

                # ====================================
                # DECODE
                # ====================================

                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
                    ):

                        torch.compiler.cudagraph_mark_step_begin()

                        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
                        )

                        # Clone into a fresh DynamicCache so
                        # the next graph replay cannot alias
                        # tensors from this iteration.
                        past_key_values = to_dynamic_cache(
                            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
                )

                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
                    "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
                    "hqq_backend": (
                        "PYTORCH_COMPILE"
                    ),

                    "torch_compile_backend": (
                        "inductor"
                    ),

                    "torch_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())