Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsStep5nemo_automodel.components.models.step5.cp

nemo_automodel.components.models.step5.cp

View as Markdown

Memory-bounded sequence transport for Step-5 context parallelism.

Module Contents

Classes

NameDescription
_GatherHaloExchange only the preceding window and return its gradients to its owners.
_GatherSequenceGather directly into one allocation; reduce-scatter the backward gradient.

Functions

NameDescription
gather_haloPrepend only the preceding causal window using point-to-point transport.
gather_metadataGather non-differentiable positions, masks or indexer projections.
gather_sequenceGather contiguous sequence shards with their autograd edges.

API

class nemo_automodel.components.models.step5.cp._GatherHalo()

Bases: Function

Exchange only the preceding window and return its gradients to its owners.

nemo_automodel.components.models.step5.cp._GatherHalo.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None, None]
staticmethod
nemo_automodel.components.models.step5.cp._GatherHalo.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    halo_size: int,
    left_padding: int
) -> torch.Tensor
staticmethod
class nemo_automodel.components.models.step5.cp._GatherSequence()

Bases: Function

Gather directly into one allocation; reduce-scatter the backward gradient.

nemo_automodel.components.models.step5.cp._GatherSequence.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None]
staticmethod
nemo_automodel.components.models.step5.cp._GatherSequence.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
staticmethod
nemo_automodel.components.models.step5.cp.gather_halo(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None,
    halo_size: int,
    left_padding: int = 0
) -> torch.Tensor

Prepend only the preceding causal window using point-to-point transport.

Parameters:

x
torch.Tensor

Tensor [batch, local_sequence, …], with equal contiguous CP shards.

group
dist.ProcessGroup | None

CP process group, or None for a single rank.

halo_size
int

Maximum number of preceding tokens, normally window - 1.

left_padding
intDefaults to 0

Local zero positions before the received halo, for aligning an attention tile with global key coordinates. No padding is sent.

Returns: torch.Tensor

Tensor [batch, left_padding + min(rank * local_sequence, halo_size) + local_sequence, …].

nemo_automodel.components.models.step5.cp.gather_metadata(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None
) -> torch.Tensor

Gather non-differentiable positions, masks or indexer projections.

Parameters:

x
torch.Tensor

Tensor of shape [batch, local_sequence, …] without gradients.

group
dist.ProcessGroup | None

CP group, or None for a single rank.

Returns: torch.Tensor

Tensor of shape [batch, global_sequence, …]. It aliases x without CP.

nemo_automodel.components.models.step5.cp.gather_sequence(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup | None
) -> torch.Tensor

Gather contiguous sequence shards with their autograd edges.

Parameters:

x
torch.Tensor

Tensor of shape [batch, local_sequence, …], contiguous by CP rank.

group
dist.ProcessGroup | None

CP group, or None for a single rank.

Returns: torch.Tensor

Tensor of shape [batch, global_sequence, …]. It aliases x without CP.