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