Один из самых частых вопросов, которые нам задают, — как мы развертываем крупные модели на чипах Cerebras пластинчатого масштаба (wafer-scale). Одиночный кристалл Cerebras WSE-3T имеет 44 ГБ памяти на чипе, в то время как передовым моделям с триллионами параметров требуются терабайты только для их весов.
В этой статье мы объясним, как мы разделяем крупную передовую модель между несколькими кристаллами WSE-3T для запуска многотриллионных моделей со скоростью 1000–2000+ токенов в секунду.
В качестве наглядного примера мы будем использовать открытую модель Kimi K3 с ~3 триллионами параметров.
К концу статьи вы поймете:
- Как мы помещаем всю модель в память на чипе
- Как мы распределяем модель по нескольким чипам
- Как мы обрабатываем большое количество одновременных запросов
Разбор на примере Kimi K3
Kimi K3 — это модель с открытым исходным кодом и 3 триллионами параметров, содержащая 69 слоев линейного внимания и 24 плотных слоя.
Для этого наглядного примера мы используем:
- Предполагаемый средний размер запроса 256k
- 4-битную квантизацию весов (как в реализации Hugging Face)
- 8-битный KV-кеш
#1: Размер каждого слоя
Чтобы рассчитать общее количество кристаллов Cerebras WSE-3T, необходимых для поддержки Kimi K3, давайте сначала вычислим размер каждого слоя модели.
Kimi K3 имеет 3 триллиона параметров. Если все 3 триллиона параметров хранятся в 4-битном формате каждый, исходный объем данных весов составляет:
Разделив этот объем данных на 93 слоя, мы получим:
Теперь, когда мы знаем, что для каждого слоя требуется около 16 ГБ, мы можем определить, сколько слоев помещается на каждом 44-гигабайтном кристалле. Ответ зависит от типа слоя:
- Плотным слоям требуется дополнительное пространство для KV-кеша, поскольку каждый плотный слой (с полным вниманием) хранит ключи и значения для каждого токена в каждом активном запросе. Поэтому мы размещаем один плотный слой на кристалл и оставляем оставшуюся SRAM для его KV-кеша.
- Слоям линейного внимания не требуется дополнительное пространство для KV-кеша. Вместо хранения ключей и значений для каждого токена, слой линейного внимания поддерживает небольшой фиксированный объем состояния для каждого запроса, независимо от длины контекста. Это оставляет место для двух линейных слоев на кристалл (2 × 16 ГБ = 32 ГБ).
Следовательно, с учетом архитектуры Kimi K3, мы получаем:
Это дает нам общее количество кристаллов, равное 35 + 24 = 59 кристаллам.
#3: Соединение кристаллов вместе
Cerebras выполняет модели в конвейере. После того как компилятор Cerebras распределяет слои модели по областям кристаллов, среда выполнения прокладывает путь для потока токенов между кристаллами. Для сверхвысокой скорости передачи данных между кристаллами необходимо передавать только активизации, и это может быть реализовано с помощью RDMA поверх Ethernet. Это занимает примерно 2 микросекунды на один переход (hop).
Важно подчеркнуть, что поскольку каждый кристалл Cerebras имеет 44 ГБ SRAM на чипе, а каждый слой составляет всего 16 ГБ, мы обычно имеем возможность размещать целые слои модели непосредственно на чипе.
Это означает, что в большинстве случаев веса модели никогда не нужно перемещать, а слои не нужно разделять между кристаллами. Только активизации переходят с одного кристалла на другой.
#4: Как мы обрабатываем 1 запрос
Теперь, когда наши кристаллы соединены вместе, а модель размечена, мы можем представить каждый кристалл как один этап конвейерной сборки: веса модели для каждого этапа остаются на месте, а промежуточные результаты запроса перемещаются по этапам слоев.
Например, Кристалл 1 выполняет слои модели 1–2 и отправляет свои активизации на кристалл 2. Кристалл 2 выполняет слои 3–4 и передает свои результаты на Кристалл 3 и так далее, пока последний кристалл не сгенерирует следующий токен.
#5: Как мы обрабатываем множество запросов
Это порождает распространенное заблуждение. Если токен должен последовательно проходить через каждый кристалл, разве каждый кристалл не простаивает в ожидании своей очереди? По сути, не будет ли система работать с размером батча, равным единице?
Нет! Если бы мы отправляли по конвейеру только один запрос за раз и ждали его завершения перед началом следующего запроса, это было бы неэффективным использованием конвейера. Вместо этого мы держим весь конвейер заполненным различными запросами. Поскольку модель уже разделена на этапы конвейера, каждый этап может одновременно работать над другим запросом.
Пока этап 3 обрабатывает токен для запроса А, этап 2 обрабатывает токен для запроса B, а этап 1 — для запроса C. В нашем примере каждый этап представляет собой слой. В продакшене Cerebras использует гораздо более мелкозернистые этапы и их значительно больше.
Загрузка плеера...
Каждый этап также может обрабатывать микробатч — небольшую группу запросов, а не один запрос. Таким образом, количество обрабатываемых в данный момент запросов примерно равно количеству этапов, умноженному на размер микробатча.
Загрузка плеера...
Если говорить о компромиссах, то меньшие микробатчи делают каждый запрос еще быстрее, в то время как большие микробатчи увеличивают общую пропускную способность.
#6: Сколько запросов мы можем обработать?
Объем памяти на чипе, выделенный под KV-кеш, определяет максимальное количество одновременных запросов. Каждый обрабатываемый запрос нуждается в собственном KV-кеше на каждом кристалле с плотным слоем, поэтому размер памяти на чипе устанавливает верхний предел для одновременных запросов. Как объяснялось выше, нам нужно рассчитывать требования к KV-кешу только для плотных слоев.
Для каждого 44-гигабайтного кристалла, на котором уже размещен один плотный слой, мы можем вычислить оставшееся пространство для KV-кеша следующим образом:
Общая емкость кристалла минус размер одного плотного слоя. В данном случае:
Расчет количества одновременных запросов
Чтобы найти количество одновременных запросов, поддерживаемых для любой конкретной конфигурации отображения модели, мы сначала рассчитываем, сколько кеша занимает один запрос на 256k:
Затем, чтобы вычислить общее количество одновременных запросов, которые поддерживает эта конфигурация, мы делим:
Общее количество одновременных запросов представляет собой количество параллельных запросов, движущихся по конвейеру в любой момент времени. В отличие от батча на GPU, эти запросы не ждут завершения друг друга, они непрерывно проходят через конвейер.
Конечно, ни параллелизм, ни размер батча сами по себе не показывают, насколько эффективно работает система инференса. При оценке пропускной способности системы инференса более значимым метрическим показателем является общая пропускная способность токенов на единицу мощности (например, токенов в секунду на мегаватт) при измерении на одинаковой скорости генерации вывода.
Масштабирование за пределы 3 триллионов параметров
Тот же подход применим и к более крупным моделям: рассчитайте память, необходимую для весов и состояния запроса, распределите модель по кристаллам и соедините их в конвейер. Каждый дополнительный кристалл добавляет вычислительную мощность и память на чипе, что позволяет нам размещать модели, превышающие емкость любого отдельного чипа.
Спасибо, что учитесь вместе с нами! Хотите развернуть модель на Cerebras или создавать решения с помощью нашего облака инференса?
- Начните работу с нашим бесплатным API: cloud.cerebras.ai
- Свяжитесь с нашим отделом продаж: sales@cerebras.ai
Дизайн: Хэлли Чанг (Halley Chang). Спасибо Атхарве Талпаде (Atharva Talpade), Чжэньвэю Гао (Zhenwei Gao), Анжеле Енг (Angela Yeung), Джессике Лю (Jessica Liu), Джойс Эр (Joyce Er), Эмме Колл (Emma Call) и Рохану Гупте (Rohan Gupta) за рецензирование и отзывы.
