nemo_automodel.components.models.qwen3_8_flash_next.cp

Contiguous context parallelism for the language-only Qwen3.8-Flash-Next model.

Module Contents

Classes

Name	Description
`Qwen3_8_FlashNextCPContext`	Per-forward metadata for Qwen3.8-Flash-Next’s contiguous CP sequence shard.

Functions

Name	Description
`_pad_right`	Right-pad a rank-two token tensor to `length`.
`_validate_right_tail_mask`	Validate and normalize a full-sequence token-validity mask.
`packed_boundaries_from_seq_lens`	Convert loader `seq_lens` metadata to physical document boundaries.
`qwen3_8_flash_next_cp_all_gather`	Gather equal contiguous sequence shards in global rank order.
`qwen3_8_flash_next_cp_left_halo`	Collect only the preceding causal boundary needed by a local operator.
`shard_batch_for_qwen3_8_flash_next_cp`	Validate, pad, and contiguously shard a Qwen3.8-Flash-Next text batch.

Data

__all__

API

class nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext(
    group: torch.distributed.ProcessGroup | None,
    rank: int,
    size: int,
    global_input_ids: torch.Tensor,
    global_padding_mask: torch.Tensor,
    local_sequence_start: int,
    local_sequence_length: int,
    global_cu_seqlens: torch.Tensor | None = None
)

Dataclass

Per-forward metadata for Qwen3.8-Flash-Next’s contiguous CP sequence shard.

global_cu_seqlens

Tensor | None = None

global_input_ids

Tensor

global_padding_mask

Tensor

global_sequence_length

int

Return the padded global physical sequence length.

global_sequence_lengths

Tensor

Return logical right-padded lengths as int64 [batch].

group

ProcessGroup | None

local_sequence_end

int

Return the exclusive global end of this rank’s sequence shard.

local_sequence_length

int

local_sequence_start

int

rank

int

size

int

nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext.__post_init__() -> None

Validate the replicated metadata and contiguous rank mapping.

nemo_automodel.components.models.qwen3_8_flash_next.cp._pad_right(
    tensor: torch.Tensor,
    length: int,
    value: int | bool
) -> torch.Tensor

Right-pad a rank-two token tensor to length.

Parameters:

tensor

torch.Tensor

Tensor of shape [batch, sequence].

length

int

Requested output sequence length, no shorter than sequence.

value

int | bool

Scalar fill value for appended positions.

Returns: torch.Tensor

Tensor of shape [batch, length]. The input is returned unchanged

nemo_automodel.components.models.qwen3_8_flash_next.cp._validate_right_tail_mask(
    mask: torch.Tensor,
    input_ids: torch.Tensor
) -> torch.Tensor

Validate and normalize a full-sequence token-validity mask.

Parameters:

mask

torch.Tensor

Binary validity tensor of shape [batch, global_sequence].

input_ids

torch.Tensor

Raw IDs of shape [batch, global_sequence] whose axes establish the expected mask shape and device.

Returns: torch.Tensor

Boolean validity tensor of shape [batch, global_sequence] on the

nemo_automodel.components.models.qwen3_8_flash_next.cp.packed_boundaries_from_seq_lens(
    seq_lens: torch.Tensor,
    total_tokens: int | None = None,
    sentinel: int = -1000
) -> torch.Tensor

Convert loader seq_lens metadata to physical document boundaries.

The THD packer concatenates documents contiguously and pads only the pack tail, so physical boundaries follow the REAL lengths (seq_lens); the loader’s seq_lens_padded is TE-specific virtual-layout metadata and must not be used for physical offsets. When total_tokens exceeds the packed length, the trailing pack padding becomes its own segment so pad tokens never join a real document.

Parameters:

seq_lens

torch.Tensor

Real per-document lengths [num_docs] or [1, num_docs].

total_tokens

int | NoneDefaults to None

Physical row length including trailing pack padding.

sentinel

intDefaults to -1000

Filler value marking unused length slots.

Returns: torch.Tensor

int64 boundaries [num_docs (+1 pad segment) + 1] starting at zero.

nemo_automodel.components.models.qwen3_8_flash_next.cp.qwen3_8_flash_next_cp_all_gather(
    tensor: torch.Tensor,
    context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext,
    sequence_dim: int = 1,
    differentiable: bool = True
) -> torch.Tensor

Gather equal contiguous sequence shards in global rank order.

Parameters:

tensor

torch.Tensor

Local tensor whose sequence_dim axis has length context.local_sequence_length (or a fixed compressed fraction of it shared by every rank). All non-sequence axes are replicated in shape across the CP group.

context

Qwen3_8_FlashNextCPContext

Qwen3.8-Flash-Next contiguous CP metadata.

sequence_dim

intDefaults to 1

Axis on which rank-ordered parts are concatenated.

differentiable

boolDefaults to True

Use PyTorch’s autograd-aware collective. Set False only for frozen routing values or integer metadata.

Returns: torch.Tensor

Tensor with the same axis order as tensor and a sequence_dim

nemo_automodel.components.models.qwen3_8_flash_next.cp.qwen3_8_flash_next_cp_left_halo(
    tensor: torch.Tensor,
    context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext,
    history: int
) -> torch.Tensor

Collect only the preceding causal boundary needed by a local operator.

Every rank contributes at most history trailing tokens. Autograd-aware All-Gather routes gradients from a later rank’s halo use back to the rank that owns those tokens. Rank zero and globally short prefixes are zero padded on the left.

Parameters:

tensor

torch.Tensor

Local sequence tensor of shape [batch, local_sequence, channels] using contiguous rank order.

context

Qwen3_8_FlashNextCPContext

Qwen3.8-Flash-Next contiguous CP metadata.

history

int

Number of immediately preceding global tokens required.

Returns: torch.Tensor

Left context of shape [batch, history, channels]. The result does

nemo_automodel.components.models.qwen3_8_flash_next.cp.shard_batch_for_qwen3_8_flash_next_cp(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: dict[str, typing.Any],
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int = 0,
    pad_multiple: int = 4
) -> tuple[collections.abc.Callable[[], contextlib.AbstractContextManager[typing.Any]], dict[str, typing.Any], nemo_automodel.components.distributed.context_parallel.sharder.ShardLayout]

Validate, pad, and contiguously shard a Qwen3.8-Flash-Next text batch.

Parameters:

cp_mesh

DeviceMesh

One-dimensional CP device mesh. Rank r owns the contiguous global interval [r * local_sequence, (r + 1) * local_sequence).

tp_mesh

DeviceMesh | None

Optional TP device mesh. Qwen3.8-Flash-Next CP requires this mesh to be absent or size one.

batch

dict[str, Any]

Mutable full-sequence batch. input_ids, labels, and optional attention_mask/padding_mask have shape [batch, global_sequence]; position_ids has shape [batch, global_sequence]. A packed THD row uses batch size one plus cu_seqlens or loader seq_lens document boundaries; packed attention and padding masks remain unsupported.

loss_mask

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, global_sequence] used by the shared sharder when labels are absent.

padding_token_id

intDefaults to 0

Raw token ID appended for CP divisibility.

pad_multiple

intDefaults to 4

Required multiple of every rank’s local sequence length. QSA requires its compression ratio, four for the released model.

Returns: Callable[[], contextlib.AbstractContextManager[Any]]

A null context factory, the mutated batch containing local token

nemo_automodel.components.models.qwen3_8_flash_next.cp.__all__ = ['Qwen3_8_FlashNextCPContext', 'packed_boundaries_from_seq_lens', 'qwen3_8_flash...