Легко добейтесь 5-кратного ускорения при развертывании трансформеров.

Введение

Экосистема для NLP и MLOps резко выросла за последние несколько лет. Существует так много невероятных инноваций, делающих развертывание моделей быстрее и проще, чем я мог себе представить. HuggingFace произвел революцию в совместном использовании моделей, а AWS SageMaker упрощает развертывание и тестирование наших моделей. Мне нравится, как легко одним щелчком мыши развернуть любую модель HuggingFace на конечной точке SageMaker. Но эти развертывания оставляют большую часть производительности на столе. Мы можем добиться 5-кратного ускорения во время вывода на SageMaker, написав всего несколько строк кода. В этом сообщении блога мы рассмотрим, как развернуть ONNX-версию модели преобразователя в SageMaker. Пост разбит на четыре части:

  1. Настройте SageMaker
  2. Разверните ванильный трансформатор
  3. ONNXify трансформатор
  4. Разверните ONNX-версию преобразователя в SageMaker.

В качестве заявления об отказе от ответственности SageMaker не является бесплатным, и его изучение потребует незначительных затрат (0,23 доллара США в час). Давайте начнем!

Настройте SageMaker

Мы предполагаем, что у вас уже есть настройка учетной записи AWS. После входа в систему нажмите на строку поиска вверху и найдите SageMaker. Затем выберите его из вариантов.

Выберите «Начать» в правом верхнем углу страницы.

Настройте домен SageMaker.

Заполните доменное имя и профиль пользователя. Это не имеет большого значения для личных проектов, просто выберите то, что вы запомните. Одно предостережение заключается в том, что они будут добавлены к ролям, и для них существует ограничение в 64 символа. Если вы выберете длинное имя, вы можете получить ошибки. В этом случае выберите более короткие имена и повторите попытку. После заполнения этих данных создайте исполнительную роль. Это даст SageMaker права доступа и использования других ресурсов.

Выберите «Создать новую роль» в раскрывающемся списке, и вы должны увидеть что-то вроде ниже. Выберите None Студия SageMaker создаст для нас корзину, которую мы будем использовать в этом уроке. Если у вас есть другое ведро, которое вы бы предпочли использовать, не стесняйтесь указывать его в поле Specifc S3 buckets.

После того, как вы создали роль исполнения, нажмите «Отправить». Настройка среды займет несколько минут, но как только это будет сделано, вы готовы к работе!

Развертывание модели Vanilla Transformer

Теперь, когда SageMaker настроен, мы готовы развернуть нашу первую модель трансформера. Я почти чувствую себя плохо, называя это ванильной моделью. Отчасти потому, что ваниль — отличный вариант, но главным образом потому, что эти модели и инфраструктура, необходимая для их развертывания, сложны. Несколько лет назад это было действительно очень тяжело. Я каждый день поражаюсь тому, какого прогресса мы добились за такое короткое время. В любом случае, в этом уроке мы собираемся развернуть модель дистиллятора, настроенную на go_emotions. Это хорошая модель для работы, потому что она немного меньше и содержит несколько меток, поэтому нам нужно немного поработать, чтобы она вернула все.

Перейдите на страницу модели на HuggingFace и нажмите кнопку Deploy в правом верхнем углу.

В раскрывающемся списке выберите Amazon SageMaker

Выберите задачу Text Classification и конфигурацию Local Machine. Затем скопируйте код, который он предоставляет.

Совершенно невероятно, что HuggingFace сгенерирует код, необходимый для развертывания SageMaker. Чтобы запустить его, обязательно установите необходимые библиотеки.

pip install sagemaker

После установки нам нужно изменить эту строку, чтобы она имела правильное имя нашей роли SageMaker, которую мы создали в предыдущем разделе:

role = iam_client.get_role(RoleName='{IAM_ROLE_WITH_SAGEMAKER_PERMISSIONS}')['Role']['Arn']

