PALM: A Efficient Performance Simulator for Tiled Accelerators with Large-scale Model Training

Based on the chosen parallelization strategy, models and input batches are distributed across workers. This makes communication and synchronization of data, like forward activation or weight/input gradients, among devices inevitable [32]. This traffic is typically formulated and processed through collective communications. Four primary collective communication operations are key contributors in DNN training [33, 34]: (i) reduce-scatter, (ii) all-gather, (iii) all-reduce, (iv) all-to-all. In Fig. 1(b), reduce-scatter operation sums all initial data in workers, resulting in each worker holding a portion of globally reduced data. The all-gather operation gathers the data initially distributed across workers, ensuring each worker possesses the complete data. All-reduce can be regarded as a combination of reduce-scatter followed by an all-gather operation. In the all-to-all pattern, each node is required to send a distinct portion of data to other nodes.

Fig. 1(c) illustrates the architecture for a tiled accelerator, which usually consists of multiple independent operating tiles. Each tile has its unique instruction queue, local memory and progresses at its own pace, which thus allows the tiled accelerators to specialize in supporting flexible dataflow and mapping. Moreover, the NoC is employed for transferring data among the tiles and synchronizing tiles at different stages throughout the program execution. Also, the NoC establishes connections among all tiles, as well as off-chip communication and memory controller blocks. As a result, each tile has access to the off-chip memory or other chips. Compared to traditional monolithic chips and single-tile SIMD GPUs, such architectures usually exhibit higher execution efficiency. Such improved efficiency comes from employing optimized dataflow strategies to spatially/temporally partition data across the tiles and fine-grained scheduling.

A sensible modeling approach is essential for simulating the training process of DL models on a substantial number of tiles, as depicted in Fig. 1(c). Real tiled accelerator systems exhibit a range of scales, from 4$\times$4 and 10$\times$12 [39, 40] to a wafer-scale architecture of 633$\times$633 [41]. A straightforward but very coarse approach is to assign each tile an independent thread or event queue. However, handling a large number of tiles using such a simulation mechanism would lead to a notable increase in software overhead. Therefore, to efficiently implement a tiled accelerator simulator for DL training tasks, it is imperative to introduce a unique scheduling mechanism among tiled accelerators.

Current simulators lack interfaces that support arbitrary parallelism strategies. Typically, users need to extract computation graphs with embedded parallelism information from established DL frameworks such as PyTorch and TensorFlow. This limitation prevents the direct iteration of parallelism strategies based on simulation results. Additionally, existing simulators lack support for various types of PP which is an important parallelism strategy of LLM, nor have they discussed the differences in bubble and capacity requirements under PP. In fact, the proposal of PP is mainly aimed at solving the storage problem of LLM, which has problems in resource utilization. The advantage of PP on tiled accelerators is that it fits the characteristics of a large number of tiles, can more evenly split the pipeline, increase the number of pipeline stages, and reduce the bubble ratio. TP and DP are two inherent parallelism strategies. In the tiled accelerators, when some tiles/cores form a tile group to execute the same operator, certain dimensions must be segmented as illustrated in Fig. 1(a). Hence, it is crucial to offer a flexible user-visible interface that supports parallelism across various dimensions.

SRAM, being faster but costlier than DRAM, is utilized to temporarily store data for computation and exchange data with DRAM. Table I indicates that the SRAM capacity per computing power unit in tiled accelerators surpasses that in traditional GPUs. Specifically, WSE’s SRAM capacity per computing unit is nearly $26\times$ that of GPU A100. Studies [31, 40] explore using SRAM to statically store frequently read data, accelerating tile computation based on dataflow. Recognizing the significant role of SRAM in computation, memory access, and communication is thus reasonable.

Efficient model training relies on DRAM with large capacity and high bandwidth. DRAM is crucial for storing extensive model parameters, intermediate activations, and optimizer states during training. Tiled accelerators, designed for high-density computing power, differ significantly from GPUs in their memory hierarchy. For example, in the WSE-2 system [41], of which the computing power is equivalent to 46 GPUs, there is no on-wafer DRAM; instead, DRAM is located off the wafer. Consequently, DRAM access in tiled accelerators becomes costly due to NoC routing, as depicted in Fig. 1(c). Therefore, modeling DRAM behavior is crucial to accurately reflect practical behaviors of tiled accelerators.

