nemo_automodel.components.models.glm5_next.cp

Contiguous packed context parallelism for GLM-5.3-Flash.

Kimi Delta Attention carries recurrent and short-convolution state from left to right, so the generic load-balanced CP permutation is not valid. This module keeps one contiguous token interval per CP rank and one global document-id map that both KDA and KPool-DSA consume.

Module Contents

Classes

Name	Description
`Glm5NextPackedContext`	Global packed-document layout for one model step.
`_AllGatherBackwardAnchor`	Return zero while retaining a backward edge to a gathered tensor.
`_AllGatherSequence`	Autograd-aware all-gather of equal contiguous sequence shards.

Functions

Name	Description
`_global_doc_ids_from_batch`	Resolve the global document map before removing packing metadata.
`_normalize_batch_axis`	Restore the placeholder batch axis used by THD VLM collaters.
`_pad_sequence_dim`	-
`all_gather_backward_anchor`	Create a zero-valued dependency that keeps gather backward collective-safe.
`all_gather_sequence`	Gather a sequence-sharded tensor while preserving K/V gradient flow.
`build_fla_cp_context`	Build FLA’s KDA context for one batch row.
`doc_ids_from_cu_seqlens`	Convert cumulative document boundaries to ids `[1, sequence]`.
`doc_ids_from_seq_lens`	Convert per-document lengths `[batch, documents]` to ids `[batch, sequence]`.
`segment_cu_seqlens`	Return boundaries for consecutive document-id runs covering the full row.
`shard_batch_for_glm5_next_cp`	Contiguously shard packed GLM-5.3 token streams.

Data

_PAD_DOC_ID

API

class nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext(
    doc_ids: torch.Tensor,
    seq_start: int = 0,
    cp_size: int = 1,
    original_seq_len: int | None = None,
    _cu_seqlens: dict[int, tuple[torch.Tensor, torch.Tensor]] = dict()
)

Dataclass

Global packed-document layout for one model step.

_cu_seqlens

dict[int, tuple[Tensor, Tensor]]

cp_enabled

bool

Return whether the sequence is split across more than one rank.

cp_size

int = 1

doc_ids

Tensor

local_doc_ids

Tensor

Return document ids [batch, local_sequence] for this rank.

local_seq_len

int

Return the padded local sequence length.

original_seq_len

int | None = None

seq_start

int = 0

nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext.row_cu_seqlens(
    row: int
) -> tuple[torch.Tensor, torch.Tensor]

Return device/CPU segment boundaries for one packed batch row.

class nemo_automodel.components.models.glm5_next.cp._AllGatherBackwardAnchor()

Bases: Function

Return zero while retaining a backward edge to a gathered tensor.

nemo_automodel.components.models.glm5_next.cp._AllGatherBackwardAnchor.backward(
    ctx,
    grad_output: torch.Tensor
) -> torch.Tensor

staticmethod

nemo_automodel.components.models.glm5_next.cp._AllGatherBackwardAnchor.forward(
    ctx,
    gathered: torch.Tensor
) -> torch.Tensor

staticmethod

class nemo_automodel.components.models.glm5_next.cp._AllGatherSequence()

Bases: Function

Autograd-aware all-gather of equal contiguous sequence shards.

nemo_automodel.components.models.glm5_next.cp._AllGatherSequence.backward(
    ctx,
    grad_output: torch.Tensor
)

staticmethod

nemo_automodel.components.models.glm5_next.cp._AllGatherSequence.forward(
    ctx,
    local_tensor: torch.Tensor,
    group: typing.Any,
    dim: int
) -> torch.Tensor

staticmethod

nemo_automodel.components.models.glm5_next.cp._global_doc_ids_from_batch(
    batch: dict[str, typing.Any],
    seq_len: int
) -> torch.Tensor

Resolve the global document map before removing packing metadata.

nemo_automodel.components.models.glm5_next.cp._normalize_batch_axis(
    batch: dict[str, typing.Any]
) -> None

Restore the placeholder batch axis used by THD VLM collaters.

nemo_automodel.components.models.glm5_next.cp._pad_sequence_dim(
    tensor: torch.Tensor,
    pad_len: int,
    value: float | int | bool
) -> torch.Tensor

nemo_automodel.components.models.glm5_next.cp.all_gather_backward_anchor(
    gathered: torch.Tensor
) -> torch.Tensor

Create a zero-valued dependency that keeps gather backward collective-safe.

nemo_automodel.components.models.glm5_next.cp.all_gather_sequence(
    tensor: torch.Tensor,
    cp_group: typing.Any,
    dim: int = 1
) -> torch.Tensor

Gather a sequence-sharded tensor while preserving K/V gradient flow.

nemo_automodel.components.models.glm5_next.cp.build_fla_cp_context(
    packed_context: nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext,
    row: int,
    cp_group: typing.Any,
    conv_kernel_size: int
)

Build FLA’s KDA context for one batch row.

Parameters:

packed_context

Glm5NextPackedContext

Global document layout.

row

int

Batch row being executed.

cp_group

Any

Context-parallel process group.

conv_kernel_size

int

Short-convolution width used for the left halo.

Returns:

FLA FLACPContext carrying segment and process-group metadata.

nemo_automodel.components.models.glm5_next.cp.doc_ids_from_cu_seqlens(
    cu_seqlens: torch.Tensor,
    seq_len: int
) -> torch.Tensor

Convert cumulative document boundaries to ids [1, sequence].

nemo_automodel.components.models.glm5_next.cp.doc_ids_from_seq_lens(
    seq_lens: torch.Tensor,
    seq_len: int,
    padding_value: int = -1000
) -> torch.Tensor

Convert per-document lengths [batch, documents] to ids [batch, sequence].

nemo_automodel.components.models.glm5_next.cp.segment_cu_seqlens(
    doc_ids_row: torch.Tensor
) -> torch.Tensor

Return boundaries for consecutive document-id runs covering the full row.

nemo_automodel.components.models.glm5_next.cp.shard_batch_for_glm5_next_cp(
    cp_mesh,
    tp_mesh,
    batch: dict[str, typing.Any],
    loss_mask = None,
    padding_token_id: int = 0,
    shard_primary: bool = False
)

Contiguously shard packed GLM-5.3 token streams.

The top-level VLM uses shard_primary=False: image features must be spliced into the full embedding sequence inside forward before that differentiable primary stream is sliced. Labels and other no-grad token streams are still sharded here.

Parameters:

cp_mesh

One-dimensional CP mesh or None.

tp_mesh

Unused tensor-parallel mesh, accepted by the sharder protocol.

batch

dict[str, Any]

Batch containing token tensors with shape [batch, sequence].

loss_mask

Defaults to None

Optional loss mask [batch, sequence].

padding_token_id

intDefaults to 0

Fill value for padded token ids.

shard_primary

boolDefaults to False

Whether to shard input_ids in this function.

Returns:

Context factory, mutated local batch, and the global shard layout.

nemo_automodel.components.models.glm5_next.cp._PAD_DOC_ID = 0