iam_client возьмет на себя роль исполнителя, если вы знаете имя. Обычно я просто копирую файл прямо из SageMaker. Чтобы скопировать его прямо из веб-интерфейса, на странице доменов щелкните свой домен.

Затем с правой стороны вы увидите свою роль исполнения.

Скопируйте арн и обновите код

role = "arn:aws:iam::000623149457:role/service-role/AmazonSageMaker-ExecutionRole-20230519T095932"

Теперь вы готовы запустить код. Моему потребовалось около 2 минут, чтобы выполнить и сделать прогноз.

Чтобы увидеть свою модель в пользовательском интерфейсе AWS, щелкните раскрывающийся список Inference слева, а затем выберите Endpoints.

Вы должны увидеть что-то вроде следующего.

Поздравляем, вы развернули модель-трансформер! Эта модель работает на сервере ml.m5.xlarge, который всегда включен и будет нести расходы (0,23 доллара США в час). Чтобы удалить конечную точку, нажмите на нее, а затем на большую кнопку удаления в правом верхнем углу, но перед этим давайте проведем несколько временных экспериментов.

В блокноте Jupyter мы можем провести очень простой временной эксперимент с возвращенным предиктором:

%%timeit
predictor.predict({
 'inputs': ["I like you. I love you"] * 128
})

Что для меня дает:

5,8 с ± 38,8 мс на цикл (среднее значение ± стандартное отклонение для 7 запусков, по 1 циклу в каждом)

Таким образом, для предсказания на 128 примерах требуется около 5,8 секунды. Что довольно быстро! Давайте сделаем это еще быстрее с помощью ONNX!

ONNX-оптимизация

ONNX — это невероятная библиотека для оптимизации, совместного использования, распространения и запуска моделей глубокого обучения. Он предоставляет множество инструментов, которые могут упростить совместное использование наших моделей и ускорить их работу на процессорах. В Torch и библиотеке трансформаторов есть несколько очень полезных методов преобразования моделей в формат ONNX.

В этом разделе мы будем использовать два основных метода для улучшения производительности наших моделей: первый — оптимизация, второй — квантование. При оптимизации мы комбинируем слои или предварительно вычисляем константы и сохраняем их, а не делаем это во время выполнения. Эти операции делают вывод быстрее. Есть много флагов и подходов к оптимизации. Рекомендую поиграться с настройками. Квантование — это стратегия уменьшения размера модели и повышения производительности путем преобразования чисел с плавающей запятой в небольшие целые числа. Целые числа занимают меньше памяти, а процессоры лучше выполняют арифметические операции с целыми числами. С помощью этих инструментов в нашем наборе инструментов мы можем значительно улучшить наши развертывания.

Давайте конвертируем эту модель эмоций в ONNX и посмотрим, какое ускорение мы можем получить! Для конвертации и оптимизации нам понадобится библиотека Optimum. Установите это с помощью:

pip install "optimum[exporters]"

Мы используем библиотеку Optimum для загрузки нашей модели преобразователя, а затем динамически квантуем ее. Существует два основных вида квантования: динамическое и статическое. Динамический делается на лету, поэтому мы можем применить его к любой модели без особых усилий, в то время как статический изучает лучшие параметры для квантования. Если вам интересно, есть хороший пост в блоге о статическом квантовании здесь.

from optimum.onnxruntime import ORTQuantizer, ORTModelForSequenceClassification
from optimum.onnxruntime.configuration import AutoQuantizationConfig

model_path = "/Users/tetracycline/models/emotions"
model_name = "jungealexander/distilbert-base-uncased-finetuned-go_emotions_20220608_1"

model = ORTModelForSequenceClassification.from_pretrained(model_name,
                                                          from_transformers=True,
                                                          export=True)

quantizer = ORTQuantizer.from_pretrained(model)

dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)

model_quantized_path = quantizer.quantize(
    save_dir=f"{model_path}/onnx",
    quantization_config=dqconfig,
)

