Flops counter issue

[jiagaoxiang]

System Info

Two issues:

1. It seems the current implementation of the flops counter is counting the total flops of all the training step which causes the tflops/s/gpu metric goes up with increasing number of training steps. (reference: flops counter error with PyTorch2.5 and 2.6 pytorch/pytorch#145947 (comment))
2. actually, it's a question: in the training config, what is the impact of the default num_freeze_layers: int = 1 on mllama model (llama3.2-11b vision model)? I am suspecting if it is causing the underestimate of total model tflops issue.

Thank you!

Information

• The official example scripts
• My own modified scripts

🐛 Describe the bug

steps to reproduce can be found in this ticket: pytorch/pytorch#145947 (comment)

Error logs

Explained in above and also in this ticket: pytorch/pytorch#145947 (comment)

Expected behavior

The Flops counting is buggy.

[jiagaoxiang]

@HamidShojanazeri Could you please take a look? Thank you!

[wukaixingxp]

@jiagaoxiang From this line, num_freeze_layers will not affect anything if freeze_layers is False (which by default it is False). From your statement The total flops count with pytorch2.5 and 2.6 is ~186TFlops; but the total flops count with PyTorch nightly (2.7) is ~2900TFlops., it looks like an issue with pytorch 2.5 and 2.6, as the issue was gone in pytorch 2.7. It could be Pytorch issue, maybe related to this FlopCounterMode ignores custom operators issue?

[jiagaoxiang]

@jiagaoxiang From this line, num_freeze_layers will not affect anything if freeze_layers is False (which by default it is False). From your statement The total flops count with pytorch2.5 and 2.6 is ~186TFlops; but the total flops count with PyTorch nightly (2.7) is ~2900TFlops., it looks like an issue with pytorch 2.5 and 2.6, as the issue was gone in pytorch 2.7. It could be Pytorch issue, maybe related to this FlopCounterMode ignores custom operators issue?

Hi @wukaixingxp , In my testing, it seems the llama-cookbook flops counter will counter the tflops of all trainnig steps. For example, if I limit the max training step to 1, even if using pytorch nightly, it will give a tflop count of around 190 (for 15 training steps, it will be around 2900Tflops). So, the problem seems to be 2 fold:

1. Llama-cook repo counts flops of all training steps (instead of counting just one step);
2. PyTorch flops counter undercounts the Tflops of a training step; i.e., the total tflops of one training step for the 11B vision model should be much more than 190TFlops with a batch size of 2 and sequence length of 4096.