scratch/check_mini_params.py

import yaml
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "model" / "scripts"))
import viu1_dense

with open(ROOT / "model" / "configs" / "viu1_dense_config.yaml", "r") as f:
    cfg = yaml.safe_load(f)

args = viu1_dense.load_args(ROOT / "model" / "configs" / "viu1_dense_config.yaml")
model = viu1_dense.Viu1Dense(args)
total, active = model.param_counts()
print(f"Total params in ViuMini-Dense-360M: {total:,} ({total/1e6:.1f}M) | Active: {active:,}")
print(f"Layers: {args.n_layers}, Dim: {args.dim}, Heads: {args.n_heads}, KV Heads: {args.n_kv_heads}")