108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
|
|
|
|
def round_intermediate(hidden_size):
|
|
return math.ceil(hidden_size * math.pi / 64) * 64
|
|
|
|
|
|
def dense_params(cfg):
|
|
h = cfg["hidden_size"]
|
|
layers = cfg["num_hidden_layers"]
|
|
vocab = cfg.get("vocab_size", 6400)
|
|
heads = cfg.get("num_attention_heads", 8)
|
|
kv_heads = cfg.get("num_key_value_heads", heads // 2)
|
|
head_dim = cfg.get("head_dim", h // heads)
|
|
intermediate = cfg.get("intermediate_size", round_intermediate(h))
|
|
attention_type = cfg.get("attention_type", "gqa")
|
|
|
|
embed = vocab * h
|
|
if attention_type == "mla":
|
|
kv_rank = cfg.get("kv_lora_rank", 128)
|
|
q_rank = cfg.get("q_lora_rank", 256)
|
|
rope_dim = cfg.get("rope_dim", head_dim // 2)
|
|
attn = (
|
|
h * q_rank
|
|
+ q_rank * heads * head_dim
|
|
+ q_rank * heads * rope_dim
|
|
+ h * kv_rank
|
|
+ kv_rank * kv_heads * head_dim
|
|
+ kv_rank * kv_heads * head_dim
|
|
+ h * rope_dim
|
|
+ heads * head_dim * h
|
|
)
|
|
kv_cache_floats = kv_rank + rope_dim
|
|
else:
|
|
q = h * heads * head_dim
|
|
k = h * kv_heads * head_dim
|
|
v = h * kv_heads * head_dim
|
|
o = heads * head_dim * h
|
|
attn = q + k + v + o
|
|
kv_cache_floats = 2 * kv_heads * head_dim
|
|
|
|
ffn = 3 * h * intermediate
|
|
norms = 4 * h
|
|
total = embed + layers * (attn + ffn + norms) + h
|
|
active = total
|
|
return {
|
|
"params": total,
|
|
"active_params": active,
|
|
"embed_params": embed,
|
|
"attn_params_per_layer": attn,
|
|
"ffn_params_per_layer": ffn,
|
|
"kv_cache_floats_per_token_layer": kv_cache_floats,
|
|
}
|
|
|
|
|
|
def fmt_params(n):
|
|
if n >= 1e9:
|
|
return f"{n / 1e9:.2f}B"
|
|
return f"{n / 1e6:.1f}M"
|
|
|
|
|
|
def print_table(profiles):
|
|
header = f"{'profile':<20}{'params':>10}{'active':>10}{'attn':>12}{'ffn':>12}{'kv_cache':>12} usage"
|
|
print(header)
|
|
print("-" * len(header))
|
|
for name, cfg in profiles.items():
|
|
if "hidden_size" not in cfg:
|
|
print(f"{name:<20}{'-':>10}{'-':>10}{'-':>12}{'-':>12}{'-':>12} {cfg.get('usage', '')}")
|
|
continue
|
|
est = dense_params(cfg)
|
|
print(
|
|
f"{name:<20}"
|
|
f"{fmt_params(est['params']):>10}"
|
|
f"{fmt_params(est['active_params']):>10}"
|
|
f"{fmt_params(est['attn_params_per_layer']):>12}"
|
|
f"{fmt_params(est['ffn_params_per_layer']):>12}"
|
|
f"{est['kv_cache_floats_per_token_layer']:>12} "
|
|
f"{cfg.get('usage', '')}"
|
|
)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Estimate SearchLM model scales without importing torch/transformers")
|
|
parser.add_argument("--profiles", default="configs/searchlm_profiles.json")
|
|
parser.add_argument("--json", action="store_true", help="Print detailed JSON")
|
|
args = parser.parse_args()
|
|
|
|
with open(args.profiles, "r", encoding="utf-8") as f:
|
|
profiles = json.load(f)
|
|
if args.json:
|
|
result = {}
|
|
for name, cfg in profiles.items():
|
|
result[name] = dict(cfg)
|
|
if "hidden_size" in cfg:
|
|
result[name].update(dense_params(cfg))
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
else:
|
|
print_table(profiles)
|
|
print()
|
|
print("Note: 2B+ full-parameter DDP is not recommended on 24GB RTX3090 cards; use LoRA/QLoRA as a baseline instead.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|