Запуск нагрузок LLM с ускорением на GPU в EKS

Рассмотрим, как запустить рабочую нагрузку на основе больших языковых моделей (Large Language Model, LLM) с использованием ускорения на GPU и Elastic Kubernetes Service (EKS).
Для демонстрации будем использовать модель Mistral AI с 7 млрд параметров и обслуживать её с помощью сервера Hugging Face text-generation-inference, который будет работать на нашем собственном EKS-кластере.

Но есть оговорка: то, что у нас получится к концу статьи, не будет полноценной службой инференса LLM для продакшена. Мы лишь дойдём до точки, где кластер EKS сможет запустить модель. Дальнейшие шаги — на ваше усмотрение

Технологии, которые используем

Не будем рассматривать процесс создания кластера EKS или установки Karpenter, но покажем, как настроить пул GPU-нод, способных запустить модель Mistral 7B на TGI сервере.


Мы будем использовать Persistent Volume Claims (PVC), поэтому убедитесь, что эти ресурсы доступны на вашем EKS-кластере. Скорее всего, вы захотите сделать это с помощью Amazon EBS или EFS CSI драйверов.

Делаем GPU доступным для подов в нашем кластере

Ускорение рабочих нагрузок в Kubernetes с помощью GPU — это не просто добавление GPU-нод в пул. Необходима конфигурация как на уровне нод, так и на уровне всего кластера. Но прежде чем мы перейдём к этому, какие именно GPU нам нужны для запуска модели Mistral 7B через наш сервер инференса?


Определение требований к GPU


Путём проб и ошибок мы поняли, что для работы нашей модели Mistral требуется Flash Attention v2. Было не сразу понятно, какие требования к GPU предъявляет Mistral, исходя из страницы модели на Hugging Face, но после более глубокого изучения нашли полезные подсказки в документации Mistral AI. Изучая Flash Attention v2, узнали, что этот алгоритм поддерживают следующие GPU:


  • NVIDIA A10
  • NVIDIA A100
  • NVIDIA H100

Изучив типы экземпляров для ускоренных вычислений, обнаружим, что G5-экземпляры используют графические процессоры NVIDIA A10G Tensor Core и имеют самую низкую почасовую ставку среди всех типов экземпляров, удовлетворяющих требованиям к GPU.


Настройка GPU-нод с Karpenter


Теперь нам потребуются два ресурса для настройки наших GPU-нод с Karpenter:


Ниже приведены манифесты для Karpenter:

---
apiVersion: karpenter.k8s.aws/v1alpha1
kind: AWSNodeTemplate
metadata:
  name: gpu-template
spec:
  subnetSelector: { ... }        # required, discovers tagged subnets to attach to instances
  securityGroupSelector: { ... } # required, discovers tagged security groups to attach to instances
---
apiVersion: karpenter.sh/v1alpha5
kind: Provisioner
metadata:
  name: gpu-a10g
spec:
  labels:
    provisioner: gpu-a10g
    zone: default
  requirements:
    - key: "karpenter.k8s.aws/instance-category"
      operator: In
      values: ["g"]
    - key: "karpenter.k8s.aws/instance-gpu-name"
      operator: In
      values: ["a10g"]
    - key: "karpenter.k8s.aws/instance-gpu-count"
      operator: Gt
      values: ["0"]
    - key: "karpenter.k8s.aws/instance-gpu-count"
      operator: Lt
      values: ["4"]
    - key: "karpenter.sh/capacity-type"
      operator: In
      values: ["spot"]
    - key: "kubernetes.io/arch"
      operator: In
      values: ["amd64"]
    - key: "kubernetes.io/os"
      operator: In
      values: ["linux"]
  limits:
    resources:
      nvidia.com/gpu: "4"
  providerRef:
    name: gpu-template
  consolidation:
    enabled: true
  kubeletConfiguration:
    clusterDNS: ["10.0.1.100"]
  taints:
    - key: nvidia.com/gpu
      value: "true"
      effect: "NoSchedule"

Выбор экземпляров


Поле requirements определяет типы экземпляров, которые Karpenter может настроить в качестве части группы нод gpu-a10g. Некоторые важные необходимые метки включают:


  • instance-category
  • instance-gpu-name и instance-gpu-count
  • arch и os

С помощью этих меток мы будем выбирать экземпляры из семейства G5, в которых установлены нужные нам GPU: NVIDIA A10G.


Karpenter автоматически выберет соответствующие AMI (Amazon Machine Images) для наших экземпляров. В нашем случае используется AMI типа amazon-eks-gpu-node, соответствующая версии EKS.


Снижение стоимости


