Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference
Authors: Rohan Baskar Prabhakar
Hengrui Zhang
David Wentzlaff
Venue: NeurIPS 2024 poster (2024)
Keywords: Neural Networks
Primary area: deep_learning_architectures
TL;DR: New Transformer architecture that allows overlapping communication collectives with compute to speed up inference.

Abstract: Large Transformer networks are increasingly used in settings where low inference latency is necessary to enable new applications and improve the end-user experience.
However, autoregressive inference is resource intensive and requires parallelism for efficiency.
Parallelism introduces collective communication that is both expensive and represents a phase when hardware resources are underutilized.
Towards mitigating this, Kraken is an evolution of the standard Transformer architecture that is designed to complement existing tensor parallelism schemes for efficient inference on multi-device systems.
By introducing a fixed degree of intra-layer model parallelism, the architecture allows collective operations to be overlapped with compute, decreasing latency and increasing hardware utilization.
When trained on OpenWebText, Kraken models reach a similar perplexity as standard Transformers while also preserving their language modeling capabilities as evaluated on the SuperGLUE benchmark.
Importantly, when tested on multi-GPU systems using TensorRT-LLM engines, Kraken speeds up Time To First Token by a mean of 35.6% across a range of model sizes, context lengths, and degrees of tensor parallelism.

Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference
Rohan Baskar Prabhakar Princeton University rohanbp @ princeton.ecu
Hengrui Zhang
Princeton University
hengrui zhang @rincetooneduu
David Wentzlaff
Princeton University
wentzlaf@ princetooopnn
Abstract
Large Transformer networks are increasingly used in settings where low infer- ence latency is necessary to enable new applications and improve the end-user experience. However, autoregressive inference resource intensive and requires parallelism for efficiency. Parallelism introduces collective communication that is both expensive and represents a phase when hardware resources are underuti- lized. Towards mitigating this, Kraken is evolution of the standard Transformer architecture that is designed to complement existing tensor parallelism schemes for efficient inference multi-device systems. By introducing fixed degree of intra-layer model parallelism, the architecture allows collective operations to be overlapped with compute, decreasing latency and increasing hardware utilization. When trained on Open WebText, Kraken models reach a similar perplexity as stan- dard Transformers while also preserving their language modeling capabilities as evaluated on the SuperGLUE benchmark. Importantly, when tested on multi-GPU systems using TensorRT-LLM engines, Kraken speeds up Time To First Token by mean of 35.6% across range of model sizes, context lengths, and degrees of tensor parallelism.
Introduction
Deep neural networks based the Transformer architecture (51) have become the prevalent choice for a variety tasks involving sequences, especially in natural language processing and computer vision (44) (7) (19). Their capabilities, particularly in language modeling, have been driven by rapid increase in parameter count (7). Today's largest language models have to a trillion parameters (21) and consequently demand more efficiency from the systems used to train and serve them. This has necessitated the need for many techniques and optimizations that focus on improving the performance of both algorithms and systems (47)
Large models are often used in interactive applications where latency is an important metric that dictates the quality of the end-user experience (28). A typical web search takes about 0.2 seconds but the Time To First Token (TTFT) for large models can be up to a few seconds (depending on context length, model size, and available hardware) (42). Additionally, because it is not always feasible to run models on local hardware, they are served to users via datacenters that use multi-device compute nodes, adding to latency constraints. Increasingly, language models also used as intermediate steps in longer processes such as augmenting web searches or presenting the results of database queries (88) The rising prevalence of such multi-step applications makes reducing inference latency even more critical.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Continuing this theme, this work focuses on reducing the latency cost of the collective operations in- troduced by tensor parallelism (49) the forward pass. In particular, it introduces Kraken, a variation of the standard Transformer architecture (7 that reduces the amount of inter-device communi- cation and allows remaining collective operators to be overlapped with compute. Kraken models have a fixed degree of innate model parallelism that allows computational graphs each device to run independently without having to wait for the results of collective operations. The architecture is designed to complement the topology of multi-device setups such nodes in typical datacenters and DGX (33 systems. By designing the model architecture to account for characteristics of the hardware, our approach increases compute utilization and allows more efficient inference.
We evaluate the improvements Kraken offers over standard Transformers in two key aspects: model quality and inference latency. For the former, we train series of Kraken models with varying degrees of parallelism and parameter count on Open WebText (3) and compare them with the GPT-2 (44) family of models on the SuperGLUE suite of benchmarks (53) We then implement Kraken using the TensorkT-LLMM library (15) and measure the Time To First Token (TTFT) given various model sizes and context lengths illustrate the efficiency gains when collective operators are longer on the critical path. We find that while maintaining the language modeling capabilities of standard Transformers, Kraken models speedup the Time To First Token (TTFT) by geomean of 35.6% when tested across a range of model sizes, context lengths, and degrees of parallelism.
2 Background
2.1 Decoder-Only Transformer models
We will briefly discuss the forward pass of decoder-only Transformer (DTransformer) models that use self-attention mechanisms to perform language modeling (39) in order to motivate our approach. Given input sequence consisting of tokens belonging to a vocabulary V, such models return a probability distribution over the vocabulary that describes what the next token in I could be i.e., the model is trained to estimate Pa=[f /] 1e[1 fI) where is the initial length of a. To compute the output logits in the forward pass, x is converted to a sequence of embeddings that incorporate information about each token and its position the sequence. These embeddings are used input to a stack of Transformer layers one of which is depicted in Figure [
LayerNorm
Time First Token with increasing parallelism for 6.7B GPT-3 model
Multi Head Attention
160
Context: 128 tobens
Contest: 2068 tokers
AllReduce
TTFT milliseconds
121.8
LayerNorm
100
Feed Forward Network
80
75,9
AllReduce
60-
38.6
GEMM: QKV
20
A GEMM: mul K
8.3
4.5
1 GPU
2 GPUs
GPUs
GPUs
Softmax
Figure 2: Increasing the degree of tensor par- allelism decreases the Time To First Token. Even when weights and KV cache fit on de- vice memory, parallelism can be worthwhile. These results for 6.7B parameter GPT-3 like model and were collected using TensorRT- LLM engines on our evaluation platform: an HGX A100 40GB system.
GEMM: A
GEMM: West
Figure 1: One layer of standard Transformer consisting of Multi-Head Attention (also shown) followed by FeedForward Network. Residual connections have been omitted.

---

Such layers comprise the bulk of the compute in the model and consist of a Multi-Head Attention (MHA) block followed by a FeedForward Network (FFN or Multi-Layer Perceptron). The FFN typically consists of two linear transformations with non-linear activation function in between. On the other hand, Multi-Head Attention (Figure I implements scaled dot-product attention i.e., each head computes Sof tmax QK3) )V where Q, K, V E RAxi, Here, h is the head dimension and is the sequence length. The General Matrix Multiply (GEMM) Wout combines the outputs of the different heads. The activations of the last layer are used as input to the unembedding operation. Henceforth, we will refer to the GPT-2.3 (7) like construction the standard Transformer architecture. This variant uses MHA, has sequential Attention and FFN blocks, expands the embedding dimension from to 4d in the hidden layer of the FFN, uses Gaussian Error Linear Unit (GELU) non- non-lineristies, and places Layer Norm operators before the MHA and FFN i.c., is a Pre-LN Transformer.
vhh
A widely used optimization during inference involves caching the Key and Value matrices (from GEMM @ V in Figure 1 of each token in a KV cache that is stored in memory. This memoization has the effect of breaking up autoregressive inference (38) into two distinct steps: 1) Prefill (when the first token and the KV cache are generated) and 2) Decode(for all subsequent tokens). Both these steps have distinct runtime characteristics with Prefill being more compute bound and Decode being more memory bandwidth bound (58). The KV cache entry for the next generated token is appended during each Decode step. Prefill, measured by TTFT, typically takes much longer than a single Decode step and will be the focus of this work.
2.2 Multi-Devic systems
Given the extensive amount of compute and memory capacity required to efficiently serve large models, most widely used systems are node-based configurations where each node has small number (between 4 and 16) of devices. These devices usually take the form of Graphics Processing Units (GPUs) or Tensor Processing Units (TPUs). The discussion in this work will focus on GPUs but we expect that our findings will also be of relevance to other choices of accelerators such as TPUs. Accessing data at levels of the system physically closer to compute cores such as scratchpad memory or caches is typically much faster and more efficient compared to accessing off-chip memory. Techniques like tiled matrix multiplication (20) and FlashAttention (I7) account for this characteristic, considerably improving the runtime performance of implementations.
Devices within node are configured in a topology and linked by interconnects such as Peripheral Component Interconnect Express (PCIe), NVLink, and NVSwitch. The different standards balance versatility and cost with performance. For instance, a topology that uses PCIe switches across some connections will have less overall bandwidth than a system that uses solely NVLink/NVSwiten Inter-device communication primitives are provided to other software by libraries like NCCL (1) and RCCL (II. Communication is comparatively expensive and represents phase in the forward pass where compute cores are mostly idle. This work strives to extend the IO-aware approach used by techniques like FlashAttention (17) towards the multi-device setting.
2.3 Tensor parallelism
Serving large models in multi-device settings requires parallelization schemes and strategies that partition the input sequence, model weights, activations, and/or incoming inference re- quests (49 (56) (35) ①I (60). In particular, model parallel (intra-operator) schemes fall into two broad categories: tensor parallelism and pipeline parallelism. Since Kraken models are designed to improve inference latency, this discussion will center around tensor parallelism.
There are several possible strategies (42) (55) to achieve distributed tensor parallelism in Transformer models but we will focus on the widely used scheme introduced by Shoeybi et al. (49) which is well-suited for multi-GPU settings. This scheme introduces two AllReduce operations per layer and takes advantage of the implicitly parallel nature of Multi-Head Attention with optimal partitioning of the FeedForward Network. In each layer, contiguous groups of Attention heads are placed across different devices and the W matrix that is used to combine the output of the different heads is partitioned across columns. The output of the MHA block is retrieved by reducing the local output of all devices; this introduces one AllReduce as shown in Figure [ Similarly, in the FFN block, the W, weight matrix partitioned across rows (using the notation where Wi is multiplied by a column-vecter of activations) and the W2 weight matrix is partitioned across columns. The output
3

---

Device 1
Standard
Device
Device
Kraken
Device
MHA Partition
MHA Partition
50
MHA
MHA
AllReduce
S1
FFN
FFN
Layer
FFN Partition
FFN Partition
S2
MHA
AllReduce
MHA
Layer
AllReduce
S3
FFN
FFN
Layer N +
MHA Partition
MHA Partition
Step 4
AllReduce
Step
FFN Partition
FFN Partition
Step
Layer N+1
Figure 3: Parallelizing two standard Transformer layers compared to executing two layers of Kraken Transformer with 2-way parallelism. Kraken Transformers have fewer AllReduce operations and these can be run concurrently with the Multi-Head Attention in the next layer. Step lengths are illustrative and not indicative of how much wall-clock time a particular operation might require.
of the FFN is computed by reducing the local output of all devices, thereby introducing another AllReduce operation. The objective of this work is to reduce and hide the runtime impact these collectives introduced by tensor parallelism.
Kraken: Model architecture
3.1 Layer construction
The primary objective of Kraken is to preserve much of the standard Transformer (GPT-2.3 (44) architecture as possible while reducing the latency cost of the collective communication introduced by tensor parallelism. To achieve this, we allow each of the individual shards of parallelized Transformer layer to behave as independent, smaller layers. More precisely, as depicted in Figure 3 instead of sharding the Multi-Head Attention and FeedForward Network blocks, each MHA and FFN block is replaced with a smaller, independent block. This introduces a fixed degree of parallelism that is chosen at the outset of training and accounts for the most common hardware deployment target. For example, if model will be mostly served on nodes with eight GPUs each, a reasonable choice would be use either 4 way 8-way parallelism (depending the size of the model). The former would be suitable for smaller models, allowing each node to serve two different inference requests at the same time.
Consequently, the only inter-device dependency is an AllReduce operation the end of every layer. This collective represents the only interchange of activations between groups of sub-layers. Its output used in the FFN block of the next layer and not in the MHA block. As shown in Figure 3 this allows for overlapping the compute in the MHA block with the AllReduce, effectively removing most inter-device communication from the critical path. Much like how positional embeddings are added to the token embeddings prior to the first layer (39), chose to use element- wise addition to combine the outputs of the various sub-layers from the previous layer. This occurs prior to the Layer Norm before the FFN. The forward pass of each sub-layer, including residual connections, is also described in A Algorithm T All pretrained models in our evaluation use the GELJU activation function in the FFN. In initial experiments that explored different constructions, we scaled the weights of the residual layer by VLLN using a similar line reasoning as that used by Radford et al. (44). Here,
L is the number of layers, and N is the degree of parallelism. We did not remove this initialization scheme in subsequent experiments even though accumulation along residual connections was limited to groups of L sub-layers.

---

Algorithm 1: Kraken Sub-Layer: Forward Pass
Input: x E Ridd
Output: y R(xd
residual =
/* The first layer replaces the AllReduce with the identity operator = AllReduce(z)
LayerNorm(z)
residual + MuliiHeeaddtteniionnl]]
residual
/* The output of the AllReduce is used only here allowing it be overlapped
with Attention
= LayerNorm(: +
7 y = residual + FeedforraardNetsorr(ss)
return y
The token and positional embeddings are shared across sub-layers of the first layer i.e., there is still one set of embeddings to maintain compatibility with weight tying (27) (3). After the last layer, we combine the outputs of the different sub-layers using linear transformation with weights W Wooneat E d xd where d is the embedding dimension and N is the degree of parallelism. The output of this transformation used as input to the unembedding. Depending on the implementation, this linear layer introduces the only blocking collective in the computational graph.
3.2 Deriving model configurations for fixed parameter budget
Increasing the degree of parallelism while keeping other hyperparameters like the embedding di- mension constant will increase the parameter count. Instead, using a configuration of standard Transformer as the basis, we make the following two modifications to derive Kraken configuration that has approximately the same number of parameters:
First, the hidden state expansion the FFN reduced from 4d to 2d where d is the embedding/model dimension.
Given the number of parameters P, degree of parallelism N, number of layers L, and vocabulary size V, we derive closed form expression for P and solve for d i.e. V V*d+E .N * ((3 (d x d) + (d x +2)+2+ (d x 2d) where the term ((3 (dx dd) comes from MHA and (d 2d)) comes from the FFN.
4 Evaluation
4.1 Model configurations and perplexity
To evaluate language modeling performance, we train series of models up to 761 million parameters large and with varying degrees of parallelism on OpenWebText (23). This allows us to compare the performance of the Kraken architecture with the GPT-2 444 family of models. Because of limited access to compute, we do not exhaustively search for hyperparameters and stop training at 150 billion tokens contrast to the about 300 billion tokens that language models of such sizes are typically trained for (4)) Table details the embedding dimensions, number of layers, and other hyperparameters for each configuration. For Kraken configurations, the number of Attention heads in each layer'n is summed across all sub-layers. The context length was set 1024 tokens. Models similar in size were trained using the same learning rate schedule and for the same number gradient steps. More details about the training setup, required compute, and how each configuration was derived can be found in Appendix I Pertinent code including the TensorRT-LLLM implementation is available at ltttt:::/// //t::/gitttt /ggggttthub.... com www.oooaon-baeeeooooooeeo kraken
Perplexity measurements for GPT-2 models are provided to add context to results with the caveat that all Kraken models were trained on WebTeet but GPT-2 models were trained on the closed-source WebText dataset. We would expect that all things equal, Kraken models will have lower perplexity unless the GPT-2 models were subsequently fine-tuned on OpenWebTexs Nonetheless.
5

---

Table Model configurations and perplexity on OpenWebText for Kraken models compared to
similarly sized GPT-2 models. Lower perplexity is better.
Model
Layers
Embedding Dimension
Attention Heads
Total Params.
Validation Perplexity
GPT-2
768
12
117M
20.64
Kraken 2-way
12
678
12
124M
18.89
Kraken 4-way
12
504
12
124.5M
18.56
Kraken 6-way
12
418
12
123.2M
19.22
GPT-2 Medium
24
1024
345M
14.87
Kraken 2-way
24
888
16
350M
14.40
Kraken 4-way
24
644
16
353.4M
14.71
GPT-2 Large
24
1280
762M
13.69
Kraken 4-way
24
960
16
761M
13.09
Table Zero-Shot performance on SuperGLUE. ReCoRD uses the F1 score the evaluation
metric. All other benchmarks use accuracy.
Model
BoolQ
RTE
CB
COPA
ReCoRD
WIC
WSC
MultiRC
Average
GPT-2
48.38
51.99
41,07
62.0
71.07
49.53
43.27
53.47
52.6
Kraken 124M
53.85
54.15
44.64
68.0
72.42
49.53
36.54
56.89
54.5
(2-way)
Kraken 124M
50.15
53.79
41.07
67.0
71.47
50.31
37.5
57.16
53.56
(4-way)
Kraken 124M
47.92
56.68
8.93
69.0
70.69
50.16
53.85
53.82
51.38
(6-way)
GPT-2
58.53
53.07
42,86
68.0
79.43
50.0
41.35
52.58
55.73
Medium
Kraken 355M
60.06
51.62
41.07
69.0
80.04
50.0
35.58
56.91
55.54
(2-way) Kraken 355M
61.68
55.23
35.71
72.0
79.01
50.0
36.54
57.2
55.92
(4-way)
GPT-2 Large
60.55
52.71
41.07
72.0
81.95
49.69
44.23 36.54
48.56 $1.65
56.34
Kraken 760M
60.58
49.1
10.71
73.0
82.04
50.0
51.7
(4-way)
even when trained over a smaller number of tokens, Kraken models reach similarly low perplexity as standard Transformers.
4.2 Performance On SuperGLUE
We used the SuperGLUE benchmark suite (53) to evaluate performance on language tasks. All benchmarks were scored on accuracy except for ReCoRD which uses the Fl Score instead. No finetuning training was performed for any combination of model and benchmark. Table containnn results for Zero-Shon performance and Table 3 presents performance in the Three-Shot setting. Scores were calculated using the Language Model Evaluation Harness with the default choice of prompts and scoring metrics for option Im-eval-Supe -uuperrrL v1. As conveyed by these results, Kraken largely preserves language modeling capabilities of the standard Transformer architecture. We expect that the gap between standard Transformers on language tasks will close further we train the models on higher quality data and using optimal choices for the various hyperparameters.
4.3 Evaluation platform
To measure improvements in inference latency, we used the TensorRT-LLLM (55) library to create engines and compare Kraken models with other widely used dense model architectures. The library provides interface to define popular Transformer models and bundles a collection of kernels,

---

Table Three-Shot performance on SuperGLUE. ReCoRD uses the Fl score as the evaluation
metric. All other benchmarks use accuracy.
Model
BoolQ
RTE
CB
COPA
ReCoRD
WIC
WSC
MultiRC
Average
GPT-2
53.76
47.65
42,86
60.0
70.12
50.0
51.92
50.62 50.56
53.37
Kraken 124M
55.14
48.74
44.64
61.0
68.83
49.69
47.12
53.22
(2-way)
Kraken 124M
59.05
44.04
48.21
64.0
69.94
47.49
51.92
54.17
54.85
(4-way)
Kraken 124M
54.28
46.93
42.86
63.0
69.17
49.37
44.23
51.96
52.72
(6-way)
GPT-2
60.58
47.65
44.64
78.35
47.65
50.0
53.22
56.26
Medium
Kraken 355M
50.73
53.07
37.5
68.0
79.01
49.06
66.35
51.34
56,88
(2-way) Kraken 355M
61.8
51.99
50.0
73.0
77.97
47.49
44.23
54.17
57.58
(4-way)
GPT-2 Large Kraken 760M
60.58
51.99
39.29
70.0
81.05
47.81
60.58
51.53
57.85
55.75
53.79
42.86
81.03
45.77
62.5
53.47
57.9
(4-way)
Table 4: Configurations for the different model engines used to compare TTFT. For models of similar sizes, hyperparameters shared for the GPT-like and Parallel Attention + FeedForward variants. 4-way denotes Kraken configurations used when evaluating tensor parallelism across devices and likewise for 8-way.
Model Layers dModel Params. Size
Attention Heads
Kraken 4-way dModel
4-way
4-way Atten- tion
Kraken 8-way dModdll
8-way Params.
8-way
Per Layer
Params. Per
Atten-
Per
tion
Layer
Heads
Layer
Heads
1.3B
2048
50.3M
1248
49.9M
12
59.0M
10
6.7B
32
4096
201.3M
2496
199.4M
24
1920
235.9M 9M
20
13B
40
5140
317.0M
40
3120
311.5M
30
2304
339.7M
32
65B
80
8192
805.3M
64
4992
797.4M
39
3648
851.7M
38
175B
96
12288
1.81B
96
7424
1.76B
58
5472
1,92B
57
plugins, and other optimizations that can be used to create efficient TensorRT engines (containing model weights) and serve them on systems with GPUs. All experiments were conducted 8 A100 GPU machine with NVSwitch and 40GB HBM memory per GPU.
4.4 Speedup in Time To First Token
For comparisons with other model architectures, we build engines for standard, GPT-like config- urations and GPT-J (54) like configurations as detailed in Table E The latter serves to contrast our approach with one that runs the FFN in parallel with the Attention (IO) (54). Such parallel layers require only one AllReduce in layer but unlike Kraken, the collective cannot be overlapped with compute. For similarly sized models, the only difference is the embedding dimension for the Kraken models. All configurations follow the convention where dModel is divided by the number of Attention heads to calculate the size of each head. Other configuration parameters such as the number of layers, maximum context length, and vocabulary size were the same. The Attention heads per sub-layer for Kraken models and parameter counts do not include biases and layer normalization.
There are two sets of engines, one for 4-way tensor parallelism and another for 8-way parallelism. We also did follow the earlier convention 3.2 of solving for the embedding dimension of Kraken model precisely. This is because the available optimizations in TensorRT-LL.V are compati- ble only with specific Attention head dimensions. dimensions that, given equivalent GPT-like configuration, have about the same number of pa- rameters while still being compatible with the available kernels. This was necessary for a fair evaluation but in practice, the chosen embedding dimension should also account for the performance
account for this, we handpicked embedding

---

Speedup in Time To First Token relative to standard Transformers: 4-way parallelism
1.7
Parallel
Bar labels are in % improvement
1.6
Relative Speedup
44.7 1.5
144
1.3
24.0
12.
16.1
14.7
15.4
13,2
1.1
6.4
1.0
1.38,
6.78,
658,
1758, 1758, Geomson
Figure 4: Speedup Time To First Token over standard Transformers on a system that uses NVSwitch and with way parallelism. x-axis labels denote the size of the model followed by the context length. Bar labels are in percentage improvement.
Speedup Time First Token relative to standard Transformers: 8-way parallelism
1.7
Parallel Attn.
FFN
Kraken
Bar labels are in % improvement
1.6
$8.4
Relative Speedup
1.5
1.4
55.2
$7.5
36.5
13:
24.9
25.4
25.4
12-
28.0
16.0
11
1.0
1.38, 128
1.38, 2048
6.78, 128
6.78.
2048
138,
128
138.
2048
658,
128
658, 2048 1758, 128 1758, 2048 Geomean
Figure 5: Speedup in Time To First Token over standard Transformers on a system that uses NVSwitch and with 8-way parallelism. x-axis labels denote the size of the model followed by the context length. Bar labels are in percentage improvement.
of the computational kernels it would map to. Anthony et al. (4) discuss this, showing how model hyperparameters can affect GEMM performance and consequently training efficiency.
Figure 4 shows the speedup in latency compared to GPT-like models on our evaluation platform and with 4-way parallelism while Figure depicts improvements for 8-way parallelism. Across range of model sizes and context lengths 128 and 2048 tokens, using Kraken models can improve inference latency by anywhere from 10.9% 63.7%. These results end-to-end and include the extra fully-connected layer required by Kraken models. All results are normalized to the equivalent latency for similar GPT-3 like model and engines Latency in terms of milliseconds reported in Appendix A.33 Importantly, our evaluation platform uses NVSwitch and consequently has considerable inter-device bandwidth (600 GB/s) which means communication can be relatively inexpensive. We expect that these gains will be more pronounced on systems with less performant interconnects.
generated with random weights and fpl6 precision.
more detailed evaluation on precisely which operators contribute to overall runtime is presented in Section 4.5 In order take advantage of the concurrent communication and compute that Kraken al- lows, we extended TensorRT-LLM with the Overlap plugin that is described Appendix A.2
Each configuration was evaluated using the gntSessionBencchmrkk that is available as part of TensorRT: L.M. We enabled the default set of optimizations which include GPT Attention plugin, Remove Input Padding, and GEMM plugin. Experiments were with batch size of but engines were

---

built with a maximum batch size of 4, vocabulary size of 51,200 tokens, maximum input context length of 2048, and maximum output length of 4096.
4.5
Runtime characterization
Forward Pass Runtime Characterization: 4-way paralielism
Norm
225
GEMM
Dverlppeed GEMM
200
Attention
175
AllReduce
Millsecoodd
1.5
125
100
75-
1.38, 128
1.38.
6.78,128
2048
138, 128
138, 2048
658. 128
2048
1758, 128 1758, 2048
Grouped (Standard,
Figure 6: Runtime characterization: 4-way parallelism. For each cluster on the x-axis, labels denote the size of model followed by the context length.
Forward Pass Runtime Characterization: 8-way parallelism
130
120
GEMM
Overlapped GEMM
Attention
Milliseconds
50
20
1.38, 128
1.38, 2048
6.78, 128
6.70, 2048
138, 120
138, 2048
658, 128
658, 2048
2048
1758, 120
1758, 2048
Grouped as (Standard, Paraliel Attn.
FFN, Kraken)
Figure 7: Runtime characterization: 8-way parallelism. For each cluster on the x-axis, labels denote the size of model followed by the context length.
To provide more context to the performance gains presented in Figures and we characterize the entire forward pass using the same experimental setting Section 4.4 The runtime profiles presented Figures [gqures/GJane [jjand] were obtained using the profiler built into gguSessionnenchhnur Across all model sizes and both context lengths, we find that a significant proportion of time is spent in the AllReduce. As expected, this proportion increases for 8-way parallelism compared to 4-way parallelism since each device works on smaller fraction of the compute.
In some cases, such for 175B models context length 2048, some of the performance gains come from the GEMMs requiring significantly more time for the GPT and GPTTJ like configurations. Despite these performance anomalies, in general, Kraken models spend much smaller fraction of runtime in inter-devics communication leading to considerable performance and efficiency im- provements. Note that the cost of the memory copies and synchronization necessary for Overlap plugin (Appendix A.2 are implicitly included in the Overlapped GEMM fraction. More precisely, Overlapped GEMM the time spent computing GEMM Wont, GEMM QKV, and other operations in the Overlap plugin.

---

5 Discussion and Limitations
Currently, our approach requires training Kraken models from scratch which can be very resource intensive. Consequently, developing techniques to distill learned weights from existing models possibly as part of the weight initialization scheme is promising area of future work. We also do not compare with more state-of-the-art Transformer training recipes or alternatives because of limited access to GPU compute. Training larger, compute-optimal (25) Kraken models on higher-quality datasets will permit evaluations on newer, more complex language modeling benchmarks. The fixed degree of model parallelism also places restrictions on the optimal choice of hardware to model For example, if we tried to deploy model with 4- 4-way parallelism system with two or six devices, we would either need to run groups of two sub-layers together or introduce more collectives. Either approach might negate most of the latency gains offered in the first place.
Need for pretraining:
Compatibility with standard Transformer improvements: Since Kraken models replicate the Attention block by the degree of tensor parallclism, the size of the KV cache will also be larger than that of an equivalent Transformer model. One way to mitigate the increased memory consumption would be to replace Multi-Head Attention with either Multi-Query Attention 47) or Grouped Query Attention (3. Similarly, many other improvements to the standard Transformer architecture also applicable to Kraken models such as RoPE embeddings (50) and R.MNNorm (57). The architecture is also compatible with "drop-in" replacements for Attention that more efficient because of time- complexity and/or sparsity (DD ()) Furthermore, we expect that the notion of parallelizing individual layers of the model will also prove useful in large deep learning models that use constructions other than Attention or MLPs (24). Incorporating a communication-awee Search (8) for Transformers is another promising area of future work.
approach to Neural Architecture
Applicability of existing optimizations: We expect that Kraken will be readily compatible with many existing system-aware techniques used in deploying standard Transformers (56) (22 II (5).
For example, it provides an extra degree of freedom to the various partitioning strategies proposed by Pope et al. (42). It is also compatible with techniques like FlashAttention (17), Speculative
Decoding (34), fusing Attention with the FFN (36), and PagedAttention (32).
Cexture-of-Expertt and hybrid models: By virtue of its architecture, Kraken evokes a comparison to Mixture-of-Experts (MoH) architectures such as the SwitchTransformer (2T) Layers in current MoE models have single Attention block and use learned Router to direct each token to one of a set of FeedForward Networks. Because the forward pass activates only a fixed fraction of the parameters, MoE models can be much larger than dense Transformers while maintaining the computational profile of inference. Nonetheless, serving them efficiently in multi-device settings is a challenge: MoEs require dynamic routing and suffer from load balancing issues. For instance, Huang et al. (26) find that the required All-to-All collectives can comprise a significant fraction of inference latency. An interesting direction of future work would be to incorporate inter-device 1O-aware approach in the construction of MoE models. In addition, recent work (52) has found that hybrid architectures that combine Multi-Head Attention, State Space Models (SSMs) (4), and FFNs in the same model perform better than models that use solely SSMs or MHA for the transformation. Similarly, Kraken layers can be combined with standard Transformer layers and pure SSM layers to develop high-quality language models that are also efficient to run on hardware.
Acknowledgements
This material is based on research sponsored by the Air Force Research Laboratory (AFRL) and the Defense Advanced Research Projects Agency (DARPA) under agreement No. FA8650-18-2- 7862. The U.S. Government authorized reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of the Air Force Research Laboratory (AFRL), the Defense Advanced Research Projects Agency (DARPA), or the U.S. Government. We thank Princeton Research Computing for their assistance in setting up and maintaining the necessary compute infrastructure, We thank Ashwinee Panda for feedback and aid in brainstorming the title. We also thank Jens Tuyls for the helpful discussions.
10

---

References
[1] Amey Agrawal, Nitin Kedia, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav Gula- vani, Alexey Tumanov, and Ramachandran Ramjee. Taming Throughput-Lalency Tradeoff in LLM Inference with Sarathi-Serve. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 117-134, Santa Clara, CA, 2024. USENIX Association.
[2] Eleuther AI. Language Model Evaluation Harness. https:///gituubb.coom.i.ecee 1m- evaluation harness/ 2022.
[3] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. GQA: Training Generalized Multi-Query Transformer Models from Multi- Head Checkpoints. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 4895-4901. Association for Computational Linguistics, December 2023.
[4] Quentin Anthony, Jacob Hatef, Deepak Narayanan, Stella Biderman, Stas Bekman, Junqi Yin, Aamir Shafi, Hari Subramoni, and Dhabaleswar Panda. The Case for Co-Designing Model Architectures with Hardware. In Proceedings of the 53rd International Conference on Parallel Processing, ICPP '24, page 84-96, New York, NY, USA, 2024, Association for Computing Machinery.
[5] Roy Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The second PASCAL recssnneco textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment, 2006.
[6] Luisa Bentivogli, Ido Dagan, Hoa Trang Dang, Danilo Giampiccolo, and Bernardo Magnini. The fifth PASCAL recognizing textual entailment challenge. In Proceedings of the Text Analysis Conference, 2009.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[8] Krishna Teja CChitty-ennatte Murali Emani, Venkatram Vishwanath, and Arun Somani. Neural Architecture Search for Transformers: A Survey. IEEE Access, 10: 10:108377-100412 2022.
[9] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking Attention with Performers. In International Conference on Learning Representations, 2021.
[10] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling Language Modeling with Pathways. Journal of Machine Learning Research, 24(240):1-111
2023.
11

---

[11] AMD Corp. ROCm Collective Communication Library. projectt/rccl//m/laatsss// 2024.
https://rocm.de amd com/
[12] NVIDIA Corp. H100 Transformer Engine Supercharges Al Training.
https://blogs
httpaaa///oooooo//-ooraraaaa traansooreer-enniiiee 2022.
[13]
NVIDIA Corp. NVIDIA DGX H100.
https:/resourree.nvviia.comm
en-us Sesssdgx-systee
2023.
[14] NVIDIA Corp. NVIDIA Collective Communications Library. httpps//ddewllpprr..rrt com/nccl 2024.
[15] NVIDIA Corp. TensorRT-LLM. attps://githur htgpChubbccc//VD//AA///sSrr----- 2024.
[16] Ido Dagan, Oren Glickman, and Bernardo Magnini. The PASCAL recognisigng textual entail- ment challenge. In Proceedings of the First International Conference on Machine Learning Challenges: Evaluating Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment, MI.CW"05, page 177-190, Berlin, Heidelberg, 2005. Springer-Vertee Verlag.
[17] Dao, Tri and Fu, Dan and Ermon, Stefano and Rudra, Atri and Ré, Christopher. FlashAtten- tion: Fast and Hemory-Efficient Exact Attention with IW-Awareness. Advances Neural Information Processing Systems, volume 35, pages 16344-16357 Curran Associates, Inc., 2022.
[18] Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The BommitmentBank: Investigating projection in naturally occurring discourse. In Proceedings of Sinn und Bedeutung 23, 2019.
[19] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x 16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.
[20] Kayvon Fatahalian, Jeremy Sugerman, and Pat Hanrahan. Understanding the Efficiency of GPU Algorithms for Matrix-Matrix Multiplication. In Proceedings of the ACM SIG- GRAPH/EUROGRAPHICC conference on Graphics hardware, pages 133-137, 2004.
[21] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[22] Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pages 1-9. Association for Computational Linguisties. 2007.
[23] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. WebTeet Corpus. http://Skyylonn007.
wit.anuttt..tttubb.oe github. ho//ppenieebTextCCrpsss 2019.
[24] Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces.
In First Conference on Language Modeling. 2024.
[25] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. Training Compute-Optimal Large Language Models. In Proceedings of the 36th International Conference Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2024. Curran Associates Inc.
[26] Haiyang Huang, Newsha Ardalani, Anna Sun, Liu Ke, Hsien-Hsin S. Lee, Anjali Sridhar, Shruti Bhosale, Carole-Jean Wu, and Benjamin Lee. Towards MoE Deployment: Mitigating Inefficiencies in Mixture-of-Expee Expert (MoE) Inference. arXiv preprint arXiv:2303.06188 2023.
12

---

[27] Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying Word Vectors and Word Classifiers: Loss Framework for Language Modeling. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017.
[28] Jean Kaddour, Joshua Harris, Maximilian Mozes, Herbie Bradley, Roberta Raileanu, and Robert McHardy, Challenges and Applications of Large Language Models. arXiv preprint arXiv:2307. 10169, 2023.
[29] Andrej Karpathy, nanoGPT. httpge///gthuub.cc/////pppttttat
2022.
[30] Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pages 252-262, 2018.
[31] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer. In International Conference on Learning Representations, 2020.
[32] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng. Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient Memory Management for Large Language Model Serving with PttentAtention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, page 61 1-626, New York, NY, USA, 2023. Association for Computing Machinery.
[33] Hector J Levesque, Ernest Davis, and Leora Morgenstern. The Winograd schema challenge. In AAA/ Spring Symposium: Logical Formalizations of Commonsense Reasoning, volume 46, page 47, 2011.
[34] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast Inference from Transformers via Speculative Decoding. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19274-19286.. PMLR, 23-29 Jul 2023.
[35] Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence Paral- lelism: Long Sequence Training from System Perspective. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2391-2404 Toronto, Canada, July 2023. Association for Computational Linguistics.
[36] Hao Liu and Pieter Abbeel. Blockwise Parallel Transformers for Large Context Models. In Advances in Neural Information Processing Systems, volume 36, pages 8828-8844, Curran Associates, Inc., 2023.
[37] Hao Liu, Matei Zaharia, and Pieter Abbeel. RingAttention with Blockwise Transformers for Near- Infinine Context. In The Twelfth International Conference on Learning Representations, 2024.
[38] Pratyush Patel, Esha Choukse, Chaojie Zhang, Aashaka Shah, fñigo Goiri, Saeed Maleki, and Ricardo Bianchini. Splitwise: Efficient Generative LLM Inference Using Phase Splitting. In 2024 ACM/EEEE 51 Anmual International Symposium on Computer Architecture (ISCA), pages 118-132, 2024.
[39] Mary Phuong and Marcus Hutter. Formal Algorithms for Transformers. arXiv preprint drXiv:2207.092. 09238, 2022.
[40] Mohammad Taher Pilehvar and Jose Camacho-Collados, WiC: The Word-in-Contex Dataset for Evaluating Context-Sensitive Meaning Representations. In Proceedings of the North American Chapter of the Association for Computational Linguistics: Human Language Technology, 2019.
[41] Adam Poliak, Aparajita Haldar, Rachel Rudinger, J. Edward Hu, Ellie Pavlick, Aaron Steven White, and Benjamin Van Durme. Collecting Diverse Natural Language Inference Problems for Sentence Representation Evaluation. In Proceedings of Empirical Methods Natural Language Processing, 2018
13

---

[42] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury. Jonathan Heek, Kefan Xiao. Shivani Agrawal, and Jeff Dean. Efficiently Scaling Transformer Inference. In Proceedings of Machine Learning and Systems, volume 5, pages 606-624, 2023.
[43] Ofir Press and Lior Wolf. Using the Output Embedding to Improve Language Models. In Conference of the European Chapter of the Association for Computational Linguistics, 2016.
[44] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and lly Sutskever. Language models are unsupervised multitask learners. OpenAl blog. 188)99, 2019.
[45] Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011.
[46] Rachel Rudinger, Jason Naradowsky, Brian Leonard, and Benjamin Van Durme. Gender Bias in Coreference Resolution. In Proceedings of the North American Chapter of the Association for Computational Linguistics: Human Language Technology, 2018.
[47] Noam Shazeer. Fast Transformer Decoding: One Write-Head All You Need. arXiv preprint arXiv: 1.9/11021500 2019.
[48] Noam Shazeer, * AAzalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of- Experts Layer. In International Conference on Learning Representattonss 2017.
[49] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LMM Training Kulti-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv: 109v00390.00803,, 2020.
[50] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 127063, 2024.
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. E ukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[52] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An Empirical Study of Mamba-based Language Models. arXiv preprint arXiv:2406.07887, 2024.
[53] Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. SuperGLUE: Stickier Benchmark for General-Purpose Language Understanding Systems. CoRR,
2019.
[54] Ben Wang and Aran Komatsuzaki, GPT-J-6B: 6 Billion Parameter Autoregressive Language Model. https: http://ggttttbub cow www./////ggoooo////oeeeoeoeeeeee http:s///gthhiiiiooo///gggg////2
htwp:////wwsssttsssteeeeeeeeeeee transf jax May 2021.
[55] Shibo Wang. Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, Sameer Kumar, Tongfei Guo, Yuanzhong Xu, and Zongwei Zhou. Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. In Proceedings of the 28th ACM International Conference on Architectrall Support for Programming Languages and Operating Systems, Volume 1, ASPLOS 2023, page 93-106, New York, NY, USA, 2022. Association for Computing Machinery.
[56] Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orea: distributed serving system for Transformer-Bann generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521-538 2022.
[57] Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization. Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
14

---

[58] Hengrui Zhang, August Ning, Rohan Baskar Prabhakar, and David Wentzlaff. LLMCompass: Enabling Efficient Hardware Design for Large Language Model Inference. In Proceedings of the 51 Annual International Symposium on Computer Architecture, 2024.
[59] Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. ReCoRD: Bridging the Gap between Human and Machine Commonsense Reading Comprchension. arXiv preprint 1810. 12885, 2018.
[60] Kan Zhu, Yilong Zhao, Liangyu Zhao, Gefei Zuo, Yile Gu, Dedong Xie, Yufei Gao, Qinyu Xu, Tian Tang, Zihao Ye, Keisuke Kamahori, Chien- Yu Lin, Stephanie Wang. Arvind Krishnamurthy, and Baris Kasikei. NanoFlow: Towards Optimal Large Language Model Serving Throughput.
arXiv preprinl arXiv:2408. / /2757, 2024.
15

---

Table Pretraining compute and setup for each Kraken configuration
Model
Layers
Embedding Dimension
Attention Heads
Total Params.
A100 GPU
Initial
Hrs
Learning
Rate
Kraken 2-way
12
678
12
124M
480
2.5c
Kraken 4-way
12
504
12
124.5M
500
2.5c 4
Kraken 6-way
12
418
12
123.2M
800
2.5e 4
Kraken 2-way
24
888
16
350M
1000
1.5e 4
Kraken 4- 4-way
24
644
16
353.4M
1750
1.5e 4
Kraken 4-way
24
960
16
761M
1800
1,5e 4
A
Appendix
A.1
Training setup and compute requirements
For all pretrained models presented in Section 4.1] we used similarly sized GPT-3 7 model's hyperparameters the basis and followed the procedure outlined in Section 3.2 to calculate the embedding dimension. We did not make an effort to optimize the codebase used for training which builds off of nanoGPT (29). It is possible to replicate pretrained models by extending nanoGPT to implement the new forward pass as described in Algorithm T The Adam optimizer was used train all models along with cosine learning rate decay with linear warmup. Initial learning rates and the approximate GPU hours required to train each configuration are presented in Table 5 All models were trained for 300, 000 gradient steps. Only the 761M parameter model was trained on node with 80GB A100 GPU machines. The other configurations were trained on 40GB A 100 machines and consequently use many more gradient accumulation steps. This is why the largest model required a similar number GPU hours as the next largest.
Weights for all pretrained GPT-2 models used when evaluating SuperGLUE perfor- mance (18) (45) (30) (59) 16) 5 (2) (6) (40) (46) (41) (33) in Section 4.2 were obtained from HuggingFace. Since the focus of this work is on illustrating the efficiency gains and language modeling capabilities, we do not implement any safeguards that will filter for biased and/or harmful content. Initial experiments that tried various variations of the model architecture required about another 1,000 hours A100 compute.
A.2 Overlap plugin implementation
We used TensorRT-LLLM version 0. /et.0.dev222007:35202 throughout the evaluation. CUDA allows kernels be launched on different streams and depending on resource availability, these kernels may be executed in parallel. However, TensorRT does not support multi-stream execution across plugins. We circumvented this limitation by implementing a Singleton that can manage a dedicated stream and global memory meant for launching collectives. This allows different instances of the plugin to launch and synchronize kernels on the same stream. Each instance the plugin can either: trigger an AllReduce op on separate low-priority CUDA stream or synchronize the stream ensure that previously launched AllReduce completes. The plugin also implements the functionality provided in the existing GEMM plugin. This allows us to perform the following within the Multi-Has Head Attention block:
1. The AllReduce is launched dedicated CUDA stream just before the GEMM used compute the Query, Key, and Value matrices (GEMM QKVD via the plugin
2. All intermediate compute is performed using existing kernels such as FttstAttention
3. The CUDA stream that the AllReduce op was placed on synchronized after the GEMM W [
also via the Plugin
This approach allowed us to overlap the collective with the computation required for Multi-Head Attention and is the only addition to TensorRT-LL.M aside from the definition of Kraken. However, it also requires two extraneous memory copy operations that can be avoided the library adds support for multi-strean execution.
16

---

Table Inference latency in milliseconds for 4-way parallelism
Model Size
Context Length
Standard
Parallel Attn. + FeedForward
Kraken
I.3B
128
3.7
3.3
3.3
1.3B
2048
17.0
13.6
13.2
6.7B
128
8.3
7.0
5.8
6.7B
2048
48.2
42.1
38.0
13B
128
13.0
11.1
10.7
13B
2048
83.7
73,8
66.9
65B
128
12.7
11.2
8.5
65B
2048
84.6
79.5
63.8
175B
128
24.6
22.5
19,9
175B
2048
243.7
230.8
158.9
Table 7: Inference latency in milliseconds for 8-way parallelism
Model Size
Context Length
Standard
Parallel Attn. + FeedForward
Kraken
1.3B
128
4.3
3.4
3.2
1.3B
2048
15.9
11.7
9.7
6.7B
128
7.1
5.7
4.7
6.7B
2048
37.3
29.2
27.6
13B
128
10.7
8.4
6.7
13B
2048
58.8
46.9
42.6
65B
128
8.7
7.2
6.2
65B
2048
55.9
48.9
42.9
175B
128
16.9
14.4
12.4
175B
2048
125.1
114.7
98.0
A.3
Time To First Token in milliseconds
Table contains the results from Figure 4 but in terms of milliseconds. Similarly, Table presents the results from Figure 5 For the 65B and 175B configurations, errors when running full sized engines because each device has only 40GB of memory. To avoid this, reduce the number of layers 1 the original number. Nonetheless, the relative latency comparison should be unaffected because the runtimes of individual layers are identical.
into CUDA Out-of-Memort
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the abstract and introduction clearly describe the work and refer only to attained experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The Discussion and Limitations section goes over the limitations of this work and contrasts with several categories of related work.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer: [NA]
Justification:
This paper does not make any theoretical claims about the efficacy of the
proposed architecture.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
All pretrained models can be replicated by using the training recipe presented in the Appendix along with the Algorithm in the Architecture section. The dataset used for training as well as the codebase are both publicly available. Similarly, the TensorRT-LLMM comparison can be reproduced by extending the library to support the new model and then implementing the Overlap plugin described in the Appendix. More importantly, we have also made pertinent code open-source.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
19

---

reproducible, they In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
Relevant code has been published online; details can be found in the Evaluation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cc
woo.iiiiissiiiiiiiioooooooliiiii for more details. Sttpi///giiids///eeeuuueee
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
7 //nips. //nips.ccppute cwtppbblcc/guidee.eeoooelpe ProdeSubmissionPooo
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All the information necessary to train the models on Open WebTeet,, perform the comparison using SuperGLUE, build TensorRT-LLLM engines, and replicate results can be found the paper and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We did not have the compute resources to train several versions of the same model configuration and report perplexity and SuperGLUE results with error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
20

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The number of GPU hours necessary train the models discussed in the Appendix along with other details about the compute resources used.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The Code of Ethics reviewed carefully for full compliance.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
21

---

Justification:
While the paper does present pretrained language models, the accompanying societal impacts and necessary safeguards are no different from that of existing language models. As such, we expect that existing mitigation strategies will be applicable to the ideas presented this paper.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: While this paper presents a few pretrained models, the focus is on the system- level efficiency gains, not the downstream performance the models. We have highlighted this in the evaluation and Appendix.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All assets such as datasets and codebases are cited and urls are provided were applicable.
22

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, our code release is documented with focus on being easy to reproduce. Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicablle, as well as details about compensation (if any)?
Answer: [NA]
Justification; The experiments in this work did not directly include data from crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
23

---

Answer:
[NA]
Justification:
None of the experiments in this work involved human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
24


Reviews:
[{"review_id": "6fTTo8HQnC", "reviewer": "Reviewer_ptrK", "summary": "This paper propose Kraken, a new evolution of the standard Transformer to reduce the communication cost of inference. Kraken overlaps the collective operations with computing, therefore achieves smaller pre-filling time cost. Kraken is specially designed for tensor parallelism on multi-device environments.", "questions": "This paper focuses on an important issue and proposes an simple yet effective solution. However, it has several limitations that prevent the paper from being accepted. \n\nFirstly, a comprehensive study about related works is lacked. It is hard for me to get the key idea of this paper. What is the advantages of Kraken compared with other works? \n\nSecondly, although the paper has preserved most of the standard Transformer architecture, it still requires to re-train a model from scratch, which is extremely costly, especially for LLMs. The necessity of altering the Transformer architecture should be clarified. Specifically, it should be determined whether Kraken could be designed as a plug-in for any Transformer-based model and can work directly or require minimal re-training.\n\nBesides, as tools like DeepSpeed, PagedAttention, and Flash Attention are widely adopted in current LLM inference, it is critical that Kraken demonstrates compatibility with these tools to enhance performance and practicality.", "limitations": "Please refer to the Weaknesses and Questions.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The focus on accelerating the pre-filling stage of inference by reducing communication overhead addresses a significant and challenging issue that merits further research.\n2. Kraken preserves much of the standard Transformer architecture while effectively reducing communication overhead.\n3. The solution of Kraken is is straightforward yet impactful.\n4. Kraken maintains the model’s accuracy despite its optimizations.", "weaknesses": "1. The paper lacks a detailed examination of related works that also focus on reducing communication costs. Overlapping techniques to accelerate inference have been widely explored in both industrial and academic contexts, as demonstrated by works such as [1, 2] and tools like NCCL. A clearer differentiation between Kraken and these existing solutions is necessary.\n\n> [1] Wang, Shibo, et al. \"Overlap communication with dependent computation via decomposition in large deep learning models.\" (ASPLOS’22)\n\n> [2] Rajbhandari, Samyam, et al. \"Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning.\" (SC’21)\n\n2. Sharing token and positional embeddings across all sub-layers might introduce limitations in how effectively the model can learn positional nuances across different parts of the input sequence. Thus, except the evaluation  in 4.1 and 4.2, the paper should clarify why the new architecture does not affect the accuracy of the model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Comparison to Related Work:**\nOur approach is readily compatible with nearly all related work such as DeepSpeed, FlashAttention, and PagedAttention. Our evaluation (Figure 4) uses FlashAttention and TensorRT-LLM’s implementation of AllReduce which is even more performant than NCCL on systems equipped with NVLink.\n\nThe fundamental contribution is to reformulate the Transformer layer and remove the true data-dependency introduced by the collectives necessary for tensor parallelism. This is illustrated by Figure 3. The closest point of comparison would be models that use parallel Attention and FeedForward which we do evaluate (Figure 4). Our model-level change is compatible with most system-level optimizations for Transformer training and inference.\n\nBy virtue of the architecture, we don’t require the techniques proposed by Wang et al.[1] unless a Kraken model is parallelized on more devices than its degree of parallelism.\n\n**Pre-training and PlugIn:**\nWhile Kraken requires training from scratch, any plug-in we develop would still need to account for the true data dependency between layer activations in the original architecture. This imposes a hard limit on efficiency. In fact, the suggested plug-in would mostly take the form of the approach proposed by Wang et al. [1]\n\n**Sharing token and positional embeddings:**\nThe absence of shared positional and token embeddings would complicate weight tying[3] given that we are computing only one probability distribution at the last layer.\n\n3: “Using the output embedding…”, Press et al. 2017", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rzrTecCrhS", "reviewer": "Reviewer_vwQM", "summary": "The paper introduces a modification to the standard Transformer architecture aimed at reducing inter-device communication during inference in multi-device systems. By predetermining the degree of model parallelism, computations on each device can operate independently, allowing collective operators to overlap with compute tasks. Evaluations conducted using the SuperGLUE benchmark suite demonstrate that the proposed models maintain similar perplexity levels compared to standard Transformers while achieving notable speed improvements.", "questions": "* Could there be additional evaluations on a more diverse set of benchmarks? Can the proposed approach be extended to applications in vision-based tasks?\n* How does the approach benefit systems with varying inter-device configurations and bandwidth capacities?", "limitations": "The authors have thoroughly discussed the limitations of their approach. From the checklist, it appears that the approach does not pose significant negative societal implications, apart from the associated costs and resources required for training large models.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* **Novelty and contribution**: The paper presents a novel approach to parallelizing the Multi-Head Attention (MHA) and Feed-Forward Network (FFN) layers, allowing them to run concurrently with the AllReduce operations. This innovative strategy shows promise and is orthogonal to existing methods like FlashAttention and Speculative Decoding.\n* **Soundness**: The design and methodology, while simple, are robust and well-justified. The authors provide clear explanations and solid rationales for their approach, enhancing the credibility of their work.\n* **Clarity**: The paper is well-structured and written in a clear, concise manner, making it easy to follow and understand.", "weaknesses": "* **Limited benchmark**: The evaluation benchmarks primarily focus on language tasks, neglecting the application of Transformers to vision tasks.\n* **Limited comparative analysis**: The paper lacks comparisons with other existing baselines, including ablation studies and sensitivity analyses. While the authors compare their approach with models running Feed-Forward Networks (FFN) in parallel with Attention, they do not explore sensitivity to varying inter-device bandwidths or different compute-to-communication ratios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Applicability to vision tasks:**\n\nYes, our approach will also translate to encoder-only Transformers including those used in vision tasks such as ViT[4]. Nonetheless, we narrowed the scope of this work to focus our resources on the decoder-only Transformer models that are typically the largest and consequently the most expensive to run inference on.\n\n4: “Image is worth 16x16 words…”, Dosovitskiy et al. 2020\n\n\n**Varying device interconnects and compute/communication ratios:**\n\nWe do present results at various compute/communication ratios because the cost of the AllReduce scales linearly with sequence length (and embedding size) but Attention itself is quadratic (Figure 4). For longer sequences, communication will occupy a smaller fraction of overall runtime.\n\nAdditionally, our evaluation platform (NVSwitch) offers the most performant interconnect between A100 GPUs that we are aware of. As such, assuming that the compute capability of each device is unchanged, the presented latency improvements are a lower bound when compared to alternative topologies where some connections traverse PCIe and consequently suffer from limited bandwidth.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QMlW47sk8m", "reviewer": "Reviewer_6qbN", "summary": "This paper proposes to modify standard transformer architectures to use less allreduce operations from tensor parallelism. Specifically, the proposal is to allow each shard of the hidden state to operate as an individual complete hidden state during attention layers and each shard of the attention layer to operate independently. Consequently, allreduce is only needed after FFN layers and can overlap in time with attention calculation. Experiments of up to 760M models show that the new architecture can achieve accuracy that is almost as good as same-sized GPT2. Latency projection suggests 1.5X speed up in context encoding for 175B-size models.", "questions": "Please see above.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Please see the summary. The proposal is simple yet makes sense. The accuracy and latency results are promising.", "weaknesses": "The biggest weakness of this paper is that accuracy results are only available up to 760M model size. There remains big uncertainty for larger models where the proposal matters. Also there seem to be consistently weak results on the ReCoRD task, is there an explanation on why?\n\nAnother weakness of this paper is that the model training is committed to a specific tensor-parallelism degree. This is not a huge issue in practice but still a weakness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Performance on ReCoRD:**\n\nPlease see the common response. This was due to a bug in the LM Evaluation Harness (Issue 1647 on GitHub).\n\n**Results only up to 760M parameters:**\n\nIndeed, it would have been great to evaluate models that are several billion parameters large. Unfortunately, we do not have access to the amount of compute necessary for such an endeavor. Nonetheless, we hope the results presented here will encourage others to try this approach given how prevalent and expensive Transformer inference can be.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for reading our work, providing helpful feedback, and finding promise in our approach.

After submission, we learned of bugs in LM Evaluation Harness that affected the scores on some SuperGLUE benchmarks such as ReCoRD. We have evaluated all models again using the most recent version of the toolkit; **updated results for Tables 2 and 3 are available in the attached PDF**.

Meta-review: This paper solved an important problem: improving the parallel efficiency for generative inference. An interesting modification is made to the standard Transformer architecture aimed at reducing inter-device communication during inference in multi-device systems. The solution is solid and has good novelty, and the experimental section is well-designed to verify the design of the proposed solution.
Decision: Accept (poster)