Synergistic Tensor and Pipeline Parallelism

Thank you for your support! We sincerely thank the reviewer for the valuable feedback and suggestions. Below, we address each point raised in your review.

Despite that the evaluation part includes several pipeline overlapping baselines, the paper omits a comparison with some recent TP overlapping frameworks (Megatron-LM, Flux). It remains unclear whether a combination of TP overlapping and ZB-V will outperform the proposed work?

Thank you for the valuable question. Megatron-LM and Flux improve TP communication primarily at the kernel level by fusing computation and communication. In contrast, our approach is implemented entirely at the software level, without relying on custom kernels, making it more flexible and broadly applicable. Moreover, in Megatron-LM, TP overlapping cannot be enabled independently—it requires coupling with sequence parallelism (SP), which introduces an additional parallel strategy. Similarly, Flux implements only several fused kernels, all-gather+gemm and gemm+reduce-scatter, which require one more parallel to fit these predefined patterns. In contrast, our method is synergistic with widely adopted pipeline and tensor parallelism strategies in large-scale model training and requires no kernel fusion, custom CUDA extensions, or modifications to core operators, enabling easier deployment across diverse architectures. Due to these constraints, we did not combine these frameworks with pipeline parallelism in our main experiments.

Here, we conducted additional experiments combining TP overlap in Megatron-LM with pipeline schedules. The throughput results are summarized below, where * indicates TP overlap kernels are enabled.

In the Qwen2-12B configuration with a sequence length of 4k, our schedule achieves a throughput of 7.27 samples/s, outperforming both 1F1B-I and ZB-V. This represents an improvement of approximately 0.14 samples/s over the best baseline. The results indicate that our synergistic schedule enables more effective cooperation between TP and PP, achieving superior performance compared to simply combining conventional TP overlap techniques with standard pipeline schedules.

Meanwhile, the throughput of Fig. 7 and 8 begins at a high y-axis baseline, which visually exaggerates the perceived speedup.

We appreciate the reviewer’s valuable feedback. We agree that starting the y-axis at a high baseline may exaggerate the visual perception of speedup. To address this, we will adjust the y-axis scale in Figures 7 and 8, ensuring a more accurate and fair visual comparison of throughput improvements.

The method increases peak activation memory compared to some baselines, which may limit applicability in memory-constrained environments if offloading is not viable.

We appreciate the reviewer’s observation. We acknowledge that our method increases activation memory compared to the baselines, as it prioritizes efficiency over memory usage. To assess this trade-off, we conducted an experiment (Table 5 in Appendix) that maximizes memory utilization under a fixed batch size and sequence length. The results show that our schedule achieves the highest throughput (2.74 samples/s), outperforming 1F1B-I (2.72 samples/s) and ZB-V (2.61 samples/s). This suggests that despite a higher memory footprint, our method remains more efficient in practical settings.

How well does the scheduling scheme integrate with other parallel overlapping schemes such as data parallelism and context parallelism? It would be valuable to understand the generality and composability of the approach.

Thank you for the thoughtful question. Our method is compatible with data parallelism and context parallelism, and we have evaluated its performance under these settings. Specifically, we conduct experiments using the Qwen2 12.1B model and report the throughput (samples per second) under different parallelism configurations:

These results demonstrate that our approach integrates well with context parallelism and data parallelism, and provides improved throughput over the baselines in both settings. This supports the generality and composability of our scheduling scheme across multiple parallelism strategies.

We thank the reviewer for the valuable feedback. Below, we provide responses to each point.

Equation 1 is not clear ... Where is $\frac{detach(X)}{t}$ coming from? Attention actually perform?

Thanks for your suggestion. We would like to note that Equation 1 is for a single TP rank.

The standard process (LayerNorm + Attention) before MLP is formulated as:

$X_{ln}=LayerNorm(X),$

$X_{attn}=Attention(X_{ln}),$

$X_{AR}=AR(X_{attn}),$

$Y=X+X_{AR},$

where $Attention$ is to calculate the attention output for a single TP rank, then $AR$ is to get the output for all TP ranks.

Due to the residual operation executing after AllReduce, it introduces a data dependency from LayerNorm, increasing engineering complexity for backward passes. Therefore, we remove this dependency by detaching $X$, which no longer calculates gradients during backward.

$\frac{detach(X)}{t}$ arises from this modification, that the residual operation is removed before AllReduce to make sure it does not stall waiting for AllReduce completion: $Y=AR(X_{attn}+\frac{detach(X)}{t})$, where $t$ is the number of TP ranks. And it keeps computational equivalence in forward as $X=AR(\frac{X}{t})$. In backward, we add $1$ gradient for $X$ to compensate for the detach operation of residual operation.

In Figure 5, ... TP bubbles illustrated clearly ...

Figure 5 is intended to illustrate the pp structure without explicitly visualizing TP bubbles. In our schedule, TP bubbles are implicitly hidden within the overlapped F&W blocks, which indeed have smaller latency than the sum of the individual F and W blocks. We chose to align all blocks in Figure 5 for visual clarity, as misaligned blocks, caused by shrinking the overlapped F&W blocks, could make it harder for readers to trace the communication patterns across stages. That said, we appreciate your feedback, and we will explore improved visualizations that more clearly reflect TP bubble reduction while maintaining readability.

Table 1 ... parameter memory ...

The methods listed in Table 1 are not bidirectional schedules, which have the same parameter memory when not considering the embedding and head layer. Thanks for the suggestion and we will update Table 1 in the final version.

In lines 202-203 where it is stated that activation offloading ...

Thank you for pointing this out. We will revise the sentence to make it more precise and reflective of the conditions under which the statement holds. We also appreciate your suggestion and will carefully review the rest of the paper to ensure similar statements are appropriately qualified.

... there’s no comparison done, for example against Flux ...

There are several works designed to improve TP communication, including Megatron-LM and Flux. Their implementations are primarily at the kernel level by fusing computation and communication. In contrast, our approach is implemented entirely at the software level, without relying on custom kernels, making it more flexible and easier to integrate into various training pipelines without requiring custom kernel development.

Moreover, in Megatron-LM, TP overlapping cannot be enabled independently—it requires coupling with sequence parallelism (SP), which introduces an additional parallel strategy. Similarly, Flux implements only several fused kernels, all-gather+gemm and gemm+reduce-scatter, which require additional parallel to fit these predefined patterns. In contrast, our method is synergistic with widely adopted pipeline and tensor parallelism strategies in large-scale model training and requires no kernel fusion, custom CUDA extensions, or modifications to core operators, enabling easier deployment across diverse architectures. Due to these constraints, we did not combine these works with baselines in our main experiments.

Due to insurmountable technical challenges in integrating Flux kernels into our experimental environment despite extensive efforts, we conducted experiments combining TP overlap in Megatron-LM with pipeline schedules, which embodies a similar overlap as Flux. The results are listed below, where * indicates TP overlap kernels are enabled.

In the Qwen2-12B with a sequence length of 4k, our schedule achieves a throughput of 7.27 samples/s, outperforming 1F1B-I and ZB-V. This represents an improvement of approximately 0.14 samples/s over the best baseline. The results indicate that our synergistic schedule enables more effective cooperation between TP and PP, achieving superior performance compared to simply combining conventional TP overlap techniques with standard pipeline schedules.

... consider a memory apples to apples comparison?

Since different schedules have varying memory footprints, we ensured that all schedules could be executed under the same configuration to enable a fair comparison. This constraint inevitably led to suboptimal memory utilization for some schedules. The primary goal of these experiments was to compare the efficiency of different schedules, rather than to assess their memory usage.

We were also fully aware of this and conducted a related experiment (Table 5) in the appendix. This experiment explored the optimal throughput that can be achieved by each schedule through maximizing the utilization of memory under the given batch size 192 and sequence length 8k. The results showed that our schedule realized the best throughput (2.74 samples/s), better than 1F1B-I (2.72 samples/s) and ZB-V (2.61 samples/s).

