Reducing Inference Times by 87% for Darwinbox’s Talent Search Engine Using AWS Inferentia

By Rony K Roy, Sr. Specialist Partner Solutions Architect, AI/ML – AWS
By Anirban De, Lead Data Analytics and AI – Minfy Technologies
By Blesson Davis, Data Scientist – Minfy Technologies
By Manish Dixit, Sr. Partner Management SA – AWS

Minfy Technologies

Darwinbox is a new-age, enterprise-ready human capital management (HCM) platform that enables organizations to automate day-to-day HR processes while simplifying human interactions and delivering actionable insights to build better workplaces.

Darwinbox is building an in-house product for finding the right talent for defined job requirements. As its user base grew into thousands, the 2.2+ seconds it took to match candidates’ resumes to job descriptions was no longer acceptable.

This led Darwinbox to seek solutions that sped up the inferencing times the platform’s PyTorch models took in processing and matching resumes against job descriptions.

In this post, we will explore how Minfy Technologies helps Darwinbox overcome these challenges and deliver an 87% reduction in inferencing times using AWS Inferentia.

Minfy is an AWS Premier Tier Services Partner and Managed Service Provider (MSP) that helps customers navigate digital journeys leveraging artificial intelligence (AI) and cloud.

Project Scope and Background

Darwinbox leverages the transformer-based PyTorch models—one model for semantic search and the other for named entity recognition—for its search and match solution to score and shortlist the best fit resume for a given job description. This solution was developed by training two HuggingFace models on proprietary datasets. However, inferencing a large number of resumes was taking a significant amount of time.

Darwinbox wanted to reduce the time to process a resume as much as possible and thus create a better user experience for end users. The team observed a pattern that reveals how the traffic fluctuates throughout the week, and wanted minimal changes in their approach and existing code while adhering to the requirements of the engineering team.

Solution Overview

Darwinbox’s existing PyTorch models were running on compute-optimized instances, but were unable to deliver the performance the organization needed. A deep learning accelerator-powered instance using GPUs or AWS Inferentia would be optimal.

Based on the customer’s specific requirements and architecture, Minfy approached the challenge in two ways:

• Run these PyTorch models on deep learning accelerator-powered instances purpose-built for natural language processing (NLP) with minimal code changes. This was achieved by compiling the existing models using the Neuron SDK and extending the pre-built Amazon SageMaker Docker containers with the required dependencies.
• The real-time endpoint was configured for both models separately and used AWS Lambda with Amazon API Gateway to invoke the endpoints. This was configured per the API design requirement of the customer’s DevOps team. To run inference on large datasets, Minfy proposed sending requests in parallel as mini batches of a few resumes. SageMaker Inference Recommender was used to determine how many concurrent calls can be made to the model per minute.

Figure 1 – High level solution architecture.

The diagram above gives a high-level overview of the architecture (development stage). Key services used are:

• Amazon SageMaker: To process the data, train the model, and deploy the model to endpoints, which are configured with autoscaling to address high traffic at peak hours.
• Amazon S3: All of the data used for training the models were stored in Amazon Simple Storage Service (Amazon S3) buckets. After training, the model artifacts were also stored in S3 buckets.
• AWS Lambda: Lambda was configured to invoke the two different endpoints in a linear manner and send the response to Amazon API Gateway.
• Amazon API Gateway: The entire solution was made available as an API using Amazon API Gateway, which made it easy to manage and monitor REST API at scale.
• AWS Inferentia : Darwinbox’s PyTorch models were compiled into Inferentia-compatible formats using the Neuron SDK.

In this post, we will be focusing on optimizing the HuggingFace models by reducing their inference time for large datasets.

HuggingFace Transformers

HuggingFace is a popular platform when it comes to transformer libraries. The HuggingFace Hub has a collection of models for different types of NLP applications like sentiment analysis, named entity recognition, and text classification to name a few. These models can be directly used in Amazon SageMaker to train them using custom data and deploy them in production.

See the documentation for blogs and resources to get started with HuggingFace on Amazon SageMaker.

Note that this use case was started before the introduction of Optimum Neuron, which simplifies the adoption of Inferentia and Trainium accelerators even further. SageMaker provides built-in container images for training and to run inference on these models. These containers are stored in a public Amazon Elastic Container Registry (Amazon ECR) repository.

For this use case, we’ll explore how to use a PyTorch container and install custom libraries and save it to our ECR repository.

Building and Deploying the Models

NLP and computer vision workloads deal with a large matrix of numbers at the fundamental level. GPUs became popular for their ability to perform these matrix calculations in parallel, thanks to the large number of threads they employ.

Using GPU-powered instances speeds up the model training process by a significant amount. Once the training is completed, one of the key requirements during the inference phase is to make these models available to millions of users at the same time.

Inferencing at Scale Using AWS Inferentia

AWS Inferentia accelerators are designed by AWS to deliver high performance at the lowest cost for most deep learning inference applications. These accelerators were built from the ground up to increase inference times by up to 2.3x and reduce inference costs by up to 70%.

Compiling the HuggingFace Models with Neuron SDK

In order to use Inferentia chips for our NLP workload, the model needs to be compiled into an Inferentia-optimized representation. AWS provides the Neuron SDK to run deep learning workloads on Inferentia instances, and it supports both PyTorch and TensorFlow frameworks.

