File size: 2,313 Bytes
13caab9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """Gate: the instantiated model must have exactly 20,453,760 parameters."""
import math
import sys
import torch
import torch.nn.functional as F
from common import EXPECTED_PARAMS, build_model
def main() -> int:
torch.manual_seed(0)
model = build_model()
cfg = model.config
m = model.model
total = sum(p.numel() for p in model.parameters())
emb = m.embed_tokens.weight.numel()
per_layer = [sum(p.numel() for p in layer.parameters()) for layer in m.layers]
final_norm = m.norm.weight.numel()
tied = model.lm_head.weight.data_ptr() == m.embed_tokens.weight.data_ptr()
print(f"params total : {total:,}")
print(f" embedding (tied) : {emb:,}")
print(f" per layer : {per_layer[0]:,} x {len(per_layer)} = {sum(per_layer):,}")
print(f" final norm : {final_norm:,}")
print(f"lm_head tied : {tied}")
print(f"attn_implementation : {cfg._attn_implementation}")
print(f"hidden_act : {cfg.hidden_act}")
print(f"rope : {getattr(cfg, 'rope_parameters', None) or getattr(cfg, 'rope_theta', None)}")
print(f"param dtype : {next(model.parameters()).dtype}")
layer0 = m.layers[0]
print("init std : "
f"embed {m.embed_tokens.weight.std():.4f} | q_proj {layer0.self_attn.q_proj.weight.std():.4f} | "
f"o_proj {layer0.self_attn.o_proj.weight.std():.4f} | gate_proj {layer0.mlp.gate_proj.weight.std():.4f} | "
f"down_proj {layer0.mlp.down_proj.weight.std():.4f} (target {0.02 / math.sqrt(2 * cfg.num_hidden_layers):.4f}) | "
f"norm mean {layer0.input_layernorm.weight.mean():.2f}")
if torch.cuda.is_available():
model.cuda()
x = torch.randint(0, cfg.vocab_size, (4, 1025), device="cuda")
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
logits = model(input_ids=x[:, :-1]).logits
loss = F.cross_entropy(logits.float().reshape(-1, cfg.vocab_size), x[:, 1:].reshape(-1))
print(f"init loss (random) : {loss.item():.4f} (ln V = {math.log(cfg.vocab_size):.4f})")
ok = total == EXPECTED_PARAMS
print(f"GATE {'PASS' if ok else 'FAIL'}: {total:,} (esperado {EXPECTED_PARAMS:,})")
return 0 if ok else 1
if __name__ == "__main__":
sys.exit(main())
|