Context-Parallel Vision Frame Sharding

View as Markdown

Overview

Context parallelism (CP) splits a model’s language sequence across multiple GPUs. In a vision-language model (VLM), the vision tower normally processes every image and video on every CP rank before the language sequence is split. This repeats the same vision computation on each rank.

Context-parallel vision frame sharding does not apply context-parallel attention inside the vision tower. Instead, it distributes independent image and video frames over the existing CP ranks. Each rank processes only its assigned frames. The resulting visual embeddings are gathered in their original order, and ordinary sequence CP then shards the assembled multimodal sequence. This reduces duplicated vision-tower computation while preserving the model’s input ordering and gradients.

Only vision encoders that compute each image or video frame independently can use this feature. The built-in integration currently supports the following model families, including dense and mixture-of-experts (MoE) variants:

ModelSupport
Dense Qwen3.5Supported
Qwen3.5-MoESupported
Dense Qwen3-VLSupported, including DeepStack features
Qwen3-VL-MoENot yet supported

Configure Vision Frame Sharding

Enable context parallelism and vision frame sharding under distributed:

1distributed:
2 cp_size: 2
3 multimodal:
4 vision:
5 frame_sharding:
6 enabled: true
7 mesh_dims: [cp]
8 min_tokens: 2048
9 min_frames: 32
10 min_local_patch_rows: 96
11 cost_alpha: auto

Vision frame sharding is disabled by default. You need a CP size greater than 1 and a supported model for the sharded path.

The following example recipes are available:

For example, run the Qwen3-VL recipe on 8 GPUs:

$uv run automodel --nproc-per-node=8 \
> examples/vlm_finetune/qwen3/qwen3_vl_8b_cp2_vision_frame_shard.yaml

Configuration Fields

The following table describes the vision frame sharding fields:

FieldDefaultDescription
enabledfalseEnables vision frame sharding across the CP group.
mesh_dims[cp]Selects the device-mesh dimensions used to distribute frames. Only [cp] is currently supported.
min_tokens2048Minimum number of merged visual tokens that can activate sharding. Set it to 0 to always satisfy the token threshold.
min_frames32Minimum number of independent image or video-frame units that can activate sharding even below min_tokens. Set it to 0 to always satisfy the frame threshold.
min_local_patch_rows96Minimum pre-merge patch rows processed by each sharded rank that owns real frames. Smaller shards append whole, same-shaped zero frames and discard their outputs before the gather. Because this counts patch rows, it accounts for both frame count and resolution. Set it to 0 to disable local padding.
cost_alphaautoControls how frame cost is estimated for load balancing. auto uses the vision hidden size and is the recommended setting. A nonnegative integer overrides the inferred value. Setting it to 0 uses a purely quadratic attention-cost estimate.

Vision frame sharding uses the replicated path only when both the merged visual-token count is below min_tokens and the image or frame count is below min_frames. The frame threshold ensures that long videos can use sharding even when reduced per-frame spatial resolution keeps their merged-token count low.

How It Works

Qwen3.5- and Qwen3-VL-style vision towers process each frame independently. Vision frame sharding uses this property to:

  1. Expand each image or video into individual frame units.
  2. Divide contiguous frames across CP ranks using an estimated per-frame compute cost.
  3. Run the vision tower once on each rank’s local frames.
  4. Gather the visual embeddings in the original frame order.
  5. Continue through the model’s existing multimodal embedding and CP sequence-sharding path.

The partitioner accounts for both attention cost and linear per-patch work. This helps balance batches that mix large images with many smaller video frames. Keeping each rank’s frames contiguous allows the gathered outputs to be concatenated without reordering.

The gather operation supports autograd, so a trainable vision tower receives the same gradient contributions as the replicated path. Built-in integrations shard across the CP dimension only, even when tensor parallelism (TP) is also enabled. This avoids over-counting gradients for vision-tower weights that are replicated across TP ranks.

If a batch contains fewer frames than CP ranks, the implementation adds minimal dummy frames so every rank participates in the vision forward and collectives. The corresponding embeddings are discarded and contribute no gradient.

Frame sharding without local padding is mathematically valid. However, enabling frame sharding can make each rank’s reduced-precision ViT input much smaller than the input used when frame sharding is disabled and the full vision tower is replicated. Tiny local shapes can select different kernel tiling and accumulation orders, increasing numerical variation relative to that replicated baseline. The default min_local_patch_rows=96 is a conservative kernel-shape floor intended to improve this sharded-versus-replicated parity.

A rank below the floor appends enough whole zero frames of the same resolution to reach it. Their outputs are removed before the differentiable gather, so they contribute no model token or gradient; they only change the local vision-kernel shape. Ranks that own only a rank-fill dummy remain minimal.

This padding is a minor parity aid, not a correctness requirement, a bitwise-equivalence guarantee, or a throughput optimization. It adds some vision work and may be unnecessary for workloads whose unpadded sharded training already matches the frame-sharding-disabled baseline closely. Set min_local_patch_rows: 0 after workload-specific convergence validation when avoiding that extra work is preferred.

Check Your Workload

The benefit depends on the amount and shape of visual input, the CP size, activation checkpointing, and communication between ranks. Small visual workloads can spend more time gathering embeddings than they save in vision computation. The default min_tokens=2048 keeps small image workloads replicated, while min_frames=32 lets long videos select sharding when their frame count still represents substantial vision-tower work.

Compare a sharded run with the replicated baseline by disabling the policy from the command line:

$uv run automodel --nproc-per-node=8 \
> examples/vlm_finetune/qwen3/qwen3_vl_8b_cp2_vision_frame_shard.yaml \
> --distributed.multimodal.vision.frame_sharding.enabled false

Use the same seed, batch, and topology for both runs. Confirm that training remains stable, then compare step time and peak memory on a representative image or video workload.

Limitations

  • Vision frame sharding applies only to vision encoders that compute each image or video frame independently. Encoders with cross-frame coupling, including cross-frame attention, require a separate model-specific parallelization path. Audio encoders are not supported.
  • The final visual embeddings are gathered on every CP rank before the language sequence is sharded. The larger intermediate vision activations are distributed, but the final embeddings remain replicated.
  • Activation checkpointing already reduces retained vision activations, so vision frame sharding can improve computation more noticeably than peak memory.
  • Additional VLM families require model-specific integration before they can use this policy.
  • CP vision frame sharding has no effect when cp_size is 1.