Measures real FLOPs for HFU
(model: GPT, x: torch.Tensor)
| 457 | |
| 458 | |
| 459 | def measure_flops(model: GPT, x: torch.Tensor) -> int: |
| 460 | """Measures real FLOPs for HFU""" |
| 461 | flop_counter = FlopCounterMode(model, display=False) |
| 462 | ctx = nullcontext() if model.training else torch.no_grad() |
| 463 | with ctx, flop_counter: |
| 464 | y = model(x) |
| 465 | if model.training: |
| 466 | y.sum().backward() |
| 467 | return flop_counter.get_total_flops() |