What about long context experiments?

We increase the sequence length to 16k and reduce the model dim and layer number to simulate this scenario. And the detailed results are as follows:

As shown in the Table, our schedule outperforms other baselines in the setup of longer context and smaller model scale, even though AR may not be fully overlapped. This indicates that our schedule is more robust to the varying sequence length.

Why are only 1F1B and ZB the baselines? ... are not compared to at all, like Hanayo.

We chose 1F1B-I and ZB-V as baselines because 1F1B-I is a classical and effective method introducing virtual stages, while ZB-V is a more recent method published at NeurIPS 2024. Other methods such as Hanayo are indeed promising; however, they are not open-sourced, and we cannot ensure faithful reproduction, especially regarding communication optimizations they incorporate.

discussion of asynchronous vs synchronous pipeline parallelism ...

We would like to correct a mistake here: ZB-V, cited in our paper, is [1], rather than the reference you mentioned. And asynchronous is not the topic in [1]. Therefore, such an issue is entirely irrelevant to our comparison. Moreover, in practice, synchronous schedules are more popular, considering the potential convergence issue in asynchronous schedules. Therefore, evaluating all methods in a synchronous setting is reasonable.

[1] Qi et al. Pipeline Parallelism with Controllable Memory. NeurIPS 2024.

the position of Fig.1; full form of LayerNorm in Eq. 1; visualization of overlapped F & B in Fig. 5; Activation Memory -> Peak Activation Memory in Tab. 1; throughput's units in Tab. 3; color annotations in Fig. 3; clear discussions in line 272-273; a pointer to Pre-Attn and Pre-MLP in Fig. 2; annotation in Fig. 4

Thanks for your valuable comments! We will revise our paper in the final version according to your suggestions.

... experiments without 2 virtual stages per device?

No! When virtual stages are disabled, our schedule is unavailable as we leverage the virtual stages to realize the overlap.

... perform with context parallelism

Our method is compatible with context parallel. Here, we conduct a corresponding experiment as follows:

As shown, our method achieves a throughput of 2.71 samples/s, outperforming both 1F1B-I (2.64 samples/s) and ZB-V (2.61 samples/s). This demonstrates that our scheduling strategy remains effective and superior even when integrated with context parallelism to handle longer sequences.

Please refer to the rebuttals for Reviewers EjFt and MuRP to get more results when it is integrated with more parallelism and strategies.

... pipeline schedule is of limited novelty from a research perspective ...

While overlapping communication and computation are common in LLM training, our method ingeniously combines TP and PP together without any custom CUDA kernels, making them synergistically work. And all experiments in the paper and appendix consistently demonstrate that our schedule is effective.

Furthermore, the value of our approach has been acknowledged by multiple reviewers. Reviewer EjFt noted that "The proposed scheduling design is practical," Reviewer P8vN highlighted its "solid contributions," stating that "the modification to the interleaved pipelining schedule is simple but effective, and the idea to overlap compute and communication for different microbatches makes sense," and Reviewer MuRP described the scheme as a "thoughtful and non-trivial approach." These comments affirm that our work offers a meaningful and well-justified advancement in the design of efficient training schedules.

Dear Reviewer nDzM

The author-reviewer discussion period will end in August 6. Please read the rebuttal, post any remaining questions or comments, and confirm the Mandatory Acknowledgement.

Best regards, Your AC

Thank you for your support! And we sincerely thank the reviewer for the valuable feedback and suggestions for improved clarity. Below, we address each point raised in your review.

Some of the writing can be made a bit more precise. For example: a) I don't think TP communication significantly grows with TP size as long as you're able to run collectives at the same bandwidth; of course compute-communication ratio decreases.

Sorry for misleading. The original statement in our paper was "The proportion of TP communications grows significantly with the increased TP size". Here, "proportion of TP communications" refers to the ratio of TP communication time to the total time, not the absolute TP communication time.

