zigzag_split_across_group_ranks

zigzag_split_across_group_ranks(data, group, seq_dim=0)

Distributes tensor data across group ranks using zigzag pattern.

Divides the input tensor along sequence dimension and distributes chunks in an alternating pattern across different ranks.

Parameters:

• data – original tensor to split across group ranks.

• group – the group to distribute the data across.

• seq_dim – the sequence/context dimension to split.

Returns:

Tensor slice for the current rank following zigzag distribution.