| 369 | |
| 370 | |
| 371 | def flops_per_param( |
| 372 | max_seq_length: int, n_layer: int, n_embd: int, n_params: int |
| 373 | ) -> int: |
| 374 | flops_per_token = ( |
| 375 | 2 * n_params |
| 376 | ) # each parameter is used for a MAC (2 FLOPS) per network operation |
| 377 | # this assumes that all samples have a fixed length equal to the block size |
| 378 | # which is most likely false during finetuning |
| 379 | flops_per_seq = flops_per_token * max_seq_length |
| 380 | attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2)) |
| 381 | return flops_per_seq + attn_flops_per_seq |
| 382 | |
| 383 | |
| 384 | def estimate_flops(model: 'GPT', training: bool) -> int: |