Экземпляры с ускорением на GPU могут быть довольно дорогими. К счастью, для достижения целей этой статьи нам не потребуется много времени работы GPU, но если не быть внимательным, можно случайно потратить много денег впустую. Есть несколько шагов, чтобы избежать растраты денег:


  • Ограничиваем Provisioner использованием нод с одной GPU. Такие ноды обычно дешевле, и нужно запустить нагрузку только в одном поде.
  • Ограничиваем Provisioner всего 4 GPU на всех нодах в пуле, чтобы снизить почасовую стоимость.
  • Включаем consolidation, которая удаляет наши GPU-нод, когда активных нагрузок нет. Мы должны удалить под с сервером инференса, когда не используем его.
  • Устанавливаем taints, чтобы предотвратить случайное планирование подов на наш GPU-промежуточный сервер. Без этого мы не сможем положиться на консолидацию.
  • Используем spot-экземпляры, чтобы платить ниже стандартной ставки за наши ноды. Это также означает, что эти ноды могут быть неожиданно завершены, что допустимо для демонстрационных нагрузок.

С учётом всех этих настроек, важно убедиться, что мы удаляем под с инференсом, когда он не нужен. Не оставляйте экземпляр EC2 стоимостью $1 и более в час в вашем кластере без дела.


Конфигурация Kubelet


Рекомендуется установить поле clusterDNS в конфигурации Provisioner, чтобы оно соответствовало фактическому CIDR-адресу службы в вашем кластере Kubernetes. Это гарантирует kubelet использование правильного IP-адреса DNS для разрешения имён служб, что поможет подам правильно разрешать имена служб. Для нашего демонстрационного примера это может не иметь значения, но лучше сразу делать всё правильно.


Предоставление доступа к GPU для подов с помощью NVIDIA K8S Device Plugin


На этом этапе мы можем планировать поды на наш пул нод gpu-a10g, но поды не смогут получить доступ к GPU до тех пор, пока мы не установим NVIDIA K8S device plugin в нашем кластере. Это daemonset, который делает NVIDIA GPU "видимыми" для Kubernetes.


Подготовка GPU-нод


Перед установкой плагина NVIDIA device plugin наши ноды должны иметь установленные и настроенные инструментарий и runtime от NVIDIA:


  • nvidia-container-toolkit
  • nvidia-container-runtime

К счастью, наши AMI amazon-eks-gpu-node уже имеют установленный и настроенный инструментарий и runtime. Если вы не можете использовать этот AMI по какой-то причине, можете обратиться к документации NVIDIA по установке инструментария и создать собственный AMI или обновить userData в ресурсе AWSNodeTemplate, который мы определили ранее.


Установка Daemonset NVIDIA Device Plugin


Есть несколько способов установить этот плагин. Например, использовать Helm, и еще NVIDIA предоставляет официальный chart. Мы можем установить плагин, используя команду:

helm repo add nvdp https://nvidia.github.io/k8s-device-plugin
helm repo update
helm upgrade -i nvdp nvdp/nvidia-device-plugin \
  --namespace nvidia-device-plugin \
  --create-namespace \
  --version 0.14.3 \
  --set nodeSelector.provisioner=gpu-a10g

Важная деталь


Одной из важных деталей является то, что мы устанавливаем nodeSelector для явного назначения только на наши GPU-ноды. Это необходимо, чтобы убедиться, что поды с нужной нагрузкой попадают именно на ноды с GPU. Более подробную информацию о данном helm chart можно найти в следующих источниках:


Быстрый тест


NVIDIA предоставляет образ контейнера, который мы можем запустить на нашем кластере, чтобы проверить, имеют ли поды доступ к GPU. Вот манифест для запуска такого пода:

apiVersion: v1
kind: Pod
metadata:
  name: gpu-test-pod
spec:
  restartPolicy: Never
  containers:
    - name: cuda-container
      image: nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda10.2
      resources:
        limits:
          nvidia.com/gpu: 1
  nodeSelector:
    provisioner: gpu-a10g
  tolerations:
    - key: nvidia.com/gpu
      operator: Exists
      effect: NoSchedule

После того как под запланирован, проверьте логи с помощью команды:

kubectl logs pod/gpu-test-pod

Убедитесь, что вы видите что-то похожее на:

[Vector addition of 50000 elements]
Copy input data from the host memory to the CUDA device
CUDA kernel launch with 196 blocks of 256 threads
Copy output data from the CUDA device to the host memory
Test PASSED
Done

Теперь мы готовы приступить к развертыванию нашего сервера инференса для генерации.


Развертывание сервера Text Generation Inference


Поскольку наша цель — доказать, что наш кластер может выполнять инференс с ускорением на GPU, сделаем конфигурацию Kubernetes для нашей нагрузки максимально простой.


Манифесты Kubernetes


Вот минимальный набор манифестов, который мы можем применить к нашему кластеру. Он создаст под, который будет запускать сервер TGI от Hugging Face и использовать модель Mistral 7B для генерации текста. Объяснение конфигурации будет следовать за фрагментом YAML.

---
apiVersion: v1
kind: Pod
metadata:
  name: text-inference
  labels:
    app: text-inference
spec:
  containers:
    - name: text-generation-inference
      image: ghcr.io/huggingface/text-generation-inference:1.3
      resources:
        limits:
          nvidia.com/gpu: 1
        requests:
          cpu: "4"
          memory: 4Gi
          nvidia.com/gpu: 1
      command:
        - "text-generation-launcher"
        - "--model-id"
        - "mistralai/Mistral-7B-v0.1"
        - "--num-shard"
        - "1"
      ports:
        - containerPort: 80
          name: http
      volumeMounts:
        - name: model
          mountPath: /data
        - name: shm
          mountPath: /dev/shm
  volumes:
    - name: model
      persistentVolumeClaim:
        claimName: text-inference-model
    - name: shm
      emptyDir:
        medium: Memory
        sizeLimit: 1Gi
  nodeSelector:
    provisioner: gpu-a10g
  tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: "NoSchedule"
  restartPolicy: Never
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: text-inference-model
spec:
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 50Gi
---
apiVersion: v1
kind: Service
metadata:
  name: text-inference
spec:
  ports:
    - port: 80
      protocol: TCP
      targetPort: http
  selector:
    app: text-inference
  type: ClusterIP

Ресурсы контейнера


Первая интересная деталь — мы используем новый тип ресурса: nvidia.com/gpu. Этот тип ресурса предоставляется плагином устройства NVIDIA для Kubernetes (NVIDIA k8s device plugin), который мы ранее настроили. Подробнее о нем можно узнать в readme плагина.


Поскольку наши ноды ограничены одной GPU, мы ожидаем, что этот под получит нод в своё распоряжение.


Команда контейнера


Здесь мы выбираем модель генеративного ИИ, с которой будем взаимодействовать. Hugging Face Hugging Face, но требования к оборудованию и ПО могут варьироваться в зависимости от модели.


Мы также устанавливаем количество шардов (частей модели) на 1, что должно соответствовать количеству GPU, запрошенных подом.


Томы пода


Наш под использует два тома:


  1. model: Содержит weights и biases(в формате safetensors).
  2. shm: Общая память (shared memory), настроенная в соответствии с рекомендациями из документации TGI.

Для тома модели используем Persistent Volume Claim (PVC). Это может уменьшить время запуска для последующих подов инференса, так как модель не нужно загружать каждый раз. Однако я считаю, что это преждевременная оптимизация, так как наша модель на 7B достаточно быстро скачивается, а служба инференса предназначена только для демонстрации и не нуждается в масштабировании.


Сервис


Использование Service в нашей демонстрации необязательно, так как мы будем использовать его только для настройки перенаправления портов (port-forward), что также можно сделать напрямую с подами. В зависимости от сети вашего кластера, вы можете пойти дальше и настроить Ingress для службы.


Применение манифестов


Теперь применим манифесты с помощью команды kubectl apply или другим предпочитаемым методом.

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

Так что подождем.


Под готов? Давайте проверим


Чтобы не обсуждать настройку Ingress, просто воспользуемся перенаправлением порта на нашу службу инференса:

kubectl port-forward svc/text-inference 8080:80

Теперь можно провести простой тест с использованием curl:

curl 127.0.0.1:8080/generate \
  -X POST -H 'Content-Type: application/json' \
  -d '{
    "inputs":"Hello? Is there anybody in there?",
    "parameters":{"max_new_tokens":20,"temperature": 0.5}
  }'

Результат может немного варьироваться в зависимости от temperature модели, у нас получилось что-то вроде:

{"generated_text":" Just nod if you can hear me. Is there anyone at home?\n\nI’m not"}

Вероятно, параметры запроса еще нужно подкорректировать, но, надеюсь, вы получили работающий сервер инференса с открытым исходным кодом, работающий на GPU в вашем кластере.

Бонус: запуск локального клиента чата

Конечно, здорово иметь модель с открытым исходным кодом, работающую в нашем кластере, но использование curl для доступа к ней кажется немного анти-кульминационным.


Hugging Face предоставляет нам несколько достойных начальных точек. Пример с Gradio:

import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient(model="http://127.0.0.1:8080")

max_tokens = 1024
system_prompt = "You are helpful AI."

def inference(message, history):
    prompt = f"System prompt: {system_prompt}\n Message: {message}."
    partial_message = ""
    for token in client.text_generation(prompt, max_new_tokens=int(max_tokens), stream=True, repetition_penalty=1.5):
        partial_message += token
        yield partial_message

gr.ChatInterface(
    inference,
    chatbot=gr.Chatbot(height=300),
    textbox=gr.Textbox(placeholder="Chat with me!", container=False, scale=7),
    description="Gradio UI consuming TGI endpoint with Mistral 7B model.",
    title="Gradio 🤝 TGI",
    examples=["Are tomatoes vegetables?"],
    retry_btn="Retry",
    undo_btn="Undo",
    clear_btn="Clear",
).queue().launch()

Есть много путей для развития. Надеемся, вы получите удовольствие от этого процесса!


И не забудьте уменьшить количество подов, когда они вам не нужны.

MLOps

Решаем бизнес-задачи, которые требуют обработки огромных объемов данных. Обучаем и интегрируем модели машинного обучения, которые помогают автоматизировать процессы вашего бизнеса с высокой степенью надежности.

Узнать подробнее
DataOps
  • Поддерживаем и ускоряем ваши процессы обработки данных. Обеспечиваем стабильность инфраструктуры для работы с данными, чтобы ваша компания оперативно получала всю необходимую информацию о внутренних процессах
Узнать подробнее