feat(vision): add Vision DP for parallel ViT computation across SP ranks

[aoshen524]

Vision Data Parallel: Distribute ViT computation across Ulysses SP ranks

Motivation

When using Ulysses Sequence Parallelism (sp_size > 1), the attention layers split across SP ranks via all-to-all, but the VisionTransformer (ViT) still processes ALL images on every rank, wasting memory proportional to total_images.

Vision DP fixes this by distributing whole images (not patches) across SP ranks:

• Before: Each of N SP ranks processes ALL images → ViT memory = O(total_images)
• After: Each rank processes total_images/N images → ViT memory = O(total_images/N)

Design

1. Image-level distribution: Avoids breaking ViT's internal cu_seqlens tracking
2. Contiguous assignment: rank 0 gets images [0,1,...], rank 1 gets next chunk → no reordering after all-gather
3. Load-balanced bin-packing: Greedy contiguous assignment that balances total patch load (not just image count) across ranks
4. Correct gradient routing: all_reduce(SUM) in backward aggregates partial sequence gradients from all SP ranks before slicing by image assignment

Key changes

File	Change
verl/utils/vision_dp.py	Core utilities: assignment, slicing, all-gather with gradient support
verl/models/transformers/monkey_patch.py	Integration: patch VisionTransformer classes with idempotency guard
tests/utils/test_vision_dp_on_cpu.py	CPU-only unit tests (21 tests, no distributed required)

Tests

python -m pytest tests/utils/test_vision_dp_on_cpu.py -v
# 21 passed

Precision Alignment: Vision DP On vs Off

Experiment Setup

Config	Value
Model	Qwen2.5-VL-7B (SFT checkpoint)
GPUs	4 × H100 (1 node)
SP size	2
Strategy	FSDP2
Train batch size	4
Rollout N	2
Max assistant turns	4
Dummy mode	full (fixed multi-image prompts with meaningful screenshots + fixed text)
Determinism	enabled (seed=42, TF32 disabled, CUBLAS_WORKSPACE_CONFIG=:4096:8)
Advantage	hardcoded to 1.0 to remove reward randomness
Training steps	1

Why fix advantage? In RL training, reward randomness from rollouts causes gradient divergence unrelated to the optimization under test. Hardcoding advantage isolates the effect of Vision DP on gradient computation.

Comparison Method

1. Online: Both runs dump per-parameter gradient tensors (as .pt files) at the optimizer step via VERL_PRECISION_DUMP_GRADS env var. Each run auto-creates a UUID subdirectory.
2. Online: Per-group SHA256 hash fingerprints are logged at 3 phases: pre_clip (raw grad), before (post-clip), after (post-step), along with L2 norm (fp64) and absmax scalars.
3. Offline: tools/compare_grads.py loads the dumped tensors and computes per-parameter max(|A-B|), mean(|A-B|), and cosine similarity, aggregated by scope (vision/language/other).

Results: Hash-Level (62 groups × 3 phases)

Phase	Param SHA256	Grad SHA256 Match	Vision Grad MM	Language Grad MM
pre_clip (raw grad)	62/62 ✅	29/62	32	0 ✅
before (post-clip)	62/62 ✅	0/62	32	29
after (post-step)	0/62	0/62	32	29

Language gradients are bitwise identical at pre_clip — Vision DP does not affect the language model's gradient computation. The before/after phase divergence is expected: global grad norm (used for clipping) includes vision gradients, so a different clip factor propagates to all parameters.

Results: Per-Parameter Element-wise Diff

Scope	Params	max_diff	mean_diff	cosine_sim
vision	390	4.70e-05	2.93e-08	0.9991
language	338	9.50e-08	1.15e-10	1.0020
other	1	9.13e-08	2.25e-13	1.0001

Top-5 worst parameters (all vision scope):

Parameter	max_diff
model.visual.patch_embed.proj.weight	4.70e-05
model.visual.blocks.7.attn.qkv.weight	1.99e-05
model.visual.blocks.1.mlp.gate_proj.weight	1.59e-05
model.visual.blocks.0.mlp.gate_proj.weight	1.18e-05
model.visual.blocks.1.norm1.weight	1.13e-05

[gemini-code-assist]

Code Review

This pull request introduces Vision Data Parallelism (DP) to enable parallel ViT computation by distributing images across sequence parallel ranks. The implementation is well-structured, with new utilities in verl/utils/vision_dp.py and corresponding monkey-patching logic. The addition of comprehensive unit tests in tests/test_vision_dp.py is commendable.

My review focuses on improving code quality, robustness, and maintainability. I've identified a few areas with high-severity issues:

