Insight Gained from Migrating a Machine Learning Model to Intelligence Processing Units

In this work, AI/ML models were tested with both GPUs and IPUs. All models were optimized for each platform architecture to guarantee that peak performance could be achieved. The objective of an ML benchmark is to measure the time required to achieve a specific level of accuracy during training, or to assess the speed and efficiency of algorithmic performance in real-world inference situations, focusing on latency and throughput. This comparison ensures fairness among various training methods by appropriately considering both throughput and accuracy. The most effective training strategy entails choosing parameters that produce both high throughput and high accuracy. Therefore, some models may be more compatible and more performant with specific hardware architectures than others, depending on how training parameters impact the model's accuracy.

We performed benchmarks on the Bow-IPU architecture using a couple of ML models. All models were implemented with the core network being the convolutional neural network (CNN), which is also the backbone for many ML tasks such as object detection and segmentation in computer vision. One of many models is the well-known ResNet [5] model which contains different skip connection layers to mainly combat the vanishing gradient problem. This model is heavily optimized for both accelerators to achieve peak performance.

On the other hand, we also analyzed the process of porting a specialized model from GPU to IPU. Besides adding simple wrappers around regular ML modules to make it run on IPU, we addressed the optimization with IPU to fully utilize the hardware for training. Additionally, several advanced options, such as pipelining and gradient accumulation, were investigated to boost IPU performance.

Theoretically, we desired to achieve linear scaling, meaning that if the number of accelerators doubles, the throughput also doubles, thereby reducing the training time by half. However, in practice, doubling the number of accelerators does not result in halving the training time. This is due to various conditions of hardware and non-compute resources, such as memory and networking. Moreover, we are much more interested in the performance of the overall hardware system than individual components. Therefore, many attributes of ML models in both IPU and GPU environments are fine-tuned to achieve the best possible optimization.

The escalating demand for safer and more efficient battery technologies has spurred a significant shift towards non-flammable solid-state batteries, and away from their lithium-ion counterparts (which can ignite or explode). While solid-state batteries have safety and energy density advantages, a notable challenge persists in the form of low power density due to non-optimal cathode microstructures that lead to high tortuosity and poor ionic conductivity [7]. High Ionic conductivity corresponds to low internal resistance, resulting in more efficient ion transport. This high ionic conductivity allows fast movement of ions within the composite cathode, enabling higher actual capacity at fast charging and discharging rates, eventually contributing to higher cell energy density, sustained stability, and consistency of performance over numerous charge and discharge cycles. To study the effect of cathode microstructure on ionic conductivity, a general workflow was proposed to predict the ionic conductivity from scanning electron microscopy (SEM) images [10]. The general workflow consists of two steps: image denoising for data preparation and effective conductivity prediction. Cross-sectional SEM images often contain striped noise from ion beam polishing, which impacts the accuracy of conductivity calculations. The denoising model, U-net, effectively eliminates striped noise from SEM images and converts them into clean binary images. The binary images have two regions: conductive (yellow) and non-conductive (black) regions, corresponding to solid electrolyte and cathode active material, respectively. The effective conductivity prediction model is a CNN implemented in PyTorch as shown in Figure 1. The model takes binary SEM images and the ionic conductivity value of the conductive region as inputs to generate predictions for effective or total ionic conductivity. The conductivity model is the model that we aim to port to the Bow-IPU.

The image datasets consist of training, validation, and testing sets, distributed in a ratio of 1216:152:152. Binary images of size 256x256 were obtained by denoising and cropping SEM images. The reference effective conductivity value of each image is determined using the finite difference method (FDM). FDM is a highly accurate method for calculating current, resistance, and effective conductivity. However, FDM accrues significant computational cost, especially in situations involving large-sized images or when processing a substantial quantity of images. In the work of Mai, et al., FDM was employed to calculate the effective conductivity value for each image, serving as the target value of the effective conductivity prediction model [9]. The acquired image datasets, along with their corresponding effective conductivity values, were utilized to train the CNN effective conductivity prediction model. The model allows a fast prediction of a composite cathode's performance from its SEM images by outputting the predicted effective ionic conductivity value.

Different hardware architectures require different techniques and workflows to improve the performance of an ML model. The specialized model is originally implemented to run on GPUs. Thus, when porting to IPUs, several modifications are required to make our model run optimally. Fortunately, the amount of work to modify our workflow is minimal and can be done quickly.