NoC acts as a physical bridge among tiles [39, 40, 42], impacting communication between pipeline stages generated by mapping and parallelism, as well as intra-stage communication. Frequent DRAM access will occupy NoC bandwidth. In Table I, various tiled accelerators exhibit different NoC hop counts to DRAM, presenting a disadvantage for on-chip access tasks. Additionally, in the same table, the Link bandwidth-to-DRAM bandwidth ratio is higher in tiled accelerators, providing an advantage for communication tasks.

In summary, it is essential to model the behavior of SRAM, DRAM, and NoC during the training process to accurately reflect the architectural characteristics of tiled accelerators.

We distinguish the concept between pipeline scheduling mechanism and pipeline parallelism. The former concerns modeling the training process effectively, while the latter involves partitioning the computation graph into stages, as discussed in the next subsection. The pipeline scheduling includes two mechanisms: pipeline execution and layer-wise execution [31]. In our modeling, layer-wise execution is treated as pipeline execution with a depth of 1. Fig. 4 illustrates the pipeline scheduling process: the computation graph is partitioned into stages (S0, S1, S2), and each stage is mapped to a tile group based on the parallelism strategy. The pipeline is divided into three processes: Forward (FD) representing the forward computation of all operators in each stage, Backward (BD) representing the backward propagation of all operators in each stage, and Gradient Update (GU) representing the gradient update process. Additionally, PALM defines Act/Grad Pass to transfer activations/gradients across stages, serving as the start signal for the next stage. In DL, a batch (mini-batch) is taken as the period for gradient updates. To reduce the pipeline bubble ratio, a batch is divided into multiple micro-batches, with one micro-batch executing FD and BD. Once all micro-batches are completed, GU is executed. Data_Fetch simulates the input data fetching of one micro-batch, representing the start of the first stage S0. In our scheduling mechanism, GPipe[19] and 1F1B[20] scheduling in Fig. 3 are supported . PALM places one of the four types of events into the Virtual Tile Executor based on the signal selected by Prior Selector. For example, in the 1F1B pipeline, priority is accorded to the execution of BD over FD. The Act/Grad Pass between different stages is accomplished through communication events on NoC. This process is primarily determined by the dependency relationships between adjacent operators in the different stages, which will be discussed in the next subsection.

Within each stage, operators are executed in the order of their dependency relationships such as op $B$ and $C$ of S1 in Fig. 3, as layer-wise execution does. Operators without dependencies are executed in the pre-order rule in the computation graph or in parallel. When tiles/cores execute the same operator, they are called a tile group. In the tile analysis level (tile analyzer in Fig. 2), PALM assumes different tiles in each tile group have the same computation and memory access cost. Therefore, each stage exclusively furnishes one or a few simulated tiles representing these tiles in tile group, denoted as virtual tiles. We have coined this modeling method as Virtual Tile Aggregation.

We assume that a single tile mainly consists of two entities: the tile internal logic unit and NoC router which have their own event queue. Additionally, we suppose the number of tiles is $N\times N$, and the number of stages $S$ is less than or equal to the number of layers $M$ in the computation graph. The naive modeling complexity is $\mathcal{O}(2N^{2})$ for all tiles, while PALM with virtual tile aggregation reduces it to $\mathcal{O}(N^{2}+M)$. By incorporating an analytical model for the NoC, the complexity is further reduced to $\mathcal{O}(M)$. Given that $M$ typically falls in the range of tens to hundreds, this significantly alleviates the modeling overhead.

In PALM, each operator also generates three types of events: forward, backward, and gradient update. Each type of event is further divided into computation, communication, and memory access tasks. Fig. 5 describes the main events during backward execution. For each operator, the backward process includes loss computation, activation re-computation, and gradient computation. Activation re-computation occurs only when there is insufficient memory capacity. Each sub-process requires accessing data from memory for computation, with non-negligible communication overhead. The next sub-process begins only after the completion of the current sub-process. For example, in Recompute sub-process, we wait for the completion of the Loss computation event before entering Gradient sub-process. During the three sub-processes, DP communication from the previous operator can overlap with the current operator’s execution. The forward process is similar to the re-computation in the backward process and is not separately listed here. The main events in the gradient update process only include full-precision weights load from DRAM and store back to DRAM, and we have omitted the accumulation computation in the gradient update process.

