OGAAA/Codes/3-2_model_training/test_chunked_attention.py

import argparse
import gc
import json
import math
import os
import random
from pathlib import Path

import torch


DEFAULT_MODEL_PATH = "/data/local/hujk/models/Llama-3.1-8B-Instruct"


def set_seed(seed: int):
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def parse_lens(raw: str):
    return [int(x) for x in raw.split(",") if x.strip()]


def dtype_from_name(name: str):
    return {
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
        "float32": torch.float32,
    }[name]


def cleanup_cuda(device):
    gc.collect()
    if device.type == "cuda":
        torch.cuda.empty_cache()


def make_keep_mask(seq_len: int, batch: int, device):
    keep = torch.ones(batch, seq_len, device=device, dtype=torch.bool)
    if seq_len > 16:
        keep[:, 7] = False
        keep[:, seq_len // 2] = False
    return keep


def make_repeated_input_ids(tokenizer, seq_len: int, device):
    ids = []
    if tokenizer.bos_token_id is not None:
        ids.append(tokenizer.bos_token_id)

    unit = tokenizer.encode(
        " context data instruction answer separator tool response",
        add_special_tokens=False,
    )
    if not unit:
        unit = [tokenizer.eos_token_id]

    while len(ids) < seq_len:
        ids.extend(unit)

    return torch.tensor([ids[:seq_len]], device=device, dtype=torch.long)


def load_model_and_tokenizer(args, device):
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        use_fast=True,
        padding_side="left",
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    common_kwargs = {
        "trust_remote_code": True,
        "low_cpu_mem_usage": True,
        "dtype": args.model_dtype,
    }
    if args.load_in_4bit:
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        )
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            quantization_config=quantization_config,
            device_map={"": 0},
            **common_kwargs,
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(args.model_path, **common_kwargs).to(device)

    model.eval()
    for param in model.parameters():
        param.requires_grad_(False)
    return model, tokenizer


def round_query_heads(query_heads: int, total_q_heads: int, total_kv_heads: int):
    group_size = total_q_heads // total_kv_heads
    query_heads = max(group_size, min(query_heads, total_q_heads))
    query_heads = int(math.ceil(query_heads / group_size) * group_size)
    query_heads = min(query_heads, total_q_heads)
    kv_heads = query_heads // group_size
    return query_heads, kv_heads


