An Architecture for Accelerated Large-Scale Inference of Transformer-Based Language Models (2021)

Abstract: This work demonstrates the development process of a machine learning architecture for inference that can scale to a large volume of requests. We used a BERT model that was fine-tuned for emotion analysis, returning a probability distribution of emotions given a paragraph. The model was deployed as a gRPC service on Kubernetes. Apache Spark was used to perform inference in batches by calling the service. We encountered some performance and concurrency challenges and created solutions to achieve faster running time. Starting with 200 successful inference requests per minute, we were able to achieve as high as 18 thousand successful requests per minute with the same batch job resource allocation. As a result, we successfully stored emotion probabilities for 95 million paragraphs within 96 hours.

## introduction as data in organizations becomes more available for analysis, it is crucial to develop efficient machine learning pipelines. previous work @xcite has highlighted the growing number of data centers and their energy and pollution repercussions. machine learning models that require less computational resources to generate accurate results reduce these externalities. on the other hand, many machine learning applications also require results in nearly real-time in order to be viable and may also require results from as many data samples as possible in order to produce accurate insights. hence, there are also opportunity costs associated with missed service-level objectives. attention-based language models such as bert @xcite are often chosen for their relative efficiency, and empirical power. compared to recurrent neural networks @xcite , each step in a transformer layer @xcite has direct access to all * work done while the author was working at wattpad. other steps and can be computed in parallel, which can make both training and inference faster. bert also easily accommodates different applications by allowing the fine-tuning of its parameters on different tasks. despite these benefits, exposing these models and communicating with them efficiently possesses some challenges. machine learning frameworks are often used to train, evaluate, and perform inference on predictive models. tensorflow @xcite has been shown to be a reliable system that can operate at a large scale. a sub-component called tensor-flow serving allows loading models as services that handle inference requests concurrently. system architectures for inference have changed over time. initial approaches favored offline settings where batch jobs make use of distributed platforms to load models and data within the same process and perform inference. for example, ijari, 2017 suggested an architecture that uses apache hadoop @xcite and apache pig for largescale data processing, where results are written to a hadoop distributed file system (hdfs) for later consumption. newer distributed platforms such as apache spark @xcite have gained prominence because of their memory optimizations and more versatile apis, compared to apache hadoop @xcite . as part of this architecture, inference services would often be reserved for applications that require faster responses. the batch-based and service-based platforms have different use cases and often run in isolation. collocating data and models in a batch job has some disadvantages. loading models in the same process as the data forces them both to scale the same way. moreover, models are forced to be implemented using the programming languages supported by the distributed data platform. ## architecture design we deployed our model as a tensorflow service in a kubernetes cluster. a sidecar service preprocessed and vectorized paragraphs and forwarded requests to this service. we used grpc to communicate with the services, foot_1 which is an efficient communication protocol on http/2. both nearly realtime and offline use cases made calls to these services. we used apache spark for batch processing, which we ran on amazon's aws emr service. 3 our batch job was developed using apache spark's python api (pyspark). the batch job fetched a dataset of relevant paragraphs, called the inference service, and stored the results. the job had two modes: a backfill mode and a daily mode, which ran on a subset of mutated and new paragraphs. this batch job was part of a data pipeline, scheduled using apache airflow foot_3 and luigi. 5 figure 1 shows the main components of this architecture. ## architecture optimization our initial approach, which used the configuration in the previous section, was resilient to failures but performed slowly at around 4 requests per second during the inference step. with a backfill target of 95 million paragraphs, running this job was intractable. our investigations concluded that the issues were rooted in a low request pressure on the backend services. thus, the sections below describe the steps taken to address these issues and speed up the inference process. ## results all steps in section 3 improve the batch job speed and results in satisfactory performance. the data pipeline is able to produce inference results for more than 95 million paragraphs in around 96 hours with an inference speed of around 300 requests per second. the semaphore value is set to 50. ## conclusion this paper discussed a successful machine learning architecture for both online and offline inference that centralizes models as services. we present solutions that use concurrency to increase the inference speed of offline batch jobs in apache spark. because of this, the majority of resources are still assigned to these services, and the batch job resources grow at a much smaller rate in comparison. we used a resource-intensive language model for emotion classification, where we demonstrated how proper tuning of tensorflow serving and kubernetes can improve the service's performance. we also showed that by parallelizing the calls made to the service in pyspark, we can significantly improve inference speed. finally, results were presented that provide useful insights into the inference performance. together, all these components resulted in a satisfactory architecture, which resulted in the emotion probabilities of 95 million paragraphs to be stored within 96 hours. we hope the architecture can be applied to other language tasks or machine learning models.