---
title: High-performance Inferencing with Transformer Models on Spark
description: "High-performance inferencing with Transformer models on Spark using PySpark, Hugging Face, and AWS GPU instances for significant speed and cost savings.\n\n===\nSummary:\nThis blog post provides a detailed tutorial on performing high-performance inferencing with Transformer models on Spark using PySpark, Hugging Face, and AWS GPU instances. It covers the setup of driver and worker instances on DataBricks, partitioning data for parallel processing, and executing inferences using the Spectre pre-trained language model. The post highlights the benefits of using GPU instances for faster and more cost-effective inferencing, managing large datasets efficiently, and the ease of scaling and collaboration in a cloud-based environment. The technique discussed allows for processing millions of rows in hours, making it feasible to handle large datasets with transformer models."
---

[AurumNova Blog](https://aurumnova.com/blog)

# [High-performance Inferencing with Transformer Models on Spark](https://aurumnova.com/blog/high-performance-inferencing-with-transformer-models-on-spark)

 Written by [Daniel Sim](https://aurumnova.com/blog/author/daniel-sim) | Nov 18, 2021, 12:00:00 PM

## A tutorial with code using PySpark, Hugging Face, and AWS GPU instances

Are you looking for up to a 100-fold increase in speed and a 50% reduction in costs with Hugging Face or Tensorflow models? By utilizing GPU instances and Spark, we can perform inferencing on anywhere from two to hundreds of GPUs simultaneously, enhancing performance with ease.

## Overview

- Setup Driver and Worker instances
- Partitioning data for parallelization
- Inferencing with Transformer Models
- Discussion

 

## Setup Driver and Worker instances

For this tutorial, we will be using [DataBricks](https://databricks.com/), and you may sign up for a [free account](https://databricks.com/try-databricks) if you still need access to one. DataBricks must connect to a cloud hosting provider like [AWS](https://aws.amazon.com/), [Google Cloud Platform](https://cloud.google.com/), or [Microsoft Azure](https://azure.microsoft.com/) to run GPU instances.

This exercise will use the AWS GPU instances of type “[g4dn.large](https://aws.amazon.com/ec2/instance-types/g4/).” However, you may still follow these instructions using Google Cloud or Microsoft Azure and select an equivalent GPU instance.

Once you have set up your DataBricks account, log in and [create a cluster](https://docs.databricks.com/getting-started/quick-start.html#step-2-create-a-cluster) with the configuration shown below:

Next, [create a notebook](https://docs.databricks.com/getting-started/quick-start.html#step-3-create-a-notebook) and attach it to the cluster by selecting it in the dropdown menu:

Now, we are all set to code.

 

## Installing Hugging Face Transformers

Firstly, let us install the [Hugging Face](https://huggingface.co/) Transformers to the cluster.

Run this in the first cell of the notebook:

```
%pip install transformers==4.2
```

Hugging Face Transformers Python library installed on the cluster

Libraries installed this way are called Notebook-scoped Python libraries. They are convenient, but they must be run at the start of a session before other code because they reset the Python interpreter.

At this point, we start with actual Python code. In the next cell, run:

Congratulations! Hugging Face Transformers will be successfully installed if the above line runs without errors.

 

## Partitioning data for parallelization

The easiest way to create data that Spark can process in parallel is by creating a Spark [DataFrame](https://spark.apache.org/docs/latest/sql-programming-guide.html). For this exercise, a DataFrame with two rows of data will suffice:

We get this DataFrame:

The Transformer model for this exercise takes in two text inputs per row. We name them “**title**” and “**abstract**” here.

For the curious, here is an excellent article by [Laurent Leturgez](https://medium.com/@laurent.leturgez) that dives into Spark partitioning strategies:  
[On Spark Performance and partitioning strategies](https://medium.com/datalex/on-spark-performance-and-partitioning-strategies-72992bbbf150)

 

## Inferencing with Transformer Models

We shall use the fantastic [Pandas UDF for PySpark](https://databricks.com/blog/2017/10/30/introducing-vectorized-udfs-for-pyspark.html) to process the Spark DataFrame in memory-efficient partitions. Each partition in the DataFrame is presented to our code as a Pandas DataFrame, which you will see below, is called “**df**” as a parameter of the function “**embed\_func**.” A Pandas DataFrame makes it convenient to process data in Python.

Code for embed\_func():

You might have noticed two things from the code above:

- The code further splits the input text found in the Pandas DataFrame into chunks of 20 as defined by the variable “**batch\_size**.”
- We use [Spectre](https://huggingface.co/allenai/specter) by AllenAI — a pre-trained language model to generate document-level embedding of documents ([pre-print here](https://arxiv.org/pdf/2004.07180.pdf).) We can easily swap this for another [Hugging Face model like BERT](https://huggingface.co/bert-base-uncased).

##### Working around GPU memory limits

When the GPU makes inferences with this Hugging Face Transformer model, the inputs and outputs are stored in memory. GPU memory is limited, especially since a large transformer model requires a lot of memory to store its parameters, leaving comparatively little memory to hold the inputs and outputs.

Hence, we control the memory usage by inferencing just 20 rows at a time. Each time 20 rows are processed, we copy the output from the GPU to a NumPy array, which resides in CPU memory (which is more abundant) with the code at Line 21 with “**.cpu().detach().numpy()**”.

##### Finally - Transformer model inferencing on GPU

As mentioned above, this is where the execution of [Pandas UDF for PySpark](https://databricks.com/blog/2017/10/30/introducing-vectorized-udfs-for-pyspark.html) happens. In this case, the Pandas UDF is “**embed\_func**” itself. Read the link above to learn more about this powerful PySpark feature.

The resulting output of this exercise is a DataFrane containing rows of floating-point arrays that represent document embeddings. Spectre creates embeddings that are vectors of 768 floats.

 

## Discussion

I hope you see how Spark, DataBricks, and GPU instances make scaling up inferencing with large transformer models relatively trivial.

The technique shown here allows for running inference on millions of rows and completing it in hours instead of days or weeks. Thus, processing big data sets on large transformer models becomes feasible in more situations.

##### Cost Savings

But wait, there is more. Despite costing 5 to 20 times a CPU instance, inferencing done by a GPU instance can be more cost-effective because it is 30 to 100 times faster.

> Since we pay for the instance on an hourly basis, **time is money** here.

##### Less time spent on plumbing

Data is easily imported into DataBricks and saved as Parquet files on AWS S3 buckets or, even better, [Data Lake](https://docs.databricks.com/delta/delta-intro.html) tables (a.k.a. Hive tables on steroids). Then, we can manipulate the data via Spark Dataframes, which, as this article shows, is trivial to parallelize for transformation and inferencing.

All data, code, and computing are accessible and managed in one place “on the cloud.” This neat solution is also intrinsically scalable as the data grows from gigabytes to terabytes, making it more “future-proof.”

##### Collaborate Seamlessly

### Being a cloud-based solution means that as the team grows, we can add more people to the project to access data and code on notebooks securely. We can create charts for reports to share with other teams with just a few clicks.

<https://medium.com/tag/hugging-face?source=post_page-----beb82e71ecc9--------------------------------------->

[View full post](https://aurumnova.com/blog/high-performance-inferencing-with-transformer-models-on-spark)

```json
{
  "@context" : "http://schema.org",
  "@type" : "BlogPosting",
  "author" : {
    "@type" : "Person",
    "name" : "Daniel Sim"
  },
  "dateModified" : "2025-02-25T09:55:33.097Z",
  "datePublished" : "2021-11-18T12:00:00Z",
  "headline" : "High-performance Inferencing with Transformer Models on Spark",
  "image" : {
    "@type" : "ImageObject",
    "height" : 600,
    "url" : "https://49121148.fs1.hubspotusercontent-na1.net/hubfs/49121148/transformers_gpu_spark_feat_img_no_logo.png",
    "width" : 2000
  },
  "mainEntityOfPage" : "https://aurumnova.com/blog/high-performance-inferencing-with-transformer-models-on-spark",
  "publisher" : {
    "@type" : "Organization",
    "logo" : {
      "@type" : "ImageObject",
      "height" : 60,
      "url" : "/hs/hsstatic/content_shared_assets/static-1.4092/img/default-amp-logo.png",
      "width" : 60
    },
    "name" : "AurumNova Blog"
  }
}
```