File size: 2,452 Bytes
f133460 ee7fc50 f133460 297c289 f133460 297c289 f133460 297c289 f133460 | 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 56 57 58 59 60 61 | """Parameter counts and FLOP estimation for the matched-budget comparisons.
Convention (documented so the numbers are reproducible):
- forward FLOPs per token per block application:
matmuls 2*(4*d^2 + 2*d*d_ff) (qkv+out proj: 4d^2; mlp: 2*d*d_ff)
attention 4*T*d (qk^T and av, T = S+2, MACs x2)
plus head 2*S*d once per forward.
- training FLOPs = 3 x forward FLOPs (backward ~2x forward).
- a recurrent model with `depth` shared blocks unrolled n_iters times pays
n_iters*depth block applications per step; its parameter count has only
`depth` block copies. Estimates ignore embeddings/layernorms/CE (<1%).
"""
from ptr.models import PtrTransformer
def count_params(model):
total = sum(p.numel() for p in model.parameters())
blocks = sum(p.numel() for b in model.blocks for p in b.parameters())
return {"total": total, "blocks": blocks}
def forward_flops_per_token(model, S, n_iters):
d, d_ff = model.d_model, model.blocks[0].mlp_in.out_features
T = S + 2
per_block = 2 * (4 * d * d + 2 * d * d_ff) + 4 * T * d
head = 2 * S * d
return n_iters * model.depth * per_block + head
def step_flops(model, S, batch, n_iters, training=True):
fwd = batch * (S + 2) * forward_flops_per_token(model, S, n_iters)
return 3 * fwd if training else fwd
def flops_budget_summary(model, S, batch, n_iters, steps):
"""Total training FLOPs for one run, plus params. `n_iters` may be an int
or a (mode, lo, hi) tuple for per-batch sampled unrolls (mean used)."""
p = count_params(model)
if isinstance(n_iters, tuple):
lo, hi = n_iters[1], n_iters[2]
mean_iters = (lo + hi) / 2.0
if len(n_iters) > 3:
mean_iters += n_iters[3]
desc = f"{n_iters[0]} U{{{lo}..{hi}}}+{n_iters[3]}"
else:
desc = f"{n_iters[0]} U{{{lo}..{hi}}}"
else:
mean_iters = n_iters
desc = str(n_iters)
total = step_flops(model, S, batch, mean_iters, training=True) * steps
return {
"params_total": p["total"],
"params_blocks": p["blocks"] + sum(q.numel() for blk in model.stem
for q in blk.parameters())
if hasattr(model, "stem") else p["blocks"],
"iters_per_step": desc,
"block_apps_per_step": model.depth * mean_iters
+ getattr(model, "stem_depth", 0),
"train_flops_total": float(total),
} |