> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.distributed.cp_vision_frame_shard

Frame-level context-parallel vision-tower sharding.

Under context parallelism (CP), a VLM must encode media and splice the resulting
features before it shards the language sequence. Without vision frame sharding,
`self.visual(...)` runs the ENTIRE vision tower on the full, un-sharded set of
images on EVERY CP rank. With `cp_size=N` that is N x redundant compute and
O(all-images) vision activations per rank.

Qwen3-VL-style vision towers attend per IMAGE/FRAME ("entry"): the forward builds
`cu_seqlens` from `grid_thw` so each entry only attends within itself.  Entries
are therefore INDEPENDENT units, and

visual(all\_entries) == concat\_r visual(entries\_owned\_by\_rank\_r)    (entry order)

holds up to numerical precision (`allclose`; all CP ranks already hold identical
vision weights after the FSDP2 all-gather -- the replicated design relies on that
too).  So we partition the entries across the CP group, run `self.visual` on each
rank's slice, and all-gather the per-entry embeddings back to the full set.  The
downstream scatter into `inputs_embeds` and the sequence shard consume numerically
equivalent embeds through the unchanged code path.

Memory/compute: vision forward + activations drop \~cp\_size x; the final gathered
embeds (small vs the ViT's internal patch activations) are reconstructed in full
on every rank, then immediately sequence-sharded by the existing CP sharder.

Sharding-group scope (deliberate, machine-checked design constraint): the caller
declares the scope of the process group it publishes via
`set_cp_vision_group(group, config=..., spans_only_cp=...)`. With a FROZEN vision tower,
frames may be sharded across the full CP x TP rank set (`spans_only_cp=False`):
the forward is numerically equivalent (`allclose`; per-frame independence,
replicated weights) and
`requires_grad=False` means no backward ever runs through the gather.  With a
TRAINABLE vision tower, only a CP-only group (`spans_only_cp=True`, the default)
is valid: the vision tower is replicated (not tensor-parallel) across TP ranks,
so gathering frames across CP x TP would make the all-gather's
reduce-scatter(SUM) backward accumulate the vision gradient tp-fold (each
TP-replicated sequence shard backprops into the single compute rank).
:func:`maybe_distribute_visual` raises when a trainable tower meets a group not
declared CP-only.  Under pure TP (`cp_size == 1`) the published group has size
1, so this sharding is not enabled.

## Module Contents

### Classes

| Name                                                                                                                      | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`CpVisionFrameShardingConfig`](#nemo_automodel-components-distributed-cp_vision_frame_shard-CpVisionFrameShardingConfig) | Declarative policy for sharding VLM vision work across CP ranks.          |
| [`_AllGatherSeqDiff`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_AllGatherSeqDiff)                     | Differentiable all-gather of equal-size per-rank shards along `seq_dim`.  |
| [`_GroupHolder`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_GroupHolder)                               | Plain get/set/reset holder for the sharding-group scope active during the |
| [`_GroupScope`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_GroupScope)                                 | The published sharding group plus the caller's declaration of its scope.  |

### Functions

| Name                                                                                                                                | Description                                                                           |
| ----------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_all_gather_var_tokens`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_all_gather_var_tokens)                     | Differentiable all-gather of per-rank `[n_r, H]` token blocks into the full           |
| [`_check_group_scope_for_trainable`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_check_group_scope_for_trainable) | Raise when a trainable vision tower is sharded over a group not declared CP-only.     |
| [`_config_value`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_config_value)                                       | -                                                                                     |
| [`_contiguous_balanced_bounds`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_contiguous_balanced_bounds)           | Partition `len(patches)` entries into `world` CONTIGUOUS groups balanced by           |
| [`_grid_for_visual`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_grid_for_visual)                                 | Move `grid_thw` (\[N, 3] long, rows of (t, h, w)) to the device `visual` expects.     |
| [`_grid_list_for_planning`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_grid_list_for_planning)                   | Return host-side grid metadata for Python partitioning.                               |
| [`_infer_vision_hidden_size`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_infer_vision_hidden_size)               | Best-effort vision-width discovery without depending on model class names.            |
| [`_raise_if_any_rank_failed`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_raise_if_any_rank_failed)               | Reach group consensus before raising so a per-rank validation failure aborts on EVERY |
| [`_vision_cost_alpha`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_vision_cost_alpha)                             | Resolve the linear term in the vision partition cost `p*(p+alpha)`.                   |
| [`cp_vision_frame_sharding_active`](#nemo_automodel-components-distributed-cp_vision_frame_shard-cp_vision_frame_sharding_active)   | Return whether the current model forward has an active CP shard group.                |
| [`maybe_distribute_visual`](#nemo_automodel-components-distributed-cp_vision_frame_shard-maybe_distribute_visual)                   | Run `visual(pixel_values, grid_thw=grid_thw, return_dict=True)` but distribute the    |
| [`reset_cp_vision_group`](#nemo_automodel-components-distributed-cp_vision_frame_shard-reset_cp_vision_group)                       | Restore the previous sharding-group scope (pair with :func:`set_cp_vision_group`).    |
| [`set_cp_vision_group`](#nemo_automodel-components-distributed-cp_vision_frame_shard-set_cp_vision_group)                           | Install the sharding process group for the current model forward.                     |

### Data

[`CpVisionGroupToken`](#nemo_automodel-components-distributed-cp_vision_frame_shard-CpVisionGroupToken)

[`_CP_VISION_GROUP`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_CP_VISION_GROUP)

[`_LOGGED_COST_ALPHAS`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_LOGGED_COST_ALPHAS)

[`_LOGGED_ONCE`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_LOGGED_ONCE)

[`_LOGGED_SMALL_FALLBACK`](#nemo_automodel-components-distributed-cp_vision_frame_shard-_LOGGED_SMALL_FALLBACK)

[`__all__`](#nemo_automodel-components-distributed-cp_vision_frame_shard-__all__)

[`logger`](#nemo_automodel-components-distributed-cp_vision_frame_shard-logger)

### API

```python
class nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig(
    enabled: bool = False,
    mesh_dims: tuple[typing.Literal['cp']] = ('cp',),
    min_tokens: int = 2048,
    cost_alpha: int | typing.Literal['auto'] | None = 'auto'
)
```

Dataclass

Declarative policy for sharding VLM vision work across CP ranks.

**Parameters:**

**`enabled`** `bool` — default: False

Enable vision frame sharding when a multi-rank CP group is published. Disabled
by default so existing CP recipes retain their replicated vision behavior.

---

**`mesh_dims`** `tuple[Literal['cp']]` — default: ('cp',)

Device-mesh dimensions across which frames are distributed. Only
`("cp",)` is currently supported.

---

**`min_tokens`** `int` — default: 2048

Minimum number of merged visual tokens required to use the sharded
path. Smaller workloads stay replicated to avoid collective overhead.

---

**`cost_alpha`** `int | Literal['auto'] | None` — default: 'auto'

Non-negative linear term in the partition cost
`p * (p + cost_alpha)`. `"auto"` infers `3 * vision_hidden_size`
and falls back to `0` when the width is unavailable. `None` remains
accepted as a backward-compatible alias for `"auto"`.

---

**`cost_alpha`** `int | Literal['auto'] | None = 'auto'`

---

**`enabled`** `bool = False`

---

**`mesh_dims`** `tuple[Literal['cp']] = ('cp',)`

---

**`min_tokens`** `int = 2048`

---

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

Validate the serialized policy fields.

```python
class nemo_automodel.components.distributed.cp_vision_frame_shard._AllGatherSeqDiff()
```

**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`.

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

staticmethod

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

staticmethod

```python
class nemo_automodel.components.distributed.cp_vision_frame_shard._GroupHolder()
```

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.

**`_value`** `_GroupScope | None = None`

---

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

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

```python
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
```

```python
class nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope(
    group: torch.distributed.ProcessGroup,
    config: nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig,
    spans_only_cp: bool
)
```

Dataclass

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.

**`config`** `CpVisionFrameShardingConfig`

---

**`group`** `ProcessGroup`

---

**`spans_only_cp`** `bool`

---

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._all_gather_var_tokens(
    local: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    world: int,
    token_counts: list[int]
) -> torch.Tensor
```

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.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._check_group_scope_for_trainable(
    visual: torch.nn.Module,
    scope: nemo_automodel.components.distributed.cp_vision_frame_shard._GroupScope
) -> None
```

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.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._config_value(
    source: object,
    name: str
) -> object | None
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._contiguous_balanced_bounds(
    patches: torch.Tensor,
    world: int,
    cost_alpha_source: object | None = None,
    config: nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig | None = None
) -> list[int] | None
```

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.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._grid_for_visual(
    grid_thw: torch.Tensor,
    pixel_values: torch.Tensor
) -> torch.Tensor
```

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:**

**`grid_thw`** `torch.Tensor`

Tensor of shape \[entries, 3] containing `(time, height, width)` rows.

---

**`pixel_values`** `torch.Tensor`

Tensor of shape \[patches, patch\_dim] whose device owns vision execution.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._grid_list_for_planning(
    grid_thw: torch.Tensor
) -> tuple[list[list[int]], torch.Tensor]
```

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:**

**`grid_thw`** `torch.Tensor`

\[N, 3] tensor, one (t, h, w) row per media entry (N = entries).

---

**Returns:** `list[list[int]]`

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

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._infer_vision_hidden_size(
    source: object | None
) -> int | None
```

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:**

**`source`** `object | None`

A model, module, or config object (or a `Mapping`) that may
carry the vision width somewhere in its attributes/children.

---

**Returns:** `int | None`

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

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._raise_if_any_rank_failed(
    local_ok: bool,
    group: torch.distributed.ProcessGroup,
    device: torch.device,
    local_detail: str
) -> None
```

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:**

**`local_ok`** `bool`

Whether THIS rank's local validation passed.

---

**`group`** `dist.ProcessGroup`

The sharding process group shared by every rank.

---

**`device`** `torch.device`

Device for the 1-element flag tensor (the compute device, so the all-reduce
matches the active backend, e.g. CUDA under NCCL).

---

**`local_detail`** `str`

Actionable message raised by a rank that itself failed.

---

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._vision_cost_alpha(
    source: object | None = None,
    config: nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig | None = None
) -> int
```

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:**

**`source`** `object | None` — default: None

Object to infer the vision width from when in `auto` mode.

---

**`config`** `CpVisionFrameShardingConfig | None` — default: None

Optional typed sharding policy containing an exact override.

---

**Returns:** `int`

The resolved non-negative `alpha`.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.cp_vision_frame_sharding_active() -> bool
```

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.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.maybe_distribute_visual(
    visual: torch.nn.Module,
    pixel_values: torch.Tensor | None,
    grid_thw: torch.Tensor | None
) -> typing.Any
```

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:**

**`visual`** `torch.nn.Module`

The vision tower; must expose `spatial_merge_size` and accept
`visual(pixel_values, grid_thw=..., return_dict=True)`.

---

**`pixel_values`** `torch.Tensor | None`

`[total_patch_rows, patch_dim]` pixel rows for ALL entries,
frame-contiguous in entry order (entry order, then frame order within each
entry) -- the exact tensor the replicated call would receive.  `None`
(no media in the batch) is forwarded to `visual` unchanged.

---

**`grid_thw`** `torch.Tensor | None`

`[N, 3]` tensor, one (t, h, w) row per media entry (N = entries;
`t * h * w` patch rows per entry).  Expected on CPU.  `None` (no media
in the batch) is forwarded to `visual` unchanged.

---

**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).

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.reset_cp_vision_group(
    token: nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionGroupToken
) -> None
```

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

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.set_cp_vision_group(
    group: torch.distributed.ProcessGroup | None,
    config: nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionFrameShardingConfig,
    spans_only_cp: bool = True
) -> nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionGroupToken
```

