nemo_automodel.components.distributed.cp_vision_frame_shard

Declarative policy for sharding VLM vision work across CP ranks.

Parameters:

nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig.__post_init__() -> None

Bases: Function

Differentiable all-gather of equal-size per-rank shards along seq_dim.

Forward concatenates every rank’s shard in rank order, producing the full sequence. Backward reduce-scatters the incoming gradient: each position may be read on any rank, so its gradient is the sum over ranks of the per-rank grad slice; reduce-scatter(SUM) hands each rank the summed gradient for the shard it owns. All shards must have equal length along seq_dim.

nemo_automodel.components.distributed.cp_vision_frame_shard._AllGatherSeqDiff.backward(
    ctx,
    grad_out: torch.Tensor
)

nemo_automodel.components.distributed.cp_vision_frame_shard._AllGatherSeqDiff.forward(
    ctx,
    x: torch.Tensor,
    group,
    seq_dim: int
) -> torch.Tensor

Plain get/set/reset holder for the sharding-group scope active during the VLM forward. Set by the VLM CP recipe around the model forward and read synchronously by maybe_distribute_visual before the vision tower runs.

A plain attribute (not a ContextVar) is sufficient: activation-checkpoint recompute re-enters vision submodules after frame selection, not this helper, so there is no cross-thread recompute read. Training steps are sequential, and token-based restore handles nested callers.

nemo_automodel.components.distributed.cp_vision_frame_shard._GroupHolder.get() -> nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope | None

nemo_automodel.components.distributed.cp_vision_frame_shard._GroupHolder.reset(
    token: nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope | None
) -> None

nemo_automodel.components.distributed.cp_vision_frame_shard._GroupHolder.set(
    value: nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope | None
) -> nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope | None

The published sharding group plus the caller’s declaration of its scope.

spans_only_cp=False marks a group that also spans replicated non-CP axes (e.g. the flattened CP x TP set); that is only safe for a frozen vision tower.

Differentiable all-gather of per-rank [n_r, H] token blocks into the full [sum(n_r), H] tensor in rank order.

Each rank holds a different token count, so we pad to the global max, all-gather equal [max_tokens, H] shards (via _AllGatherSeqDiff: forward all-gather + cat, backward reduce-scatter SUM), then slice each rank’s block back to its true token_counts[r] and concat. Backward = SUM is correct: every rank reassembles the full embeds and keeps only its sequence shard, so an entry’s gradient is produced on the shard-owner rank and reduce-scatter routes it back to the compute-rank (summing the straddling case); padding rows are never used downstream so they contribute zero.

Raise when a trainable vision tower is sharded over a group not declared CP-only.

The ViT is replicated (not tensor-parallel) across TP ranks, so gathering frames across a CP x TP group makes the all-gather’s reduce-scatter(SUM) backward accumulate the vision gradient tp-fold. This check is deterministic across ranks (requires_grad and the declaration are identical everywhere), so it raises on every rank before any collective.

Partition len(patches) entries into world CONTIGUOUS groups balanced by approximate vision-attention cost, with >=1 entry per group.

patches[i] is entry i’s pixel-row count (= grid_thw[i].prod()), a proxy for one frame-unit’s attention sequence length. Since the ViT attends within each frame/window, the hot path scales closer to patches[i] ** 2 than to patch count. Returns cut points cuts of length world+1 where rank r owns entries [cuts[r], cuts[r+1]); returns None when num_entries < world (caller falls back to the replicate path so every rank still runs the ViT once and FSDP collectives stay uniform).

Contiguous (not round-robin) so the gathered per-rank blocks concatenate back in the original entry order with no reshuffle. Deterministic: patches is identical on every rank, so all ranks compute the same partition (and thus the same per-rank token counts they need to unpack the all-gather).