❶ PP. PP partitions operators of the computation graph into different stages to minimize the pipeline bubble. The ideal execution time in the pipeline training scenario can be evaluated using Eq. (1).

where $ET$ is executing time, $B$ is the batch size and $b$ is the micro-batch size. In fact, on the tiled accelerator, the execution time is influenced by the spatial position of the physical tiles corresponding to the stages. We will further discuss this phenomenon with experiments in Section V-B2. PALM takes into account that PP results in differences in memory capacity requirements, as discussed in [41]. Considering a training pipeline with $S$ stages, activations from each stage are stored in the FD process, until they are consumed for GU in the BD process. For example, the first stage should store $S$ times the activation in 1F1B, and $B$ times the activation in GPipe as illustrated in Fig. 3. Incorporating the aforementioned considerations into PP modeling, PALM supports users to bind stages based on tile IDs and op IDs with Adaptive Parallelism Interface in Fig. 2, and provides a default way for DL models to allocate stages based on computing power requirements.

❷ TP and DP. We analyzed the communication size of all-reduce generated by TP and DP strategies in common operators, as shown in Table III. PALM partitions mapped physical tile groups into communication groups, automatically inserting collective communication events into the tile group event queue. Taking the simple linear operator as an example: The linear operator $Y=WX^{T}$ has four dimensions $(B,M,N,K)$, where $B$ represents batch size, $K$ represents the reduce dimension, $M,N$ represents the output dimension, $X^{(N\times K)}$represents input, $W^{(M\times K)}$ represents weights, and $Y^{(M\times N)}$ represents output. The dimensions $(b,m,n,k)$ represent the parallelism degree for each corresponding dimension. If we map the operator onto 16 tiles from 0 to 15, it is essential to ensure that $b\times m\times n\times k=16$. The parallelism strategy can be configured by the user as $(2,2,2,2)$ or $(4,4,1,1)$, and so on. Further, corresponding communication groups are automatically generated. During the FD, BD, and GU processes, there is a need for collective communication in the corresponding tile groups. The parallelism of other operators like Conv2 and Pool are the same. For simplicity, we assume that the input shape of Conv2 or Pool is $(B,C,I,I)$, the shape of weight is $(W,W,K)$, and the shape of output is $(B,K,O,O)$. Specially, $K$ is equal to 1 in Pool operator. The communication size of all-reduce is also represented in Table III. For transformer operator, it is a combination of a series of linear operators. And the shapes of input and output are $(B,S,H)$. We support both DP ($N_{d}$) and TP ($N_{m}$) as described by Megatron[20]. The communication size generated by splitting these linear operators is accumulated. These parallelism dimensions such as $(b,m,n,k)$ and $(N_{d},N_{m})$ can also be configured by the user with the interface in Fig. 2.

❶ SRAM allocation. PALM holds the view that SRAM primarily influences DRAM access. Alg. 1 explains the main modeling idea: Operators are split by the parallelism strategy and their corresponding tiles are taken as input to obtain the corresponding SRAM strategy and DRAM access size in forward process. Strategies $S_{WSG\_ACT}$, $S_{WSG}$, $S_{ACT}$, $S_{PTY}$ respectively represent weights, optimizer states, weight gradients (WSG) and input/output activation ($ACT_{IN},ACT_{OUT}$) either statically stored in on-chip SRAM, one of them stored in on-chip SRAM, or none stored on-chip. It is worth noting that when WSG and ACT cannot be retained in SRAM for a long time, PALM adopts a penalty strategy $S_{PTY}$, modeling extra DRAM accesses for WSG and ACT. When $ACT\geq W$, we use input stationary (IS), otherwise, we use weight stationary (WS). PALM considers storage differences brought about by the optimizer. For optimizer Adam, it requires storage for first-order and second-order moments related to weights, and gradients of backward activations, significantly increasing storage requirements. If optimizer SGD is used, there is no overhead for optimizer states. During inference, there is no storage overhead for gradients. Alg. 1 only lists the DRAM access size for the forward process. The analysis for the backward and gradient update process follows the same methodology, thus being neglected here.

