Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference

Authors: rohan baskar prabhakar
hengrui  zhang
david  wentzlaff
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 19:19:16

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.

Main paper: 
 introduction: 
 Deep neural networks based on the Transformer architecture (51) have become the prevalent choice for a variety of tasks involving sequences, especially in natural language processing and computer vision (44) (7) (19) . Their capabilities, particularly in language modeling, have been driven by a rapid increase in parameter count (7) . Today's largest language models have up 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) (34) (17) (37) (12) .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 are also used as intermediate steps in longer processes such as augmenting web searches or presenting the results of database queries (28) . The rising prevalence of such multi-step applications makes reducing inference latency even more critical.Continuing this theme, this work focuses on reducing the latency cost of the collective operations introduced by tensor parallelism (49) in the forward pass. In particular, it introduces Kraken, a variation of the standard Transformer architecture (44) (7) that reduces the amount of inter-device communication and allows remaining collective operators to be overlapped with compute. Kraken models have a fixed degree of innate model parallelism that allows computational graphs on 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 as nodes in typical datacenters and DGX (13) 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 a series of Kraken models with varying degrees of parallelism and parameter count on OpenWebText (23) and compare them with the GPT-2 (44) family of models on the SuperGLUE suite of benchmarks (53) . We then implement Kraken using the TensorRT-LLM library (15) and measure the Time To First Token (TTFT) given various model sizes and context lengths to illustrate the efficiency gains when collective operators are no 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 a geomean of 35.6% when tested across a range of model sizes, context lengths, and degrees of parallelism.For all pretrained models presented in Section 4.1, we used a similarly sized GPT-3 (7) model's hyperparameters as 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 1. The Adam optimizer was used to train all models along with a 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 a node with 80GB A100 GPU machines. The other configurations were trained on 40GB A100 machines and consequently use many more gradient accumulation steps. This is why the largest model required a similar number of GPU hours as the next largest.Weights for all pretrained GPT-2 models used when evaluating SuperGLUE performance (18) (45) (30) (59) (16) (5) (22) (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 of A100 compute.We used TensorRT-LLM version 0.12.0.dev2024073000 throughout the evaluation. CUDA allows kernels to 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 of the plugin can either: trigger an AllReduce op on a separate low-priority CUDA stream or synchronize the stream to ensure that a 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-Head Attention block:1. The AllReduce is launched on a dedicated CUDA stream just before the GEMM used to compute the Query, Key, and Value matrices (GEMM QKV 1) via the plugin 2. All intermediate compute is performed using existing kernels such as FlashAttention 3. The CUDA stream that the AllReduce op was placed on is synchronized after the GEMM W out 1, also via the Plugin This approach allowed us to overlap the collective with all the computation required for Multi-Head Attention and is the only addition to TensorRT-LLM aside from the definition of Kraken. However, it also requires two extraneous memory copy operations that can be avoided if the library adds support for multi-stream execution. Table 6 contains the results from Figure 4 but in terms of milliseconds. Similarly, Table 7 presents the results from Figure 5 . For the 65B and 175B configurations, we ran into CUDA Out-of-Memory errors when running full sized engines because each device has only 40GB of memory. To avoid this, we reduce the number of layers to 1 4 the original number. Nonetheless, the relative latency comparison should be unaffected because the runtimes of individual layers are identical.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.• 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 in 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 as long as it is clear that these goals are not attained by the paper. 
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 an input sequence x consisting of tokens belonging to a vocabulary V , such models return a probability distribution over the vocabulary that describes what the next token in x could be i.e., the model is trained to estimate P (x[ℓ + 1] | x[1 : ℓ]) where ℓ is the initial length of x. 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 in the sequence. These embeddings are used as input to a stack of Transformer layers one of which is depicted in Figure 1 . Even when weights and KV cache fit on device memory, parallelism can be worthwhile. These results are for a 6.7B parameter GPT-3 like model and were collected using TensorRT-LLM engines on our evaluation platform: an HGX A100 40GB system.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 a non-linear activation function in between. On the other hand, Multi-Head Attention (Figure 1 ) implements scaled dot-product attention i.e., each head computes Softmax(QK T √ h )V where Q, K, V ∈ R h×l .Here, h is the head dimension and l is the sequence length. The General Matrix Multiply (GEMM) W out 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 (44)(7) like construction as the standard Transformer architecture. This variant uses MHA, has sequential Attention and FFN blocks, expands the embedding dimension from d to 4d in the hidden layer of the FFN, uses Gaussian Error Linear Unit (GELU) non-linearities, and places Layer Norm operators before the MHA and FFN i.e., is a Pre-LN Transformer.A widely used optimization during inference involves caching the Key and Value matrices (from GEMM QKV 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) (58) 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.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 a 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 (17) account for this characteristic, considerably improving the runtime performance of implementations.Devices within a 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/NVSwitch. Inter-device communication primitives are provided to other software by libraries like NCCL (14) and RCCL (11) . Communication is comparatively expensive and represents a 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.Serving large models in multi-device settings requires parallelization schemes and strategies that partition the input sequence, model weights, activations, and/or incoming inference requests (49) (56) (35) (1) (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 out 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 1 . Similarly, in the FFN block, the W 1 weight matrix is partitioned across rows (using the notation where W 1 is multiplied by a column-vector of activations) and the W 2 weight matrix is partitioned across columns. The output Step 4Step 5Step 6Standard Kraken 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 of these collectives introduced by tensor parallelism.3 Kraken: Model architectureQuestion: 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 it with several categories of related work.Guidelines:• The answer NA means that the paper has no 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 (e.g., independence assumptions, noiseless settings, model well-specification, 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
layer construction: 
 The primary objective of Kraken is to preserve as much of the standard Transformer (GPT-2,3 like)(44)(7) 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 a 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 a model will be mostly served on nodes with eight GPUs each, a reasonable choice would be to use either 4-way or 8-way parallelism (depending on 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 at the end of every layer. This collective represents the only interchange of activations between groups of sub-layers. Its output is 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) , we 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 Algorithm 1. All pretrained models in our evaluation use the GELU activation function in the FFN. In initial experiments that explored different constructions, we scaled the weights of the residual layer by 1 √ L * N using a similar line of 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. The token and positional embeddings are shared across all sub-layers of the first layer i.e., there is still one set of embeddings to maintain compatibility with weight tying ( 27) (43) . After the last layer, we combine the outputs of the different sub-layers using a linear transformation with weights W concat ∈ R d * N ×d where d is the embedding dimension and N is the degree of parallelism. The output of this transformation is used as input to the unembedding. Depending on the implementation, this linear layer introduces the only blocking collective in the computational graph.Increasing the degree of parallelism while keeping other hyperparameters like the embedding dimension constant will increase the parameter count. Instead, using a configuration of a standard Transformer as the basis, we make the following two modifications to derive a Kraken configuration that has approximately the same number of parameters:• First, the hidden state expansion in the FFN is 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 a closed form expression for P and solve for d i. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
model configurations and perplexity: 
 To evaluate language modeling performance, we train a 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 (44) family of models. Because of limited access to compute, we do not exhaustively search for hyperparameters and stop training at 150 billion tokens in contrast to the about 300 billion tokens that language models of such sizes are typically trained for (44) (24) . Table 1 details the embedding dimensions, number of layers, and other hyperparameters for each configuration. For Kraken configurations, the number of Attention heads in each layer is summed across all sub-layers. The context length was set at 1024 tokens. Models similar in size were trained using the same learning rate schedule and for the same number of gradient steps. More details about the training setup, required compute, and how each configuration was derived can be found in Appendix A.1. Pertinent code including the TensorRT-LLM implementation is available at https://github.com/rohan-bp/kraken.Perplexity measurements for GPT-2 models are provided to add context to the results with the caveat that all Kraken models were trained on OpenWebText 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 OpenWebText. Nonetheless, even when trained over a smaller number of tokens, Kraken models reach a similarly low perplexity as standard Transformers.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 F1 Score instead. No finetuning or training was performed for any combination of model and benchmark. Table 2 contains results for Zero-Shot performance and Table 3 presents performance in the Three-Shot setting. Scores were calculated using the Language Model Evaluation Harness (2) with the default choice of prompts and scoring metrics for option lm-eval-SuperGLUE v1. As conveyed by these results, Kraken largely preserves the language modeling capabilities of the standard Transformer architecture. We expect that the gap between standard Transformers on language tasks will close further if we train the models on higher quality data and using optimal choices for the various hyperparameters.To measure improvements in inference latency, we used the TensorRT-LLM (15) library to create engines and compare Kraken models with other widely used dense model architectures. The library provides an interface to define popular Transformer models and bundles a collection of kernels, 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 on a 8 x A100 GPU machine with NVSwitch and 40GB of HBM memory per GPU.For comparisons with other model architectures, we build engines for standard, GPT-like configurations and GPT-J (54) like configurations as detailed in Table 4 . The latter serves to contrast our approach with one that runs the FFN in parallel with the Attention (10) (54) . Such parallel layers require only one AllReduce in a 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 are 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 not follow the earlier convention 3.2 of solving for the embedding dimension of a Kraken model precisely. This is because the available optimizations in TensorRT-LLM are compatible only with specific Attention head dimensions. To account for this, we handpicked embedding dimensions that, given an equivalent GPT-like configuration, have about the same number of parameters 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 Figure 4 shows the speedup in latency compared to GPT-like models on our evaluation platform and with 4-way parallelism while Figure 5 depicts improvements for 8-way parallelism. Across a range of model sizes and context lengths of 128 and 2048 tokens, using Kraken models can improve inference latency by anywhere from 10.9% to 63.7%. These results are end-to-end and include the extra fully-connected layer required by Kraken models. All results are normalized to the equivalent latency for a similar GPT-3 like model and engines were generated with random weights and fp16 precision. Latency in terms of milliseconds is reported in Appendix A.3. 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.A more detailed evaluation on precisely which operators contribute to overall runtime is presented in Section 4.5. In order to take advantage of the concurrent communication and compute that Kraken allows, we extended TensorRT-LLM with the Overlap plugin that is described in Appendix A.2.Each configuration was evaluated using the gptSessionBenchmark that is available as part of TensorRT-LLM. We enabled the default set of optimizations which include GPT Attention plugin, Remove Input Padding, and GEMM plugin. Experiments were run with a batch size of 1 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. To provide more context to the performance gains presented in Figures 4 and 5 , we characterize the entire forward pass using the same experimental setting as Section 4.4. The runtime profiles presented in Figures 6 and 7 were obtained using the profiler built into gptSessionBenchmark. 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 a smaller fraction of the compute.In some cases, such as for 175B models at context length 2048, some of the performance gains come from the GEMMs requiring significantly more time for the GPT and GPT-J like configurations. Despite these performance anomalies, in general, Kraken models spend a much smaller fraction of runtime in inter-device communication leading to considerable performance and efficiency improvements. Note that the cost of the memory copies and synchronization necessary for the Overlap plugin (Appendix A.2) are implicitly included in the Overlapped GEMM fraction. More precisely, Overlapped GEMM is the time spent computing GEMM W out , GEMM QKV , and other operations in the Overlap plugin. 
discussion and limitations: 
 Need for pretraining: 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 a 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 run a model on. For example, if we tried to deploy a model with 4-way parallelism on a 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.Compatibility with standard Transformer improvements: Since Kraken models replicate the Attention block by the degree of tensor parallelism, 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 are also applicable to Kraken models such as RoPE embeddings (50) and RMSNorm (57) . The architecture is also compatible with "drop-in" replacements for Attention that are more efficient because of timecomplexity and/or sparsity (31) (9) . 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-aware approach to Neural Architecture Search (8) for Transformers is another promising area of future work.We expect that Kraken will be readily compatible with many existing system-aware techniques used in deploying standard Transformers (56) (32) (1) (35) . 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) .Mixture-of-Experts and hybrid models: By virtue of its architecture, Kraken evokes a comparison to Mixture-of-Experts (MoE) architectures such as the SwitchTransformer (21) (48) . Layers in current MoE models have a single Attention block and use a 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 an inter-device IO-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) (24) , and FFNs in the same model perform better than models that use solely SSMs or MHA for the sequence-to-sequence 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.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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 is authorized to 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. A.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 OpenWebText, perform the comparison using SuperGLUE, build TensorRT-LLM engines, and replicate results can be found in 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, or as supplemental material. 
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.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at 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. Justification: The number of GPU hours necessary to train the models is 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 runs 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). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: The Code of Ethics was reviewed carefully for full compliance.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [No] 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 in this paper.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 out 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 or 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). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes] Justification: While this paper presents a few pretrained models, the focus is on the systemlevel efficiency gains, not the downstream performance of 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 a 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 we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: All assets such as datasets and codebases are cited and urls are provided were applicable.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 a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's 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 a 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 is used. 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

