mstar.model.components.distributed.sequence_parallel

mstar.model.components.distributed.sequence_parallel#

Ulysses sequence parallelism — model-agnostic primitives.

Sequence parallelism (SP) shards the token sequence across an sp_group and replicates the weights, the mirror image of tensor parallelism (which shards the weights and replicates the sequence). Attention is the only operator that needs the whole sequence, so Ulysses brackets it with two all-to-alls: the first turns a sequence-sharded [seq/P, heads, dim] tensor into a head-sharded [seq, heads/P, dim] one (full sequence, this rank’s slice of heads), the attention runs locally, and the second converts back. Net effect: attention runs exactly as if at tensor-parallel degree tp*sp — the kernel (run_attention) is unchanged — while every pointwise op (norms, MLP, residuals) runs on the seq/P shard for free.

These helpers take an sp_group (a CommGroup over the SP axis of the mesh; trivial when sp_size == 1) and are model-independent: any attention that projects to [tokens, heads, head_dim] and calls a run_attention(q, k, v) can use ulysses_attention(). seq_sizes is the per-rank token count along the sequence (need not be equal — sequences indivisible by sp are handled without padding); the scatter/gather of the residual stream at the model boundary uses the same split.

Functions

gather_sequence(sp_group, x_shard, seq_sizes)

All-gather the sequence shards back into the full tensor (rank order).

scatter_sequence(sp_group, x_full, seq_sizes)

Return this rank's contiguous slice of a sequence-replicated tensor.

sp_head_gather(sp_group, x)

Reassemble full TP-local heads from per-rank head-groups (the inverse of sp_head_slice()): [T, H/P, D] -> [T, H, D] via an all-gather over heads, restoring the layout the row-parallel output projection expects.

sp_head_slice(sp_group, x)

Keep this SP rank's contiguous head-group: [T, H, D] -> [T, H/P, D].

sp_seq_split(total, world_size)

Even per-rank token counts summing to total (earlier ranks get the remainder).

ulysses_attention(sp_group, q, k, v, ...[, ...])

Run attention under Ulysses SP.
mstar.model.components.distributed.sequence_parallel.gather_sequence(sp_group, x_shard, seq_sizes, dim=0)[source]#

All-gather the sequence shards back into the full tensor (rank order).

Pads each shard to the max per-rank length so the underlying collective is a native equal-size all-gather, then trims — robust to sequences not divisible by the group size. For the common even split it is a plain all-gather.

Parameters:
Return type:

Tensor

mstar.model.components.distributed.sequence_parallel.scatter_sequence(sp_group, x_full, seq_sizes, dim=0)[source]#

Return this rank’s contiguous slice of a sequence-replicated tensor. No communication — every SP rank holds the identical x_full (the denoise latent is replicated), so each simply narrows to its seq_sizes window.

Parameters:
Return type:

Tensor

mstar.model.components.distributed.sequence_parallel.sp_head_gather(sp_group, x)[source]#

Reassemble full TP-local heads from per-rank head-groups (the inverse of sp_head_slice()): [T, H/P, D] -> [T, H, D] via an all-gather over heads, restoring the layout the row-parallel output projection expects.

Parameters:

  • sp_group (CommGroup)

  • x (Tensor)

Return type:

Tensor

mstar.model.components.distributed.sequence_parallel.sp_head_slice(sp_group, x)[source]#

Keep this SP rank’s contiguous head-group: [T, H, D] -> [T, H/P, D].

The selected heads [rank*H/P : (rank+1)*H/P] are exactly those that ulysses_attention()’s all-to-all routes to this rank, so a tensor sliced here (e.g. the replicated UND prefix K/V) lands on the same head partition as the sequence-parallel GEN attention.

Parameters:

  • sp_group (CommGroup)

  • x (Tensor)

Return type:

Tensor

mstar.model.components.distributed.sequence_parallel.sp_seq_split(total, world_size)[source]#

Even per-rank token counts summing to total (earlier ranks get the remainder). The sequence need not be divisible by world_size.

Parameters:

  • total (int)

  • world_size (int)

Return type:

list[int]

mstar.model.components.distributed.sequence_parallel.ulysses_attention(sp_group, q, k, v, run_attention, seq_sizes, prefer_all_gather=False)[source]#

Run attention under Ulysses SP.

q: [seq/P, Hq, D]; k/v: [seq/P, Hkv, D] (this rank’s tokens, TP-local heads). Returns [seq/P, Hq, D]. The head counts must be divisible by sp (Ulysses shards heads). When the group is trivial this is a passthrough — byte-identical to the non-SP path.

prefer_all_gather selects the all-gather collective instead of the all-to-all (see _ulysses_attention_via_all_gather()). The caller sets it on the denoise forward that the CUDA graph captures: the all-to-all is grouped point-to-point send/recv and does not replay from a captured graph, whereas all-gather (a true collective, like the TP all-reduce) does. It must be set consistently across warmup, capture and replay so the all-gather kernels are compiled and autotuned during eager warmup, not mid-capture (autotuning synchronizes, which is illegal while a graph is recording). Eager paths (video, uncaptured resolutions) leave it off for the lighter all-to-all. Both produce identical results.

Parameters:
Return type:

Tensor