| 422 | |
| 423 | |
| 424 | def flops_per_param(config: Config, n_params: int) -> int: |
| 425 | flops_per_token = ( |
| 426 | 2 * n_params |
| 427 | ) # each parameter is used for a MAC (2 FLOPS) per network operation |
| 428 | # this assumes that all samples have a fixed length equal to the block size |
| 429 | # which is most likely false during finetuning |
| 430 | flops_per_seq = flops_per_token * config.block_size |
| 431 | attn_flops_per_seq = ( |
| 432 | config.n_layer * 2 * 2 * (config.n_embd * (config.block_size**2)) |
| 433 | ) |
| 434 | return flops_per_seq + attn_flops_per_seq |
| 435 | |
| 436 | |
| 437 | def estimate_flops(model: GPT) -> int: |