| """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), |
| } |
