nemo_automodel.components.distributed.context_parallel.mamba

Context parallelism for Mamba/SSM layers using a hidden-parallel strategy.

Instead of splitting the sequence across CP ranks (as attention CP does), this module uses an all-to-all redistribution so that each CP rank processes the full sequence but only a subset of heads (d_inner / cp_size). The data flow is::

[B, L_local, D] —> all-to-all —> [B, L_global, D/cp] —> conv1d + SSM kernel —> [B, L_global, D/cp] —> all-to-all —> [B, L_local, D]

This module is intentionally not a subclass of nn.Module because it owns no trainable parameters. It holds references to the Mamba mixer’s parameters and slices them in the forward path so that gradients flow back to the full (unsliced) parameters.

Module Contents

Classes

Name	Description
`MambaContextParallel`	Hidden-parallel context parallelism for a Mamba2 mixer layer.
`_AllToAll`	Autograd wrapper around `torch.distributed.all_to_all_single`.

Functions

Name	Description
`_all_to_all`	Functional entry-point for the autograd-aware all-to-all.
`_all_to_all_cp2hp`	Transform from sequence-sharded to hidden-sharded layout (batch-first).
`_all_to_all_hp2cp`	Transform from hidden-sharded to sequence-sharded layout (batch-first).
`_deinterleave_packed_seqs`	Rearrange tokens from rank-major to sequence-major order after all-to-all.
`_redo_attention_load_balancing`	Reorder from sequential back to DualChunkSwap for attention.
`_reinterleave_packed_seqs`	Inverse of `_deinterleave_packed_seqs`.
`_reorder_chunks`	Reorder equal-sized chunks of a tensor according to order.
`_undo_attention_load_balancing`	Reorder from DualChunkSwap to sequential for SSM processing.

API

class nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel(
    cp_group: torch.distributed.ProcessGroup,
    num_heads: int,
    head_dim: int,
    n_groups: int,
    d_state: int,
    mixer: torch.nn.Module
)

Hidden-parallel context parallelism for a Mamba2 mixer layer.

This class does not own trainable parameters. It stores a reference to the mixer module and accesses its parameters (conv1d, dt_bias, A_log, D) on the fly so that gradients propagate to the original (full) parameters and FSDP-managed DTensor replacements are picked up correctly.

DualChunkSwap reordering is always undone before the SSM kernel and redone after, because both TE CP (p2p) and PyTorch’s context_parallel(allgather) reorder sequence chunks for load balancing.

Parameters:

cp_group

torch.distributed.ProcessGroup

Context-parallel process group.

num_heads

int

Total number of SSM heads (before any parallelism).

head_dim

int

Dimension per head.

n_groups

int

Number of SSM groups (for grouped B/C states).

d_state

int

SSM state dimension.

mixer

nn.Module

Reference to the Mamba mixer module (owns conv1d, dt_bias, A_log, D).

cp_rank

= cp_group.rank()

cp_size

= cp_group.size()

d_inner

= num_heads * head_dim

d_inner_local

= self.num_heads_local * head_dim

group_repeat_count

= self.cp_size // n_groups

n_groups_local

= 1

num_heads_local

= num_heads // self.cp_size

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel._repeat_group_state(
    state: torch.Tensor
) -> torch.Tensor

Repeat group states for CP ranks when n_groups < cp_size.

[B, L, n_groups * d_state] -> [B, L, n_groups * repeat * d_state] Also supports THD 2D input [T, n_groups * d_state].

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel._slice_conv_param(
    param: torch.Tensor
) -> torch.Tensor

Slice a conv1d parameter (weight or bias) along its channel dimension.

Parameter slicing is done in the forward path so that gradients backpropagate to the original (full) parameters.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel._slice_vector_param(
    param: torch.Tensor
) -> torch.Tensor

Slice a per-head vector parameter for the current CP rank.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.get_A_log(
    A_log: torch.Tensor | None = None
) -> torch.Tensor

Slice A_log for the current CP rank.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.get_D(
    D: torch.Tensor | None = None
) -> torch.Tensor

Slice D for the current CP rank.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.get_conv1d_bias() -> torch.Tensor

Slice conv1d.bias for the current CP rank.

Bias shape: [conv_dim]. Returns [conv_dim_local].

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.get_conv1d_weight() -> torch.Tensor

Slice conv1d.weight for the current CP rank.

Weight shape: [conv_dim, 1, kernel_size] where conv_dim = d_inner + 2 * n_groups * d_state. Returns [conv_dim_local, kernel_size] (squeezed for causal_conv1d kernel).

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.get_dt_bias(
    dt_bias: torch.Tensor | None = None
) -> torch.Tensor