We agree that if the bandwidth remains constant, the absolute communication time may not grow significantly. However, as TP size increases, the compute-to-communication ratio decreases, leading to TP communication becoming a larger fraction of the overall runtime, which is we intended to highlight.

To avoid misunderstanding, we will revise the sentence to clarify this point in the final version.

b) What does "at the kernel or hardware levels" mean when discussing communication overlapping with computation?

Thanks for your question. The statement of "at the kernel or hardware levels" means that the overlapped communication is implemented based on the well-designed and complex CUDA kernels, such as all-gather+gemm kernel [1] (kernel level), or on some specific hardware, such as track and trigger in [2] (hardware level).

Thanks for your feedback! We will revise this sentence in our final version to make it clearer.

[1] Chang, Li-Wen, et al. Flux: Fast software-based communication overlap on gpus through kernel fusion. arxiv. [2] Pati, Suchita, et al. T3: Transparent tracking & triggering for fine-grained overlap of compute & collectives. ASPLOS 2024.

c) Clarify hardware setup in statements like "where limited inter-device bandwidth can severely degrade training efficiency".

Thank you for the insightful suggestion. We clarify the sentence as follows:

"where limited inter-device bandwidth (e.g., systems without NVLink) can severely degrade training efficiency."

I'm not sure how "synergistic" the modification to the interleaved pipelining schedule is, can't the proposed memory improvements be used for pipeline parallelism in isolation?

Yes! The proposed memory improvements can be used for pipeline parallelism in isolation, but they are not synergistic with the interleaved pipelining schedule actually. The modification of memory is just designed for coping with limited memory scenarios. "Synergistic" here is to reflect in the execution timeline of overlapped F&B block in the schedule, where we interleave the forward and backward passes in the same PP stage to overlap the TP communication in forward and backward passes.

Figure 3 doesn't make full sense to me. For example, in Figure 3a, is the green AR in the communication stream for a different microbatch than the green "Attn (F)"? If they are for the same microbatch, how does "Attn (F)" complete before its communication?

Sorry for the misunderstanding here. In Figure 3a, the green "Attn (F)" and the green AR both correspond to the same microbatch. "Attn (F)" refers solely to the computation of the attention forward pass, while AllReduce (AR) represents the associated communication operation. And those two constitute a complete attention operation. The green AR is initiated only after the corresponding computation (Attn (F)) is completed. Therefore, the forward pass of the attention module is only considered fully complete once both the computation and the communication have finished.

