GPGPU Russia

Уже влезли в долги, чтобы наскрести на подписку ChatGPT за 500$?

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

Забавно, что изначально Cerebras в первую очередь пытались решить задачу обучения нейросетей.

И, если верить довольно жёсткому разбору Irrational Analysis, именно с training всё получилось не очень хорошо. Архитектура получилась сильно перекошенной: внутри самого WSE пропускная способность огромная, но канал обмена с внешней памятью и другими устройствами относительно узкий. Поэтому для обучения Cerebras приходилось дополнять WSE внешними системами MemoryX и SwarmX, а ограниченный внешний I/O становился одним из главных bottleneck'ов.

А что вообще такое WSE, MemoryX и SwarmX и как всё это устроено, читайте в отчёте Irrational Analysis.

А вот с инференсом всё оказалось гораздо интереснее.

Cerebras довольно самобытная архитектура. Вместо относительно классической связки GPU с расположенной рядом HBM они буквально делают процессор размером почти с целую кремниевую пластину, Wafer-Scale Engine. WSE-3 содержит около 900 тысяч маленьких ядер, у каждого из которых есть локальный SRAM и простые SIMD-блоки. В сумме получается примерно 44 GB SRAM прямо на wafer.

И вот для inference это оказалось очень интересной комбинацией. Обычный GPU при генерации токенов во многом упирается в необходимость постоянно читать веса модели из HBM.

У Cerebras веса и KV-cache можно держать прямо в SRAM рядом с вычислительными ядрами, а слои модели распределять по WSE. Если одного WSE не хватает, модель распределяется между несколькими. Каждый хранит свою часть модели и передаёт следующему только промежуточный результат вычислений, а не сами веса. Благодаря этому веса остаются в быстром локальном SRAM, а между WSE нужно передавать сравнительно небольшой объём данных. Отсюда и очень высокая скорость генерации.

То есть архитектура, которая довольно странно выглядела как конкурент Nvidia в training, неожиданно нашла свою сильную сторону именно там, где сегодня особенно ценится latency: ультрабыстрый LLM inference.

Но silver bullet из Cerebras всё равно не получился. Вопрос, как оптимальнее всего строить железо для LLM inference, сегодня всё ещё открыт. Одна из главных проблем Cerebras в том, что объём SRAM нельзя наращивать так же легко, как внешнюю память, а большие модели и длинный контекст требуют всё больше памяти. В итоге приходится использовать больше WSE, и впечатляющая скорость начинает обходиться всё дороже с точки зрения масштабирования всей системы. Сам отчёт отдельно отмечает, что KV-cache и веса конкурируют за один и тот же ограниченный SRAM, а с ростом размера модели экономика такого подхода становится всё сложнее.

Но сам инженерный подход всё равно совершенно безумный и очень интересный.

Кстати, если вы забыли, какого размера «чип» у Cerebras, приложу фотку в комментариях.

P.S. Пост чуть больше чем полностью построен на отчёте Irrational Analysis. Обязательно ознакомьтесь с ним, невероятно увлекательный и познавательный разбор.