Deploy pytorch model with AWS inferentia

Dec 22, 2023

Deploy pytorch model with AWS inferentia

What is AWS inferentia?

AWS Inferentia 및 AWS Inferentia2는 저렴한 비용으로 높은 성능의 추론을 제공하도록 AWS가 설계 및 구축한 기계 학습 추론 액셀러레이터입니다. 각 AWS Inferentia 액셀러레이터에는 4개의 1세대 NeuronCore가 있으며 FP16, BF16 및 INT8 데이터 유형을 지원합니다. 각 AWS Inferentia2 액셀러레이터에는 2개의 2세대 NeuronCores가 있으며 FP32, TF32 및 구성 가능한 새로운 FP8(cFP8) 데이터 유형을 추가로 지원합니다.

https://aws.amazon.com/ko/ec2/instance-types/inf1/

해당 글에서는 pytorch model을 AWS inferentia instance에 배포하고, 비슷한 사양의 CPU instance와 추론 속도를 비교하여 inferentia를 통해 얼마나 큰 폭의 비용절감과 속도 향상이 있는지 알아봅니다.

Pytorch model for inferentia

Create inferentia instance

inferentia instance를 생성합니다. aws neuron SDK는 inferentia instance에서만 정상 작동 하기 때문에, 앞으로의 과정은 해당 instance 내에서 진행해야 합니다.

https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/EC2_GetStarted.html#ec2-launch-instance

instance: inf1.2xlarge

storage: 16GB~ 권장

security rule: default

instance를 생성할 때 발급했거나 이미 존재하는 key로 instance에 접속합니다.

(접속이 불가능할 시 security rule에서 0.0.0.0 ip가 허용되어 있는지 확인.)

ssh -i <key>.pem ubuntu@<server_ip>

Install pytorch-neuron dependencies (Inferentia1)

아래의 명령어들을 실행하여 inferentia에서 model을 추론하기 위해 필요한 환경을 구축합니다.

install drivers

# Configure Linux for Neuron repository updates
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb <https://apt.repos.neuron.amazonaws.com> ${VERSION_CODENAME} main
EOF
wget -qO - <https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB> | sudo apt-key add -

# Update OS packages 
sudo apt-get update -y

# Install OS headers 
sudo apt-get install linux-headers-$(uname -r) -y

# Install git 
sudo apt-get install git -y

# install Neuron Driver
sudo apt-get install aws-neuronx-dkms=2.* -y

# Install Neuron Tools 
sudo apt-get install aws-neuronx-tools=2.* -y

# Add PATH
export PATH=/opt/aws/neuron/bin:$PATH

install pytorch-neuron

# Install Python venv 
sudo apt-get install -y python3.8-venv g++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_pytorch_inf1 

# Activate Python venv 
source aws_neuron_venv_pytorch_inf1/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_pytorch_inf1 --display-name "Python (torch-neuron)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url <https://pip.repos.neuron.amazonaws.com>

# Install PyTorch Neuron
python -m pip install torch-neuron neuron-cc[tensorflow] "protobuf" torchvision

Compile model for inferentia

torch model을 neuron을 사용하여 compile해야 inferentia에 알맞게 최적화 할 수 있습니다. (fp16, bf16, int8 지원)

python virtual env가 활성화 되어있지 않다면, source aws_neuron_venv_pytorch_inf1/bin/activate 명령어를 통해 활성화 합니다.

이후 아래와 같이 model을 불러오고, 알맞은 dummy input을 생성하여 neuron을 통해 compile을 진행합니다.

import torch
import torch_neuron

model = Model() # Init your model

model.eval()

data = torch.rand(1, 3, 1088, 1920) # Create dummy data

torch.neuron.analyze_model(model, example_inputs=[data])

model_neuron = torch.neuron.trace(model, example_inputs=[data])

model_neuron.save("simpleunet_neuron.pt")

아래와 같이 Compiler status PASS 가 출력되면 compile을 성공적으로 마친것입니다.

설정한 경로에 pt 파일이 생성된 것을 확인할 수 있습니다.

INFO:Neuron:The following operations are currently supported in torch-neuron for this model:
INFO:Neuron:prim::Constant
INFO:Neuron:aten::_convolution
INFO:Neuron:aten::add
INFO:Neuron:prim::ListConstruct
INFO:Neuron:aten::relu_
INFO:Neuron:aten::upsample_bilinear2d
INFO:Neuron:100.00% of all operations (including primitives) (656 of 656) are supported
INFO:Neuron:100.00% of arithmetic operations (63 of 63) are supported
INFO:Neuron:All operators are compiled by neuron-cc (this does not guarantee that neuron-cc will successfully compile)
INFO:Neuron:Number of arithmetic operators (pre-compilation) before = 63, fused = 63, percent fused = 100.0%
...
...
Compiler status PASS
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 63, compiled = 63, percent compiled = 100.0%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
...

Define pre/post processing functions

필요 시 data 전/후 처리 pipeline을 구성합니다.

예제에서는 간단하게 dummy input으로 batching 해주는 함수만 구현 하였습니다.

import numpy as np

