benchmark_lora_kernels.py

import logging
import time
from contextlib import contextmanager

from axolotl.monkeypatch.lora_kernels import (
    apply_lora_kernel_patches,
    patch_self_attn_lora,
)
from axolotl.utils.dict import DictDefault
from axolotl.utils.distributed import is_distributed
import numpy as np
import torch
from peft import PeftModelForCausalLM, get_peft_config
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

if is_distributed():
    print("distributed.")


def cleanup():
    """Force cleanup between model loads"""
    import gc

    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()


@contextmanager
def timer(name):
    start = time.perf_counter()
    yield
    end = time.perf_counter()
    logger.info(f"{name}: {(end - start)*1000:.2f}ms")


def profile_forward_backward(model, batch_size=1, seq_len=512, n_samples=30):
    """Profile forward and backward passes with multiple samples. Returns averaged stats."""
    device = next(model.parameters()).device
    input_ids = torch.randint(0, 32000, (batch_size, seq_len), device=device)

    # Warmup
    for _ in range(10):
        with torch.no_grad():
            out = model(input_ids)
            loss = out.logits.mean()
            if loss.requires_grad:
                loss.backward()
                model.zero_grad()

    forward_times = []
    forward_memories = []
    backward_times = []
    backward_memories = []

    for _ in range(n_samples):
        # Profile forward pass
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.empty_cache()

        forward_start = time.perf_counter()
        out = model(input_ids)
        loss = out.logits.mean()
        forward_times.append((time.perf_counter() - forward_start) * 1000)  # ms
        forward_memories.append(torch.cuda.max_memory_allocated() / 1024**3)

        # Profile backward pass
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.empty_cache()

        backward_start = time.perf_counter()
        loss.backward()
        backward_times.append((time.perf_counter() - backward_start) * 1000)  # ms
        backward_memories.append(torch.cuda.max_memory_allocated() / 1024**3)

        model.zero_grad()

    # Calculate averages and standard deviations
    forward_duration = np.mean(forward_times)
    forward_duration_std = np.std(forward_times)
    forward_memory = np.mean(forward_memories)
    forward_memory_std = np.std(forward_memories)

    backward_duration = np.mean(backward_times)
    backward_duration_std = np.std(backward_times)
    backward_memory = np.mean(backward_memories)
    backward_memory_std = np.std(backward_memories)

    logger.info(
        f"Forward pass: {forward_duration:.2f}±{forward_duration_std:.2f}ms, "
        f"{forward_memory:.2f}±{forward_memory_std:.2f}GB"
    )
    logger.info(
        f"Backward pass: {backward_duration:.2f}±{backward_duration_std:.2f}ms, "
        f"{backward_memory:.2f}±{backward_memory_std:.2f}GB"
    )

    return {
        "forward": {
            "duration_ms": forward_duration,
            "duration_std": forward_duration_std,
            "memory_gb": forward_memory,
            "memory_std": forward_memory_std,
            "tokens_per_second": (batch_size * seq_len) / (forward_duration / 1000),
        },
        "backward": {
            "duration_ms": backward_duration,
            "duration_std": backward_duration_std,
            "memory_gb": backward_memory,
            "memory_std": backward_memory_std,
            "tokens_per_second": (batch_size * seq_len) / (backward_duration / 1000),
        },
    }


MODEL = "HuggingFaceTB/SmolLM2-135M"
# MODEL = "HuggingFaceTB/SmolLM2-1.7B"
# MODEL = "meta-llama/Llama-3.2-3B"
# MODEL = "mhenrichsen/gemma-2b"


def get_fresh_model(quantize=True):
    bnb_config = None
    if quantize:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float32,
            bnb_4bit_use_double_quant=True,
        )

    return AutoModelForCausalLM.from_pretrained(
        MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.float32,
    )


# Profiling parameters
quantizes = [False, True]
ranks = [8]
batch_sizes = [1]
seq_lens = [1024, 2048, 4096, 8096]

target_modules = [
    "gate_proj",
    "up_proj",
    "down_proj",
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
]  # MLP + QKV + O params

cfg = DictDefault(
    base_model=MODEL,
    lora_mlp_kernel=True,
    lora_qkv_kernel=True,
    lora_o_kernel=True,
)

# Collect results for both pre and post patch
results = {"pre": {}, "post": {}}

# Pre-patch profiling
logger.info("\n=== Pre-patch Performance ===")
for quantize in quantizes:
    for rank in ranks:
        cleanup()

        # Apply LoRA
        peft_config = get_peft_config(
            {
                "peft_type": "LORA",
                "task_type": "CAUSAL_LM",
                "r": rank,
                "lora_alpha": 16,
                "target_modules": target_modules,
                "lora_dropout": 0,
                "bias": "none",
            }
        )
        model = get_fresh_model(quantize)
        model = PeftModelForCausalLM(model, peft_config)

        for batch_size in batch_sizes:
            for seq_len in seq_lens:
                logger.info(
                    f"\n=== Profiling with quantize={quantize}, rank={rank}, batch_size={batch_size}, seq_len={seq_len} ==="
                )
                results["pre"][(quantize, rank, batch_size, seq_len)] = (
                    profile_forward_backward(model, batch_size, seq_len)
                )

        del model

