nemo_automodel.components.models.inkling2.context_parallel

Contiguous Inkling2 context shards, differentiable KV and convolution halos.

Module Contents

Functions

Name	Description
`convolution_history`	Prepend causal history to masked [batch, local_sequence, channels].
`gather_sequence`	Gather equal-length [batch, local_sequence, …] tensors in CP rank order.
`prepare_cp_inputs`	Return a sharder for packed or unpacked global [batch, sequence] text/image inputs.
`unpack_row_layout`	Normalize shared THD packs to document positions and validity, in place.
`validate_document_positions`	Require unit increments inside documents and zero at a new document.

Data

CPBatch

API

nemo_automodel.components.models.inkling2.context_parallel.convolution_history(
    masked: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh | None,
    width: int
) -> torch.Tensor

Prepend causal history to masked [batch, local_sequence, channels].

Parameters:

masked

torch.Tensor

This rank’s shard [batch, local_sequence, channels] with invalid positions already zeroed.

mesh

DeviceMesh | None

One-dimensional CP mesh; None or a single rank pads with zeros only.

width

int

Convolution taps; width - 1 history positions are prepended.

Returns: torch.Tensor

New [batch, local_sequence + width - 1, channels] tensor. Rank zero and positions

nemo_automodel.components.models.inkling2.context_parallel.gather_sequence(
    value: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh | None,
    differentiable: bool
) -> torch.Tensor

Gather equal-length [batch, local_sequence, …] tensors in CP rank order.

Parameters:

value

torch.Tensor

This rank’s shard [batch, local_sequence, …]; every rank supplies the same shape and dtype.

mesh

DeviceMesh | None

One-dimensional CP mesh; None or a single rank returns the input itself.

differentiable

bool

Route gradients back to each shard’s source rank.

Returns: torch.Tensor

[batch, global_sequence, …] with the shards concatenated in rank order, newly

nemo_automodel.components.models.inkling2.context_parallel.prepare_cp_inputs(
    batch: nemo_automodel.components.models.inkling2.context_parallel.CPBatch,
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]

Return a sharder for packed or unpacked global [batch, sequence] text/image inputs.

IDs, image masks and images remain global for per-microbatch image splicing inside forward. Labels, positions and padding become local. The shared recipes own target alignment and loss normalization; no weights are touched.

nemo_automodel.components.models.inkling2.context_parallel.unpack_row_layout(
    batch: nemo_automodel.components.models.inkling2.context_parallel.CPBatch
) -> None

Normalize shared THD packs to document positions and validity, in place.

The shared packer supplies real/padded lengths [batch, documents], padded with -1000. Padding is determined from those lengths, never from token IDs: a real EOS remains valid even when the tokenizer also uses it as PAD.

Parameters:

batch

CPBatch

Mutable mapping with input_ids [batch, sequence], seq_lens and seq_lens_padded [batch, documents], optional position_ids and labels [batch, sequence]. Replaces metadata with boolean padding_mask and integer position_ids [batch, sequence]; token IDs and labels are unchanged.

nemo_automodel.components.models.inkling2.context_parallel.validate_document_positions(
    keep: torch.Tensor,
    positions: torch.Tensor
) -> None

Require unit increments inside documents and zero at a new document.

A contiguous CP shard may begin partway through a document. Later resets explicitly identify new documents; a gap cannot silently resume an old one.

Parameters:

keep

torch.Tensor

Boolean validity [batch, sequence].

positions

torch.Tensor

Integer logical positions [batch, sequence]. The first real span may start at a nonzero offset, including after left padding.

nemo_automodel.components.models.inkling2.context_parallel.CPBatch: TypeAlias = dict[str, torch.Tensor | list[torch.Tensor] | str | None]