IPU Options

The first modification that we perform is to add Poptorch options to our program. Graphcore provides a placeholder for many options that can be defined for the hardware; thus all options are wrapped in a unique and central variable. IPU options define how IPUs control the workflow of the training and inference of a model. One of the many important options is the device iterations option. The option organizes batches of data on a queue that can be used directly by the IPU, which significantly reduces the time spent on data transferring between CPUs and IPUs. Another advantageous attribute involves instructing the IPU to execute convolution and matrix multiplication in float16, thereby reducing the computational workload, as these two operations account for the majority of computations.

Dataloader

For implementation on IPUs, we opt for a dataloader specifically designed for IPUs, as opposed to utilizing PyTorch's dataloader. The IPU dataloader provides the flexibility to incorporate additional options for modifying data processing methods. With the predefined Poptorch options mentioned above, we simply add the unique option variable to the IPU dataloader. As a result, we achieve an efficient data pipeline for IPU training and inference processes.

Pipelining

Pipelining is an effective method that we use to improve the performance of our model. Pipelining involves dividing the entire model into several computational stages, with each stage's output being fed into the next. Each stage is allocated an IPU to perform its forward and backward propagation. By distributing the stages across various IPUs and ensuring an adequate number of mini-batches, all stages will concurrently process a data batch after an initial "ramp-up" phase. These stages execute concurrently across multiple IPUs, which enables a more efficient utilization of resources compared to relying solely on sharding. As a result, this pipeline maximizes the parallel processing capabilities of all IPUs involved, leading to enhancements in processing efficiency, throughput, and latency performance.

When employing pipelining with a model on the IPU, it is necessary to specify the gradient accumulation option which enables concurrent processing of several mini-batches through the pipeline, enabling IPUs to execute multiple stages simultaneously. Particularly, in the backpropagation, gradients accumulate across multiple mini-batches. After processing the predefined number of accumulation, the model parameters are updated with the average value of gradients. With gradient accumulation enabled, the global batch size is defined to be equal to the multiplication of the mini-batch size, the number of replicas, and the number of gradient accumulation.

Graphcore provides several methods to pipeline a model. All of them are equally efficient and the choice of which method to use depends on the user's preference. Using the provided ‘poptorch.BeginBlock’ wrapper, we are able to control the partitioning of the model's granularly. The process of organizing blocks into different stages is straightforward, which only requires access to components of an ML model.

To further boost the performance of multiple IPUs via pipelining, the amount of workload and memory of each partition should be equally distributed. It is essential to consider that these partitions will be allocated and executed across multiple IPUs. Thus, the way in which we segment a model will have a direct influence on both memory usage and performance at each stage. Subsequent compromises must be taken into account:

• Stages must be accommodated within memory, taking into account both active and intermittently active memory. A report from the PopVision tool can assist with identifying which Ops and tensors utilize the most memory.
  
• Stages should exhibit comparable execution durations to prevent IPUs from remaining inactive while others finish the computation. The time taken by the longest stage dictates the duration of a single step. Again, the PopVision execution trace can aid in identifying stages with imbalances.
  
• Communication between IPUs incurs a slower rate compared to accessing memory within a single IPU and should be kept to a minimum.
  

Even though Graphcore provides convenient methods to pipeline a model, the task still requires manually partitioning the original blocks of a model into a corresponding pipelined block. This step requires a good understanding of all components of the model. The conductivity model is built upon many layers. Based on the model implementation, we group all layers into 10 different blocks, each of which contains a different number of parameters. After analyzing the memory usage and computational workload during pipelining with PopVision, we obtain configurations for pipelining with a different number of IPUs as shown in Table 1.

As can be seen from Table 1, there is no single IPU model. The reason behind this is that the model cannot fit into a single Bow-IPU with an input size of 256x256. This is due to the limited memory of each IPU in the system. To overcome this problem, pipelining is the method Graphcore recommends. Therefore, we came up with three pipelined models:

The number in the second column of Table 1 indicates how many layers are being placed in multiple IPUs. For instance, in the case of pipelining with two IPUs, the first six layers are placed on the first IPU with the index being 0, while the remaining four layers are placed in the second IPU with the index being 1. The placement of layers is dictated by the number of parameters at each layer. The number of parameters gradually increases from the first layer and peaks at the eighth layer before drastically decreasing with the last two layers. To further confirm our partitioning, PopVision is used to check the workload of each stage. Results from PopVision are in agreement with our pipeline configurations to achieve peak performance for IPUs.