Slice dt_bias for the current CP rank.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.post_conv_ssm(
    output: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor

Redistribute SSM output from hidden-sharded back to sequence-sharded layout.

nemo_automodel.components.distributed.context_parallel.mamba.MambaContextParallel.pre_conv_ssm(
    projected_states: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor

Redistribute from sequence-sharded to hidden-sharded layout, undoing DualChunkSwap.

class nemo_automodel.components.distributed.context_parallel.mamba._AllToAll()

Bases: Function

Autograd wrapper around torch.distributed.all_to_all_single.

For equal-sized splits the all-to-all operation is its own inverse, so the backward pass is simply another all-to-all on the same group.

nemo_automodel.components.distributed.context_parallel.mamba._AllToAll.backward(
    ctx,
    grad_output: torch.Tensor
)

staticmethod

nemo_automodel.components.distributed.context_parallel.mamba._AllToAll.forward(
    ctx,
    input_: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor

staticmethod

nemo_automodel.components.distributed.context_parallel.mamba._all_to_all(
    input_: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor

Functional entry-point for the autograd-aware all-to-all.

nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_cp2hp(
    input_: torch.Tensor,
    cp_group: torch.distributed.ProcessGroup,
    batch_size: int
) -> torch.Tensor

Transform from sequence-sharded to hidden-sharded layout (batch-first).

Parameters:

input_

torch.Tensor

Tensor of shape [B, L_local, H] (BSHD) or [T, H] (THD) where H is the full hidden dimension on this rank.

cp_group

torch.distributed.ProcessGroup

Context-parallel process group.

batch_size

int

Batch size B (needed to recover dimensions after reshape).

Returns: torch.Tensor

Tensor of shape [B, L_global, H / cp_size] (BSHD) or [T, H / cp_size] (THD).

nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_hp2cp(
    input_: torch.Tensor,
    cp_group: torch.distributed.ProcessGroup,
    batch_size: int
) -> torch.Tensor

Transform from hidden-sharded to sequence-sharded layout (batch-first).

This is the inverse of _all_to_all_cp2hp.

Parameters:

input_

torch.Tensor

Tensor of shape [B, L_global, H_local] (BSHD) or [T, H_local] (THD) where H_local = H / cp_size.

cp_group

torch.distributed.ProcessGroup

Context-parallel process group.

batch_size

int

Batch size B.

Returns: torch.Tensor

Tensor of shape [B, L_local, H] (BSHD) or [T, H] (THD)

nemo_automodel.components.distributed.context_parallel.mamba._deinterleave_packed_seqs(
    input_: torch.Tensor,
    cu_seqlens: torch.Tensor,
    cp_size: int
) -> torch.Tensor

Rearrange tokens from rank-major to sequence-major order after all-to-all.

After _all_to_all_cp2hp on packed 2-D data the token layout along the sequence dimension is::

This function rearranges to::

so that each sequence’s tokens are contiguous (required by the _undo_attention_load_balancing reorder that follows).

Parameters:

input_

torch.Tensor

Tensor of shape [global_tokens, local_hidden], containing CP ranks’ local token blocks consecutively after all-to-all.

cu_seqlens

torch.Tensor

Integer tensor of shape [sequences + 1], containing positive-length local padded sequence boundaries. Starts at zero and ends at global_tokens // cp_size. May reside on CPU or the input device.

cp_size

int

Context-parallel world size; must divide global_tokens.

Returns: torch.Tensor

Tensor of shape [global_tokens, local_hidden], in sequence-major order

nemo_automodel.components.distributed.context_parallel.mamba._redo_attention_load_balancing(
    input_: torch.Tensor,
    cp_size: int,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor

Reorder from sequential back to DualChunkSwap for attention.

Inverse of _undo_attention_load_balancing.

Parameters:

input_

torch.Tensor

Tensor [batch, sequence, hidden], or packed [tokens, hidden].

cp_size

int

Context-parallel world size.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional integer tensor [sequences + 1] of global padded THD boundaries on CPU or the input device. Each document length must be positive and divisible by 2 * cp_size.

Returns: torch.Tensor

Tensor with the input shape, dtype, and device, in DualChunkSwap order

nemo_automodel.components.distributed.context_parallel.mamba._reinterleave_packed_seqs(
    input_: torch.Tensor,
    cu_seqlens: torch.Tensor,
    cp_size: int
) -> torch.Tensor

Inverse of _deinterleave_packed_seqs.

Rearranges from sequence-major back to rank-major order before the inverse all-to-all in post_conv_ssm.

Parameters:

input_

torch.Tensor

Tensor of shape [global_tokens, local_hidden], containing each packed sequence’s CP rank blocks consecutively.

cu_seqlens

torch.Tensor

Integer tensor of shape [sequences + 1], containing positive-length local padded sequence boundaries. Starts at zero and ends at global_tokens // cp_size. May reside on CPU or the input device.

cp_size

int

Context-parallel world size; must divide global_tokens.

Returns: torch.Tensor

Tensor of shape [global_tokens, local_hidden], with all of each CP

nemo_automodel.components.distributed.context_parallel.mamba._reorder_chunks(
    input_: torch.Tensor,
    order: list[int] | torch.Tensor,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor

Reorder equal-sized chunks of a tensor according to order.

Parameters:

input_

torch.Tensor

Tensor of shape [batch, sequence, hidden] (BSHD), or [tokens, hidden] (packed THD) when cu_seqlens is provided.

order

list[int] | torch.Tensor

Permutation of equal-sized chunk indices. For packed THD, also accepts an integer tensor [num_chunks] on the input device to avoid uploading a Python list on each call. BSHD uses a list.

cu_seqlens

torch.Tensor | NoneDefaults to None

Integer tensor of shape [sequences + 1] containing the global padded THD boundaries, starting at zero and ending at tokens. Each sequence length must be positive and divisible by len(order). May reside on CPU or the input device.

Returns: torch.Tensor

Tensor with the input shape, dtype, and device, with chunks permuted

nemo_automodel.components.distributed.context_parallel.mamba._undo_attention_load_balancing(
    input_: torch.Tensor,
    cp_size: int,
    cu_seqlens: torch.Tensor | None = None
) -> torch.Tensor

Reorder from DualChunkSwap to sequential for SSM processing.

Parameters:

input_

torch.Tensor

Tensor [batch, sequence, hidden], or packed [tokens, hidden].

cp_size

int

Context-parallel world size.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional integer tensor [sequences + 1] of global padded THD boundaries on CPU or the input device. Each document length must be positive and divisible by 2 * cp_size.

Returns: torch.Tensor

Tensor with the input shape, dtype, and device, in sequential order