Install the sharding process group for the current model forward.

**Parameters:**

**`group`** `dist.ProcessGroup | None`

Process group to shard vision frames across, or `None` to disable
sharding for the call.

---

**`config`** `CpVisionFrameShardingConfig`

Declarative sharding policy resolved from the recipe configuration.

---

**`spans_only_cp`** `bool` — default: True

Declaration of the group's scope.  `True` (default) states
that `group` spans context-parallel ranks only -- always
gradient-correct.  Pass `False` only for a group that also spans
replicated non-CP axes (e.g. the flattened CP x TP rank set); that is
valid solely for a FROZEN vision tower, and
:func:`maybe_distribute_visual` raises when a trainable tower meets a
group not declared CP-only (the gather's reduce-scatter(SUM) backward
would otherwise accumulate the vision gradient tp-fold).

---

**Returns:** `CpVisionGroupToken`

A token to pass to :func:`reset_cp_vision_group`.

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.CpVisionGroupToken = _GroupScope | None
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._CP_VISION_GROUP = _GroupHolder()
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._LOGGED_COST_ALPHAS: set[tuple[str, int, str]] = set()
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._LOGGED_ONCE = False
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard._LOGGED_SMALL_FALLBACK = False
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.__all__ = ['CpVisionFrameShardingConfig', 'cp_vision_frame_sharding_active', 'maybe_distri...
```

```python
nemo_automodel.components.distributed.cp_vision_frame_shard.logger = logging.getLogger(__name__)
```