efreelancer 14 сен 2024 в 19:07

Распределённый инференс llama.cpp через RPC

Средний

9 мин

9.8K

Машинное обучение*DevOps*Искусственный интеллект

Туториал

Приветствую, хабровчане!

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

Погуглив некоторое время узнал, что проект LocalAI уже относительно давно поддерживает такую возможность, недолго думая я раскатал на нескольких компьютерах данный проект, после чего выполнил все необходимые настройки связав все инстансы в единую систему и, мягко говоря, был разочарован, уж слишком "фатально-недостаточным" оказалось данное решение, Docker-образ собран неоптимально, он был огромный по весу и только под amd64, неотключаемый веб-интерфейс шел в комплекте с проектом, скупой выбор моделей, некоторые из доступных LLM не работали в режиме RPC, все эмбеддинговые модели тоже отказывались запускаться в таком режиме, и так далее и тому подобное.

Повозившись ещё немного, полез в исходники и обнаружил упоминание проекта llama.cpp, затем нашёл вызов бинарника rpc-server. И вот я оказался на странице llama.cpp/examples/rpc и всё заверте...

Краткий(?) обзор

Давайте для начала спросим GigaChat о том, что такое протокол RPC:

Протокол RPC (Remote Procedure Call) позволяет программам вызывать функции или процедуры в другом адресном пространстве, на удаленных узлах или в независимых системах на том же узле. Он включает в себя сетевой протокол для обмена данными в режиме клиент-сервер и язык сериализации объектов для кодирования данных при их передаче через сеть.
Существуют различные реализации RPC, включая SOA, CORBA и DCOM. Для транспортного уровня часто используются протоколы TCP и UDP, но также существуют реализации на основе HTTP. Примерами реализации RPC являются XML-RPC, который использует XML для кодирования сообщений и HTTP в качестве транспортного механизма, и gRPC, использующий HTTP/2 и Protocol Buffers для описания интерфейсов. RPC широко применяется в различных сетевых сервисах, включая NFS.

В проекте llama.cpp данный протокол реализован в формате клиент-сервер, при этом в роли RPC-клиентов выступают утилиты навроде llama-server, llama-cli, llama-embedding и так далее, а в роли RPC-серверов специализированные бинарники rpc-server.

Если очень кратко расписать как всё это работает получается следующее:

Некий RPC-клиент, скажем llama-server, в момент запуска получает через аргументы командной строки список RPC-серверов и модель;
RPC-клиент считывает модель затем "нарезает" её слои таким образом, чтобы они были равномерно распределены между всеми RPC-серверами;
Далее RPC-клиент разливает слои по серверам и запускает инференс.

В общих чертах вся эта схема будет выглядеть следующим образом:

При этом rpc-server может быть собран под разные бэкенды, это могут быть разные архитектуры процессоров, с поддержкой тех или иных функций, скажем можно собрать один RPC-сервер под x86_64 с поддержкой CUDA, а второй - под x86_64 без CUDA, ну а третий - скажем под ARM64 чтобы на RepkaPi 3 запустить и... RPC-клиент сможет с ними всеми прекрасно работать и выполнять инференс.

Сборка бинарников

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

llama-cli - утилита командной стройки, которая позволяет запускать инференс LLM;
llama-embedding - утилита командной стройки, которая позволяет запускать инференс эмбеддинговых моделей;
llama-server - это очень простой API-сервер который может работать как в режиме инференса LLM так и в режиме инференса эмбеддинговых моделей;
rpc-server - бинарник который будет запускаться на удалённых машинах и выполнять всю работу по инференсу.

Ну так вот, если очень кратко сборку llama.cpp можно выполнить в три простых шага.

Ставим пакеты необходимые для сборки:

apt install -fyq bash wget git make g++

Клонируем репозиторий к себе на хост и переходит в директорию с исходниками:

git clone https://github.com/ggerganov/llama.cpp.git
cd llama.cpp

Запускаем компиляцию (в инструкции приводится пример через cmake, но мне больше нравится make):

GGML_RPC=ON make llama-server llama-cli llama-embedding rpc-server libggml.so libllama.so

Важно перед перед make прописать переменную окружения GGML_RPC=ON (можно и через export, но мне удобнее в inline формате) данная переменная позволяет включить в инструкциях по сборке блоки кода добавляющие поддержку RPC.

По завершению компиляции в директории появятся перечисленные после make исполняемые бинарные файлы.

Сборка Docker-образов

Возможность компилировать бинарники под разные архитектуры конечно штука полезная, но что делать если у нас имеется скажем десяток компьютеров и виртуальных машин или скажем кластер в Kubernetes, не будем же мы на каждом узле запускать компиляцию? Конечно нет! Вместо этого мы воспользуемся DevOps практиками и соберём бинарники в Docker-образы.

