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', training: bool)
| 382 | |
| 383 | |
| 384 | def estimate_flops(model: 'GPT', training: bool) -> int: |
| 385 | """Measures estimated FLOPs for MFU. |
| 386 | |
| 387 | Refs: |
| 388 | * https://ar5iv.labs.arxiv.org/html/2205.05198#A1 |
| 389 | * https://ar5iv.labs.arxiv.org/html/2204.02311#A2 |
| 390 | """ |
| 391 | # using all parameters for this is a naive over estimation because not all model parameters actually contribute to |
| 392 | # this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage |
| 393 | # (~10%) compared to the measured FLOPs, making those lower but more realistic. |
| 394 | # For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper. |
| 395 | n_trainable_params = num_parameters(model, requires_grad=True) |
| 396 | trainable_flops = flops_per_param( |
| 397 | model.max_seq_length, |
| 398 | model.config.n_layer, |
| 399 | model.config.n_embd, |
| 400 | n_trainable_params, |
| 401 | ) |
| 402 | # forward + backward + gradients (assumes no gradient accumulation) |
| 403 | ops_per_step = 3 if training else 1 |
| 404 | n_frozen_params = num_parameters(model, requires_grad=False) |
| 405 | frozen_flops = flops_per_param( |
| 406 | model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params |
| 407 | ) |
| 408 | # forward + backward |
| 409 | frozen_ops_per_step = 2 if training else 1 |
| 410 | return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops |
| 411 | |
| 412 | |
| 413 | class CycleIterator: |
nothing calls this directly
no test coverage detected