Развертывание Open Assistant 12B на Amazon SageMaker с Hugging Face LLM Inference DLC

Введение: Зачем развертывать Open Assistant 12B на SageMaker?

Open Assistant 12B построена на архитектуре Pythia 12B и обучена сообществом LAION как открытая альтернатива коммерческим диалоговым моделям. Модель умеет отвечать на вопросы, писать код, суммаризировать тексты и поддерживать длинный контекст диалога. Для команд, которым нужен собственный контур обработки данных без передачи во внешние API, это рабочий вариант с коммерчески допустимой лицензией.

Amazon SageMaker снимает с разработчика рутинную работу: управление кластером, мониторинг, автомасштабирование. Hugging Face LLM Inference DLC добавляет поверх этого преднастроенный контейнер с Text Generation Inference (TGI). Внутри уже собраны оптимизации для популярных архитектур, поэтому деплой сводится к указанию идентификатора модели и параметров инстанса. В статье разберём полный цикл: настройка окружения, загрузка DLC, конфигурация модели на g5.12xlarge, инференс с параметрами генерации и запуск чат-бота на Gradio.

Если вы уже работали с SageMaker, но не использовали Hugging Face DLC, этот материал покажет, насколько быстрее проходит запуск open-source LLM по сравнению с ручной сборкой контейнера. Если вы только начинаете, получите пошаговую инструкцию с кодом, которую можно адаптировать под свою модель из Hugging Face Hub.

Предварительные требования и настройка окружения

Для выполнения шагов нужны: AWS аккаунт с правами на создание ресурсов SageMaker, Python 3.9 или новее, установленный SageMaker SDK. Понадобится IAM роль с разрешениями на запуск инстансов и обращение к ECR. Квоты на инстансы g5.12xlarge в выбранном регионе должны быть больше нуля.

Установка SageMaker SDK и настройка IAM роли

Установите SDK через pip:

pip install sagemaker boto3

Создайте IAM роль с политикой AmazonSageMakerFullAccess. Если вы работаете в SageMaker Studio или Notebook Instance, роль уже создана, её ARN можно получить через SDK:

import sagemaker
from sagemaker import get_execution_role

role = get_execution_role()
print(role)

При работе вне SageMaker создайте роль в консоли IAM и укажите её ARN вручную. Роль должна доверять сервису sagemaker.amazonaws.com.

Проверка квот на инстансы g5.12xlarge

Инстансы g5 не всегда доступны в новых аккаунтах. Проверьте лимит в AWS Service Quotas: откройте консоль, найдите сервис SageMaker, выберите квоту для ml.g5.12xlarge для training-job или endpoint usage. Если значение равно 0, запросите увеличение до 1 или больше. Обработка запроса обычно занимает от нескольких часов до суток.

Загрузка Hugging Face LLM Inference DLC через SageMaker SDK

Hugging Face LLM Inference DLC доступен как готовый образ в ECR. SageMaker SDK предоставляет функцию для получения URI образа под нужный регион и версию контейнера:

from sagemaker.huggingface import get_huggingface_llm_image_uri

image_uri = get_huggingface_llm_image_uri(
    backend="huggingface",
    region="us-east-1"
)
print(image_uri)

Функция возвращает полный URI, например 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu22.04. Версия контейнера зависит от региона и даты обращения, поэтому всегда берите URI через SDK, а не хардкодьте его.

Этот контейнер уже включает TGI, PyTorch, CUDA и все зависимости для запуска моделей с Hugging Face Hub. Вам не нужно писать Dockerfile и собирать образ вручную.

Конфигурация модели и деплой на инстансе g5.12xlarge

Деплой выполняется через класс HuggingFaceModel. Он принимает URI образа, IAM роль и переменные окружения, которые управляют поведением TGI:

from sagemaker.huggingface import HuggingFaceModel

model = HuggingFaceModel(
    image_uri=image_uri,
    role=role,
    env={
        "HF_MODEL_ID": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
        "HF_TASK": "text-generation",
        "SM_NUM_GPUS": "4",
        "MAX_INPUT_LENGTH": "2048",
        "MAX_TOTAL_TOKENS": "4096",
        "HF_MODEL_QUANTIZE": "bitsandbytes"
    }
)

predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.12xlarge",
    endpoint_name="open-assistant-12b"
)

Деплой занимает 5-10 минут: SageMaker поднимает инстанс, скачивает контейнер, затем TGI загружает модель из Hugging Face Hub в память GPU.

Настройка переменных окружения для TGI

Ключевые переменные:

• HF_MODEL_ID - идентификатор модели на Hugging Face Hub. Для Open Assistant 12B используйте OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5.
• SM_NUM_GPUS - количество GPU, выделяемых под модель. Для тензорного параллелизма укажите больше 1.
• MAX_INPUT_LENGTH - максимальная длина входного промпта в токенах.
• MAX_TOTAL_TOKENS - суммарный лимит входных и выходных токенов.
• HF_MODEL_QUANTIZE - метод квантизации: bitsandbytes или gptq. Уменьшает потребление видеопамяти.

Для 12B модели в FP16 требуется около 24 ГБ видеопамяти. Один GPU A10G на g5.12xlarge имеет 24 ГБ, поэтому модель помещается на один GPU. Указание SM_NUM_GPUS=4 включает тензорный параллелизм и распределяет веса по четырём GPU, что ускоряет обработку длинных последовательностей и увеличивает пропускную способность.

Выбор инстанса: почему g5.12xlarge?

ml.g5.12xlarge оснащён четырьмя GPU NVIDIA A10G по 24 ГБ, 48 vCPU и 192 ГБ оперативной памяти. Для 12B модели это избыточно по памяти, но оправдано по скорости: четыре GPU позволяют TGI использовать тензорный параллелизм и обрабатывать несколько запросов параллельно через динамическое пакетирование.

Сравнение с g4dn.12xlarge: у g4dn четыре GPU T4 по 16 ГБ, чего недостаточно для 12B модели в FP16 без квантизации. g5.12xlarge стоит дороже, но даёт запас по памяти и более быстрые вычисления на A10G. Если бюджет ограничен, можно выбрать g5.2xlarge с одним A10G и включить квантизацию bitsandbytes, тогда модель займёт около 12 ГБ.

Выполнение инференса с настройкой параметров генерации

После деплоя отправляйте запросы через predictor.predict. Параметры генерации передаются в поле parameters:

payload = {
    "inputs": "Объясни разницу между тензорным параллелизмом и пайплайнным параллелизмом",
    "parameters": {
        "temperature": 0.7,
        "top_p": 0.9,
        "max_new_tokens": 256,
        "repetition_penalty": 1.1
    }
}

response = predictor.predict(payload)
print(response[0]["generated_text"])

Ответ приходит в формате списка словарей. TGI возвращает сгенерированный текст с учётом заданных параметров.

Основные параметры генерации и их влияние

Для чат-ботов рабочие значения: temperature 0.6–0.8, top_p 0.9, repetition_penalty 1.05–1.15. Для фактологических ответов снижайте temperature до 0.2–0.4, чтобы уменьшить галлюцинации.

Создание чат-бота на Gradio с интеграцией SageMaker API

Gradio позволяет поднять веб-интерфейс за несколько строк кода. Установите библиотеку:

pip install gradio

Создайте функцию, которая принимает текст пользователя, отправляет его в SageMaker endpoint и возвращает ответ:

import gradio as gr
from sagemaker.huggingface import HuggingFacePredictor

predictor = HuggingFacePredictor(
    endpoint_name="open-assistant-12b"
)

def generate_response(user_input):
    payload = {
        "inputs": user_input,
        "parameters": {
            "temperature": 0.7,
            "top_p": 0.9,
            "max_new_tokens": 256,
            "repetition_penalty": 1.1
        }
    }
    response = predictor.predict(payload)
    return response[0]["generated_text"]

iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=3, placeholder="Введите сообщение..."),
    outputs=gr.Textbox(),
    title="Open Assistant 12B",
    description="Чат-бот на базе Open Assistant 12B, развёрнутый на Amazon SageMaker"
)

iface.launch(share=True)

