nemo_automodel.components.models.mistral_large4.context_parallel

Contiguous CP batch preparation and differentiable key transport for Large 4.

Module Contents

Functions

Name	Description
`gather_tokens`	Gather [local_tokens, …] tensors to [CP * local_tokens, …] in rank order.
`prepare_cp_inputs`	Prepare contiguous CP shards with positions computed before partitioning.
`prepare_packed_cp_inputs`	Preserve packed documents while preparing contiguous model-owned CP.
`shifted_cp_loss`	Compute HF shifted CE across shard boundaries, scaled for FSDP CP averaging.

Data

CPBatch

API

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

Gather [local_tokens, …] tensors to [CP * local_tokens, …] in rank order.

Local token counts must agree. Differentiable transport sums remote-query contributions back to each key owner; parameter averaging remains FSDP-owned.

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

Prepare contiguous CP shards with positions computed before partitioning.

Parameters:

batch

CPBatch

Global BSHD token tensors [batch, sequence, …] and optional complete image tensors [channels, height, width]. THD uses one packed row per microbatch, with seq_lens/seq_lens_padded document lengths.

num_chunks

intDefaults to 1

Positive number of PP microbatches along the batch axis.

shard_primary

boolDefaults to True

Whether to shard IDs/embeddings here. False keeps them global for image splicing and per-microbatch sharding in forward.

Returns: dict[str, ContextParallelSharder]

Mapping containing the model-owned ContextParallelSharder. Labels,

nemo_automodel.components.models.mistral_large4.context_parallel.prepare_packed_cp_inputs(
    batch: nemo_automodel.components.models.mistral_large4.context_parallel.CPBatch,
    num_chunks: int,
    shard_primary: bool
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]

Preserve packed documents while preparing contiguous model-owned CP.

Parameters:

batch

CPBatch

Global IDs/labels/positions [chunks, sequence], document lengths [chunks, documents], and optional complete images [channels, height, width]. Exactly one packed row belongs to each PP microbatch.

num_chunks

int

Positive number of independent pipeline microbatches; one without PP.

shard_primary

bool

False retains global IDs until complete-image splicing and/or stage-zero embedding; True shards text IDs here.

Returns: dict[str, ContextParallelSharder]

Sharder yielding local IDs/labels/positions/padding/document IDs

nemo_automodel.components.models.mistral_large4.context_parallel.shifted_cp_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    valid: torch.Tensor,
    mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.Tensor

Compute HF shifted CE across shard boundaries, scaled for FSDP CP averaging.

Parameters:

logits

torch.Tensor

FP32 capped logits [batch, local_sequence, vocab].

labels

torch.Tensor

Integer unshifted labels [batch, local_sequence].

valid

torch.Tensor

Boolean flattened local [batch * local_sequence] keep mask.

mesh

DeviceMesh

CP mesh; peers must hold shards of the same padded global batch.

Returns: torch.Tensor

Local loss sum times CP / global valid-target count, a differentiable

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