• Redundant computation in prepare_local_vision_inputs.
• Fragile logic for determining hidden_size in the DP wrapper, which could lead to runtime errors with different models.
• Significant code duplication in monkey_patch.py that should be refactored.

Addressing these points will make the new functionality more robust and easier to maintain in the future.

[gemini-code-assist]

This block of code for monkey-patching the vision transformer is repeated multiple times in this file (here for Qwen2/2.5-VL, and again for Qwen3 models). This duplication makes the code harder to read and maintain. Please consider refactoring this logic into a helper function that takes the module path and class name as arguments. This would significantly reduce code duplication and improve maintainability.

For example, a helper could look like this:

def _patch_vision_model_for_dp(module_path, class_name, ulysses_sp_size):
    try:
        module = __import__(module_path, fromlist=[class_name])
        model_class = getattr(module, class_name)
        original_forward = model_class.forward
        model_class.forward = create_dp_vision_forward(original_forward)
        print(f"Monkey patch {class_name}.forward for Vision DP (dp_size={ulysses_sp_size})")
    except (ImportError, AttributeError) as e:
        print(f"Warning: Could not patch {class_name} for Vision DP: {e}")

Additionally, the step numbering is inconsistent ("Step 4" here, but "Step 3" for the Qwen3 models). A refactor would help resolve such inconsistencies.

[gemini-code-assist]

The patch_counts are recalculated here, but they have already been computed in the calling function create_dp_vision_forward. This leads to redundant computation. You can improve efficiency and code clarity by passing patch_counts as an argument to prepare_local_vision_inputs and removing this line.

[gemini-code-assist]

The logic to determine hidden_size for ranks with no images is complex and relies on a series of hasattr checks for specific model attributes. This approach is brittle and may break with new or refactored models, leading to RuntimeError. A more robust approach would be to make this explicit. For example, create_dp_vision_forward could accept the hidden_size or a hidden_size_getter function as an argument. This would be provided at the patching site where the model context is clear, making the wrapper more reliable and easier to maintain.

[aoshen524]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces Vision Data Parallelism (Vision DP) to distribute Vision Transformer (ViT) computation across sequence parallel ranks, which is a valuable feature for reducing memory consumption. The implementation includes core utilities for image assignment, a custom autograd function for correct gradient aggregation, and integration via monkey-patching. The changes are accompanied by comprehensive unit tests and a detailed precision analysis in the description. My review identified an opportunity to refactor duplicated code in the monkey-patching logic to improve maintainability.

[gemini-code-assist]

There is significant code duplication in this block and in the qwen3_vl block below (lines 487-522) for applying the Vision DP monkey patch. This repetitive logic can be refactored into a helper function to improve maintainability and reduce redundancy.

A helper function could encapsulate the try...except block, the import, the idempotency check, and the patching logic itself. This would make the code cleaner and easier to extend to other vision models in the future.

For example, you could define a helper function like this:

def _patch_vision_model_for_dp(model_class_name: str, module_path: str, ulysses_sp_size: int):
    """Applies the Vision DP monkey patch to a given model class."""
    try:
        from verl.utils.vision_dp import create_dp_vision_forward
        module = __import__(module_path, fromlist=[model_class_name])
        model_class = getattr(module, model_class_name)

        if not getattr(model_class, "_vision_dp_patched", False):
            original_forward = model_class.forward
            model_class.forward = create_dp_vision_forward(original_forward)
            model_class._vision_dp_patched = True
            print(
                f"Monkey patch {model_class_name}.forward for Vision DP (dp_size={ulysses_sp_size})"
            )
    except (ImportError, AttributeError) as e:
        print(f"Warning: Could not patch {model_class_name} for Vision DP: {e}")

And then call it for each model:

# Step 4: patch VisionTransformer for Vision DP (image-level distribution)
if ulysses_sp_size > 1:
    if vision_dp:
        _patch_vision_model_for_dp(
            "Qwen2VisionTransformerPretrainedModel",
            "transformers.models.qwen2_vl.modeling_qwen2_vl",
            ulysses_sp_size
        )
        _patch_vision_model_for_dp(
            "Qwen2_5_VisionTransformerPretrainedModel",
            "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl",
            ulysses_sp_size
        )
    else:
        print(
            f"Vision DP disabled (vision_dp=False). "
            f"ViT runs replicated on all {ulysses_sp_size} SP ranks."
        )

This refactoring would also apply to the qwen3_vl block.

[aoshen524]

@ISEEKYAN Hi, could you please help review this PR? Thanks!