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

[aoshen524]

Summary

• Adds Vision Data Parallel (DP) to distribute whole images across Ulysses SP ranks for parallelized ViT computation
• Hugely reduces ViT memory overhead: without this, every SP rank redundantly processes ALL images through the VisionTransformer. With Vision DP, each rank only processes 1/sp_size of the images, reducing ViT peak memory by ~sp_sizex (e.g. SP=4 → ~4x ViT memory reduction)
• When ulysses_sp_size > 1, each rank processes a subset of images independently, then all-gathers embeddings once at the end
• This avoids breaking cu_seqlens semantics (which would happen if patches within images were split across ranks)
• Model-agnostic create_dp_vision_forward() wrapper supports any VisionTransformer with forward(self, hidden_states, grid_thw) signature
• Supports Qwen2-VL, Qwen2.5-VL, Qwen3-VL, and Qwen3-VL-MoE VisionTransformers
• Includes GatherVisionEmbeddings custom autograd function with proper gradient scaling for FSDP/DDP compatibility

Why this matters

In VLM RL training with Ulysses SP, the ViT (VisionTransformer) is a major memory bottleneck. Text SP splits the sequence across ranks at each attention layer, but the ViT runs on the full set of images on every rank — meaning ViT memory usage is completely unaffected by SP. For scenarios with many images (e.g. multi-turn GUI agent training with screenshots), ViT activation memory can dominate.

Vision DP solves this by distributing images at the ViT level:

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

Key design choices

• Image-level distribution (not patch-level): avoids breaking ViT's internal cu_seqlens tracking
• Contiguous assignment: rank 0 gets images [0,1,...], rank 1 gets next chunk, etc. — no reordering needed after all-gather
• Gradient scaling: backward pass scales gradients by dp_size to compensate for partial image processing before FSDP/DDP reduction

Test plan

• 17 unit tests covering all utility functions (get_image_patch_counts, assign_images_to_dp_ranks, prepare_local_vision_inputs)
• Integration tests for full workflow with varying image sizes
• Edge cases: empty inputs, fewer images than ranks, single rank
• Multi-GPU integration test with actual VLM model

🤖 Generated with Claude Code

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