How much is computation slowed down by the overlapped communication (since communication probably uses some SMs that can't be used by GEMMs)?

Thank you for your insightful question. We understand your concern that overlapped communication might occupy certain GPU resources (e.g., SMs), potentially slowing down GEMM computations. To examine the issue, we conduct experiments to measure GEMM performance with and without overlapping AllReduce (AR) operation. The execution time of GEMM is set to be larger than that of AllReduce in Experiment 1, which is opposite in Experiment 2. The execution times are reported as follows:

In Experiment 1, the communication is effectively hidden behind computation, resulting in minimal overhead (approximately 7.5%) for GEMM and substantial time savings compared to sequential execution (GEMM+AR, 11.969ms).

In Experiment 2, although the GEMM operation completes earlier, leaving a portion of the communication time unhidden, the overall execution time remains close to the theoretical lower bound (i.e., the communication durations) and is significantly shorter than sequential execution (1.977 ms).

These results demonstrate that overlap is beneficial in both scenarios, with minimal interference and improved overall efficiency.

Is there a way to apply these ideas to PP=1?

The idea of leveraging the backward passes to overlap TP communication can be valid to a certain extent. Specifically, when PP=1, our method degenerates into pure TP training with gradient accumulation. In this case, where gradient accumulation is enabled, the backward pass for microbatch 1 can be interleaved with the forward pass for microbatch 2. Therefore, the overlap is valid for the majority of microbatches.

What are the throughputs in Figure in TFLOP/s/GPU? That is, how fast on an absolute scale are the baselines and "Ours"?

Thank you for the question. Due to the absence of specific numbers, we present some results in TFLOP/s/GPU of Figures 7 and 8.

As shown in the table, in the TP=4, Seq=6144, PP=4, bs=192 setup, our approach achieves an absolute improvement of 8.7 TFLOPS/s/GPU over 1F1B-I. Furthermore, in the TP=8, Seq=4096, PP=4, bs=256 configuration, we achieve an absolute improvement of 10.1 TFLOPS/s/GPU compared to ZB-V. These results underscore the consistent and significant gains in computational efficiency delivered by our scheduling method across diverse parallel and workload settings.

Dear Reviewer P8vN

The author-reviewer discussion period will end in August 6. Please read the rebuttal, post any remaining questions or comments, and confirm the Mandatory Acknowledgement.

Best regards, Your AC

Thank you for your support! And we sincerely thank the reviewer for the valuable feedback. Below, we address each point raised in your review.

I am not very familiar with pipeline parallelism, so I found some figures (e.g. Figure 3) a bit challenging to fully understand.

Sorry for that. Figure 3 presents the execution timelines of the computation units. "F" and "B" denote the forward and backward passes of these units. Taking the Attn unit as an example, it completes computation in the computation stream and then immediately activates the communication stream to execute the AllReduce (AR) operation. The forward computation in the MLP unit must wait for the completion of that AllReduce due to the inherent data dependency. Therefore, the computation stream would remain idle during the AR execution in the naive implementation.

In our schedule, we interleave the forward and backward passes of the units to reduce idle time. And that AR operation is naturally overlapped with the computation of backward passes in Figure 3 (a). Figure 3 (b) presents the timeline in which the weight gradient computation is separated from the full backward pass. We will revise Figure 3 to make it easy to understand in the final version.

How would the proposed scheduling handle highly imbalanced stages, such as those caused by MoE layers or custom modules?

Thanks for your valuable question. The imbalance among stages indeed negatively affects performance, especially for multimodal models, custom modules, and so on. In this work, we indeed try to address this issue, although it is not emphasized in the paper, as our focus is primarily on the implementation of the proposed PP and TP schedules.

Specifically, we first profile the execution time of the forward and backward passes of each unit, as well as the memory footprint. Then, we design a simple algorithm to split the consecutive unit sequence into balanced PP stages by ensuring that the total forward and backward time per stage is approximately equal and ensuring that the memory footprint of each PP stage stays within memory limits.

Is the system compatible with activation checkpointing strategies commonly used in large-scale models?

Yes! Our schedule is compatible with activation checkpointing strategies. We conduct an experiment on Qwen2-12.1B using activation checkpointing (AC) with a batch size of 128 and a sequence length of 6k, and report the results in the following table:

As shown, disabling AC yields the highest throughput of 4.79 samples per second but requires a peak memory footprint of 56.0 GB. Enabling AC selectively on the MLP modules reduces peak memory by 20.2% to 44.7 GB, with a corresponding throughput decrease of 12.5% to 4.19 samples/s. Extending AC to both Attention and MLP modules further reduces peak memory to 41.5 GB (a 25.9% reduction), while throughput declines to 3.94 samples/s. The most aggressive configuration, applying AC to all modules, achieves the greatest memory savings, reducing peak memory by 35.2% to 36.3 GB, at the cost of a 21.7% reduction in throughput, resulting in 3.75 samples/s. These results confirm that our scheduling framework is fully compatible with activation checkpointing strategies.

Dear Reviewer EjFt,

The author-reviewer discussion period will end in August 6. Please read the rebuttal, post any remaining questions or comments, and confirm the Mandatory Acknowledgement.

Best regards, Your AC

Dear Reviewer EjFt,

Could you please review the rebuttal and engage in a discussion with the authors?

Best regards, Your AC