Measures estimated FLOPs for MFU. Refs: * https://ar5iv.labs.arxiv.org/html/2205.05198#A1 * https://ar5iv.labs.arxiv.org/html/2204.02311#A2
(model: GPT)
| 435 | |
| 436 | |
| 437 | def estimate_flops(model: GPT) -> int: |
| 438 | """Measures estimated FLOPs for MFU. |
| 439 | |
| 440 | Refs: |
| 441 | * https://ar5iv.labs.arxiv.org/html/2205.05198#A1 |
| 442 | * https://ar5iv.labs.arxiv.org/html/2204.02311#A2 |
| 443 | """ |
| 444 | # using all parameters for this is a naive over estimation because not all model parameters actually contribute to |
| 445 | # this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage |
| 446 | # (~10%) compared to the measured FLOPs, making those lower but more realistic. |
| 447 | # For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper. |
| 448 | n_trainable_params = num_parameters(model, requires_grad=True) |
| 449 | trainable_flops = flops_per_param(model.config, n_trainable_params) |
| 450 | # forward + backward + gradients (assumes no gradient accumulation) |
| 451 | ops_per_step = 3 if model.training else 1 |
| 452 | n_frozen_params = num_parameters(model, requires_grad=False) |
| 453 | frozen_flops = flops_per_param(model.config, n_frozen_params) |
| 454 | # forward + backward |
| 455 | frozen_ops_per_step = 2 if model.training else 1 |
| 456 | return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops |
| 457 | |
| 458 | |
| 459 | def measure_flops(model: GPT, x: torch.Tensor) -> int: |
no test coverage detected