AI5 августа в 01:01 · 5 мин

Смесь котят: как Cursor открыла ядро, ускоряющее обучение ИИ до 2.4 раза

Разработка больших языковых моделей достигла этапа, когда узким местом стали не вычисления, а обмен данными между видеокартами. Компания Cursor преодолела эту проблему, представив в открытый доступ Mixture-of-Kittens (MoK) — специализированное ядро для работы с архитектурами Mixture of Experts (MoE) на мощнейших стойках Nvidia GB300 NVL72. Это решение позволяет сократить время обучения моделей в разы, используя уникальный подход к организации коммуникации между графическими процессорами.

# «Смесь котят»: как Cursor создало открытое ядро, ускоряющее обучение ИИ

В мире развития искусственного интеллекта часто говорят о проблеме вычислительной мощности, но инженеры компании Cursor подсознали, что истинным ограничением стала архитектура обмена данными. Для обучения моделей следующего поколения, особенно использующих архитектуру Mixture of Experts (MoE), критически важным становится не столько скорость выполнения математики внутри чипа, сколько скорость пересылки данных между сотнями графических процессоров в одной стойке.

Кудабыло для этого была представлена платформа, которая позволяет обрабатывать данные в 2.37 раза быстрее лучших публичных решений на текущем этапе разработки. Это достигается за счет внедрения ядра MoK, полностью перепрофилированного под специфику оборудования Nvidia GB300 NVL72.

Архитектура вычислений и проблема узких мест

MoE-модели представляют собой трансформеры, где для обработки каждого отдельного токена активация выбирает лишь небольшую часть «экспертов» из огромного набора тысяч параметров. Это позволяет модели быть компактной при обучении, но при реальной работе требует сложной координации.

В традиционной настройке веса всех экспертов не помещаются на одном видеокарту. Они распределены по dozens (десяткам) GPU. Процесс обучения включает в себя два ключевых этапа: dispatch (отправка токенов на карты, где лежат нужные эксперты) и combine (сбор результатов обратно). В классическом подходе эти операции часто выполняются синхронно или с частыми синхронизациями с центральными процессорами (CPU), что создает задержки.

Стойки Nvidia GB300 NVL72 представляют собой массив из 72 GPU, объединенных в единый домен NVLink, где обмен данными происходит практически мгновенно. Однако инженеры обнаружили, что встроенные в эту стойку процессоры ARMGrace оказались значительно медленнее видеокарт. GPU простаивали, ожидая команды от CPU, даже для выполнения незначительных задач, таких как отправка метрик обучения на сервер. Это создало «бутылочное горлышко», которое тормозило весь кластер.

Решение Cursor было радикальным: полностью исключить CPU из критического цикла обучения. Вся коммуникация и вычисления были перенесены непосредственно на видеокарты, что позволило создать единое гигантское вычислительное ядро.

Инновации в пересылке данных

Ключевым элементом оптимизации стала переосмысление механизма передачи данных между видеокартами. Индустриальный стандарт, например библиотека DeepEP, использует подход «push»: видеокарта-источник самостоятельно отправляет данные всем получателям. Однако на специфике NVL72 эффективнее оказался обратный подход — «pull» (по «смеси котят»). В этой схеме карта-получатель сама инициирует запрос и «вытягивает» нужные токены.

Такой подход позволил упростить логическое расписание операций: вместо сложных трехколоночных схем с сортировками, система работает по двум основным колонкам. Главное преимущество заключается в отсутствии сигналов о завершении операций между картами. Не нужно ждать подтверждения от 71 соседней карты перед продолжением работы. Микробенчмарки показали, что устранение этих сигналов ускорило процесс более чем в пять раз по сравнению с традиционными методами, повысив утилизацию пропускной способности NVLink до 29% при неравномерной нагрузке.

Дополнительным техническим трюком стало использование кольцевого буфера для токенов. Раньше системы либо выделяли огромный резервный буфер, либо жертвовали качеством обучения, отбрасывая излишки данных. MoK использует фиксированное кольцевое окно в несколько сотен мегабайт. Токены циркулируют по кольцу: слоты заполняются данными при передаче, а освобождаются после вычислений. Это полностью устраняет необходимость обращения к процессору и позволяет поддерживать непрерывный поток данных без потери информации.

Производительность и масштабируемость

Результаты тестирования на реальных конфигурациях демонстрируют впечатляющий прирост эффективности. При прямом проходе обучения в формате MXFP8 скорость работы MoK оказалась на 2.37 раза выше лучших публичных стеков, включая рекомендованные связки от Nvidia. В режиме обратного прохода (для дообучения и RL) ускорение составило до 1.78 раза.

В условиях реальной производственной среды эти изменения переросли в повышение сквозной скорости обучения на 41%. Если раньше видеокарта могла обработать в среднем 761 токен в секунду, то с использованием MoK этот показатель вырос до 1070 токенов в секунду на одну карту. На базе этого ядро уже используются для обучения модели Composer в масштабах десятков тысяч GPU.

Важным аспектом является детерминизм вычислений. Новый подход гарантирует, что при одинаковых входных данных результат будет бит-в-бит идентичным, независимо от того, как именно железо распланировало выполнение инструкций. Это критически важно для воспроизводимости экспериментов и работы с алгоритмами обучения с подкреплением (RL).

Роль искусственного интеллекта в создании инфраструктуры

Особый интерес вызывает и методология создания самого ядра. Название Mixture-of-Kittens является отсылкой к фреймворку ThunderKittens из лаборатории Hazy Research при Стэнфордском университете. Согласно заявлениям команды, значительная часть кода однооператорных ядер теперь пишется с помощью ИИ-агентов. Эти агенты, получив архитектурные ориентиры, выдают решения уровня state-of-the-art с первой попытки.

Год назад для написания подобных компонентов требовался специальный фреймворк-прослойка, упрощающий работу человеку. В этот раз команда обошлась без него и разобрала всю сложность напрямую, вместе с агентами. Получается ядро для обучения ИИ, в значительной части написанное ИИ. Это подчеркивает тенденцию к тому, что инструменты создания инфраструктуры для ИИ сами начинают разрабатываться с помощью ИИ.

Открытость и переносимость решения

Хотя ядро жестко заточено под архитектуру GB300 NVL72 — стойку стоимостью в миллионы долларов, доступную лишь у ограниченного числа компаний — подход авторов остается открытым. Весь код, алгоритмы планирования и бенчмарк-скрипты доступны для публичного доступа.

Авторы подчеркивают, что писали код так, чтобы его было легко перенести на другое железо. Идеи вроде pull-пересылки и кольцевых буферов универсальны и применимы к другим архитектурам, не только к NVL72. Это снимает скепсис относительно заявленного «снижения барьера для ИИ-исследований», делая передовые методы доступными и для менее мощных конфигураций.

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

Первоисточники

Habr AI
← Вернуться в эфир