DSP: Dynamic Sequence Parallelism for Multi-Dimensional Transformers

We sincerely thank the reviewer o73Q for the valuable questions and comments. For the concerns and questions, here are our responses:

Q1: Performance gains are not convincingly demonstrated for long sequences

A1: We thank the reviewer for this important question and have conducted additional experiments to robustly demonstrate DSP’s performance on long sequences.

Experiment1: Communication latency comparison.

Analysis1: DSP consistently reduces communication latency across sequence lengths.

Experiment2: End-to-end speedup for 8M long sequecnes with two seuqence shapes A:[8192, 1024] and B:[8, 1048576].

Analysis2: DSP significantly reduces communication time for both shapes. However, end-to-end speedup varies: shape A benefits more due to lower computation time, while shape B’s longer single dimension increases attention computation, dominating total time. This shows DSP’s communication efficiency, though overall gains depend on computation complexity.

Improvement plan: We have added a detailed discussion of these results in the Appendix of the revised manuscript.

Conclusion:

1. DSP delivers consistent communication reductions across sequence.
2. End-to-end speedup depends on computation time, which varies with sequence dimensionality.

Q2: Is it possible to extend it to other models other than multi-dimensional transformers?

A2: DSP’s simple yet general design extends beyond multi-dimensional transformers to various multi-dimensional neural networks, including conv+attention models (e.g., 2D-UNet, 2D-VAE), AlphaFold, and pure conv.

For example, in 2D-UNet, DSP enables efficient resharding between computations along different dimensions:

$x -> conv\\_dim(1) -> attn\\_dim(1) -> DSP\\ reshard -> conv\\_dim(2) -> attn\\_dim(2) -> y$

Q3: The paper should include more comprehensive experiments on diverse multi-dimensional transformer architectures.

A3: We thank the reviewer for raising this point. Main stream multi-dimensional transformer architectures typically fall into two categories:

1. Dimension-by-dimension processing as shown in our paper:

$x -> transformer\\_dim1(x) -> transformer\\_dim2(x) -> y$

2. Dimension-by-dimension with cross-dimensional interactions:

$x -> transformer\\_dim1(x) -> transformer\\_dim2(x) -> transformer\\_cross\\_dims(x) -> y$

Experiment: End-to-end performance on the second architectures on Transformer-2D 3B with 8 GPUs.

There many be some changes within the transformer (e.g., layernorm, positional embedding) but they don't affect much for the results.

Improvement plan: We have added a detailed discussion of these results in the Appendix of the revised manuscript.

Conclusion: Our method is able to deliver stable performance on main-stream multi-dimensional transformer architectures.

Q4: The authors should provide a more thorough analysis of the trade-off between performance and accuracy.

A4: DSP is designed to preserve accuracy strictly, introducing only negligible errors (1e-7 to 1e-8) due to communication. For example, in video generation tasks, DSP outputs are identical to the originals.

Q5: How can DSP be further optimized to provide more substantial speedups in scenarios with small overhead from baseline methods?

A5: While DSP already minimizes communication volume effectively, further speedups are possible by overlapping communication with computation. DSP’s lightweight communication and decoupling from computation make it ideal for asynchronous execution, hiding latency behind computation.

Q6: Essential references not discussed (GSPMD)

A6: We appreciate the reviewer highlighting this point.

Improvement plan: In the revised manuscript, we have added a detailed discussion of GSPMD and other relevant works in the related work section, ensuring a more comprehensive literature review.

We sincerely thank the reviewer qVhs for the valuable questions and comments, especially about the evalutions which are indeed not detailed enough.

We have conducted the following experiments to make our evalutions more comprehensive and have added them to the experiment section or appendix or in the latest version.

Q1: DSP's scalability according to the model size.

A1: We conduct an end-to-end analysis for DSP with 3B, 7B and 13B using 1M tokens.

Analysis:

• When scale from 3B to 7B, although the communication cost increase, the performance grows because the compuation density increases more.
• When scale from 7B to 13B, the benefits of denser computation is marginal. The total throughput decrease because of more communication cost, but the loss is still marginal.

