| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from transformers import AutoTokenizer |
|
|
|
|
| PRESETS = { |
| "large": {"d": 1024, "layers": 24, "heads": 16, "rank": 128}, |
| "agillm4_floor": {"d": 1280, "layers": 28, "heads": 20, "rank": 160}, |
| "agillm4_main": {"d": 1536, "layers": 32, "heads": 24, "rank": 192}, |
| "agillm4_big": {"d": 1792, "layers": 36, "heads": 28, "rank": 224}, |
| } |
|
|
|
|
| def estimate(vocab: int, d: int, layers: int, heads: int, rank: int, tie_weights: bool = False) -> dict: |
| dk = d // heads |
| embed = vocab * d |
| attn = 3 * d * d + d * d + dk * rank |
| ff = d * (4 * d) + (4 * d) + (4 * d) * d + d |
| norms = 4 * d |
| block = attn + ff + norms |
| core = embed + layers * block + 2 * d |
| ar = 0 if tie_weights else vocab * d + vocab |
| nat = vocab * d + vocab |
| sat = vocab * d + vocab + 2 * d + 2 |
| total = core + ar + sat + nat |
| return { |
| "vocab": vocab, |
| "d_model": d, |
| "layers": layers, |
| "heads": heads, |
| "rank": rank, |
| "tie_weights": tie_weights, |
| "embedding_params": embed, |
| "block_params_each": block, |
| "core_params": core, |
| "ar_head_params": ar, |
| "nat_head_params": nat, |
| "sat_head_params": sat, |
| "total_params": total, |
| "tokens_at_100_to_1": total * 100, |
| } |
|
|
|
|
| def fmt(n: int) -> str: |
| if n >= 1_000_000_000: |
| return f"{n / 1_000_000_000:.3f}B" |
| if n >= 1_000_000: |
| return f"{n / 1_000_000:.1f}M" |
| return f"{n:,}" |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description="Estimate AGILLM-4 parameter and token targets") |
| parser.add_argument("--tokenizer", default="deepseek-ai/DeepSeek-V4-Pro") |
| parser.add_argument("--preset", default="agillm4_main", choices=PRESETS) |
| parser.add_argument("--tie_weights", action="store_true") |
| parser.add_argument("--json_out", default="") |
| args = parser.parse_args() |
|
|
| tok = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True, trust_remote_code=True) |
| vocab = max(tok.get_vocab().values()) + 1 |
| result = estimate(vocab=vocab, tie_weights=args.tie_weights, **PRESETS[args.preset]) |
| print(f"preset={args.preset}") |
| print(f"total_params={fmt(result['total_params'])} ({result['total_params']:,})") |
| print(f"target_tokens_100_to_1={fmt(result['tokens_at_100_to_1'])} ({result['tokens_at_100_to_1']:,})") |
| print(json.dumps(result, indent=2, sort_keys=True)) |
| if args.json_out: |
| Path(args.json_out).write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
