Moonshot представила Checkpoint Engine: дообучение LLM с триллионом параметров за 20 секунд

Китайская компания Moonshot, ответственная за выпуск передовой большой языковой модели Kimi-K2, представила новый новый метод дообучения ИИ-моделей — Checkpoint Engine. Разработчики заявляют, что система смогла выполнить обучение с подкреплением нейросети Kimi-K2 c 1 триллионом параметров всего за 20 секунд.

Checkpoint-engine представляет собой оптимизированное промежуточное программное решение, разработанное специально для обновления весов моделей в механизмах вывода. Основная логика обновления весов реализована через систему ParameterServer — сервис, развертываемый совместно с механизмами вывода, который предлагает два различных метода обновления:

• Широковещательный метод: применяется, когда требуется синхронное обновление весов в множестве клиентов инференса LLM. Это наиболее быстрая реализация, используемая по умолчанию. 
• P2P-передача: используется при динамическом добавлении новых клиентов инференса, в то время как существующие экземпляры LLM продолжают обработку пользовательских запросов. Для избежания влияния на производительность активных экземпляров LLM используется движок mooncake-transfer-engine для прямой P2P-передачи весов с GPU существующих экземпляров на GPU новых. 

Помимо, система поддерживает технологию оптимизированной трансляции весов, благодаря которой механизм Checkpoint-engine хранит ссылки на сегментированные веса в памяти CPU, эффективно передавая их GPU-кластеру для инференса LLM. Передача происходит в три этапа:

• H2D-передача: перемещение весов в память GPU с диска или из систем обучения.
• Фаза широковещания: обмен данными между обработчиками контрольных точек с использованием буферов CUDA IPC.
• Фаза перезагрузки: механизмы вывода идентифицируют и копируют соответствующие подмножества весов из опубликованных данных.

Другими словами, Checkpoint-engine оркестрирует весь процесс передачи данных. Изначально он собирает необходимые метаданные для формирования плана передачи, включая определение оптимальных размеров буферов. Затем он выполняет передачу данных, управляя механизмами вывода через коммуникацию ZeroMQ-сокетов. Для максимальной производительности он организует передачу данных по конвейеру с перекрывающимися операциями связи и копирования, как подробно описано в техническом отчете Kimi-K2. Такой конвейерный подход требует дополнительной памяти GPU, а при недостаточном объеме памяти checkpoint-engine автоматически переходит в режим последовательного выполнения.

Компания Moonshot активно меняет игру в области разработки искусственного интеллекта — они появились из ниоткуда, выпустив свой главный бенгер Kimi-K2, который сразу же занял топы лучших открытых языковых моделей, а теперь стартап создает инновационные системы, максимально оптимизирующие дообучение ИИ по методу обучения с подкреплением. Checkpoint-engine уже доступен для скачивания на GitHub, а провести тестирование революционного подхода можно через движок vLLM.