Ladder Residual: Redefining Tensor Parallelism in Transformers for Accelerated Inference

(as a reference, even using A100 GPU can reach the regime of 2000 tokens/s/GPU, while the highest throughput in this paper is in Figure 3.(1), slightly above 600 tokens/s/GPU). Besides, it would be beneficial to also report the baseline absolute number, since all other latency numbers are reported in the form of a relative improvement.

We previously did all the benchmarks under CudaGraph + flash attention, since with torch compile the communication and computation will not be launched into different streams. Recently we resolved this issue with torch compile new features (which were released several days ago - by setting torch._inductor.config.reorder_for_compute_comm_overlap = True) and were able to benchmark under torch compile. With batch size 64, we are able to obtain 2463 token/sec on one H100 80G GPU. We redo all benchmarking under torch compile and found that while the overall trend is the same, the speedup relative to the standard Transformer is even larger for all methods due to increased latency portion from the communication. The result for 8B can be found below:

Note that in all experiments below, we switch to 512 decoded tokens to make better token/sec comparison with https://lmsys.org/blog/2024-07-25-sglang-llama3/.

---

NVLink=True

See https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md for results under more batch size and tp size, as well as results under NVLink=False

---

These findings show that despite we were previously using a slower benchmarking setting, as the overall trend doesn’t change when we switch to a faster setting, our conclusion in the paper still holds.

The model size used in Figure 2 (8 billion parameters) is very small, compared to the capacity of hardware (up to 8xH100 GPUs). In this case, running the decode operation with a small batch size makes the computation very sparse, and thus the evaluation setup is not representative. …(split into another response below)... It would make the result more convincing to measure the latency/throughput of a larger model (e.g. 65B) or batch size (e.g. 256 or 512)

We did benchmark 70B size in table 1 and we agree with the reviewer that tensor parallelism is more interesting to study on larger models. To provide more comprehensive results, we benchmarked 70B, 405B with various batch sizes and TP sizes below:

---

70B results, NVLink=True:

---

For 405B, even loading the model in bf16 requires > 800GB GPU memory, therefore we only benchmark under TP=16 (2 nodes, each with 8 H100 GPUs) for various batch sizes. Below are results with NVLink=True:

See https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md for results under more batch size, tp size, as well as results without NVLink.

---

Our method (Ladder Residual) still achieves significant speedup for both 70B and 405B and consistently outperforms the parallel attn-mlp baseline on both speed and model quality (shown on smaller size in the paper due to computation constraints). Note that the relative improvement of 70B is smaller than 8B size, since the computation scales faster than communication as model size increases (bits to be communicated scale linearly, while computation can scale quadratically). However, when we use TP=16, where the communication needs to happen across nodes, the improvement is larger as communication is more expensive. We expect these heterogeneous networking settings to become prevalent in the future as models get larger and inference hardware become more diverse.

In Figure 2, for the case Without NVLink, how to explain the reason of a throughput drop of batch size 4 and 16, when the TP world size shifts from 2 to 4?

We reran the code and generated an updated version of Fig. 2(2) with a larger sequence length (512) to confirm the observed performance degradation, as shown in https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/8B%20-%20Without%20NVLink%20-%201024x512%20-%20no_compile.PNG. The performance degradation happens similarly as shown in the paper when batch_size = 16. We hypothesize the following reasons for the degradation when tp-size = 4:

1. Increased Communication Overhead: A larger TP size leads to more communication between GPUs, which becomes costly, especially without NVLink, as seen in the Fig. 2(2) results. This lack of NVLink results in higher latency and more expensive data transfer.
2. GPU Latency Bound: A larger TP size reduces memory I/O and compute load per GPU. Lower utilization prevents GPUs and optimized kernels from operating at their full potential. As a result, the performance benefit from reduced computation volume and latency with tp-size = 4 is outweighed by the increased communication overhead. This leads to diminished overall performance gains.

How to understand the correlation of the 3 columns in Table 2?

We agree with your opinion on the relationship between prefill, decode, and end-to-end speedup and we thank the reviewer for pointing out the weird inconsistency in that entry. After verifying the results, we found that we swapped the data of prefill and decode somehow, the log file output as below (gpt-dense is the standard Transformer): For gpt-dense, we got: prefill: 15.89ms, decode: 172.39ms, token/sec: 169.53; For upper-bound we got: prefill: 10.01ms, decode: 122.94ms, token/sec: 240.16; Thus in the paper, the prefill improvement is 37.00(%) and the decode improvement is 28.69(%).

