unsloth/scripts/benchmarks/cb_vs_vllm_generation.py

"""Standalone generation microbenchmark: vLLM vs transformers CB vs Unsloth.

Backends (one per process; all engines are GPU-greedy):
- `vllm`             : Unsloth `fast_inference=True` (vLLM colocated).
- `tpaged`           : `model.generate_batch` on paged HF + `--attn_impl`.
- `unsloth_fi_false` : Unsloth `fast_inference=False` with the custom HF
                       inference kernels (cached fp16 LoRA).

LoRA: pass `--lora_adapter PATH` to activate a PEFT-style rank-32 adapter on
both vLLM (`LoRARequest`) and the HF paths (`peft.PeftModel.from_pretrained`,
or for `unsloth_fi_false` `FastLanguageModel.get_peft_model` pointed at the
same weights).

Equivalence-friendly sampling defaults (`--temperature 0.1 --top_p 0.97
--min_p 0.5 --top_k 5`) keep rollouts comparable across backends for the KL /
reward diff checks done in Phase 2.

Usage:
    CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
        --backend vllm --stats_path logs/lora_vllm_gen.json \
        --lora_adapter outputs/lora_rank32_fresh
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import torch  # noqa: E402

# FA4 shim for the `tpaged` backend with paged_attention. No-op for vLLM /
# unsloth_fi_false.
import flash_attn_fa4_shim  # noqa: E402

flash_attn_fa4_shim.apply()


def build_prompts(tokenizer, n_prompts):
    from unsloth_grpo_common import (
        apply_chat_template_to_tokenizer,
        SYSTEM_PROMPT,
    )
    from datasets import load_dataset

    apply_chat_template_to_tokenizer(tokenizer)
    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train")
    ds = ds.shuffle(seed=3407).select(range(n_prompts))
    messages = [
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": x["prompt"]},
        ]
        for x in ds
    ]
    prompts_text = [
        tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
        for m in messages
    ]
    prompt_ids = [
        tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=True)
        for m in messages
    ]
    return prompts_text, prompt_ids


def run_vllm(args):
    os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        load_in_4bit=False,
        fast_inference=True,
        max_lora_rank=32,
        gpu_memory_utilization=args.gpu_memory_utilization,
    )
    prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)

    lora_request = None
    if args.lora_adapter:
        from vllm.lora.request import LoRARequest
        lora_request = LoRARequest("fresh", 1, str(Path(args.lora_adapter).resolve()))

    from vllm import SamplingParams
    sp = SamplingParams(
        temperature=args.temperature,
        top_p=args.top_p,
        min_p=args.min_p,
        top_k=args.top_k,
        seed=3407,
        max_tokens=args.max_new_tokens,
        stop=[tokenizer.eos_token],
        include_stop_str_in_output=True,
    )

    # Warmup on 16 prompts then discard.
    warmup_text = prompts_text[:16]
    _ = model.fast_generate(warmup_text, sampling_params=sp, lora_request=lora_request)
    torch.cuda.synchronize()

    n_prompt_tokens = sum(len(p) for p in prompt_ids)
    wall_times = []
    total_decoded = None
    last_outputs = None
    for _ in range(args.n_rounds):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        outputs = model.fast_generate(
            prompts_text, sampling_params=sp, lora_request=lora_request
        )
        torch.cuda.synchronize()
        wall_times.append(time.perf_counter() - t0)
        total_decoded = sum(len(o.outputs[0].token_ids) for o in outputs)
        last_outputs = outputs

    med = sorted(wall_times)[len(wall_times) // 2]
    sample_texts = [o.outputs[0].text[:200] for o in (last_outputs[:3] or [])] if last_outputs else []
    return {
        "backend": "vllm",
        "lora_adapter": args.lora_adapter,
        "n_prompts": args.n_prompts,
        "n_prompt_tokens": n_prompt_tokens,
        "n_decoded_tokens": total_decoded,
        "wall_times_s": wall_times,
        "median_wall_s": med,
        "prompt_tps": n_prompt_tokens / med,
        "decode_tps": (total_decoded or 0) / med,
        "max_new_tokens": args.max_new_tokens,
        "sample_completions": sample_texts,
    }


def run_tpaged(args):
    """Vanilla HF + paged cache.

    Unsloth's Qwen3Attention monkey-patch does not compose with the
    `paged|<impl>` functional attention interface, so we use plain HF.
    """
    from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name,
        dtype=torch.bfloat16,
        attn_implementation=args.attn_impl,
    ).to("cuda")
    model.eval()

    if args.lora_adapter:
        from peft import PeftModel
        # NOTE: no merge_adapter -- we measure LoRA-active inference.
        model = PeftModel.from_pretrained(
            model, str(Path(args.lora_adapter).resolve()), is_trainable=False
        )
        model.eval()

    if args.persistent_cb:
        from persistent_cb import install_for_model  # noqa: WPS433
    prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)

    gen_config = GenerationConfig(
        max_new_tokens=args.max_new_tokens,
        do_sample=True,
        temperature=args.temperature,
        top_p=args.top_p,
        min_p=args.min_p,
        top_k=args.top_k,
        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=True,
    )
    gen_config.max_batch_tokens = args.max_batch_tokens
    gen_config.num_blocks = args.num_blocks

    if args.persistent_cb:
        install_for_model(model, gen_config)

    warmup_ids = prompt_ids[:16]
    with torch.inference_mode():
        _ = model.generate_batch(warmup_ids, generation_config=gen_config, progress_bar=False)
    torch.cuda.synchronize()

    n_prompt_tokens = sum(len(p) for p in prompt_ids)
    wall_times = []
    total_decoded = None
    last_outputs = None
    for _ in range(args.n_rounds):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.inference_mode():
            outputs = model.generate_batch(
                prompt_ids, generation_config=gen_config, progress_bar=False
            )
        torch.cuda.synchronize()
        wall_times.append(time.perf_counter() - t0)
        total_decoded = sum(len(v.generated_tokens) for v in outputs.values())
        last_outputs = outputs

    # Sample completions for coherence sanity check.
    sample_texts = []
    if last_outputs is not None:
        for k in list(last_outputs.keys())[:3]:
            toks = last_outputs[k].generated_tokens
            sample_texts.append(tokenizer.decode(toks, skip_special_tokens=False)[:200])

    med = sorted(wall_times)[len(wall_times) // 2]
    return {
        "backend": "tpaged",
        "lora_adapter": args.lora_adapter,
        "attn_impl": args.attn_impl,
        "persistent_cb": args.persistent_cb,
        "n_prompts": args.n_prompts,
        "n_prompt_tokens": n_prompt_tokens,
        "n_decoded_tokens": total_decoded,
        "wall_times_s": wall_times,
        "median_wall_s": med,
        "prompt_tps": n_prompt_tokens / med,
        "decode_tps": (total_decoded or 0) / med,
        "max_new_tokens": args.max_new_tokens,
        "sample_completions": sample_texts,
    }


def run_unsloth_fi_false(args):
    """Unsloth `fast_inference=False` path with custom HF inference kernels.

    This is the path that backs regular Unsloth training's sampling loop
    (Triton RMSNorm/RoPE, cached fp16 LoRA copies in `fast_linear_forward`).
    Previously only exercised through full GRPO runs -- isolating it lets us
    compare it head-to-head against vLLM on the same workload.

    LoRA is attached via `FastLanguageModel.get_peft_model`; if a PEFT adapter
    path is provided we re-load its weights into the Unsloth-wrapped model so
    every backend uses the *same* weights.
    """
    os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        load_in_4bit=False,
        fast_inference=False,
        max_lora_rank=32,
    )
    # Attach LoRA rank 32 the same way the GRPO notebook does.
    model = FastLanguageModel.get_peft_model(
        model,
        r=32,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_alpha=64,
        use_gradient_checkpointing="unsloth",
        random_state=3407,
    )

    # Optional: overlay a shared adapter so weights match other backends.
    if args.lora_adapter:
        from safetensors import safe_open
        adapter_file = Path(args.lora_adapter).resolve() / "adapter_model.safetensors"
        loaded_tensors = {}
        with safe_open(str(adapter_file), framework="pt") as f:
            for key in f.keys():
                loaded_tensors[key] = f.get_tensor(key)
        # PEFT saves with keys like `base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight`.
        # Unsloth's `get_peft_model` produces the same key shape.
        own_state = {n: p for n, p in model.named_parameters() if "lora_" in n}
        # Re-key by the suffix after `base_model.model.`.
        matched = 0
        with torch.no_grad():
            for name, tensor in loaded_tensors.items():
                # Try direct + strip `base_model.model.` prefix variants.
                candidates = [name, name.replace("base_model.model.", ""),
                              "base_model.model." + name]
                for cand in candidates:
                    # PEFT sometimes inserts `.default.` between module and lora_A.
                    variants = [cand, cand.replace(".default.", ".")]
                    for v in variants:
                        # own_state keys typically have `.default.weight` suffix
                        for own_name, own in own_state.items():
                            if own_name.endswith(v.split("base_model.model.")[-1]) \
                               or v.endswith(own_name.split("base_model.model.")[-1]):
                                if own.shape == tensor.shape:
                                    own.data.copy_(tensor.to(own.device, own.dtype))
                                    matched += 1
                                    break
        print(f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors "
              f"(out of {len(loaded_tensors)} adapter entries).")

    FastLanguageModel.for_inference(model)

    prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)

    # `model.generate` accepts batched input_ids; pad to max length.
    from transformers import GenerationConfig
    if tokenizer.padding_side != "left":
        tokenizer.padding_side = "left"  # decoder needs left padding
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    gen_config = GenerationConfig(
        max_new_tokens=args.max_new_tokens,
        do_sample=True,
        temperature=args.temperature,
        top_p=args.top_p,
        min_p=args.min_p,
        top_k=args.top_k,
        pad_token_id=tokenizer.pad_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=True,
    )

    def _batched_generate(texts):
        batch = tokenizer(texts, return_tensors="pt", padding=True).to("cuda")
        with torch.inference_mode():
            out = model.generate(**batch, generation_config=gen_config)
        prompt_len = batch["input_ids"].shape[1]
        return out, prompt_len

    # Warmup on 16 prompts.
    _ = _batched_generate(prompts_text[:16])
    torch.cuda.synchronize()

    n_prompt_tokens = sum(len(p) for p in prompt_ids)
    wall_times = []
    total_decoded = None
    last_out_ids = None
    last_prompt_len = None
    for _ in range(args.n_rounds):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        out_ids, prompt_len = _batched_generate(prompts_text)
        torch.cuda.synchronize()
        wall_times.append(time.perf_counter() - t0)
        # Count generated tokens past prompt_len per sequence (subtract any
        # trailing pad-only tail by comparing against EOS).
        total_decoded = int((out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item())
        last_out_ids = out_ids
        last_prompt_len = prompt_len

    med = sorted(wall_times)[len(wall_times) // 2]
    sample_texts = []
    if last_out_ids is not None:
        for i in range(min(3, last_out_ids.shape[0])):
            sample_texts.append(tokenizer.decode(
                last_out_ids[i, last_prompt_len:], skip_special_tokens=False)[:200])

    return {
        "backend": "unsloth_fi_false",
        "lora_adapter": args.lora_adapter,
        "n_prompts": args.n_prompts,
        "n_prompt_tokens": n_prompt_tokens,
        "n_decoded_tokens": total_decoded,
        "wall_times_s": wall_times,
        "median_wall_s": med,
        "prompt_tps": n_prompt_tokens / med,
        "decode_tps": (total_decoded or 0) / med,
        "max_new_tokens": args.max_new_tokens,
        "sample_completions": sample_texts,
    }


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--backend", choices=["vllm", "tpaged", "unsloth_fi_false"], required=True)
    p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
    p.add_argument("--max_seq_length", type=int, default=2048)
    p.add_argument("--n_prompts", type=int, default=32)
    p.add_argument("--n_rounds", type=int, default=2)
    p.add_argument("--max_new_tokens", type=int, default=512)
    p.add_argument("--gpu_memory_utilization", type=float, default=0.8)
    p.add_argument("--attn_impl", default="sdpa")
    p.add_argument("--max_batch_tokens", type=int, default=8192)
    p.add_argument("--num_blocks", type=int, default=16384)
    p.add_argument("--persistent_cb", action="store_true")
    p.add_argument("--lora_adapter", default=None,
                   help="Path to a PEFT adapter (rank 32) applied in every backend.")
    p.add_argument("--temperature", type=float, default=0.1)
    p.add_argument("--top_p", type=float, default=0.97)
    p.add_argument("--min_p", type=float, default=0.5)
    p.add_argument("--top_k", type=int, default=5)
    p.add_argument("--stats_path", required=True)
    return p.parse_args()


def main():
    args = parse_args()
    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True)

    torch.cuda.reset_peak_memory_stats()
    if args.backend == "vllm":
        out = run_vllm(args)
    elif args.backend == "unsloth_fi_false":
        out = run_unsloth_fi_false(args)
    else:
        out = run_tpaged(args)

    out["peak_memory_gb"] = torch.cuda.max_memory_allocated() / 1024**3
    out["sampling"] = {
        "temperature": args.temperature,
        "top_p": args.top_p,
        "min_p": args.min_p,
        "top_k": args.top_k,
    }
    with open(args.stats_path, "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))
    os._exit(0)


if __name__ == "__main__":
    main()