В качестве базового образа с целью унификации была выбрана библиотечная Ubuntu 22.04 LTS, так как она же используется в базовых контейнерах nvidia/cuda.

Для реализации проекта решил использовать multi-stage сборку разделённую на два этапа.

На первом этапе пусть выполняется загрузка всего необходимого для компиляции и собественно сама компиляция:

FROM ubuntu:22.04 AS builder
WORKDIR /app

ARG LLAMACPP_REPO="https://github.com/ggerganov/llama.cpp.git"
ARG LLAMACPP_VERSION="master"

# Install dependencies
RUN apt update -q \
 && apt install -fyq bash wget git make g++ \
 && apt clean

# Clone repo
RUN git clone --branch "$LLAMACPP_VERSION" --depth 1 "$LLAMACPP_REPO"

# Build binaries
WORKDIR /app/llama.cpp
RUN GGML_RPC=ON make -j$(nproc) llama-server llama-cli llama-embedding rpc-server libggml.so libllama.so

А на втором этапе пусть копируются собранные бинарные файлы в чистый базовый образ:

FROM ubuntu:22.04
WORKDIR /app

# Install basic dependencies
RUN apt update -q \
 && apt install -fyq libgomp1 \
 && apt clean

# Create folders
RUN mkdir -pv /app/models

# Copy compiled tools  
COPY --from=builder /app/llama.cpp/libllama.so /usr/lib/x86_64-linux-gnu
COPY --from=builder /app/llama.cpp/libggml.so /usr/lib/x86_64-linux-gnu
COPY --from=builder /app/llama.cpp/rpc-server .
COPY --from=builder /app/llama.cpp/llama-cli .
COPY --from=builder /app/llama.cpp/llama-embedding .
COPY --from=builder /app/llama.cpp/llama-server .

# Init entrypoint  
ADD entrypoint.sh .  
ENTRYPOINT ["/app/entrypoint.sh"]

Полный код Dockerfile в репозитории на GitHub.

Сборка Docker-образов с поддержкой CUDA

Принципиальных отличий от Dockerfile основанном на библиотечной ubuntu нет, разве что на первом этапе сборки используется контейнер nvidia/cuda:devel:

# Stage 1
FROM nvidia/cuda:12.5.1-devel-ubuntu22.04 AS builder

Ну а команда сборки бинарников с поддержкой CUDA выглядит следующим образом:

GGML_CUDA=ON GGML_RPC=ON make llama-server llama-cli llama-embedding rpc-server libggml.so libllama.so

Как видно помимо GGML_RPC добавлена ещё и переменная GGML_CUDA.

На втором этапе используется nvidia/cuda:runtime:

# Stage 2
FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04

Полный код Dockerfile.cuda в репозитории на GitHub.

Про entrypoint.sh

Поскольку мне хотелось собрать универсальный контейнер который можно использовать в различных режимах потребовалось реализовать специальный entrypoint.sh скрипт, который будет выполнять каждый раз при запуске контейнера.

По плану контейнер будет работать в следующих режимах:

backend
Режим в котором запускается rpc-server, команда его запуска сервера выглядит следующим образом:

rpc-server --host "0.0.0.0" --port "50052" --mem "1024"

Тут видно, что есть некая странная опция --mem она позволяет указать какое количество оперативной памяти (в Мегабайтах) может использовать данный RPC-сервер, если rpc-server собран под CUDA то этот параметр отвечает за количество VRAM (видепамяти), если без поддержки CUDA то за количество RAM (системной оперативной памяти).

server
Режим в котором запускается llama-server, представляющий из себя простейший API-сервер предоставляющий возможность интерактивного взаимодействия с большими (и малыми) языковыми и эмбеддинговыми моделями, команда запуска выглядит следующим образом:

llama-server --host "0.0.0.0" --port "8080" --model "/app/models/TinyLlama-1.1B-q4_0.gguf" --gpu-layers 99 --rpc backend01:50052,backend02:50052

Тут важно обратить внимание на опцию --gpu-layers при обычных обстаятельствах она указывает на то сколько слоёв максимум можно выгрузить в память видеокарты, однако, в случае если указана опция --rpc, её поведение меняется и она указывает сколько слоёв можно выгрузить на RPC-серверы.

С опцией --rpc в ней мы через запятую перечисляем хосты и порты RPC-серверов, к которым RPC-клиент будет подключаться.

none
Специальный режим, который запускает команду sleep inf, чтобы можно было подключиться к контейнеру и вручную запустить llama-cli или скажем llama-embedding.

Если собрать всё это в рамках одного скрипта то получится универсальный entrypoint.sh.

Кросс-платформенная сборка Docker-образов

