Skip to content

[perf] fix: correct Qwen-VL head_dim consistency and merger token count (#6903)#6944

Open
chethanuk wants to merge 1 commit into
verl-project:mainfrom
chethanuk:fix/issue-6903-flops-counter-qwenvl
Open

[perf] fix: correct Qwen-VL head_dim consistency and merger token count (#6903)#6944
chethanuk wants to merge 1 commit into
verl-project:mainfrom
chethanuk:fix/issue-6903-flops-counter-qwenvl

Conversation

@chethanuk

@chethanuk chethanuk commented Jul 6, 2026

Copy link
Copy Markdown

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).

  1. head_dim sourcing (_estimate_qwen3_vl_flops) — read head_dim from config.text_config via getattr(config.text_config, "head_dim", hidden_size // num_attention_heads), matching the sibling estimators (_estimate_qwen3_vl_moe_flops and the other Qwen estimators). Qwen3-VL variants may configure a head_dim that differs from hidden_size // num_attention_heads; deriving it only from hidden_size under-counts attention FLOPs in that case.
  2. Merger token count (_estimate_qwen3_vit_flop) — the vision patch merger and the deepstack mergers run after spatial merging, so their linear layers process tokens_sum // spatial_merge_size**2 tokens, not the full tokens_sum. Multiplying merger parameters by the full tokens_sum over-counted merger FLOPs by a factor of spatial_merge_size**2 (= 4).

Scope is deliberately limited to these two objective claims. Claims 2 (lm-head * 2) and 3 (seqlen**2 attention accounting under packing) from the issue are excluded — they were part of upstream PR #4929, which was reverted by #4937.

Checklist Before Starting

Test

torch/ray are not installed in this environment, so pytest cannot collect tests/utils/test_flops_counter.py directly. The arithmetic was instead validated with a standalone torch-free replica of both estimators that:

  1. reproduces the old goldens exactly (confirming the replica matches the pre-change code), and
  2. reproduces the new goldens exactly (confirming the correction).

Both configs and both batches pass:

qwen3_vl     b0: old=195379819708416  new=192792370348032  delta=-2587449360384
qwen3_vl     b1: old=709446422495232  new=700575167545344  delta=-8871254949888
qwen3_vl_moe b0: old= 92975451340800  new= 90388001980416  delta=-2587449360384
qwen3_vl_moe b1: old=367622860308480  new=358751605358592  delta=-8871254949888

The delta on every case is exactly 6 * merger_total_N * (tokens_sum//4 - tokens_sum), i.e. only the merger term moved. The head_dim change produces no golden change because head_dim == hidden_size // num_attention_heads for the test configs. python -m py_compile passes on both files.

Correctness of the underlying architecture was cross-checked against the HuggingFace transformers Qwen3-VL source: Qwen3VLVisionPatchMerger.hidden_size = config.hidden_size * spatial_merge_size**2 and its forward reshapes with x.view(-1, self.hidden_size), so each merger emits tokens_sum // spatial_merge_size**2 tokens; the deepstack mergers are the same Qwen3VLVisionPatchMerger module (use_postshuffle_norm=True), so the same merger_tokens applies to both.

API and Usage Example

No API change. estimate_flops / FlopsCounter signatures are unchanged; only the internal FLOPs value for Qwen3-VL / Qwen3-VL-MoE models is corrected.

Design & Code Changes

  • _estimate_qwen3_vl_flops: head_dim now sourced via getattr(config.text_config, "head_dim", ...).
  • _estimate_qwen3_vit_flop: mergers pulled out of dense_N and multiplied by merger_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_vl and qwen3_vl_moe goldens 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 --> E
Loading

AI 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.

  • Read the Contribute Guide.
  • Apply pre-commit checks.
  • Add / Update the documentation. N/A — internal FLOPs-accounting fix, no user-facing docs.
  • Add unit or end-to-end test(s) to cover the code — existing tests/utils/test_flops_counter.py goldens updated to lock in the corrected values.

…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).

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@wuxibin89 wuxibin89 changed the title fix(flops_counter): correct Qwen-VL head_dim consistency and merger token count (#6903) [perf] fix: correct Qwen-VL head_dim consistency and merger token count (#6903) Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant