Optimizing Transformers for GPUs with Optimum

%pip install "optimum[onnxruntime-gpu]==1.3.0" --upgrade

from onnxruntime import get_available_providers, get_device
import onnxruntime
 
# check available providers
assert 'CUDAExecutionProvider' in get_available_providers(), "ONNX Runtime GPU provider not found. Make sure onnxruntime-gpu is installed and onnxruntime is uninstalled."
assert "GPU" == get_device()
 
# asser version due to bug in 1.11.1
assert onnxruntime.__version__ > "1.11.1", "you need a newer version of ONNX Runtime"

from optimum.onnxruntime import ORTModelForQuestionAnswering
from transformers import AutoTokenizer
from pathlib import Path
 
 
model_id="distilbert-base-cased-distilled-squad"
onnx_path = Path("onnx")
 
# load vanilla transformers and convert to onnx
model = ORTModelForQuestionAnswering.from_pretrained(model_id, from_transformers=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
 
# save onnx checkpoint and tokenizer
model.save_pretrained(onnx_path)
tokenizer.save_pretrained(onnx_path)

context="Hello, my name is Philipp and I live in Nuremberg, Germany. Currently I am working as a Technical Lead at Hugging Face to democratize artificial intelligence through open source and open science. In the past I designed and implemented cloud-native machine learning architectures for fin-tech and insurance companies. I found my passion for cloud concepts and machine learning 5 years ago. Since then I never stopped learning. Currently, I am focusing myself in the area NLP and how to leverage models like BERT, Roberta, T5, ViT, and GPT2 to generate business value."
question="As what is Philipp working?"

from transformers import pipeline
 
vanilla_qa = pipeline("question-answering", model=model, tokenizer=tokenizer, device=0)
 
print(f"pipeline is loaded on device {vanilla_qa.model.device}")
print(vanilla_qa(question=question,context=context))
#    pipeline is loaded on device cuda:0
#    {'score': 0.6575328707695007, 'start': 88, 'end': 102, 'answer': 'Technical Lead'}

from optimum.onnxruntime import ORTOptimizer
from optimum.onnxruntime.configuration import OptimizationConfig
 
# create ORTOptimizer and define optimization configuration
optimizer = ORTOptimizer.from_pretrained(model_id, feature=model.pipeline_task)
optimization_config = OptimizationConfig(optimization_level=99,
                                         optimize_for_gpu=True,
                                         fp16=True
                                         )
 
# apply the optimization configuration to the model
optimizer.export(
    onnx_model_path=onnx_path / "model.onnx",
    onnx_optimized_model_output_path=onnx_path / "model-optimized.onnx",
    optimization_config=optimization_config,
)

from transformers import pipeline
 
# load optimized model
model = ORTModelForQuestionAnswering.from_pretrained(onnx_path, file_name="model-optimized.onnx")
 
# create optimized pipeline
optimized_qa = pipeline("question-answering", model=model, tokenizer=tokenizer, device=0)
print(optimized_qa(question=question,context=context))

from transformers import pipeline
 
trfs_qa = pipeline("question-answering", model=model_id, device=0)

from datasets import load_metric,load_dataset
 
metric = load_metric("squad_v2")
metric = load_metric("squad")
eval_dataset = load_dataset("squad")["validation"]
 
# creating a subset for faster evaluation
# COMMENT IN to run evaluation on a subset of the dataset
# eval_dataset = eval_dataset.select(range(1000))

def evaluate(example):
  default = vanilla_qa(question=example["question"], context=example["context"])
  optimized = optimized_qa(question=example["question"], context=example["context"])
  return {
      'reference': {'id': example['id'], 'answers': example['answers']},
      'default': {'id': example['id'],'prediction_text': default['answer']},
      'optimized': {'id': example['id'],'prediction_text': optimized['answer']},
      }
 
result = eval_dataset.map(evaluate)
 
default_acc = metric.compute(predictions=result["default"], references=result["reference"])
optimized = metric.compute(predictions=result["optimized"], references=result["reference"])
 
print(f"vanilla model: f1={default_acc['f1']}%")
print(f"optimized model: f1={optimized['f1']}%")
print(f"The optimized model achieves {round(optimized['f1']/default_acc['f1'],2)*100:.2f}% accuracy of the fp32 model")
 
# vanilla model: f1=86.84859514665654%
# optimized model: f1=86.8536859246896%
# The optimized model achieves 100.00% accuracy of the fp32 model
 

from time import perf_counter
import numpy as np
 
context="Hello, my name is Philipp and I live in Nuremberg, Germany. Currently I am working as a Technical Lead at Hugging Face to democratize artificial intelligence through open source and open science. In the past I designed and implemented cloud-native machine learning architectures for fin-tech and insurance companies. I found my passion for cloud concepts and machine learning 5 years ago. Since then I never stopped learning. Currently, I am focusing myself in the area NLP and how to leverage models like BERT, Roberta, T5, ViT, and GPT2 to generate business value."
question="As what is Philipp working?"
 
def measure_latency(pipe):
    latencies = []
    # warm up
    for _ in range(10):
        _ = pipe(question=question,context=context)
    # Timed run
    for _ in range(300):
        start_time = perf_counter()
        _ =  pipe(question=question,context=context)
        latency = perf_counter() - start_time
        latencies.append(latency)
    # Compute run statistics
    time_avg_ms = 1000 * np.mean(latencies)
    time_std_ms = 1000 * np.std(latencies)
    time_p95_ms = 1000 * np.percentile(latencies,95)
    return f"P95 latency (ms) - {time_p95_ms}; Average latency (ms) - {time_avg_ms:.2f} +\- {time_std_ms:.2f};", time_p95_ms
 
vanilla_model=measure_latency(vanilla_qa)
optimized_model=measure_latency(optimized_qa)
 
print(f"Vanilla model: {vanilla_model[0]}")
print(f"Optimized model: {optimized_model[0]}")
print(f"Improvement through optimization: {round(vanilla_model[1]/optimized_model[1],2)}x")
 
# Vanilla model: P95 latency (ms) - 7.784631400227273; Average latency (ms) - 6.87 +\- 1.20;
# Optimized model: P95 latency (ms) - 3.392388850079442; Average latency (ms) - 3.32 +\- 0.03;
# Improvement through optimization: 2.29x