Benchmarking on GPUs is performed on NVIDIA A100 (40GB, PCle) [13] on the NSF FASTER cluster at Texas A&M University. These GPUs are composed over the PCIe fabric to dual-socket nodes powered by 64-core processors (two 32-core Intel Xeon 8352Y Ice Lake processors)

On the other hand, IPU calculations are performed on Graphcore Bow-IPU compute nodes provided by ACES. Each compute node includes 16 IPUs. Graphcore provides their own implementation of both Tensorflow and PyTorch through the Poplar software development toolkit. Graphcore also provides benchmarking tools and documentation for widely-used AI/ML models, and comprehensive guidance on ML workload environments [3, 6]. It is imperative to emphasize that benchmarking measurements for training are conducted starting from the second epoch onward. This is because ML models undergo compilation during the initial epoch, a procedure that usually demands a significant amount of time. Consequently, it is more precise to assess benchmarking criteria from the second epoch onward.

Graphcore IPUs allow training and inference to be performed in parallel with single or multiple hosts. The number of hosts (number of instances) can be specified using the command ‘poprun’ to launch different IPU nodes. Furthermore, the number of replicas (model copies) can also be defined to create multiple copies of the same graph of a model. Based on these two parameters, ‘poprun’ automatically allocates the correct number of IPUs for the given task. The total number of IPUs is a result of the multiplication between the number of instances and the number of replicas. Each host (instance) manages the input/output of data, but it does not directly affect the computation of IPUs. As a result, this parameter is specified based on the workload of data transferring between IPUs and hosts. In general, the number of instances is equal to or divisible by the number of replicas.

For benchmarking on IPUs with computer vision models, we used the same PyTorch ResNet50 model that was used in the previous paper [11, 12]. ResNet model [5] is a popular CNN model for its innovation of skip connections to address the gradient degradation problem. It also serves as the backbone for many models of image classification, object detection, semantic and instance segmentation, etc.

Figure 2 shows the image throughput of PyTorch ResNet50 varies with the number of epochs on ACES Bow and Colossus IPU systems respectively across 8 IPUs. Initially, for both Bow and Colossus IPU systems, the first epoch exhibits lower performance primarily because of graph compilation overhead. However, the performance improves and stabilizes after the second epoch. Therefore, the performance values reported in this study are obtained after 2 epochs. Also, we can see the performance comparison of the two systems is evident. On average, the Bow IPU system outperforms the Colossus system by approximately 42%. This is consistent with the comparison of the IPU frequencies. The Bow IPU frequency is 1.85 GHz. The Colossus IPU frequency is 1.325 GHz. And the difference is around 40%.

In deep learning models, batch size can impact how efficiently hardware resources are utilized and different batch sizes may lead to different levels of memory usage and computational efficiency. We study the image throughput variation of the PyTorch ResNet50 model with the micro-batch sizes across 8 IPUs (Figure 3). As shown in the figure, there is a clear positive correlation between the micro-batch size and the throughput. At smaller micro-batch sizes, the performance of throughput increases, indicating higher efficiency in processing images at a higher rate with micro-batch size. However, at larger micro-batch sizes, the increase slows down because of the memory and computation being almost saturated.

Monitoring the training loss and accuracy of deep learning models over epochs can help assess whether hte models are converging. Figure 4 depicts the training loss and accuracy of the PyTorch ResNet50 model with the number of epochs on both ACES Bow and Colossus IPU systems. We can see that the training loss and accuracy on the two systems are very close. The accuracy gradually increases from approximately 1.2% to nearly 50%, while the training loss steadily decreases from around 6.5 to about 3.2 over 10 epochs. In the final epochs, the trends start to plateau.

Figure 5 shows the performance of Meta's Large Language Model for Text Generation (LLaMa2-7B) with 7 billion parameters on the ACES Bow IPU system. This evaluation was conducted with micro-batch size of 2 and a sequence length of 1024. The model is obtained from the Gradient Hugging Face GitHub repository. From the figure, we can see that as the number of IPUs increases, the throughput (measured in generated tokens/sec) increases almost linearly, indicating enhanced processing efficiency when more IPUs are involved, and the latency decreases, indicating a quicker response. This proves the scalability and efficiency improvements achievable by allocating more IPUs to LLaMa2-7B model inference.