Q2: Breakdown analysis of communication overhead.

A2: We breakdown for communication overhead/actual computation for DSP under weak scaling condition using Transformer-2D 3B:

Analysis: The communication costs grows significantly for inter-node communication even for our method. Highlight the necessity of efficient parallelism.

If you need more experiments for each parallel or model, please tell us!

Q3: Cost of layout change/reshape/transpose.

A3: Let's analyze from memory and speed:

For memory, this usually will not cause much fragmentation as we show in Figure 9 (Line 447) because of the following reasons:

• The sequence have already been splitted. Therefore there will be only a part of sequence on one device, reducing the overhead to layout change.
• DSP changes actually less frequent than any other methods, lead to less overhead.
• Duing training, the layout of a sinlge sequence does not affect overall memory cost because parameter and activation is much larger than it.

To further clarify this, we make experiments about the memory overhead of our method:

For speed, we conduct an experiment to evaluate the time of changing layout and communication:

Analysis:

• The layout changing time is much less than the communication time. But still takes 5.2% of total time.

• As the sequence and gpu number get larger, the layout change time become less as communication takes more time.

• To fully elimiate the layout change's time, we can overlap communication with layout change.

• We use activation checkpointing in training.

Conclusion:

• Memory: The layout changes have little impact on memory cost.
• Speed: The layout changes will occupy 4.1%-5.2% of communication time, but can be avoided by overlapping. And it will reduce as sequence and gpu becomes larger.

We sincerely thank the reviewer Ux6h for the valuable questions and comments. For the concerns and questions, here are our responses.

Q1: The experiments only use H100 GPUs, leaving performance on other GPUs unknown.

A1: Thanks for pointing out this! We test the performance on A100 GPUs.

Experiment1: End-to-end performance on A100 on Transformer-2D 3B with 8 GPUs.

Experiment2: Runtime breakdown for A100 and H100 on Transformer-2D 3B with 16 GPUs with DSP.

Analysis:

1. The performance improvement slightly decrease for A100 (from 23.5% to 22.7% compared with Deepspeed-Ulysses).
2. The change is because A100 takes less time ratio to communicate than H100. Because H100 significantly improves computation capbility, therefore taking more time ratio to communicate than A100, making our method more effective.

Improvement plan: We have added theroritcal analysis above to the Appendix of our work.

Conclusion: For various GPUs, the more time communication takes, the more speedup our work can achieve.

Q2: Lacks theoretical analysis of complexity, memory, and communication costs.

A2: Thanks for pointing out this problem.

Communication costs:

Therefore, the communication volume for each parallel method is

Memory cost:

For memory, our method usually will not cause much extra memory cost as we show in Figure 9 (Line 447) because of the following reasons:

• The sequence have already been splitted. Therefore there will be only a part of sequence on one device, reducing the overhead to layout change.
• DSP changes actually less frequent than any other methods, lead to less overhead.
• Duing training, the layout of a sinlge sequence does not affect overall memory cost because parameter and activation is much larger than it.
• We use activation checkpointing during training.

Experimental analysis of memory cost:

Theoretical analysis of memory cost:

In training, everytime we create/change a tensor, it will make a copy, which causes extra memory cost. There are two reshape before and after communication. So let the tensor size be M, we will have activation memory cost of 4M for one communication.

In 2D-Transformer, every layer has activation memory for communication of 4M size as mentioned above, and other components of 42M. If we consider activation memory per layer only, the memory overhead is 8%.

But since we use activation checkpoint, the actual memory cost of our method during training will be Communication Activation / (Communication Activation + Other Activation + Activation Checkpointing * Num Layers + Parameter + Gradient + Optimizer), which will be approximately less than 1%. Taking the settings of Transformer-2D 720M into this formula, the memory overhead is 0.67%.

Improvement plan: We have added theroritcal analysis above to the Appendix of our work.

We sincerely thank the reviewer 4hHH for the valuable comments and acknowledgement of our work.