Интерфейс запускается на локальном хосте, параметр share=True создаёт публичную ссылку для демонстрации коллегам.

Настройка интерфейса и обработка запросов

Для продакшен-сценариев добавьте обработку таймаутов и повторных попыток. SageMaker endpoint может отвечать несколько секунд при длинных ответах, поэтому установите таймаут в predictor:

from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer

predictor = HuggingFacePredictor(
    endpoint_name="open-assistant-12b",
    serializer=JSONSerializer(),
    deserializer=JSONDeserializer(),
    invocation_timeout=60
)

Если endpoint не отвечает за 60 секунд, вы получите исключение, которое можно перехватить и показать пользователю сообщение об ошибке вместо бесконечной загрузки.

Оптимизации TGI: тензорный параллелизм, динамическое пакетирование, квантизация

TGI включает три ключевые оптимизации для продакшен-инференса. Тензорный параллелизм разбивает матрицы весов по GPU и выполняет вычисления параллельно. Для 12B модели на четырёх A10G это даёт прирост пропускной способности в 2–3 раза по сравнению с одним GPU.

Динамическое пакетирование объединяет несколько входящих запросов в один батч без ожидания заполнения фиксированного размера. Это увеличивает утилизацию GPU и снижает задержку для каждого отдельного запроса при потоковой нагрузке.

Квантизация уменьшает точность весов с FP16 до INT8 или INT4. Для Open Assistant 12B квантизация bitsandbytes снижает потребление видеопамяти с 24 ГБ до 12 ГБ, что позволяет запускать модель на инстансе g5.2xlarge с одним GPU. Компромисс: незначительное снижение качества ответов, особенно на сложных рассуждениях.

Включение квантизации для уменьшения потребления памяти

Установите переменную HF_MODEL_QUANTIZE в конфигурации модели:

env={
    "HF_MODEL_ID": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
    "HF_TASK": "text-generation",
    "SM_NUM_GPUS": "1",
    "HF_MODEL_QUANTIZE": "bitsandbytes",
    "MAX_INPUT_LENGTH": "1024",
    "MAX_TOTAL_TOKENS": "2048"
}

С квантизацией модель помещается на один GPU A10G. Это снижает стоимость инференса в четыре раза по сравнению с g5.12xlarge. Для чат-бота с умеренной нагрузкой такой конфигурации достаточно.

Оценка стоимости и очистка ресурсов

Стоимость ml.g5.12xlarge в регионе us-east-1 составляет около $5.672 в час по он-деманд тарифу. При непрерывной работе это $136 в сутки и $4080 в месяц. Для тестирования и разработки это дорого, поэтому удаляйте endpoint после завершения работы:

predictor.delete_endpoint()

Если endpoint нужен периодически, рассмотрите SageMaker Serverless Inference или переключение на g5.2xlarge с квантизацией. Serverless платит за фактическое время обработки запросов, а не за простой инстанса. Для моделей с нерегулярной нагрузкой это снижает затраты в 5–10 раз.

Связанные материалы по теме: партнёрство Hugging Face и AWS объясняет, почему DLC стал стандартом для деплоя open-source моделей на SageMaker. Руководство по оптимизации инференса в SageMaker дополняет эту статью бенчмарками и сравнением LMI и vLLM. Если рассматриваете альтернативы, обзор Hugging Face Inference Endpoints показывает разницу между управляемым сервисом и самостоятельным деплоем.

Заключение: Итоги и дальнейшие шаги

Развертывание Open Assistant 12B на SageMaker с Hugging Face LLM Inference DLC занимает 15–20 минут: установка SDK, получение URI контейнера, конфигурация HuggingFaceModel, вызов deploy. Дальше модель доступна через API для инференса и может быть подключена к Gradio-интерфейсу.

Для углубления изучите документацию TGI по параметрам окружения и попробуйте другие модели из Hugging Face Hub с тем же контейнером. Для продакшен-нагрузки настройте автомасштабирование endpoint и мониторинг метрик в CloudWatch. Если нужен закрытый контур для внутренних данных, адаптируйте этот подход под свою модель и VPC-конфигурацию.

Подписаться на канал