Cost model: p*(p + alpha). A configured cost_alpha wins; otherwise alpha is inferred as 3 * vision_hidden_size from cost_alpha_source. The quadratic term is per-frame attention; alpha*p adds the LINEAR per-patch work (qkv/MLP projections) that the pure quadratic ignores — without it, packs mixing big image frames with many small video frames heap the small frames onto few ranks (attention-cost-”balanced” but frame-count- and wall-clock- imbalanced). Unknown architectures fall back to alpha=0. Partition-only choice: forward/grad are identical for any cuts.

Move grid_thw ([N, 3] long, rows of (t, h, w)) to the device visual expects.

Vision forwards may use the grid both as host-readable shape metadata and as indices into device-resident position embeddings. Keep it on pixel_values’s device; callers that need Python values can synchronize through tolist(), while a CPU grid cannot index a CUDA embedding table. Attention backend alone does not determine every grid consumer’s device requirement.

Parameters:

Returns: torch.Tensor

Tensor of shape [entries, 3] on pixel_values’s device. The input is returned

Return host-side grid metadata for Python partitioning.

The generic VLM recipe moves batch tensors, including grid metadata, to the model device. Make the required host copy explicit instead of hiding the synchronization behind tolist(). Callers that already keep the grid on CPU avoid the copy.

Parameters:

Returns: list[list[int]]

Tuple of grid_thw.tolist() (a list of [t, h, w] ints) and the CPU long

Best-effort vision-width discovery without depending on model class names.

Supported VLMs expose the width in different places: Qwen vision towers use visual.config.hidden_size, while Nemotron Omni’s multimodal config uses vit_hidden_size and RADIO-style towers may expose embed_dim. Prefer explicitly vision-named fields and nested vision configs over a root hidden_size so a text-model width is not accidentally selected.

Parameters:

Returns: int | None

The discovered positive width, or None when no vision width is found.

Reach group consensus before raising so a per-rank validation failure aborts on EVERY rank instead of deadlocking peers on the next collective.

A per-rank raise placed BEFORE a collective hangs the group: the failing rank raises while its peers block forever in all_gather. Instead, each rank contributes 0 (ok) or 1 (failed) and an all-reduce(MAX) over group makes the outcome identical everywhere. When any rank failed, the offending rank(s) raise their own local_detail and the rest raise a group-level message, so all ranks raise and none is left blocking.

Parameters:

Resolve the linear term in the vision partition cost p*(p+alpha).

A configured non-negative integer is an exact override (0 selects the legacy pure-quadratic model). With cost_alpha="auto" or the backward-compatible None alias, use 3 * vision_hidden_size: the three Q/K/V projections are a portable proxy for linear per-patch ViT work. Unknown architectures safely retain the legacy alpha=0 behavior.

Parameters:

Returns: int

The resolved non-negative alpha.

Return whether the current model forward has an active CP shard group.

This intentionally shares :func:maybe_distribute_visual’s typed policy. Model-specific multimodal implementations can use it to leave their ordinary (replicated / CP-off) multimodal forward completely untouched.

Run visual(pixel_values, grid_thw=grid_thw, return_dict=True) but distribute the forward across the CP group when enabled, returning an output object whose pooler_output (and deepstack_features list, if any) are the FULL gathered embeds in original entry order — a drop-in for the direct visual(...) call.

Falls back to the plain replicated call (exact pre-sharding behaviour) when sharding is disabled, no CP group is active, cp_size <= 1, there are no media inputs, or the visual workload is below the minimum sharding size.

Parameters:

Returns: Any

The vision tower’s output object. pooler_output is the full gathered

Raises:

• ValueError: When the vision tower has trainable parameters but the published group was declared spans_only_cp=False (see :func:set_cp_vision_group), or when any rank’s local visual output does not match its planned token count. The token-count mismatch is reduced across the sharding group before raising, so a single diverging rank makes EVERY rank raise (rather than deadlocking peers on the gather).

Restore the previous sharding-group scope (pair with :func:set_cp_vision_group).

Install the sharding process group for the current model forward.

Parameters:

Returns: CpVisionGroupToken

A token to pass to :func:reset_cp_vision_group.