def preprocess(batch_size=1, num_neuron_cores=1):
		image = torch.rand(1, 3, 1088, 1920)

    # Create a "batched" image with enough images to go on each of the available NeuronCores
    # batch_size is the per-core batch size
    # num_neuron_cores is the number of NeuronCores being used
    batch_image = image
    for i in range(batch_size * num_neuron_cores - 1):
        batch_image = torch.cat([batch_image, image], 0)
     
    return batch_image

Run inference using neuron model

속도 테스트를 할 수 있는 전체 코드 예시는 아래와 같습니다.

# inference.py
from time import time

import torch
import torch_neuron
import numpy as np

def preprocess(batch_size=1, num_neuron_cores=1):
    image = torch.rand(1, 3, 1088, 1920)

    # Create a "batched" image with enough images to go on each of the available NeuronCores
    # batch_size is the per-core batch size
    # num_neuron_cores is the number of NeuronCores being used
    batch_image = image
    for i in range(batch_size * num_neuron_cores - 1):
        batch_image = torch.cat([batch_image, image], 0)
     
    return batch_image

def benchmark(model, image):
    print('Input image shape is {}'.format(list(image.shape)))

    # The first inference loads the model so exclude it from timing
    results = model(image)

    # Collect throughput and latency metrics
    latency = []
    throughput = []

    # Run inference for 100 iterations and calculate metrics
    num_infers = 100
    for _ in range(num_infers):
        delta_start = time()
        results = model(image)
        delta = time() - delta_start
        latency.append(delta)
        throughput.append(image.size(0)/delta)

    # Calculate and print the model throughput and latency
    print("Avg. Throughput: {:.0f}, Max Throughput: {:.0f}".format(np.mean(throughput), np.max(throughput)))
    print("Latency P50: {:.0f}".format(np.percentile(latency, 50)*1000.0))
    print("Latency P90: {:.0f}".format(np.percentile(latency, 90)*1000.0))
    print("Latency P95: {:.0f}".format(np.percentile(latency, 95)*1000.0))
    print("Latency P99: {:.0f}\\n".format(np.percentile(latency, 99)*1000.0))

if __name__ == "__main__":
    image = preprocess()
    model = torch.jit.load("<neuron_model_name>")

    benchmark(model, image)

위의 코드를 실행하면, 100번의 추론에 대한 평균 throughput과 latency(s) 가 출력됩니다.

Run inference using DataParallel (Optional)

inferentia에는 여러개의 neuron core가 존재합니다. (inf1.2xlarge의 경우 4개)

DataParallel을 이용하면, input batch를 slice 하여 각각의 neuron core로 보내 연산하기 때문에, throughput의 향상을 가져올 수 있습니다.

neuron core는 아래와 같이 torch.neuron.DataParallel 을 사용하여 사용할 수 있습니다.

# inference.py
...
...
if __name__ == "__main__":
    model = torch.jit.load("simpleunet_neuron.pt")
    model_parallel = torch.neuron.DataParallel(model)

    batch_size = 1
    num_neuron_cores = 4

    image = preprocess(batch_size, num_neuron_cores)

    benchmark(model_parallel, image)

Compile with different batch sizes (Optional)

Small batch size로 compile 된 모델로 dynamic batching을 사용하는것은 slicing, transfer 등의 작업 때문에 sub-optimal 한 throughput을 초래할 수 있습니다.

따라서 compile시에 더 큰 batch size를 설정하여 compile하면 dynamic batching을 통해 더 효율적으로 throughput을 개선할 수 있습니다.

batch size를 설정하는 법은 위에서 compile 할 때 model에 넘겨준 dummy input의 batch size를 키우기만 하면 됩니다.

...
data = torch.rand(<batch_size>, 3, 1088, 1920) # Create dummy data
...

Results

같은 Unet model을 동등한 spec의 c5.2xlarge instance에서 속도를 비교하였습니다.

1 batch size에서 각각 1.53fps, 4.23fps로 큰 차이가 있을 뿐 아니라,

CPU instance에서는 batch size를 높여도 아무런 이득이 없는 반면, Inferentia에서는 큰 개선이 있는것을 확인할 수 있습니다.

c5.2xlarge

OpenVino: 0.8259968497245238 -> 1.21 fps
Eager:  0.8111959931666711 -> 1.23 fps
Script:  0.6788944289999361 -> 1.49 fps
Compile:  0.6576376958332983 -> 1.53 fps
Intel-extension:  0.6574252737221337 -> 1.53 fps
Quantized script:  0.5683428937777788 -> 1.78 fps

inf1.2xlarge

#Input image shape is [1, 3, 1088, 1920]
Avg. Throughput: 4.23, Max Throughput: 4.25
Latency P50: 0.2357
Latency P90: 0.2360
Latency P95: 0.2364
Latency P99: 0.2558


#Input image shape is [24, 3, 1088, 1920]
Avg. Throughput: 15.60, Max Throughput: 15.81
Latency P50: 1.5380
Latency P90: 1.5540
Latency P95: 1.5585
Latency P99: 1.5689

Ref

https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/inference-torch-neuron.html#inference-torch-neuron