> [1] Wang, Shibo, et al. "Overlap communication with dependent computation via decomposition in large deep learning models." (ASPLOS’22)

> [2] Rajbhandari, Samyam, et al. "Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning." (SC’21)

2. 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.
questions: 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. 

Firstly, 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? 

Secondly, 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.

Besides, 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: limitations: Please refer to the Weaknesses and Questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: * **Limited benchmark**: The evaluation benchmarks primarily focus on language tasks, neglecting the application of Transformers to vision tasks.
* **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.
questions: 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?
* How does the approach benefit systems with varying inter-device configurations and bandwidth capacities?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

Another 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.
questions: questions: Please see above.
limitations: limitations: Yes

Comments: title: Energy Usage and Infrastructure Allocation For Training vs Inference
Thanks for the response
comment: We thank the reviewer for appreciating how our approach is compatible with existing techniques. Yes, training from scratch is indeed very expensive. However, as the usage of Transformers grows, inference will require a much larger fraction of the available power and infrastructure. As an example, take the case study presented by Wu et al. [1]; here, the authors find that 70% of available power at a major tech company goes to ML inference (Figure 3) as opposed to 20% for training. Moreover, this work[1] and data is from 2022 which is a little before the current "Gen AI" trend that has taken over the industry.

Sustainable AI: Environmental Implications, Challenges, and Opportunities, Wu et al., MLSYS 2022.
Thank you for the explanation. I now clearly understand Kraken's key idea and how it differs from related works such as [1, 2]. It's a notable strength that Kraken is compatible with many existing tools like TensorRT. However, the requirement to train from scratch significantly impacts its application in LLMs due to the high cost of re-training. Considering both its strengths and weaknesses, I’d like to adjust my rating to 5.
Thank you to the authors for the response. I will maintain my current rating.
Thanks for the response. My ratings remain the same.