# Patch LlamaAttention and variants
patch_self_attn_lora(cfg)

# Post-patch profiling
logger.info("\n=== Post-patch Performance ===")
for quantize in quantizes:
    for rank in ranks:
        cleanup()

        # Apply LoRA
        peft_config = get_peft_config(
            {
                "peft_type": "LORA",
                "task_type": "CAUSAL_LM",
                "r": rank,
                "lora_alpha": 16,
                "target_modules": target_modules,
                "lora_dropout": 0,
                "bias": "none",
            }
        )
        model = get_fresh_model(quantize)
        model = PeftModelForCausalLM(model, peft_config)

        # Apply patches
        apply_lora_kernel_patches(model=model, cfg=cfg)

        for batch_size in batch_sizes:
            for seq_len in seq_lens:
                logger.info(
                    f"\n=== Profiling with quantize={quantize}, rank={rank}, batch_size={batch_size}, seq_len={seq_len} ==="
                )
                results["post"][(quantize, rank, batch_size, seq_len)] = (
                    profile_forward_backward(model, batch_size, seq_len)
                )

        del model

# Print comparison
for pass_type in ["forward", "backward"]:
    logger.info(f"\n=== {pass_type.title()} Pass Performance Comparison ===")
    logger.info(
        f"{'Quantize':^6} | {'Rank':^6} | {'Batch':^6} | {'Seq':^5} | {'Pre-Time':^10} | {'Post-Time':^10} | {'Speedup':^8} | "
        f"{'Pre-Mem':^8} | {'Post-Mem':^8} | {'Mem Save':^8} | {'Pre-TPS':^10} | {'Post-TPS':^10} | {'TPS Gain':^8}"
    )
    logger.info("-" * 110)

    for quantize, rank, batch_size, seq_len in sorted(results["pre"].keys()):
        pre = results["pre"][(quantize, rank, batch_size, seq_len)][pass_type]
        post = results["post"][(quantize, rank, batch_size, seq_len)][pass_type]

        speedup = pre["duration_ms"] / post["duration_ms"]
        mem_save_pct = (pre["memory_gb"] - post["memory_gb"]) / pre["memory_gb"] * 100
        tps_gain_pct = (
            (post["tokens_per_second"] - pre["tokens_per_second"])
            / pre["tokens_per_second"]
            * 100
        )

        logger.info(
            f"{quantize} | {rank} | {batch_size:^6d} | {seq_len:^5d} | "
            f"{pre['duration_ms']:^10.2f} | {post['duration_ms']:^10.2f} | {speedup:^8.2f}x | "
            f"{pre['memory_gb']:^8.2f} | {post['memory_gb']:^8.2f} | {mem_save_pct:^7.1f}% | "
            f"{pre['tokens_per_second']:^10.0f} | {post['tokens_per_second']:^10.0f} | {tps_gain_pct:^7.1f}%"
        )

# Summary statistics
logger.info("\n=== Summary Statistics ===")
for pass_type in ["forward", "backward"]:
    avg_speedup = np.mean(
        [
            results["pre"][k][pass_type]["duration_ms"]
            / results["post"][k][pass_type]["duration_ms"]
            for k in results["pre"].keys()
        ]
    )
    avg_mem_save = np.mean(
        [
            (
                results["pre"][k][pass_type]["memory_gb"]
                - results["post"][k][pass_type]["memory_gb"]
            )
            / results["pre"][k][pass_type]["memory_gb"]
            * 100
            for k in results["pre"].keys()
        ]
    )
    avg_tps_gain = np.mean(
        [
            (
                results["post"][k][pass_type]["tokens_per_second"]
                - results["pre"][k][pass_type]["tokens_per_second"]
            )
            / results["pre"][k][pass_type]["tokens_per_second"]
            * 100
            for k in results["pre"].keys()
        ]
    )

    logger.info(f"\n{pass_type.title()} Pass:")
    logger.info(f"Average speedup: {avg_speedup:.2f}x")
    logger.info(f"Average memory savings: {avg_mem_save:.1f}%")
    logger.info(f"Average throughput gain: {avg_tps_gain:.1f}%")

# Best configurations
logger.info("\n=== Best Configurations ===")
for pass_type in ["forward", "backward"]:
    pre_best = max(
        results["pre"].items(), key=lambda x: x[1][pass_type]["tokens_per_second"]
    )
    post_best = max(
        results["post"].items(), key=lambda x: x[1][pass_type]["tokens_per_second"]
    )

    logger.info(f"\n{pass_type.title()} Pass:")
    logger.info(
        f"Pre-patch best: batch_size={pre_best[0][0]}, seq_len={pre_best[0][1]}, "
        f"throughput={pre_best[1][pass_type]['tokens_per_second']:.0f} tokens/sec"
    )
    logger.info(
        f"Post-patch best: batch_size={post_best[0][0]}, seq_len={post_best[0][1]}, "
        f"throughput={post_best[1][pass_type]['tokens_per_second']:.0f} tokens/sec"
    )