To double-verify, we rerun NVLINK-UpperBound-Llama-8B with two settings (a. Cuda_Graph - Flash_Attention which is the same as the paper, b. Torch.compile), and found that both results agree with the formula the reviewer concludes. Again, we appreciate the detailed reading of the reviewer that helped us detect this subtle mistake.

It would be better to have a breakup of computation (both MLP and Attention) and communication (both with and without NVL) to help understand the importance of the overlapping, as well as how much overlapped can be, and is achieved with Ladder Transformer.

We agree with the reviewer that providing detailed speedup can provide a better landscape on where the Ladder Residual gained its largest setup, and such results can potentially shape the architecture design. We will do more detailed analysis on this in the future.

When NVLink is disabled, is NVLS still turned on?

Yes, NVLS (SHARP) is still turned on. But based on our experiments, the results are almost the same for NCCL_NVLS_ENABLE=0 and NCCL_NVLS_ENABLE=0.

What is the precision used during the inference?

We use bf16 precision for all the experiments.

What is the memory of a single H100 GPU? For Figure 2, when Batch Size = 64, why there is no number reported for TP=1 and 2

For batch size 64, we run into OOM with the models in tp size = 1, 2 cause we were using cuda_graph + flash attention for benchmarking. When we generate static cuda graph for our 8B model, we observe high memory consumption and OOM as below:

[rank0]:     static_next_token = prefill(model, static_x.view(batch_size, -1), static_input_pos, **sampling_kwargs)
[rank0]:     return func(*args, **kwargs)
[rank0]:     logits = model(x, input_pos)
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 15.66 GiB. GPU 0 has a total capacity of 79.11 GiB of which 4.97 GiB is free. Including non-PyTorch memory, this process has 74.13 GiB memory in use. Of the allocated memory 41.05 GiB is allocated by PyTorch, with 13.52 GiB allocated in private pools (e.g., CUDA Graphs), and 32.02 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  

With our new setup under torch.compile, we can run BS=64 with all the TP sizes, as expected from the reviewer’s calculation. Results can be found at https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/8B%20-%20With%20NVLink%20-%201024x512%20-%20torch_compile.PNG

For results in Table 3, it seems like Parallel Transformer is catching up with Ladder Transformer (Average from -0.17 to +0.12, Wikitext PPL from -0.53 to -0.06), as the number of parameters scales up from 1B to 3B, this raises concerns on the scalability of Ladder Transformer. Providing additional experiments with a larger model size or an explanation for why this trend might not continue could be helpful in supporting the significance of the Ladder Transformer.

It is difficult to conclude the scaling trend from two data points (1B and 3B) and we agree with the reviewer that pretrain more models would make this clear. However, just the 3B pretraining run already costs 1228 H100 hours for us, and pretraining a larger model is beyond our computation budgets at this point.

Although Ladder Transformer is friendly to Tensor Parallel, it is not friendly to Pipeline Parallel, which is also necessary during the training process and the inference of large models, because both prev_attn_out and prev_mlp_out should be sent from a pipeline stage to the other. This prevents the Ladder Transformer from scaling up to a larger size. It would make the contribution more solid if the authors can discuss how to address this concern or why it may not be critical.

