Is this page helpful?

core.ssm.gdp_context_parallel#

Context parallel support for Gated Delta Product (GDP) with num_householder > 1.

The key difference from GDNContextParallel (which assumes a single copy of V/K/b) is that GDP has num_householder copies of V, K, and b (beta). The in_proj output layout is:

[z(d_inner), V(d_inner*M), K(ngroups*d_state*M), Q(ngroups*d_state), b(nheads*M), a(nheads)]

where M = num_householder. Similarly, the conv1d operates on:

[V(d_inner*M), K(ngroups*d_state*M), Q(ngroups*d_state)]

The all-to-all communication and parameter slicing must account for this.

Strategy for householder-multiplied tensors (V, K, b): We fold the M (householder) dimension into the batch dimension before calling the standard all-to-all, then unfold afterward. This ensures each householder copy is independently partitioned by heads across CP ranks.

Module Contents#

Classes#

GDPContextParallel

Context parallel support for Gated Delta Product (GDP) models with num_householder >= 1.

API#

class core.ssm.gdp_context_parallel.GDPContextParallel(
cp_group: torch.distributed.ProcessGroup,
d_inner_local_tp: int,
nheads_local_tp: int,
ngroups_local_tp: int,
d_state: int,
num_householder: int,
headdim: int,
conv1d_cp1: torch.nn.Conv1d,
dt_bias_cp1: torch.Tensor,
A_log_cp1: torch.Tensor,
D_cp1: torch.Tensor,
D_has_hdim: bool,
)#

Context parallel support for Gated Delta Product (GDP) models with num_householder >= 1.

Handles the “all-to-all” CP strategy where heads are partitioned across CP ranks and each rank processes the full sequence for its head partition. Correctly handles the num_householder multiplier on V, K, and beta projections.

Parameters:

  • cp_group – The process group for context parallel.

  • d_inner_local_tp – d_inner on the current TP rank.

  • nheads_local_tp – nheads on the current TP rank.

  • ngroups_local_tp – ngroups on the current TP rank.

  • d_state – SSM state dimension.

  • num_householder – Number of householder reflections (M).

  • headdim – Dimension per head.

  • conv1d_cp1 – The conv1d module for cp_size=1.

  • dt_bias_cp1 – The dt_bias parameter for cp_size=1.

  • A_log_cp1 – The A_log parameter for cp_size=1.

  • D_cp1 – The D parameter for cp_size=1 (can be None).

  • D_has_hdim – Whether D is sized to the hidden dimension.

Initialization

pre_conv_ssm(
input_: torch.Tensor,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
) → torch.Tensor#

All-to-all from sequence-partitioned to head-partitioned layout, before conv + SSM.

Input layout (last dim): [z, V, K, Q, b, a] with sizes [d_inner, d_innerM, ngroupsd_stateM, ngroupsd_state, nheads*M, nheads]

Output layout (last dim, after head partitioning): [z, V, K, Q, b, a] with sizes [d_inner/cp, d_inner/cpM, ngroups_cpd_stateM, ngroups_cpd_state, nheads/cp*M, nheads/cp]

packed_seq_params must be passed for THD/SFT input — without it the post-all-to-all undo uses the non-packed zigzag pattern, which scrambles token order across pack boundaries.

post_conv_ssm(
input_: torch.Tensor,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
) → torch.Tensor#

Method to be applied after the conv + SSM (on y and z, which have no M dim).

conv1d(input_: torch.Tensor) → torch.Tensor#

Performs conv1d using sliced weights for the current CP rank.

conv1d_channels()#

Number of conv channels on the current CP rank.

get_conv1d_weight() → torch.Tensor#

Returns sliced conv1d weight for the current CP rank.

get_conv1d_bias() → torch.Tensor#

Returns sliced conv1d bias for the current CP rank.

get_dt_bias() → torch.Tensor#

Returns sliced dt_bias for the current CP rank.

get_A_log() → torch.Tensor#

Returns sliced A_log for the current CP rank.

get_D() → torch.Tensor#

Returns sliced D for the current CP rank.

_slice_conv_param(param: torch.Tensor) → torch.Tensor#

Slices a cp_size=1 conv1d parameter along the channel dimension, returning the channels needed on the current CP rank.

Conv param layout (dim 0): [V(d_inner * M), K(ngroups * d_state * M), Q(ngroups * d_state)]

For V and K (which have M copies), we reshape to (M, per_copy_channels, …), slice the per-copy channels for this CP rank, then flatten back.

_slice_vector_param(
param: torch.Tensor,
has_hdim: bool = False,
) → torch.Tensor#

Slices a per-head vector parameter (dt_bias, A_log, D) for the current CP rank. These are single-copy (no householder dimension).