feat(models): shard vision encoder across Ulysses SP ranks

[aoshen524]

Summary

• Vision DP distributes whole images across Ulysses SP ranks for ViT computation, then all-gathers embeddings. Reduces ViT memory from O(total_images) to O(total_images/N) when sp_size > 1.
• Gated behind vision_dp: bool in TrainEngineConfig — opt-in, default behavior unchanged.
• Supports Qwen2-VL, Qwen2.5-VL, Qwen3-VL (including Qwen3-VL deepstack).

Design

Aspect	Detail
Distribution unit	Whole images (not patches) — preserves ViT cu_seqlens semantics
Assignment	Load-balanced contiguous: greedy bin-packing by patch count, preserving image order
Gather	GatherVisionEmbeddings autograd Function; embedding counts computed locally (no all_gather)
Gradient sync	all_reduce(SUM) across SP group in backward before slicing — recovers complete gradient for each image from all sequence shards
Integration point	apply_monkey_patch() in ulyssess_patch.py

Files Changed

File	Change
areal/utils/vision_dp.py	Core Vision DP utilities (new)
areal/models/transformers/ulyssess_patch.py	Gate apply_vision_dp_patch() behind vision_dp flag
areal/api/cli_args.py	Add vision_dp: bool to TrainEngineConfig
areal/engine/fsdp_engine.py	Pass vision_dp=self.config.vision_dp to apply_monkey_patch()
areal/tests/test_vision_dp.py	21 CPU-only unit tests (new)

Usage

--train_engine.vision_dp true --train_engine.ulysses_sp_size 2

Precision Alignment (verl reference experiment)

Validated in verl-project/verl#5230 under controlled conditions (same algorithm shared across frameworks):

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

• Vision DP is numerically lossless: all differences within bf16 precision (~1e-05 max)
• Language gradients are bitwise identical at pre-clip phase
• Root cause: all_reduce(SUM) changes FP accumulation order in vision backward

Adapted from verl PR #5230.

🤖 Generated with Claude Code

[gemini-code-assist]

Summary of Changes

Hello @aoshen524, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces Vision Data Parallelism (Vision DP) to significantly optimize memory consumption for Vision Transformers (ViT) within Visual Language Models (VLMs) when operating with Ulysses Sequence Parallelism. By distributing images across SP ranks rather than processing all images on every rank, it effectively scales down ViT memory requirements. The integration is seamless for supported VLMs and leverages existing FSDP mechanisms for gradient synchronization.

Highlights

• Memory Optimization: Implemented Vision Data Parallel (Vision DP) to distribute whole images across Ulysses Sequence Parallel (SP) ranks, reducing Vision Transformer (ViT) memory usage from O(total_images) to O(total_images/N) when sp_size > 1.
• Automatic Activation: Vision DP is automatically activated for Visual Language Models (VLMs) such as Qwen2-VL, Qwen2.5-VL, and Qwen3-VL through the existing apply_monkey_patch() flow when sp_size > 1.
• Gradient Synchronization: Eliminated the need for explicit gradient synchronization as AReaL's FSDP utilizes a dp_sp submesh, which inherently covers both Data Parallel (DP) and Sequence Parallel (SP) dimensions, handling gradient reduction automatically.
• Image Distribution Strategy: Designed Vision DP to distribute whole images contiguously across ranks, preserving ViT cu_seqlens semantics and simplifying the all-gather process without requiring reordering.
• Gradient Scaling: Incorporated a GatherVisionEmbeddings autograd Function that includes gradient scaling in the backward pass to correctly compensate for partial image processing on each rank.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• areal/models/transformers/ulyssess_patch.py
  • Integrated the apply_vision_dp_patch function to enable Vision DP for supported VLM models.
• areal/tests/test_vision_dp.py
  • Added comprehensive CPU-only unit tests for the Vision Data Parallel utilities, covering image patch counting, embedding counting, image assignment to ranks, and local input preparation.
• areal/utils/vision_dp.py
  • Added a new module containing core Vision Data Parallel utilities, including functions for calculating image patch and embedding counts, assigning images to distributed ranks, preparing local vision inputs, and a custom GatherVisionEmbeddings autograd function for distributed embedding collection with gradient scaling.

Activity

• Initial implementation adapted from verl PR #5230.
• Automated generation of parts of the code using Claude Code.
• Ruff linting and formatting checks have passed.
• Unit tests for areal/tests/test_vision_dp.py are pending execution in the CI container.
• Multi-GPU correctness tests with VLM models are pending.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces Vision Data Parallelism (Vision DP) to distribute Vision Transformer (ViT) computation across sequence parallelism ranks, which is a great feature for reducing memory usage in Vision-Language Models. The implementation is well-structured, with the core logic encapsulated in areal/utils/vision_dp.py, integration via monkey-patching, and comprehensive unit tests. The code is clean and well-documented. I have one suggestion to improve the performance and readability of the patch extraction logic in prepare_local_vision_inputs by leveraging tensorized operations.

[rchardx]

Thanks for the clean implementation! The design is well-documented and the contiguous assignment approach is elegant. I have a few questions and concerns I'd like to discuss — particularly around gradient flow in the backward pass. Please see inline comments.

[rchardx]

LGTM