It is possible to use Pipeline Parallel (PP) with the Ladder architecture. There are 2 ways in which PP can be used for Ladder Transformer model: Just before the pipeline boundary, we wait for the async AllReduces to complete, and send 3 tensors to the next pipeline stage: residual tensor, current_mlp_out tensor and current_attention_out tensor. It should be noted that this is still pretty cheap since generally the P2P communication is latency bound and can be implemented easily using the batch_isend_irecv API (https://pytorch.org/docs/stable/distributed.html#torch.distributed.batch_isend_irecv).

No PP:

def forward(
   self,
   previous_attention_out: Tensor,
   previous_mlp_out: Tensor,
   residual: Tensor,
   attention_handle,
   mlp_handle,
) -> Tensor:
   attention_handle.wait()
   residual = residual + previous_attention_out


   current_attention_out = self.attention(self.attention_norm(residual))
   current_attention_out = all_reduce(current_attention_out, async_op=True)


   mlp_handle.wait()
   residual = residual + previous_mlp_out


   current_mlp_out = self.feed_forward(self.ffn_norm(residual))
   current_mlp_out, mlp_handle = all_reduce(current_mlp_out, async_op=True)


   return current_attention_out, current_mlp_out, residual, attention_handle, mlp_handle

PP:

def forward(
   self,
   previous_attention_out: Tensor,
   previous_mlp_out: Tensor,
   residual: Tensor,
   attention_handle,
   mlp_handle,
) -> Tensor:
   attention_handle.wait()
   residual = residual + previous_attention_out


   current_attention_out = self.attention(self.attention_norm(residual))
   current_attention_out = all_reduce(current_attention_out, async_op=True)


   mlp_handle.wait()
   residual = residual + previous_mlp_out


   current_mlp_out = self.feed_forward(self.ffn_norm(residual))
   current_mlp_out, mlp_handle = all_reduce(current_mlp_out, async_op=True)


   if is_last_layer_on_pp_stage:
       attention_handle.wait()
       mlp_handle.wait()


   return current_attention_out, current_mlp_out, residual, attention_handle, mlp_handle

We thank the reviewer's time for responding to our rebuttals. Here are our responses for further addressing the concerns:

The file https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md does not provide me any data. The raw format only has three paragraphs of pure text;

The anonymous link we linked to is an anonymous GitHub repo, you should be able to find other files in it. Either on the left-hand side or the top side there should be a list of files in the repo and those are our additional benchmarking results. Please let us know if you still cannot find it though!

I assume that bs=4,tp=1 is actually bs=4,tp=4 because that pattern applies for all other batch sizes, and tp=1 does not make sense to any speedup (since you don't have communication in that case);

Yes bs=4, tp=4 is correct, thanks for catching that.

My concern on using too many GPUs for a small model and small batch size is still unsolved;

Using too many GPUs for a small model can generally come into play when the small model (8B for example) is being used as a speculator for a larger model (say 405B). In this scenario, since the larger model is already running on a large number of GPUs, it makes sense to run the small model on a large number of GPUs as well. In the meantime, we also show speedup on larger models.

Is there a range of setup including (model size, batch size, TP, ...), where LadderTransformer is significantly better than Parallel Attn? If so, how to prove the importance of that domain;

From https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/70B%20-%20With%20NVLink%20-%201024x512%20-%20torch_compile.PNG, we are consistently better than the parallel attn/mlp baseline on 70B model with TP=8 from bs=1 to bs=32. Indeed we observe a smaller gap as the batch size increases (9.05% more speedup than parallel attn/mlp with bs=1, to 4.81% with bs=32), this is a current limitation of our method, but in theory we shouldn’t observe this much degradation. We are planning to investigate this in the future.

Adding experiments for (8B, bs=128,256...) could help with my concern on the fact that computation is too sparse;

We are running these experiments. From https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/8B%20-%20With%20NVLink%20-%201024x512%20-%20torch_compile.PNG we can also infer that at bs > 128, ladder-residual is likely to only have similar speedup as parallel attn/mlp, but we will provide results for completeness.

We want to note that parallel attn/mlp as an alternative approach isn’t what the community has adopted, and in our experiment, we found that ladder-residual performs better in terms of model accuracy on 1B size and similar on 3B size compare with the parallel attn/mlp architecture. Ladder Residual shows significant speedup against the standard Transformer, and a modest gain over an alternative approach. It also has the potential to be directly converted from a pre-trained standard Transformer. We believe these results make Ladder Residual a better alternative worth being recognized.

It could be better to have a discussion and comparison to the model comparison scheme, which can also serve the LLM efficiently with high accuracy.

Lots of post-training model compression techniques have been proposed for LLM, for example N:M sparsity, quantization, or depth-wise pruning. These techniques are orthogonal to our proposed methods and can be applied in combination. Our contribution is a modification of computation flow that makes tensor parallelism more efficient, and does not negatively interfere with other techniques. We leave for the future work to study the performance impact of applying compression techniques on our Ladder Residual architecture.

I suggest having some discussion and comparison to quantization. For example, the int4 weight quantization can make a 70B model running on a single H100 GPU, without losing too much accuracy.

Given the ever increasing model capacity, quantization is a very viable option however, even with quantization, it’s impossible to run larger models like Llama 405B on a single H100 80GB GPU. Using tensor parallelism to run large models is still an approach the community actively uses and our contribution is to make it faster. Our work is orthogonal to quantization and it should be noted that our proposed model architecture can be easily combined with current SOTA quantization approaches.

Having a comparison to the pipeline parallelism can also better demonstrate the contribution of this paper. Note some recent studies have shown good performance of pipeline parallelism than the pure TP (https://blog.vllm.ai/2024/07/23/llama31.html).

In the post, a combination of pipeline parallelism (PP) and tensor parallelism (TP) are used to run 405B models on 16 GPUs. Our method makes the TP portion faster while still being compatible with PP. Due to the expensive cross-node communication, TP is usually not used across nodes, but our methods provide an opportunity to reconsider the trade-off and potentially scale TP to multiple nodes. The best practice of multi-dimensional parallelism under different scenarios is still a resolved question for researchers and practitioners, and as long as TP is still being used, our method can provide efficiency gain.

It uses up to 8B model for the motivation illustration and evaluation. However, it is less likely to use 8 H100 GPUs to inference the small 8B model in practice.

We have benchmarked 70B size in table 1 under the setup of 1024 prompt length, 256 generated token, TP=8 and batch size 1, and observed 17% tokens/second improvement and 30% without NVLink.

In the response below, we benchmarked 70B and 405B more thoroughly to demonstrate the effectiveness of Ladder Residual at larger model size.

Why not conduct the experiment on 70B models? I believe this can be a better model that worth the multi-GPU serving. But given the high compute-intensity of the 70B model, the ratio of the AllReduce overhead can be smaller than 8B, which can be a concern to the motivation of this paper.

To provide a more comprehensive results, beyond the one setup in Table 1 that shows the effectiveness of our method on 70B model, we additionally benchmarked it under various batch sizes . As shown in https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/70B%20-%20With%20NVLink%20-%201024x32%20-%20no_compile.PNG, With NVLink, the speedup on 70B is lower due to higher compute-intensity as the reviewer points out, but is still significant and our method outperforms the baseline. For example, the improvement goes from 1.297x in 8B to 1.188x in the batch size 1 case with 32 decoded tokens.

However, if we increase decode tokens to 512 and disabled NVLink P2P communication, the communication volume is larger and sometimes outweighs the downward trend going from 8B to 70B due to increasing compute-intensity, leading to higher speedup. Below for batch size 4, at TP=4 8B observes a larger speedup while at TP=8 70B has a larger speedup.

---

---

To enhance our experimental results, we also provide benchmarking results with 512 decoding tokens on 70B and 405B with various batch sizes and TP sizes as below. Here we benchmark with torch.compile as it is more memory efficient than our previous benchmarking setting (there is a recent feature that allows us to benchmark with compile). The compile setup brings higher speedup due to reduced non-communication overhead but the overall trend stays the same (please refer to our response to Reviewer Jbio for more detail).

---

70B results, NVLink=True:

---

For 405B, even loading the model in bf16 requires > 800GB GPU memory, therefore we only benchmark under TP=16 (2 nodes, each with 8 H100 GPUs) for various batch sizes. Below are results with NVLink=True:

See https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md for results under more batch size, tp size, as well as results without NVLink.

---

Our method (Ladder Residual) still achieves significant speedup for both 70B and 405B and consistently outperforms the parallel attn-mlp baseline on both speed and model quality (shown on smaller size in the paper due to computation constraints). Note that when we use TP=16, where the communication needs to happen across nodes, the speedup is larger as communication is more expensive, and this leads to much larger speedup despite the increased compute-intensity from 70B to 405B. As we have seen that models are getting larger and larger, we expect that these techniques that decouple computation and communication will become more important.

why are the results consistently better without nvlink than with nvlink, across all three baselines? I would assume Megatron-style TP would perform better with NVlink.

We are reporting the relative speedup instead of the absolute speed in the paper. It’s correct that without NVLink, we would observe worse absolute speed. Since without NVLink the communication is slower, our technique that overlaps communication with computation is expected to lead to a larger relative improvement than with NVLink as we reported.

In Fig 2(2), why is there a performance degradation when TP world size =4?

We reran the code and generated an updated version of Fig. 2(2) with a larger sequence length (512) to confirm the observed performance degradation, as shown in https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/8B%20-%20Without%20NVLink%20-%201024x512%20-%20no_compile.PNG. The performance degradation happens similarly as shown in the paper when batch_size = 16. We hypothesize the following reasons for the degradation when tp-size = 4:

1. Increased Communication Overhead: A larger TP size leads to more communication between GPUs, which becomes costly, especially without NVLink, as seen in the Fig. 2(2) results. This lack of NVLink results in higher latency and more expensive data transfer.
2. GPU Latency Bound: A larger TP size reduces memory I/O and compute load per GPU. Lower utilization prevents GPUs and optimized kernels from operating at their full potential. As a result, the performance benefit from reduced computation volume and latency with tp-size = 4 is outweighed by the increased communication overhead. This leads to diminished overall performance gains.

Is the baseline "standard transformer" using data parallelism or Megatron-style tensor parallelism?

The baseline transformer uses standard Megatron-style Tensor Parallelism.

The evaluation sections are restricted to models with relatively small sizes (up to 8B).

For efficiency, we did benchmark 70B size in table 1 under one setup and we agree with the reviewer that tensor parallelism is more interesting to study on larger models. To provide more comprehensive results, we benchmarked 70B, 405B with various batch sizes and TP sizes below, all with NVLink:

---

70B speedup results:

---

For 405B, even loading the model in bf16 requires > 800GB GPU memory, therefore we only benchmark under TP=16 (2 nodes, each with 8 H100 GPUs) for various batch sizes.

See https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md for results under more batch size, tp size, as well as results without NVLink.

---

Our method (Ladder Residual) still achieves significant speedup for both 70B and 405B and consistently outperforms the parallel attn-mlp baseline on both speed and model quality (shown on smaller size in the paper due to computation constraints). Note that the relative improvement of 70B is smaller than 8B size, since the computation scales faster than communication as model size increases (bits to be communicated scale linearly, while computation can scale quadratically). However, when we use TP=16, where the communication needs to happen across nodes, the improvement is larger as communication is more expensive. As we have seen that models are getting larger and larger, we expect that these techniques that decouple computation and communication will become more important.

We are currently running the experiment of adapting llama-3.1-70B-Instruct to Ladder Residual and will update the results if we are able to finish it within the rebuttal period.

Implementing Ladder Residual may require substantial retraining or adaptation efforts for existing models, which may be challenging for larger models;

With only 2 epochs of supervised finetuning on 1.6 billion tokens (1-2 days on an H100 node with an off-the-shelf finetuning library), we can recover the performance of llama-3.1-8b-instruct while gaining 14.5% speedup, which makes it more efficient serve the model for months. This is significantly cheaper than the trillion-tokens scale pre-training cost. (we could not find the report on post-training token count for Llama-3, but it is likely more costly than our setup, and we don’t require human preference data or RL training). Given that inference is starting to account for a significant fraction of compute, we think this is a favorable tradeoff. More importantly, the new architecture offers much more flexibility in terms of hardware to serve the model as we have decoupled the computation and communication (NVLink). We are excited to see how this changes the networking required for LLM inference.

What are the considerations when deciding to apply ladder-residual on the upper half (or later half) of the transformer layers?

We found that it’s difficult to apply ladder-residual on all the layers for a pre-trained LLM, and used the zero-shot result to decide which layers to apply. Our hypothesis is that a lot of knowledge is stored within the lower half of the layer, and with a light-weight retraining (1.6B token in our case), we can’t expect the fine-tuning dataset to contain the lost knowledge. As the field as a whole develops more understanding on how knowledge is stored in these models, we expect it would be easier to apply these architectural changes to pretrained models.

First, the paper only evaluates with TP size up to 8. However, the paper has no evaluation when scale up to more than 8 GPUs. It is thus obscure how the performance of the proposed technique when scaling up to multiple GPU nodes.

We only benchmarked TP size up to 8 in the paper since cross-node TP becomes expensive and we didn’t get time to evaluate that. Since our method hides the communication latency behind computation, it allows us to reconsider cross-node TP as a potential approach for a very large model.

For 405B, even loading the model in bf16 requires > 800GB GPU memory, therefore > 8 GPU is needed. Here we benchmarked llama-3.1-405B under TP=16 (2 nodes, each with 8 H100 GPUs) and obtain a larger speedup compare with 8B or 70B as the saving from communication becomes a larger portion of the end-to-end latency. Our method still outperforms the parallel attn/mlp baseline across batch size while the model accuracy is higher on the smaller models we tested (training a 405B model is out of our computation budget).

---

---

See https://anonymous.4open.science/r/ICLR2025_rebuttal-B81D/README.md for more details and compare with 8B and 70B results

Second, the paper has no consideration how the proposed technique behave when using other parallelisms. While the proposed architecture should work orthogonally with other parallelisms, it would be much speedup the proposed method could provide when used in conjunction with other parallelisms.

One of the advantages of our proposed architecture is that it solely modifies the input/output of each module and therefore can be combined with other parallelism as using them on a standard transformer. Pipeline Parallelism needs a bit special treatment, which can be combined with our proposed model architecture by synchronizing before the pipeline boundary and sending the tensors to the next pipeline stage (please refer to our response to Reviewer Jbio for the detail). The expected speedup should be the same as if we apply them on a standard Transformer.