> 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.context_parallel.sharder

Context-parallel batch-sharding contract.

Every CP backend is a :class:`ContextParallelSharder`. A model that owns its CP batch
sharding and attention transport returns one from
`prepare_model_inputs_for_cp` under the `"cp_sharder"` batch key; the
framework constructs its own for the remaining backends (torch
`context_parallel` round-robin, TE/THD, MagiAttention) so
constructing :class:`ContextParallelSharder` resolves and configures the backend;
callers then invoke `sharder.shard(batch)`. This replaces the retired private batch keys
(`_cp_make_batch_fn`, `_cp_metadata_seq_dims`, `_cp_metadata_pad_values`,
`_cp_full_logits_grad_touch`).

The contract is a closed verb set with open implementations: the dataclass
slots are plain callables filled with functions from the owning model's
directory or from the framework. `local_token_global_indices` — the global
position of every local token — is the universal layout coordinate system: the
default token-tensor shard/gather are synthesized from it, so a sharder only
overrides them when it has a cheaper communication pattern. Layouts that are a
pure function of `(cp_mesh, padded_seq_len)` (contiguous, round-robin)
provide it at construction; data-dependent layouts (THD `cu_seqlens`
partitioning, magi's dispatch solver) start as None and report the partition
they computed as :class:`ShardLayout`, which the dispatch stores on the
sharder — sharders are built per resolution/hook call, so the layout never
leaks across steps. Before the first `shard_batch` their token verbs
raise.
The sharder carries no backend tag: nothing may branch on which backend
produced it.

This module also hosts the framework's `shard_batch` implementations: the
shared contiguous-shard batch prep used by models whose CP ranks own contiguous
sequence slices (Gemma4, DeepSeek V4), and the torch `context_parallel`
round-robin load-balanced prep with its index map. The TE/THD and magi preps
live with their dependencies (`context_parallel.utils`, `context_parallel.magi`); the
dispatcher wraps them into sharders at resolution time.

## Module Contents

### Classes

| Name                                                                                                               | Description                                                                 |
| ------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------- |
| [`ContextParallelSharder`](#nemo_automodel-components-distributed-context_parallel-sharder-ContextParallelSharder) | CP backend description: how a batch is sharded and where local tokens live. |
| [`ShardLayout`](#nemo_automodel-components-distributed-context_parallel-sharder-ShardLayout)                       | What a `shard_batch` learned about the layout it just applied.              |

### Functions

| Name                                                                                                                                               | Description                                                                         |
| -------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`_cp_rank`](#nemo_automodel-components-distributed-context_parallel-sharder-_cp_rank)                                                             | Resolve this rank's index within the CP submesh (0 without distributed).            |
| [`_normalize_cp_primary_and_positions`](#nemo_automodel-components-distributed-context_parallel-sharder-_normalize_cp_primary_and_positions)       | Resolve the primary stream and normalize `position_ids` for a CP shard.             |
| [`_pad_cp_buffers_to_divisor`](#nemo_automodel-components-distributed-context_parallel-sharder-_pad_cp_buffers_to_divisor)                         | Right-pad every CP buffer on its sequence dim to `cp_divisor` in place.             |
| [`_pad_position_ids_seq_dim_`](#nemo_automodel-components-distributed-context_parallel-sharder-_pad_position_ids_seq_dim_)                         | -                                                                                   |
| [`_pad_tensor_seq_dim_`](#nemo_automodel-components-distributed-context_parallel-sharder-_pad_tensor_seq_dim_)                                     | -                                                                                   |
| [`_reorder_gathered_token_tensor`](#nemo_automodel-components-distributed-context_parallel-sharder-_reorder_gathered_token_tensor)                 | Reassemble rank-order gathered shards into global sequence order.                   |
| [`contiguous_local_indices`](#nemo_automodel-components-distributed-context_parallel-sharder-contiguous_local_indices)                             | Global positions owned by this rank under contiguous rank-ordered sharding.         |
| [`convert_attention_mask_to_padding_mask`](#nemo_automodel-components-distributed-context_parallel-sharder-convert_attention_mask_to_padding_mask) | Pop `attention_mask` and derive `padding_mask` (True == pad) in place.              |
| [`gather_token_tensor_by_indices`](#nemo_automodel-components-distributed-context_parallel-sharder-gather_token_tensor_by_indices)                 | Differentiably gather a token-aligned local shard back to the full sequence.        |
| [`identity_local_indices`](#nemo_automodel-components-distributed-context_parallel-sharder-identity_local_indices)                                 | Global positions owned by this rank when CP is inactive: all of them.               |
| [`round_robin_local_indices`](#nemo_automodel-components-distributed-context_parallel-sharder-round_robin_local_indices)                           | Global positions owned by this rank under torch `context_parallel` load balancing.  |
| [`shard_batch_aux_only`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_batch_aux_only)                                     | Round-robin CP shard of the no-grad AUXILIARY streams only.                         |
| [`shard_batch_contiguous`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_batch_contiguous)                                 | Prepare and contiguously shard a batch for model-owned CP.                          |
| [`shard_batch_identity`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_batch_identity)                                     | `shard_batch` of the identity sharder: a no-op that returns a trivial shard layout. |
| [`shard_batch_load_balanced`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_batch_load_balanced)                           | Shard a batch with torch `context_parallel` round-robin load balancing.             |
| [`shard_sequence_for_cp_contiguous`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_sequence_for_cp_contiguous)             | Pad and contiguously shard a full-length sequence tensor inside a model forward.    |
| [`shard_sequence_for_cp_round_robin`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_sequence_for_cp_round_robin)           | Pad and round-robin shard a full-length sequence tensor inside a model forward.     |
| [`shard_token_tensor_by_indices`](#nemo_automodel-components-distributed-context_parallel-sharder-shard_token_tensor_by_indices)                   | Keep this rank's slice of a full-length token-aligned tensor.                       |

### Data

[`_NO_SHARD_LAYOUT`](#nemo_automodel-components-distributed-context_parallel-sharder-_NO_SHARD_LAYOUT)

[`_ROUND_ROBIN_PAD_FILL`](#nemo_automodel-components-distributed-context_parallel-sharder-_ROUND_ROBIN_PAD_FILL)

### API

```python
class nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder(
    model: torch.nn.Module | None = None,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    batch: dict[str, typing.Any] | None = None,
    shard_batch: collections.abc.Callable[..., tuple[collections.abc.Callable, dict[str, typing.Any], 'ShardLayout | None']] | None = None,
    local_token_global_indices: collections.abc.Callable[..., torch.Tensor] | None = None,
    shard_layout: 'ShardLayout | None' = None,
    padding_token_id: int = 0,
    num_chunks: int = 1,
    loss_mask: torch.Tensor | None = None,
    invoke_pre_embed: bool = True,
    extra_seq_buffers: dict[str, int] | None = None
)
```

CP backend description: how a batch is sharded and where local tokens live.

```python
nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder._indices(
    padded_seq_len: int,
    device
) -> torch.Tensor
```

```python
nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.gather_token_tensor(
    tensor: torch.Tensor,
    seq_dim: int = 1,
    trim: bool = False,
    fill: float | int | None = None
) -> torch.Tensor
```

Differentiably gather a token-aligned local shard to global order.

With `trim=True` the result is returned in the caller's original
coordinates using the shard layout: sliced back to `original_seq_len`,
un-flattened to `input_row_shape` (THD), or mapped through the
reported position map (`fill` for input positions whose tokens were
dropped, e.g. re-padded pack slots). Raises when no layout is present
(nothing to trim to).

```python
nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.shard(
    batch: dict[str, typing.Any]
) -> tuple[collections.abc.Callable, dict[str, typing.Any]]
```

Shard a batch and retain its layout for token-aligned tensors.

```python
nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder.shard_token_tensor(
    tensor: torch.Tensor,
    seq_dim: int = 1,
    fill: float | int | None = None
) -> torch.Tensor
```

Shard a full-length token-aligned tensor exactly like the model inputs.

When shard layout are present (after the first `shard_batch`), the
caller may pass tensors in its own coordinates and the verb applies the
same transform the batch went through:

* `[B, S_in]` tensors on a repositioned-row layout (reported position
  map, e.g. DSV4 packed repad) are scattered into the padded rows,
  `fill` filling the pad slots;
* tensors matching the reported pre-flatten `input_row_shape` on a
  flat-stream (THD) layout are flattened first (the returned shard is
  in the model's local stream coordinate);
* tensors of `original_seq_len` are right-padded to
  `padded_seq_len` with the explicit `fill` value;
* tensors already at `padded_seq_len` shard directly.

Any other length raises instead of silently sharding the wrong slice.

```python
class nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout(
    local_token_global_indices: torch.Tensor | None = None,
    original_seq_len: int | None = None,
    padded_seq_len: int | None = None,
    input_row_shape: tuple[int, ...] | None = None,
    input_token_stream_positions: torch.Tensor | None = None
)
```

Dataclass

What a `shard_batch` learned about the layout it just applied.

Returned as the third element of `shard_batch` and stored on the sharder
by the dispatch (`sharder.shard_layout = layout`), so the token verbs can
accept and return tensors in the CALLER's coordinates: pad is an internal
detail of the CP layout.

```python
nemo_automodel.components.distributed.context_parallel.sharder._cp_rank(
    cp_mesh
) -> int
```

Resolve this rank's index within the CP submesh (0 without distributed).

```python
nemo_automodel.components.distributed.context_parallel.sharder._normalize_cp_primary_and_positions(
    batch
) -> tuple[str, torch.Tensor, int, int, torch.Tensor, int]
```

Resolve the primary stream and normalize `position_ids` for a CP shard.

Pops nothing; determines whether the batch carries `inputs_embeds`
(`[batch, sequence, hidden]`) or `input_ids` (`[batch, sequence]`),
injects a 1-D `position_ids` arange when absent, and expands a shared
`[1, sequence]` position row to the batch size.

**Parameters:**

The full-sequence batch; `position_ids` is added/expanded in place.

**Returns:** `str`

\`\`(primary\_key, primary\_seq\_tensor, seq\_len, batch\_size, position\_ids,

```python
nemo_automodel.components.distributed.context_parallel.sharder._pad_cp_buffers_to_divisor(
    cp_buffers,
    cp_seq_dims,
    cp_no_restore_buffers,
    batch_buffer_keys,
    seq_len,
    cp_divisor,
    batch
)
```

Right-pad every CP buffer on its sequence dim to `cp_divisor` in place.

Pads sequence length to be divisible by `2 * cp_size` (required by
`context_parallel` load balancing); the inputs\_embeds path can hit
arbitrary seq lengths from the VLM collator, so padding happens here rather
than relying on dataset-side padding. Each padded, batch-sourced buffer is
mirrored back into `batch` so downstream dict readers see the padded shape.

**Parameters:**

Sequence-aligned tensors, each `[..., sequence, ...]`,
padded and replaced in place.

Per-buffer sequence axis.

Set of buffers the CP context must not restore;
rebuilt to reference the padded tensors.

Map from buffer index to its batch key (buffers not
sourced from the batch are omitted).

The pre-pad sequence length.

`2 * cp_size`.

The batch dict; padded batch-sourced buffers are mirrored back.

**Returns:**

`(cp_buffers, cp_no_restore_buffers)` referencing the padded tensors.

```python
nemo_automodel.components.distributed.context_parallel.sharder._pad_position_ids_seq_dim_(
    position_ids: torch.Tensor,
    seq_dim: int,
    pad_len: int
) -> torch.Tensor
```

```python
nemo_automodel.components.distributed.context_parallel.sharder._pad_tensor_seq_dim_(
    tensor: torch.Tensor,
    seq_dim: int,
    pad_len: int,
    value: float | int = 0
) -> torch.Tensor
```

```python
nemo_automodel.components.distributed.context_parallel.sharder._reorder_gathered_token_tensor(
    parts: list[torch.Tensor],
    index_parts: list[torch.Tensor],
    seq_dim: int = 1
) -> torch.Tensor
```

Reassemble rank-order gathered shards into global sequence order.

**Parameters:**

Per-rank local shards, in rank order.

Per-rank global token positions, in the same order.

The sequence dimension of the shards.

**Returns:** `torch.Tensor`

The full-sequence tensor with position `i` holding the token whose

```python
nemo_automodel.components.distributed.context_parallel.sharder.contiguous_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor
```

Global positions owned by this rank under contiguous rank-ordered sharding.

Rank `r` owns `[r * L, (r + 1) * L)` with `L = padded_seq_len // cp_size`.

**Parameters:**

The context-parallel device (sub)mesh.

Global sequence length after CP padding; must be
divisible by the CP size.

Device for the returned index tensor.

**Returns:** `torch.Tensor`

A 1-D int64 tensor of this rank's global token positions.

```python
nemo_automodel.components.distributed.context_parallel.sharder.convert_attention_mask_to_padding_mask(
    batch: dict
) -> None
```

Pop `attention_mask` and derive `padding_mask` (True == pad) in place.

Preserves padding semantics for modules such as MoE routers after the
attention mask is stripped for CP. Idempotent: a no-op when
`padding_mask` already exists or there is no `attention_mask`.

```python
nemo_automodel.components.distributed.context_parallel.sharder.gather_token_tensor_by_indices(
    cp_mesh,
    tensor: torch.Tensor,
    local_indices: torch.Tensor,
    seq_dim: int = 1
) -> torch.Tensor
```

Differentiably gather a token-aligned local shard back to the full sequence.

Uses `torch.distributed.nn.functional.all_gather` so gradients route back
to the owning ranks, then restores global sequence order from the gathered
per-rank indices. Identity when CP is inactive.

**Parameters:**

The context-parallel device (sub)mesh.

This rank's local shard, e.g. `[B, S/cp]` token logprobs.

This rank's global token positions.

The sequence dimension of `tensor`.

**Returns:** `torch.Tensor`

The full-sequence tensor in global order (padding not trimmed).

```python
nemo_automodel.components.distributed.context_parallel.sharder.identity_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor
```

Global positions owned by this rank when CP is inactive: all of them.

Index map of the identity sharder, so consumers run the same
token-verb code path at cp\_size \<= 1 (shard/gather become identities).

```python
nemo_automodel.components.distributed.context_parallel.sharder.round_robin_local_indices(
    cp_mesh,
    padded_seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor
```

Global positions owned by this rank under torch `context_parallel` load balancing.

The sequence splits into `2 * cp_size` chunks and rank `r` owns chunks
`r` and `2 * cp_size - 1 - r` (head-tail pairing, so causal-attention
work stays even across ranks).

**Parameters:**

The context-parallel device (sub)mesh.

Global sequence length after CP padding; must be
divisible by `2 * cp_size`.

Device for the returned index tensor.

**Returns:** `torch.Tensor`

A 1-D int64 tensor of this rank's global token positions, in local

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_aux_only(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_buffers: dict[str, int] | None = None
)
```

Round-robin CP shard of the no-grad AUXILIARY streams only.

The framework-owned peer of :func:`shard_batch_load_balanced` for models that
embed and sequence-shard their primary stream (`input_ids` /
`inputs_embeds` and any `pixel_*` keys) inside their own forward
(Megatron-style per-microbatch CP). Pads and round-robin-shards
`labels`/`position_ids`/`loss_mask`/`padding_mask` (and any
`extra_seq_buffers`) exactly like the load-balanced path, installs the same
ring-SDPA `context_parallel` context, and leaves the primary stream
full-length in the batch for the model to embed and shard via
:func:`shard_sequence_for_cp_round_robin`.

Because the primary stream never enters the `context_parallel` buffer list,
the grad-carrying-buffer constraint (`resize_` rejects tensors requiring
grad) does not apply: the model shards its embeddings with a differentiable
`index_select` and gradients route to the embeddings / vision tower.

**Parameters:**

The context-parallel device (sub)mesh (size > 1).

The tensor-parallel device (sub)mesh (or None); unused here.

The full-sequence batch; aux tensors are padded/mirrored in place
and the primary stream (`input_ids`/`inputs_embeds`) is left
full-length. `labels` is required.

Optional per-token loss mask `[batch, sequence]`, sharded
alongside the batch.

Accepted per the contract but unused (the primary
stream is not sharded here).

Additional batch keys mapped to their sequence axis,
padded and sharded alongside the aux tensors.

**Returns:**

`(ctx_factory, batch, ShardLayout)` where entering `ctx_factory()`

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_contiguous(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    pad_multiple: int = 1,
    extra_seq_keys: dict[str, int] | None = None,
    extra_pad_values: dict[str, typing.Any] | None = None,
    shard_primary: bool = True
)
```

Prepare and contiguously shard a batch for model-owned CP.

Normalizes the batch (mask conversion, position\_ids, labels), pads the
sequence to `cp_size * max(pad_multiple, 2)`, then keeps one contiguous
`seq_start:seq_end` slice per CP rank.

`shard_primary` is the only behavioral switch. When True (default) the
primary stream (`input_ids` / `inputs_embeds`) is padded and sliced with
the aux streams — the dispatch-level shard (e.g. DSV4). When False the primary
and pixel streams are left FULL-length for a model that embeds and slices them
inside its own forward per microbatch (e.g. Gemma4; see
:func:`shard_sequence_for_cp_contiguous`). The padded layout is identical
either way, so the two paths stay bit-for-bit slice-equivalent.

**Parameters:**

The context-parallel device (sub)mesh.

The tensor-parallel device (sub)mesh (or None).

The full-sequence batch; mutated and sharded in place.

Optional per-token loss mask; used as labels when the batch
has none, otherwise sharded alongside the batch.

Pad sentinel for `input_ids`.

Required per-CP-rank shard length multiple (e.g. DSV4's
compress-ratio LCM). The effective divisor is
`cp_size * max(pad_multiple, 2)`.

Model-specific per-token batch keys to pad and shard,
mapped to their sequence dim (e.g. Gemma4 vision group ids).

Pad sentinels for `extra_seq_keys` (default 0).

When False, leave the primary/pixel streams full-length for
in-forward slicing (aux-only shard); `ShardLayout.padded_seq_len` is
then the length the model must pad its primary to before slicing.

**Returns:**

`(contextlib.nullcontext, batch, ShardLayout)` — transport lives in the

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_identity(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0
)
```

`shard_batch` of the identity sharder: a no-op that returns a trivial shard layout.

Returned by the dispatch when no CP prep applies, so callers hold working
token verbs at every cp\_size (nothing was padded: original == padded).

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_batch_load_balanced(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_buffers: dict[str, int] | None = None
)
```

Shard a batch with torch `context_parallel` round-robin load balancing.

`ContextParallelSharder.shard_batch` implementation for the default framework-owned CP
path (layout `"round_robin"`, indices from
:func:`round_robin_local_indices`). Assumes an active CP mesh (size > 1).
`padding_token_id` is accepted per the contract but unused: CP-pad slots
are zero-filled, sit after every real token under the causal mask, and
carry -100 labels.

**Returns:**

`(ctx_factory, batch, ShardLayout)` where entering `ctx_factory()`

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_sequence_for_cp_contiguous(
    cp_mesh,
    tensor: torch.Tensor,
    seq_dim: int = 1,
    pad_value: float | int = 0,
    pad_multiple: int = 1
) -> tuple[torch.Tensor, torch.Tensor, int]
```

Pad and contiguously shard a full-length sequence tensor inside a model forward.

The contiguous peer of :func:`shard_sequence_for_cp_round_robin`: a model that keeps one
contiguous `seq_start:seq_end` slice per CP rank (model-owned p2p ring, e.g.
Gemma4) and leaves its primary stream full-length (see
:func:`shard_batch_contiguous` with `shard_primary=False`) calls this on
its embedded, spliced hidden states (and its 4-D `per_layer_inputs`) to keep this rank's
contiguous shard. It pads the sequence axis to `cp_size * max(pad_multiple, 2)`
— Gemma4's `pad_multiple` divisor, NOT the round-robin `2 * cp_size` — so
the slice aligns with the aux streams the contiguous aux-only sharder sliced.
Differentiable (`index_select`), so gradients route to the embeddings /
vision tower.

**Parameters:**

The context-parallel (sub)mesh; None or size \<= 1 is an identity.

Full-length sequence tensor, e.g. `inputs_embeds` `[B, S, H]`
or `per_layer_inputs` `[B, S, L, H]` (seq axis 1).

The sequence axis of `tensor`.

Fill for the CP-padding slots appended on `seq_dim`.

Per-CP-rank shard length multiple (the effective divisor is
`cp_size * max(pad_multiple, 2)`), matching
:func:`shard_batch_contiguous`.

**Returns:** `torch.Tensor`

`(local, local_indices, padded_seq_len)`: this rank's contiguous shard

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_sequence_for_cp_round_robin(
    cp_mesh,
    tensor: torch.Tensor,
    seq_dim: int = 1,
    pad_value: float | int = 0
) -> tuple[torch.Tensor, torch.Tensor, int]
```

Pad and round-robin shard a full-length sequence tensor inside a model forward.

The in-forward peer of :func:`shard_batch_load_balanced`: a model that leaves
its primary stream full-length (see :func:`shard_batch_aux_only`) calls this
on its embedded, spliced hidden states to keep this CP rank's head-tail chunk
pair, matching the layout the ring-SDPA context and the aux streams were
sharded to. Differentiable (`index_select`), so gradients route to the
embeddings / vision tower.

**Parameters:**

The context-parallel device (sub)mesh; None or size \<= 1 makes
this an identity (no pad, arange indices).

Full-length sequence tensor, e.g. `inputs_embeds` of shape
`[batch, sequence, hidden]`; the axis given by `seq_dim` is padded
to `2 * cp_size` and sharded.

The sequence axis of `tensor`.

Fill for the CP-padding slots appended on `seq_dim` (0 for
embeddings, matching the zero-padding the load-balanced path applies
to floating buffers).

**Returns:** `torch.Tensor`

`(local, local_indices, padded_seq_len)`: this rank's shard laid out

```python
nemo_automodel.components.distributed.context_parallel.sharder.shard_token_tensor_by_indices(
    tensor: torch.Tensor,
    local_indices: torch.Tensor,
    seq_dim: int = 1
) -> torch.Tensor
```

Keep this rank's slice of a full-length token-aligned tensor.

**Parameters:**

Full (padded) sequence tensor, e.g. `[B, S]` advantages.

This rank's global token positions.

The sequence dimension of `tensor`.

**Returns:** `torch.Tensor`

The local shard, laid out identically to the model inputs.

```python
nemo_automodel.components.distributed.context_parallel.sharder._NO_SHARD_LAYOUT = ShardLayout()
```

```python
nemo_automodel.components.distributed.context_parallel.sharder._ROUND_ROBIN_PAD_FILL = {'labels': -100, 'padding_mask': True, 'attention_mask': False}
```