def extract_llama_qkv(model, tokenizer, seq_len: int, layer_idx: int, query_heads: int, device, qkv_dtype):
    from transformers.models.llama.modeling_llama import apply_rotary_pos_emb

    if layer_idx != 0:
        raise ValueError("This smoke test extracts q/k/v from layer 0 only to avoid a full model forward.")

    input_ids = make_repeated_input_ids(tokenizer, seq_len, device)
    position_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0)
    layer = model.model.layers[layer_idx]
    attn = layer.self_attn
    cfg = model.config
    total_q_heads = int(cfg.num_attention_heads)
    total_kv_heads = int(cfg.num_key_value_heads)
    head_dim = int(getattr(attn, "head_dim", cfg.hidden_size // total_q_heads))
    query_heads, kv_heads = round_query_heads(query_heads, total_q_heads, total_kv_heads)

    with torch.no_grad():
        hidden_states = model.model.embed_tokens(input_ids)
        hidden_states = layer.input_layernorm(hidden_states)
        batch, seq, _ = hidden_states.shape

        q = attn.q_proj(hidden_states).view(batch, seq, total_q_heads, head_dim).transpose(1, 2)
        k = attn.k_proj(hidden_states).view(batch, seq, total_kv_heads, head_dim).transpose(1, 2)
        v = attn.v_proj(hidden_states).view(batch, seq, total_kv_heads, head_dim).transpose(1, 2)

        rotary_emb = getattr(attn, "rotary_emb", None)
        if rotary_emb is None:
            rotary_emb = model.model.rotary_emb
        cos, sin = rotary_emb(hidden_states, position_ids)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        q = q[:, :query_heads].contiguous().detach().to(dtype=qkv_dtype)
        k = k[:, :kv_heads].contiguous().detach().to(dtype=qkv_dtype)
        v = v[:, :kv_heads].contiguous().detach().to(dtype=qkv_dtype)

    return q, k, v, {
        "layer_idx": layer_idx,
        "query_heads": query_heads,
        "kv_heads": kv_heads,
        "head_dim": head_dim,
        "seq_len": seq_len,
    }


def kv_indices_for_gqa(q_heads: int, kv_heads: int, device):
    if q_heads % kv_heads != 0:
        raise ValueError(f"q_heads must be divisible by kv_heads: {q_heads} vs {kv_heads}")
    group_size = q_heads // kv_heads
    return torch.arange(q_heads, device=device, dtype=torch.long) // group_size


def apply_keep_mask(logits, keep, start: int = 0, end: int | None = None):
    mask = keep[:, start:end] if end is not None else keep
    return logits.masked_fill(~mask[:, None, :], torch.finfo(logits.dtype).min)


def full_last_attention_gqa(q, k, v, keep):
    scale = 1.0 / math.sqrt(q.shape[-1])
    kv_indices = kv_indices_for_gqa(q.shape[1], k.shape[1], q.device)
    k_expanded = k.index_select(1, kv_indices)
    v_expanded = v.index_select(1, kv_indices)

    logits = torch.matmul(q[:, :, -1:, :].float(), k_expanded.float().transpose(-2, -1))
    logits = logits.squeeze(2) * scale
    logits = apply_keep_mask(logits, keep)
    probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
    out = torch.matmul(probs[:, :, None, :], v_expanded.float()).squeeze(2)
    return out, probs


def chunked_last_attention_gqa(q, k, v, keep, chunk_size: int):
    scale = 1.0 / math.sqrt(q.shape[-1])
    seq_len = k.shape[2]
    kv_indices = kv_indices_for_gqa(q.shape[1], k.shape[1], q.device)
    q_last = q[:, :, -1, :].float()
    max_scores = None

    for start in range(0, seq_len, chunk_size):
        end = min(start + chunk_size, seq_len)
        k_chunk = k.index_select(1, kv_indices)[:, :, start:end, :].float()
        logits = (q_last[:, :, None, :] * k_chunk).sum(dim=-1) * scale
        logits = apply_keep_mask(logits, keep, start, end)
        cur_max = logits.max(dim=-1).values
        max_scores = cur_max if max_scores is None else torch.maximum(max_scores, cur_max)

    denom = torch.zeros_like(max_scores, dtype=torch.float32)
    for start in range(0, seq_len, chunk_size):
        end = min(start + chunk_size, seq_len)
        k_chunk = k.index_select(1, kv_indices)[:, :, start:end, :].float()
        logits = (q_last[:, :, None, :] * k_chunk).sum(dim=-1) * scale
        logits = apply_keep_mask(logits, keep, start, end)
        denom = denom + torch.exp(logits - max_scores[:, :, None]).sum(dim=-1)
    denom = denom.clamp_min(1e-20)

    out = torch.zeros(q.shape[0], q.shape[1], v.shape[-1], device=q.device, dtype=torch.float32)
    prob_sums = torch.zeros_like(max_scores, dtype=torch.float32)
    v_expanded = v.index_select(1, kv_indices)
    for start in range(0, seq_len, chunk_size):
        end = min(start + chunk_size, seq_len)
        k_chunk = k.index_select(1, kv_indices)[:, :, start:end, :].float()
        v_chunk = v_expanded[:, :, start:end, :].float()
        logits = (q_last[:, :, None, :] * k_chunk).sum(dim=-1) * scale
        logits = apply_keep_mask(logits, keep, start, end)
        probs = torch.exp(logits - max_scores[:, :, None]) / denom[:, :, None]
        out = out + (probs[:, :, :, None] * v_chunk).sum(dim=2)
        prob_sums = prob_sums + probs.sum(dim=-1)
    return out, prob_sums


def compare_short_case(args, model, tokenizer, device):
    set_seed(args.seed)
    q_src, k_src, v_src, qkv_info = extract_llama_qkv(
        model,
        tokenizer,
        args.short_len,
        args.layer_idx,
        args.query_heads,
        device,
        args.qkv_dtype,
    )
    keep = make_keep_mask(args.short_len, q_src.shape[0], device)
    target = torch.randn(q_src.shape[0], q_src.shape[1], q_src.shape[-1], device=device, dtype=torch.float32)

    q_full = q_src.clone().detach().requires_grad_()
    k_full = k_src.clone().detach().requires_grad_()
    v_full = v_src.clone().detach().requires_grad_()
    q_chunk = q_src.clone().detach().requires_grad_()
    k_chunk = k_src.clone().detach().requires_grad_()
    v_chunk = v_src.clone().detach().requires_grad_()

    full_out, full_probs = full_last_attention_gqa(q_full, k_full, v_full, keep)
    chunk_out, chunk_prob_sums = chunked_last_attention_gqa(
        q_chunk, k_chunk, v_chunk, keep, args.chunk_size
    )

    full_loss = (full_out * target).sum()
    chunk_loss = (chunk_out * target).sum()
    full_loss.backward()
    chunk_loss.backward()

    grad_stats = {}
    for name, a, b in (
        ("q", q_full.grad, q_chunk.grad),
        ("k", k_full.grad, k_chunk.grad),
        ("v", v_full.grad, v_chunk.grad),
    ):
        diff = (a.float() - b.float()).abs()
        grad_stats[name] = {
            "max_abs": float(diff.max().item()),
            "mean_abs": float(diff.mean().item()),
            "full_norm": float(a.float().norm().item()),
            "chunk_norm": float(b.float().norm().item()),
        }

    out_diff = (full_out - chunk_out).abs()
    prob_sum_diff = (chunk_prob_sums - 1.0).abs()
    return {
        **qkv_info,
        "chunk_size": args.chunk_size,
        "attn_out_max_abs": float(out_diff.max().item()),
        "attn_out_mean_abs": float(out_diff.mean().item()),
        "prob_full_sum_max_abs": float((full_probs.sum(dim=-1) - 1.0).abs().max().item()),
        "prob_chunk_sum_max_abs": float(prob_sum_diff.max().item()),
        "loss_full": float(full_loss.item()),
        "loss_chunk": float(chunk_loss.item()),
        "loss_abs_diff": float(abs(full_loss.item() - chunk_loss.item())),
        "grad": grad_stats,
    }


def long_len_probe(args, model, tokenizer, device):
    results = []
    for seq_len in args.long_lens:
        cleanup_cuda(device)
        set_seed(args.seed + seq_len)
        if device.type == "cuda":
            torch.cuda.reset_peak_memory_stats(device)
        try:
            q_src, k_src, v_src, qkv_info = extract_llama_qkv(
                model,
                tokenizer,
                seq_len,
                args.layer_idx,
                args.query_heads,
                device,
                args.qkv_dtype,
            )
            keep = make_keep_mask(seq_len, q_src.shape[0], device)
            target = torch.randn(
                q_src.shape[0], q_src.shape[1], q_src.shape[-1], device=device, dtype=torch.float32
            )
            q = q_src.clone().detach().requires_grad_()
            k = k_src.clone().detach().requires_grad_()
            v = v_src.clone().detach().requires_grad_()

            out, prob_sums = chunked_last_attention_gqa(q, k, v, keep, args.chunk_size)
            loss = (out * target).sum()
            loss.backward()
            if device.type == "cuda":
                torch.cuda.synchronize(device)
                peak_gb = torch.cuda.max_memory_allocated(device) / (1024**3)
                allocated_gb = torch.cuda.memory_allocated(device) / (1024**3)
            else:
                peak_gb = 0.0
                allocated_gb = 0.0
            grad_ok = q.grad is not None and k.grad is not None and v.grad is not None
            results.append(
                {
                    **qkv_info,
                    "status": "ok",
                    "loss": float(loss.item()),
                    "grad_ok": bool(grad_ok),
                    "prob_sum_max_abs": float((prob_sums - 1.0).abs().max().item()),
                    "peak_allocated_gb": float(peak_gb),
                    "allocated_after_backward_gb": float(allocated_gb),
                }
            )
            del q_src, k_src, v_src, q, k, v, keep, target, out, prob_sums, loss
        except torch.cuda.OutOfMemoryError as exc:
            if device.type == "cuda":
                torch.cuda.empty_cache()
            results.append({"seq_len": seq_len, "status": "oom", "error": str(exc).splitlines()[0]})
        finally:
            cleanup_cuda(device)
    return results


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH)
    parser.add_argument("--short-len", type=int, default=512)
    parser.add_argument("--long-lens", type=parse_lens, default=parse_lens("1024,2048,4096,8192,16384"))
    parser.add_argument("--chunk-size", type=int, default=1024)
    parser.add_argument("--query-heads", type=int, default=8)
    parser.add_argument("--layer-idx", type=int, default=0)
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument("--model-dtype", choices=["float16", "bfloat16", "float32"], default="float16")
    parser.add_argument("--qkv-dtype", choices=["float16", "bfloat16", "float32"], default="float16")
    parser.add_argument("--load-in-4bit", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--output", default="")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this test.")
    device = torch.device("cuda:0")
    args.model_dtype = dtype_from_name(args.model_dtype)
    args.qkv_dtype = dtype_from_name(args.qkv_dtype)

    set_seed(args.seed)
    model, tokenizer = load_model_and_tokenizer(args, device)
    cleanup_cuda(device)

    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
    short = compare_short_case(args, model, tokenizer, device)
    long = long_len_probe(args, model, tokenizer, device)

    result = {
        "model_path": args.model_path,
        "visible_device": os.environ.get("CUDA_VISIBLE_DEVICES", ""),
        "cuda_name": torch.cuda.get_device_name(device),
        "load_in_4bit": bool(args.load_in_4bit),
        "model_dtype": str(args.model_dtype),
        "qkv_dtype": str(args.qkv_dtype),
        "short_compare": short,
        "long_probe": long,
    }

    text = json.dumps(result, indent=2, ensure_ascii=False)
    print(text)
    if args.output:
        out = Path(args.output)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(text + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()