[perf] fix: correct Qwen-VL head_dim consistency and merger token count (#6903)#6944
[perf] fix: correct Qwen-VL head_dim consistency and merger token count (#6903)#6944chethanuk wants to merge 1 commit into
Conversation
…oken count (verl-project#6903) Two objective fixes to the Qwen3-VL FLOPs estimators in verl/utils/flops_counter.py: - _estimate_qwen3_vl_flops now sources head_dim via getattr(config.text_config, "head_dim", hidden_size // num_attention_heads), matching the sibling estimators. Some Qwen3-VL variants set a head_dim distinct from hidden_size//num_attention_heads. - _estimate_qwen3_vit_flop counts merger and deepstack-merger FLOPs on the post-spatial-merge token count (tokens_sum // spatial_merge_size**2), not the full tokens_sum, fixing a ~spatial_merge_size**2 (4x) over-count of the merger term. Test goldens for qwen3_vl and qwen3_vl_moe updated accordingly; verified with a torch-free numeric replica that reproduces both the old and new goldens exactly. Addresses verl-project#6903 (claims 1 and 4 only).
There was a problem hiding this comment.
Code Review
This pull request corrects the FLOPs estimation for Qwen3 VL models by properly accounting for spatial merger downsampling. Specifically, merger FLOPs are now calculated using downsampled tokens (tokens_sum // spatial_merge_size**2) rather than the full token count, and the expected FLOPs in the test suite have been updated accordingly. Additionally, the head_dim retrieval in _estimate_qwen3_vl_flops has been made more robust by attempting to fetch it directly from the configuration before falling back to a default calculation. There are no review comments, and I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
What does this PR do?
Corrects two objective inaccuracies in the Qwen3-VL FLOPs estimators in
verl/utils/flops_counter.py, and refreshes the affected test goldens. Addresses #6903 (claims 1 and 4 only).head_dimsourcing (_estimate_qwen3_vl_flops) — readhead_dimfromconfig.text_configviagetattr(config.text_config, "head_dim", hidden_size // num_attention_heads), matching the sibling estimators (_estimate_qwen3_vl_moe_flopsand the other Qwen estimators). Qwen3-VL variants may configure ahead_dimthat differs fromhidden_size // num_attention_heads; deriving it only fromhidden_sizeunder-counts attention FLOPs in that case._estimate_qwen3_vit_flop) — the vision patch merger and the deepstack mergers run after spatial merging, so their linear layers processtokens_sum // spatial_merge_size**2tokens, not the fulltokens_sum. Multiplying merger parameters by the fulltokens_sumover-counted merger FLOPs by a factor ofspatial_merge_size**2(= 4).Scope is deliberately limited to these two objective claims. Claims 2 (lm-head
* 2) and 3 (seqlen**2attention accounting under packing) from the issue are excluded — they were part of upstream PR #4929, which was reverted by #4937.Checklist Before Starting
is:pr 6903 in:body,flops_counter qwen-vlfix(flops_counter): ...referencing the issue.Test
torch/rayare not installed in this environment, sopytestcannot collecttests/utils/test_flops_counter.pydirectly. The arithmetic was instead validated with a standalone torch-free replica of both estimators that:Both configs and both batches pass:
The delta on every case is exactly
6 * merger_total_N * (tokens_sum//4 - tokens_sum), i.e. only the merger term moved. Thehead_dimchange produces no golden change becausehead_dim == hidden_size // num_attention_headsfor the test configs.python -m py_compilepasses on both files.Correctness of the underlying architecture was cross-checked against the HuggingFace
transformersQwen3-VL source:Qwen3VLVisionPatchMerger.hidden_size = config.hidden_size * spatial_merge_size**2and itsforwardreshapes withx.view(-1, self.hidden_size), so each merger emitstokens_sum // spatial_merge_size**2tokens; the deepstack mergers are the sameQwen3VLVisionPatchMergermodule (use_postshuffle_norm=True), so the samemerger_tokensapplies to both.API and Usage Example
No API change.
estimate_flops/FlopsCountersignatures are unchanged; only the internal FLOPs value for Qwen3-VL / Qwen3-VL-MoE models is corrected.Design & Code Changes
_estimate_qwen3_vl_flops:head_dimnow sourced viagetattr(config.text_config, "head_dim", ...)._estimate_qwen3_vit_flop: mergers pulled out ofdense_Nand multiplied bymerger_tokens = tokens_sum // spatial_merge_size**2; a code comment documents why (spatial-merge downsampling, shared by the deepstack mergers).tests/utils/test_flops_counter.py:qwen3_vlandqwen3_vl_moegoldens updated with the derived deltas; derivation comments made self-contained.The single non-obvious point is that dense layers and mergers see different token counts:
flowchart TD A["tokens_sum = sum(images_seqlens)"] --> B["dense layers (patch_embed, attn, mlp)<br/>run on full tokens_sum"] A --> C["mergers downsample by spatial_merge_size**2"] C --> D["merger_tokens = tokens_sum // spatial_merge_size**2"] B --> E["dense_N_flops = 6·dense_N·tokens_sum<br/>+ 6·merger_total_N·merger_tokens"] D --> EAI assistance was used to prepare this change; the arithmetic and the architecture claims above were verified end-to-end.
Checklist Before Submitting
Important
Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.
tests/utils/test_flops_counter.pygoldens updated to lock in the corrected values.