188 lines
7.0 KiB
Python
188 lines
7.0 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
|
|
|
|
|
PRESETS = {
|
|
"tiny": {
|
|
"hidden_size": 64,
|
|
"num_hidden_layers": 2,
|
|
"num_attention_heads": 4,
|
|
"num_key_value_heads": 2,
|
|
"vocab_size": 6400,
|
|
"max_position_embeddings": 128,
|
|
},
|
|
"minimind3": {
|
|
"hidden_size": 768,
|
|
"num_hidden_layers": 8,
|
|
"num_attention_heads": 8,
|
|
"num_key_value_heads": 4,
|
|
"vocab_size": 6400,
|
|
"max_position_embeddings": 32768,
|
|
},
|
|
}
|
|
|
|
|
|
def fmt_params(value):
|
|
return f"{value:,} ({value / 1_000_000:.3f}M)"
|
|
|
|
|
|
def count_params(module):
|
|
return sum(p.numel() for p in module.parameters())
|
|
|
|
|
|
def build_config(args, overrides=None):
|
|
values = dict(PRESETS[args.preset])
|
|
values.update(
|
|
{
|
|
"hidden_size": args.hidden_size or values["hidden_size"],
|
|
"num_hidden_layers": args.num_hidden_layers or values["num_hidden_layers"],
|
|
"num_attention_heads": args.num_attention_heads or values["num_attention_heads"],
|
|
"num_key_value_heads": args.num_key_value_heads or values["num_key_value_heads"],
|
|
"vocab_size": args.vocab_size or values["vocab_size"],
|
|
"max_position_embeddings": args.max_position_embeddings or values["max_position_embeddings"],
|
|
"use_moe": bool(args.use_moe),
|
|
"tie_word_embeddings": bool(args.tie_word_embeddings),
|
|
"dropout": 0.0,
|
|
}
|
|
)
|
|
if overrides:
|
|
values.update(overrides)
|
|
return MiniMindConfig(**values)
|
|
|
|
|
|
def make_model(config, device):
|
|
if device == "meta":
|
|
with torch.device("meta"):
|
|
return MiniMindForCausalLM(config)
|
|
return MiniMindForCausalLM(config)
|
|
|
|
|
|
def collect_row(name, config, device):
|
|
model = make_model(config, device)
|
|
block0 = model.model.layers[0]
|
|
token_matrix = model.lm_head.weight.numel()
|
|
all_blocks = count_params(model.model.layers)
|
|
return {
|
|
"name": name,
|
|
"hidden": config.hidden_size,
|
|
"layers": config.num_hidden_layers,
|
|
"vocab": config.vocab_size,
|
|
"heads": config.num_attention_heads,
|
|
"kv_heads": config.num_key_value_heads,
|
|
"head_dim": config.head_dim,
|
|
"tied": model.model.embed_tokens.weight is model.lm_head.weight,
|
|
"total": count_params(model),
|
|
"token": token_matrix,
|
|
"block0": count_params(block0),
|
|
"attn0": count_params(block0.self_attn),
|
|
"mlp0": count_params(block0.mlp),
|
|
"blocks": all_blocks,
|
|
}
|
|
|
|
|
|
def print_single_report(config, device):
|
|
model = make_model(config, device)
|
|
block0 = model.model.layers[0]
|
|
tied = model.model.embed_tokens.weight is model.lm_head.weight
|
|
token_matrix_params = model.lm_head.weight.numel()
|
|
|
|
print("[Config]")
|
|
print(f"hidden_size={config.hidden_size}")
|
|
print(f"num_hidden_layers={config.num_hidden_layers}")
|
|
print(f"vocab_size={config.vocab_size}")
|
|
print(f"num_attention_heads={config.num_attention_heads}")
|
|
print(f"num_key_value_heads={config.num_key_value_heads}")
|
|
print(f"head_dim={config.head_dim}")
|
|
print(f"intermediate_size={config.intermediate_size}")
|
|
print(f"use_moe={config.use_moe}")
|
|
print(f"tie_word_embeddings={config.tie_word_embeddings}")
|
|
print()
|
|
|
|
print("[Important Shapes]")
|
|
print(f"embed_tokens.weight={tuple(model.model.embed_tokens.weight.shape)}")
|
|
print(f"lm_head.weight={tuple(model.lm_head.weight.shape)}")
|
|
print(f"embedding_and_lm_head_are_same_parameter={tied}")
|
|
print(f"hidden_states=[batch_size, seq_len, {config.hidden_size}]")
|
|
print(f"attention q heads={config.num_attention_heads}, kv heads={config.num_key_value_heads}, head_dim={config.head_dim}")
|
|
print()
|
|
|
|
print("[Parameter Breakdown]")
|
|
print(f"total_unique_params={fmt_params(count_params(model))}")
|
|
print(f"token_matrix_params_once={fmt_params(token_matrix_params)}")
|
|
print(f"all_blocks={fmt_params(count_params(model.model.layers))}")
|
|
print(f"block0_total={fmt_params(count_params(block0))}")
|
|
print(f"block0.self_attn={fmt_params(count_params(block0.self_attn))}")
|
|
print(f"block0.mlp={fmt_params(count_params(block0.mlp))}")
|
|
print(f"final_norm={fmt_params(count_params(model.model.norm))}")
|
|
print()
|
|
|
|
print("[Simple Formulas]")
|
|
print(f"vocab_size * hidden_size = {config.vocab_size} * {config.hidden_size} = {fmt_params(config.vocab_size * config.hidden_size)}")
|
|
print(f"num_hidden_layers * block0_total = {config.num_hidden_layers} * {count_params(block0):,} = {fmt_params(config.num_hidden_layers * count_params(block0))}")
|
|
print(f"q projection output dim = num_attention_heads * head_dim = {config.num_attention_heads} * {config.head_dim} = {config.num_attention_heads * config.head_dim}")
|
|
|
|
|
|
def print_compare(args):
|
|
base = build_config(args)
|
|
variants = [
|
|
("base", {}),
|
|
("wider_hidden", {"hidden_size": base.hidden_size + 32}),
|
|
("deeper_layers", {"num_hidden_layers": base.num_hidden_layers + 2}),
|
|
("larger_vocab", {"vocab_size": base.vocab_size + 1600}),
|
|
("untied_lm_head", {"tie_word_embeddings": False}),
|
|
]
|
|
rows = [collect_row(name, build_config(args, overrides), args.device) for name, overrides in variants]
|
|
|
|
print("name | hidden | layers | vocab | heads/kv/head_dim | tied | total | token | block0 | attn0 | mlp0")
|
|
print("-" * 128)
|
|
for row in rows:
|
|
print(
|
|
f"{row['name']:15s} | "
|
|
f"{row['hidden']:6d} | "
|
|
f"{row['layers']:6d} | "
|
|
f"{row['vocab']:5d} | "
|
|
f"{row['heads']}/{row['kv_heads']}/{row['head_dim']:<6d} | "
|
|
f"{str(row['tied']):5s} | "
|
|
f"{row['total']:10d} | "
|
|
f"{row['token']:9d} | "
|
|
f"{row['block0']:9d} | "
|
|
f"{row['attn0']:9d} | "
|
|
f"{row['mlp0']:9d}"
|
|
)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Inspect MiniMind config, tensor shapes, and parameter counts.")
|
|
parser.add_argument("--preset", choices=sorted(PRESETS), default="tiny")
|
|
parser.add_argument("--device", choices=["cpu", "meta"], default="cpu")
|
|
parser.add_argument("--hidden_size", type=int)
|
|
parser.add_argument("--num_hidden_layers", type=int)
|
|
parser.add_argument("--num_attention_heads", type=int)
|
|
parser.add_argument("--num_key_value_heads", type=int)
|
|
parser.add_argument("--vocab_size", type=int)
|
|
parser.add_argument("--max_position_embeddings", type=int)
|
|
parser.add_argument("--use_moe", type=int, choices=[0, 1], default=0)
|
|
parser.add_argument("--tie_word_embeddings", type=int, choices=[0, 1], default=1)
|
|
parser.add_argument("--compare", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
if args.compare:
|
|
print_compare(args)
|
|
else:
|
|
print_single_report(build_config(args), args.device)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
main()
|