❷ Detailed NoC model. The ideal communication latency of the NoC can be obtained using Eq. (2), where $Link\_Time$ represents single hop link delay and $Hops$ represents the total number of hops in the communication path. However, the analytical model [38] does not consider whether all links are idle at a given moment in a transmission path. Hence, the specific latency of contention_delay can not be obtained by the analytical method. In the presence of congestion, the communication time may degrade to Eq. (3) in the analytical model, which means a hop-by-hop data transmission, without forming a pipeline transmission along the link. But it is equivalent to reducing the bandwidth of the NoC by $Hops$ times. Even the modeling of the latter cannot guarantee that the single hop transmission is not occupied by other tasks.

PALM considers NoC congestion, treating the link as an exclusive resource during execution. When a link is occupied by the current task, the execution time can be obtained by Eq. (2). Communication tasks can only be executed when needed link are not occupied. Otherwise, they will wait for the release of resources.

❸ Detailed DRAM model. Through the analysis of SRAM, the size of DRAM memory access has been determined, and ideally, the memory access latency $Access\_Time$ can be obtained using Eq. (4). However, in the tiled accelerator, the DRAM is shared among tiles. Due to the varying distances of different tiles from DRAM and the different times they initiate memory access requests, the understanding of whether the bandwidth ($BW_{DRAM}$) is occupied at a particular moment is not clear enough. Eq. (4) cannot accurately represent memory access latency.

Based on the above equation, PALM constructs a memory access model for edge-shared DRAM in tiled accelerators. PALM considers DRAM bandwidth as a resource that is occupied during execution like the NoC model. The data transmission time, denoted as $NoC\_Time$, through the NoC has been taken into account. Therefore, the total DRAM access time $DRAM\_Time$ of a tile can be obtained using Eq. (5).

We explore the influence of wafer-scale architecture on the optimal parallelism of LLM. Based on PALM, we build a wafer-scale architecture with specific parameters, as shown in Table VI. The overall system consists of a $5\times 4$ tile array with $4\times 4$ core per tile, communicated with tile-to-tile and core-to-core NoC. We have selected models T-18B, T-76B, and T-145B as the baseline in Table VII, with (TP=8, DP=2, PP=20). The performance of the baseline is close to the result presented in Table IV.

Due to the bandwidth limitations of the GPU cluster architecture, there is only a single choice for its communication strategy [51]. In wafer-scale systems, close intra- and inter-bandwidth can support different communication strategies to minimize costs. Adapter tiles[37] are the tiles within the destination group receiving data from the source tile group.

Two communication strategies for inter-tile groups are depicted in Fig 11. The first involves all-reduce within the source group, data transmission to the destination, and broadcast within the destination. The second reduces the source based on adapters, performs inter-tile transmission, and conducts all-reduce and broadcast in the destination.

Strategy 1’s inter-tile communication time is shown by Formula 7, while Strategy 2’s is shown by Formula 8. In the formulas SG represents the source tile group, DG represents the destination tile group, AR represents all-reduce, R represents reduce, and B represents broadcast.

Based on BERT-base model, we assess the performance of two communication strategies. The first set of experiments compares 12 tile source and destination groups under ring shape all-reduce, while the second adds a tile to disrupt ring formation and reassessing performance.

In Fig 12a, when a ring structure is formed in the source tile group, strategy 1 outperforms strategy 2 in inter-communication performance. This is due to the smaller overall latency of ring all-reduce, resulting in a smaller communication time compared to strategy 2. Moreover, with more adapters participating in inter-communication, the performance of strategy 1 gradually improves by reducing broadcast time in the destination tile group. In Fig 12b, when a ring structure cannot be formed, strategy 2 shows better communication performance. In this case, the total time of the reduce and the all-reduce in strategy 2 is smaller than the all-reduce time in the source group of strategy 1. Additionally, the performance of strategy 2 initially improves and then declines as the number of adapters increases, due to the trade-off between the reduce cost and the all-reduce time among adapters.

According to the result, it is evident that inter-tile communication in ring shape configurations exhibits superior performance under strategy 1, leading to 3.08$\times$ performance gap over strategy 2. Conversely, non-ring shapes are more suitable for the adoption of strategy 2, with a performance increase of approximately 1.23$\times$ compared with strategy 1.