Potentially incorrect attention flop calculation due to wrong head_dim?

[gau-nernst]

Bug description

torchtitan/torchtitan/models/utils.py

Lines 432 to 437 in a8899e4

	l, h, q, t = (
	model_args.n_layers,
	model_args.n_heads,
	model_args.dim // model_args.n_heads,
	seq_len,
	)

However, head_dim is not necessarily equal to dim / n_heads

e.g. Qwen3-4B, dim=2560, n_heads=32, head_dim=128

Versions

latest main

[tianyu-l]

It seems we need to revisit the formula of attention flops computation, especially for Qwen3 and DeepSeek V3.

Maybe we can separate num params computation from total flops calculation -- the former can be shared across models, the latter may need to be customized for each model due to different attention details.

Overall I think attention part may not matter much in flop computation due to the its scale compared with the matmul part.

cc @shuhuayu @wwwjn

[gau-nernst]

Also just noticed that, due to causal attention, shouldn't the attention FLOPs be further divided by 2?

[tianyu-l]

Also just noticed that, due to causal attention, shouldn't the attention FLOPs be further divided by 2?

See

torchtitan/torchtitan/models/utils.py

Line 384 in a8899e4

# 4. we follow the convention and do not account for sparsity in causal attention

[shuhuayu]

Proposed a fix in this pr: #1923. Following the previously mentioned convention, the sparsity introduced by fixed_block_size (attention window) in local layers of Llama 4 models is also ignored.