After training the model on custom data, we can deploy the model on Inferentia. To do this, we’ll trace and compile the model with the Neuron SDK, which optimizes the model for Inferentia accelerators at the lowest cost. It’s not directly possible to deploy the model on Inferentia instances unless we trace and compile the model using Neuron SDK.

Here, we’ll use a pre-trained HuggingFace model to demonstrate how to compile a PyTorch model to deploy, trace, and compile the model. You may access the entire code used for this blog in this GitHub repository.

Step 1: Install and Import All the Libraries

The following code installs transformers and torch-neuron packages:

!pip install --upgrade --no-cache-dir torch-neuron neuron-cc[tensorflow] torchvision torch --extra-index-url=https://pip.repos.neuron.amazonaws.com
!pip install --upgrade --no-cache-dir 'transformers==4.6.0'
!pip install transformers

import torch
import torch_neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig

# Build tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("gchhablani/bert-base-cased-finetuned-mrpc")
model = AutoModelForSequenceClassification.from_pretrained("gchhablani/bert-base-cased-finetuned-mrpc", return_dict=False)

# Setup some example inputs
sequence_0 = "The company HuggingFace is based in New York City"
sequence_1 = "Apples are especially bad for your health"
sequence_2 = "HuggingFace's headquarters are situated in Manhattan"

max_length=128
paraphrase = tokenizer.encode_plus(sequence_0, sequence_2, max_length=max_length, padding='max_length', truncation=True, return_tensors="pt")
not_paraphrase = tokenizer.encode_plus(sequence_0, sequence_1, max_length=max_length, padding='max_length', truncation=True, return_tensors="pt")

# Run the original PyTorch model on compilation example
paraphrase_classification_logits = model(**paraphrase)[0]

# Convert example inputs to a format that is compatible with TorchScript tracing
example_inputs_paraphrase = paraphrase['input_ids'], paraphrase['attention_mask'], paraphrase['token_type_ids']
example_inputs_not_paraphrase = not_paraphrase['input_ids'], not_paraphrase['attention_mask'], not_paraphrase['token_type_ids']

# Run torch.neuron.trace to generate a TorchScript that is optimized by AWS Neuron
model_neuron = torch.neuron.trace(model, example_inputs_paraphrase, verbose=1, compiler_workdir='./compilation_artifacts')

# Save the TorchScript for later use
model_neuron.save('neuron_compiled_model.pt')

# Create a model.tar.gz file to be used by SageMaker endpoint
!tar -czvf model.tar.gz neuron_compiled_model.pt

import boto3
import time
from sagemaker.utils import name_from_base
import sagemaker

# upload model to S3
role = sagemaker.get_execution_role()
sess=sagemaker.Session()
region=sess.boto_region_name
bucket=sess.default_bucket()
sm_client=boto3.client('sagemaker')
model_key = '{}/model/model.tar.gz'.format('inf1_compiled_model')
model_path = 's3://{}/{}'.format(bucket, model_key)
boto3.resource('s3').Bucket(bucket).upload_file('model.tar.gz', model_key)
print("Uploaded model to S3:")
print(model_path)

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference-neuron:1.7.1-neuron-py36-ubuntu18.04

# Install packages
RUN pip install "transformers==4.7.0"

%%sh

# The name of our algorithm
algorithm_name=neuron-py36-inference
cd container
account=$(aws sts get-caller-identity --query Account --output text)

# Get the region defined in the current configuration (default to us-west-2 if none defined)
region=$(aws configure get region)
region=${region:-ap-south-1}
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# If the repository does not exist in ECR, create it.
aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1

if [ $? -ne 0 ]
then
    aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null
fi

# Get the login command from ECR in order to pull the SageMaker PyTorch image
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

%%sh

# Build the docker image locally with the image name and then push it to ECR with the full name.
docker build  -t ${algorithm_name} . --build-arg REGION=${region}
docker tag ${algorithm_name} ${fullname}

%%sh 
# The name of our algorithm
algorithm_name=neuron-py36-inference cd container account=$(aws sts get-caller-identity --query Account --output text)

# Get the region defined in the current configuration (default to us-west-2 if none defined)
region=$(aws configure get region)
region=${region:-ap-south-1}
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# If the repository does not exist in ECR, create it.
aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1

if [ $? -ne 0 ]
then     aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null fi 
# Get the login command from ECR in order to pull the SageMaker PyTorch image
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com 
# Build the docker image locally with the image name and then push it to ECR with the full name.
docker build  -t ${algorithm_name} . --build-arg REGION=${region}
docker tag ${algorithm_name} ${fullname}

# Get the login command from ECR and execute it directly
aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com docker push ${fullname}

# Create a PyTorch Estimator using the inference code and the ECR image
from sagemaker.pytorch.model import PyTorchModel

pytorch_model = PyTorchModel(
    model_data=pretrained_model_data,
    role=role,
    source_dir="code",
    framework_version="1.7.1",
    entry_point="inference.py",
    image_uri=ecr_image
)

# indicate SageMaker that the model has been compiled using neuron-cc. 
pytorch_model._is_compiled_model = True

# Deploy the model
predictor = pytorch_model.deploy(initial_instance_count=1, instance_type="ml.inf1.2xlarge")