Этот код сохранит квантованную модель ONNX в папке onnx на определенном model_pathwe. Эта полезная функция сохраняет все, что нам нужно для загрузки и запуска этой модели, и почти все, что нам нужно, чтобы загрузить ее в SageMaker.

Мы можем загрузить нашу новую квантованную модель и запустить ее с помощью следующего кода:

tokenizer = AutoTokenizer.from_pretrained(f"{model_path}/onnx")
model = ORTModelForSequenceClassification.from_pretrained(f"{model_path}/onnx")
inputs = tokenizer("What am I using?", "Using DistilBERT with ONNX Runtime!", return_tensors="pt")
outputs = model(**inputs)

Мы проквантовали и оптимизировали нашу первую модель! Эта модель меньше (170 МБ против 500 МБ) и обладает некоторыми хорошими свойствами для развертывания.

Создайте обработчик вывода

Чтобы развернуть нашу модель ONNX в SageMaker, нам нужно сообщить ему, как делать прогнозы и обрабатывать ввод. SageMaker использует скрипт под названием inference.py для обработки этих входных данных. Мы создадим тот, который работает с ONNX. Наш сценарий должен реализовать две функции: model_fn, которая загружает нашу модель, и tranform_fn, которая применяет нашу модель к входящим данным.

import json
from typing import Any, Dict, List
from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForSequenceClassification

def model_fn(model_dir=None):
    pass

def transform_fn(model, input_data, content_type, accept) -> List[List[Dict[str, Any]]]:
    pass

model_fn выбирает путь от SageMaker к местоположению артефактов нашей модели. Все, что нам нужно сделать здесь, это взять этот путь и передать его в соответствующие определения конвейера из библиотеки optimum.onnxruntime. Этот код использует этот путь для создания экземпляра сохраненного токенизатора и модели ONNX, а затем передает их в конвейер преобразователей. Вы заметите аргумент top_k в конвейере. Это устанавливает конвейер для возврата прогнозов для всех классов. Мне нравится устанавливать этот параметр при работе с моделями, потому что он позволяет мне выбирать, как работать с выходными данными модели во внешнем интерфейсе. Если вам просто нужен лучший прогноз, вы можете удалить это.