Одна из любопытных особенностей библиотечного образа ubuntu является то, что она он поддерживает множество процессорных архитектур, но мне в первую очередь было важно amd64, arm64 и arm/v7, первая понятно почему, а вот последние две мне нужны чтобы иметь возможность запускать RPC-сервер на микрокомпьютерах, а вот контейнер nvidia/cuda поставляется только под архитектуры amd64 и arm64.

Сама же сборка будет выполняться при помощи docker buildx специального плагина, расширяющего базовый функционал Docker, в нашем же случае интересно только лишь возможность кросс-компиляции контейнеров, так как сборку под ARM64 планируется выполнять на x86_64 процессоре.

И так, для начала создадим сборщик buildx, назовём его скажем my_builder.

docker buildx create --name my_builder --driver=docker-container

Далее, предположим что файл Dockerfile и entrypoint.sh находятся в директории под названием llama.cpp:

docker buildx build --builder=my_builder --platform=linux/amd64,linux/arm64,linux/arm/v7 --build-arg LLAMACPP_VERSION=master ./llama.cpp/

Тут видим, что сборка происходит под три архитектуры, в качестве версии используется HEAD из master ветки репозитория llama.cpp.

Добавив опции --tag=${owner}/${repo}:${tag} и --push мы сможем тегировать образы и выгружать их в регистри.

Полный пример сборки и публикации контейнеров при помощи GitHub Actions.

Запускаем через Docker Compose

И так, предположим мы собрали несколько контейнеров, запушили их на Docker Hub и теперь хотим запустить всё это добро на своём железе, предположим у нас имеется два сервера, на одном мы можем использовать видеокарту, но при этом только 1Гб VRAM, а на втором у нас нет видеокарты и можно использовать только 2Гб RAM. Мы планируем запустить на них модель TinyLlama 1.1B таким образом чтобы пользователь взаимодействовал с API-сервером.

В общем виде такая схема будет выглядеть следующим образом:

Схема из двух RPC-серверов и одного RPC-клиента

В результате у нас получится следующего вида docker-compose.yml

version: "3.9"

services:

  main:
    image: evilfreelancer/llama.cpp-rpc:latest
    restart: unless-stopped
    volumes:
      - ./models:/app/models
    environment:
      APP_MODE: server
      APP_MODEL: /app/models/TinyLlama-1.1B-q4_0.gguf
      APP_RPC_BACKENDS: backend-cuda:50052,backend-cpu:50052
    ports:
      - "127.0.0.1:8080:8080"

  backend-cpu:
    image: evilfreelancer/llama.cpp-rpc:latest
    restart: unless-stopped
    environment:
      APP_MODE: backend
      APP_MEM: 2048

  backend-cuda:
    image: evilfreelancer/llama.cpp-rpc:latest-cuda
    restart: "unless-stopped"
    environment:
      APP_MODE: backend
      APP_MEM: 1024
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [ gpu ]

Далее потребуется рядом с docker-compose.yml создать директорию models и скачать в неё файл TinyLlama-1.1B-q4_0.gguf.

Запускаем композицию командой:

docker compose up -d

Далее ждём некоторое время и после того как композиция запустится можем попробовать через curl выполнить инференс:

curl \
    --request POST \
    --url http://localhost:8080/completion \
    --header "Content-Type: application/json" \
    --data '{"prompt": "Building a website can be done in 10 simple steps:"}'

В ответе будет что-то вроде этого:

Что дальше?

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

Из любопытного на что я обратил бы ваше внимание это вот этот небольшой PR в проект ollama (который на момент публикации данной статьи ещё висел в несмердженных) и вот это обсуждение, всё там же, в тикетах проекта ollama. Если кратко, то разработчики хотят добавить возможность выполнять распределённый инференс на RPC-бэкендах по типу того, что был продемонстрирован в данной публикации. Так что в дальнейшем я планирую попробовать подружить ollama с моими Docker-контейнерами.

Ещё я планирую использовать данные контейнеры в Kubernetes, поэтому скорее всего в ближайшем будущем подготовлю k8s operator или просто deployment в формате Helm-чарта дабы упростить процедуру развёртывания серверов по нодам.

А ещё у меня на антрисолях есть немало микрокомпьютеров, а также две специальные материнские платы под названием TuringPi v1 для кластеризации RaspberryPi CM3, на них я тоже планирую проводить эксперименты в будущем и именно поэтому среди всех перечисленных архитектур контейнеров ниличествует arm/v7.

В общем планов грамодьё, было бы время...

За сим откланиваюсь, спасибо что дочитали статью до конца, если вас заинтересовало будущее данного проекта приглашаю ко мне на канал @evilfreelancer в Телеграме.

Ссылки

Прочее:

Теги:

Хабы: