Repository navigation
Fix FLOPs overcount for models with untied word embeddings - #7450
Merged
albertvillanova merged 2 commits intoSep 30, 2026
Merged
Conversation
Fixes huggingface#6708 compute_flops_per_token charged the input embedding as a 2*V*h matmul and dropped the LM head when word embeddings are tied. The embedding is a lookup, and the LM head is a 2*V*h matmul whether or not its weight is tied. Drop the embedding term and always count the LM head. Tied models are unchanged; untied models no longer get an extra 2*V*h per token (+5.4% for a Llama-3-8B shaped config at seq_len 8192). Tests: test_tied_vs_untied_lm_head now asserts tied == untied, and test_vocab_size_scaling checks that vocab size adds 3 * 2 * dV * h. Both fail on main and pass here; TestComputeFlopsPerToken and the rest of tests/test_utils.py pass; ruff check/format clean. AI assistance (Claude) was used; I reviewed the change and ran the tests. Co-Authored-By: Claude Opus 5.5 <[email protected]>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
albertvillanova
left a comment
albertvillanova
left a comment
Member
There was a problem hiding this comment.
Thanks, the fix is correct: the input embedding is a lookup, and the LM head costs 2 * V * h whether or not its weight is tied, so untied models were overcounted by exactly one vocabulary projection. Pinning tie_word_embeddings = False in test_vocab_size_scaling is the right call, since on a tied config the old code already gave the correct vocabulary delta and the test would not catch the regression.
One nit: could you make the code comments shorter?
albertvillanova
approved these changes
Sep 30, 2026
albertvillanova
changed the title
5 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #6708.
compute_flops_per_token charges the input embedding as a 2·V·h matmul, but it's a lookup with no FLOPs, and it charges the LM head as 0 when weights are tied, but the LM head is a 2·V·h matmul either way. So tied models come out right by accident; untied models are overcounted by 2·V·h forward FLOPs per token: +5.44% MFU on a Llama-3-8B-shaped config at seq_len 8192, +3.49% on the tiny Qwen3 test model.
Fix: drop the embedding term, always count the LM head once (removes the tie_word_embeddings branch). Tied-model numbers remain unchanged.
Tests: test_tied_vs_untied_lm_head now asserts tied == untied; new test_vocab_size_scaling checks +ΔV adds exactly 3·2·ΔV·h. Both fail on main, pass with the fix.
Matches how nanochat counts it (raised on the original MFU PR #5698).
Before submitting
AI writing disclosure
We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.
Note
Low Risk
Changes only analytical FLOP/MFU helpers and unit tests; no training forward/backward or model behavior is modified.
Overview
Corrects per-token training FLOP estimates in
compute_flops_per_tokenused for MFU reporting.Embedding is no longer counted as a
2·V·hmatmul (it's a lookup with zero FLOPs). LM head is always counted once as2·V·h, regardless oftie_word_embeddings—weight tying saves memory, not the projection compute.Untied models were previously overcounted by one embedding-sized term; tied totals were accidentally right. Reported MFU for untied configs will increase slightly because the denominator shrinks.
Tests now expect tied and untied FLOPs to match, and add
test_vocab_size_scalingso vocab-only changes affect FLOPs by3·2·ΔV·h(fwd+bwd).Reviewed by Cursor Bugbot for commit 597541f. Bugbot is set up for automated code reviews on this repo. Configure here.