def model_fn(model_dir=None):
    """
    Load in the model from the container and prepare it for predictions.

    Sagemaker will pass this the base directory created from unziping model.tar.gz.
    This happens once on the startup of the service.

    :param model_dir: The directory containing model.onnx and associated files. This is
        the path created from SageMaker unziping model.tar.gz.
    :return pipe: The transformers pipeline with the ONNX model.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = ORTModelForSequenceClassification.from_pretrained(
        model_dir, file_name="model.onnx"
    )
    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None)
    return pipe

Когда SageMaker запускает новый экземпляр, он сначала запускает функцию model_fn для загрузки модели. Он будет передавать выходные данные model_fn в качестве именованного параметра model последующим функциям. Когда мы попадаем в конечную точку SageMaker, SageMaker вызывает функцию transform_fn. Это основная рабочая лошадка, которая управляет вводом и вызовом модели.

def transform_fn(model, input_data, content_type, accept) -> List[List[Dict[str, Any]]]:
    """Load in the data and run predictions.

    :param model: The output of model_fn which is the transformers pipeline.
    :param input_data: The incoming payload in the format specified by `content_type`.
    :param content_type: The format of the incoming data.
    :return: The predictions from the model returned to the user.
    """
    data = json.loads(input_data)
    return model(data["inputs"])

Разверните модель

Чтобы выполнить развертывание, нам нужно упаковать весь этот код и нашу модель в один файл с именем model.tar.gz и загрузить его на S3. Папка onnx, которую мы создали ранее в этом уроке, содержит почти все, что нам нужно. В нем отсутствует папка с именем code, которая сообщает SageMaker, как обрабатывать наши операции с пользовательскими моделями. Создайте папку code в папке onnx ранее и добавьте в нее inference.py. Этот скрипт SageMaker будет использовать для загрузки и запуска наших моделей на конечной точке. Нам также нужно создать файл requirements.txt со следующей строкой:

optimum[onnxruntime]

При запуске конечной точки, если присутствует файл requirements.txt, SageMaker установит необходимые библиотеки. В этом случае мы используем optimum для запуска нашей модели ONNX, и поэтому нам нужно установить эту библиотеку. Окончательная папка должна иметь следующую структуру:

model/
    |- config.json
    |- model.onnx
    |- ort_config.json 
    |- special_tokens_map.json
    |- tokenizer_config.json
    |- tokenizer.json
    |- vocab.txt
    |- code
      |- inference.py
      |- requirements.txt

Заархивируйте папку в файл с именем model.tar.gz, и мы готовы к развертыванию. Для этого вы можете использовать следующий Python:

import tarfile
import os

tar_dir = "/path/to/your/folder"
output_file = "model.tar.gz"
with tarfile.open(output_file, "w:gz") as tar:
    tar.add(tar_dir, arcname=os.path.sep)

Загрузите этот файл model.tar.gz в корзину s3, созданную SageMaker, и скопируйте путь. Для меня это был s3://nb-sagemaker-dev/model.tar.gz. С ролью сверху мы можем развернуть модель так же, как и раньше, с одной небольшой поправкой. Добавьте поле model_data, которое указывает путь к нашему файлу model.tar.gz.

from sagemaker.serverless import ServerlessInferenceConfig
from sagemaker.huggingface import HuggingFaceModel

role = "arn:aws:iam::000623149457:role/service-role/AmazonSageMaker-ExecutionRole-20230519T095932" 

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    model_data="s3://nb-sagemaker-dev/model.tar.gz",  # path to your trained SageMaker model
    role=role,  # IAM role with permissions to create an endpoint
    transformers_version="4.17.0",  # Transformers version used
    pytorch_version="1.10.2",  # PyTorch version used
    py_version="py38",
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    endpoint_name="emotions-onnx",
    initial_instance_count=1,
    instance_type="ml.m5.xlarge"
)

Это должно занять около четырех минут для развертывания. Как только это будет сделано. Теперь протестируйте наше новое развертывание так же, как и раньше, в записной книжке Jupyter:

%%timeit
predictor.predict({
 'inputs': ["I like you. I love you"] * 128
})

Результаты приведенного выше кода для меня:

1,2 с ± 194 мс на цикл (среднее значение ± стандартное отклонение для 7 запусков, по 1 циклу в каждом)

Нашей модели с чистыми трансформерами потребовалось 5,8 для выполнения 128 прогнозов, а этой модели, оптимизированной для ONNX, требуется всего 1,2 с! Это почти в 5 раз быстрее!!!

Заключение

Мы можем реализовать 5-кратное ускорение логического вывода с очень небольшой дополнительной работой, используя ONNX. Для получения дополнительной информации о развертывании моделей HuggingFace в SageMaker см. это пошаговое руководство.

Примечание

В приведенной выше статье я в основном использовал класс predictor, который возвращается при развертывании этих моделей. Это удобно, но не очень полезно, когда вы хотите вызвать модель откуда-то не из того места, где вы развернули модель. Чтобы вызвать модель без класса предиктора, вам необходимо:

  1. Создайте экземпляр клиента sagemaker-runtime boto3 в регионе, где вы развернули свою модель.
  2. Создайте словарь полезной нагрузки с правильными ключами и значениями. Для нас наш inference.py ожидает объект json с ключом под названием inputs
  3. используйте invoke_endpoint на клиенте с соответствующей полезной нагрузкой
  4. Прочитайте результаты.

Приведенный ниже код позаботится об этом за вас.

import json
import boto3

client = boto3.client("sagemaker-runtime", region_name="us-west-2")

endpoint_name = "emotions-onnx"
payload = {
    "inputs": ["I like you, I love you."] * 128,
}

response = client.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType="application/json",
    Accept="application/json",
    Body=json.dumps(payload),
)
print(response["Body"].read().decode())

Удачного строительства :).