Potentially incorrect attention flop calculation due to wrong head_dim? #1920
Description
Bug description
torchtitan/torchtitan/models/utils.py
Lines 432 to 437 in a8899e4
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