| 56 | |
| 57 | |
| 58 | def parse_args(): |
| 59 | p = argparse.ArgumentParser(description="GLM-4.7-Flash MLA -> GQLA via PCA + absorption") |
| 60 | p.add_argument("--model_path", required=True) |
| 61 | p.add_argument("--save_path", required=True) |
| 62 | p.add_argument("--num_kv_heads", type=int, required=True, |
| 63 | help="GQA group count; must divide num_attention_heads.") |
| 64 | p.add_argument("--dtype", choices=list(_DTYPES), default="bf16") |
| 65 | p.add_argument("--device_map", default="auto", |
| 66 | help="HF device_map ('auto' shards across visible GPUs).") |
| 67 | p.add_argument("--cal_dataset", choices=["wikitext2", "pg19", "alpaca"], default="wikitext2") |
| 68 | p.add_argument("--cal_nsamples", type=int, default=128) |
| 69 | p.add_argument("--cal_seqlen", type=int, default=512) |
| 70 | p.add_argument("--seed", type=int, default=42) |
| 71 | p.add_argument("--head_grouping", choices=["neighbor", "similarity"], default="neighbor", |
| 72 | help="Per-layer head grouping for PCA. 'neighbor' (default): heads " |
| 73 | "0..gs-1, gs..2gs-1, ... — byte-identical to legacy. " |
| 74 | "'similarity': full all-pair K/V cov + trace-normalised nuclear-norm " |
| 75 | "head similarity + seed-and-grow greedy balanced grouping. Costs an " |
| 76 | "extra O(H^2 d^2) per-layer cov collection but typically reduces PCA " |
| 77 | "truncation error noticeably.") |
| 78 | p.add_argument("--sim_w_k", type=float, default=1.0, |
| 79 | help="K-side weight in nuclear-norm head similarity. w_k=1, w_v=0 groups " |
| 80 | "by K subspace only (default 1.0).") |
| 81 | p.add_argument("--sim_w_v", type=float, default=1.0, |
| 82 | help="V-side weight in nuclear-norm head similarity (default 1.0).") |
| 83 | p.add_argument("--hessian_pca", action="store_true", default=False, |
| 84 | help="Weight the calibration cov by per-token NLL from the teacher's " |
| 85 | "next-token prediction (SparseGPT/OBS-style Hessian-weighted PCA): " |
| 86 | "rare / surprising tokens contribute more than easy / predictable " |
| 87 | "ones. Requires one extra norm + lm_head forward over the cached " |
| 88 | "final hidden states; OFF path is byte-identical to legacy.") |
| 89 | p.add_argument("--hessian_mode", choices=["nll"], default="nll", |
| 90 | help="Weighting scheme when --hessian_pca is set. 'nll' (default): per-" |
| 91 | "token NLL of the next token. Other modes from the research codebase " |
| 92 | "are not exposed in opensource.") |
| 93 | p.add_argument("--eval_ppl", action="store_true") |
| 94 | p.add_argument("--eval_ppl_dataset", choices=["wikitext2", "pg19", "alpaca"], default="wikitext2") |
| 95 | p.add_argument("--eval_ppl_seqlen", type=int, default=2048) |
| 96 | p.add_argument("--eval_ppl_batch_size", type=int, default=1) |
| 97 | p.add_argument("--eval_ppl_max_chunks", type=int, default=None) |
| 98 | return p.parse_args() |
| 99 | |
| 100 | |
| 101 | def main(): |