Kog Laneformer 2B генерирует 3000 токенов в секунду за счет отложенного тензорного параллелизма

Компания Kog выпустила модель Laneformer 2B с 2,3 млрд параметров, спроектированную с приоритетом на скорость вывода, а не на максимальную точность в тестах. В отличие от стандартного подхода, где ускорение достигается на этапе развертывания, здесь архитектура модели изначально адаптирована под «железо» через механизм отложенного тензорного параллелизма (DTP). Это позволило достичь скорости генерации 3 000 токенов в секунду на одном запросе при использовании 8 AMD MI300X, что является рекордом для моделей такого класса.

Архитектурные решения и принцип работы

Разработчики отказались от традиционной схемы, где каждый слой модели требует синхронизации между видеокартами. Вместо этого они внедрили систему из 8 «полос» (lanes), которые позволяют скрывать задержки коммуникации. Пока одни части модели вычисляют данные, другие передают результаты, что устраняет простои процессора.

Ключевые технические особенности:

• Отложенный тензорный параллелизм (DTP): Синхронизация между GPU происходит не после каждого слоя, а с задержкой в 2 шага, что маскирует время передачи данных.
• Внимание со скользящим окном (SWA): Используется в 10 из 15 слоев модели для оптимизации работы с кэшем и предотвращения задержек при потоковой передаче.
• Группировка запросов (GQA): 32 головы запросов и 16 голов ключей/значений, равномерно распределенных по 8 полосам.

Важный нюанс: Максимальная скорость достигается только при использовании проприетарного движка Kog Inference Engine. Стандартные реализации в библиотеке Hugging Face не смогут полностью раскрыть потенциал архитектуры DTP.

Обучение и специализация данных

Модель обучалась с нуля на 4 трлн токенов общих данных, после чего прошла дообучение на 2 трлн токенов, сфокусированных на коде и логике. Финальная настройка под инструкции (SFT) заняла 210 млн токенов.

Процесс обучения включал три этапа:

1. Предварительное обучение: Создание базовой модели на широком наборе данных (по рецепту NVIDIA Nemotron).
2. Среднее обучение (Mid-training): Резкий сдвиг в сторону кода и рассуждений. Это привело к снижению разнообразия токенов, но значительно улучшило специализированные навыки.
3. Финальная настройка: Легковесная донастройка для работы с инструкциями.

Результаты на бенчмарках:

• HumanEval+: 45,1% (greedy decoding).
• MBPP+: 51,6% (greedy decoding).

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

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

Тренировка проводилась на 192 видеокартах NVIDIA H100 (24 узла по 8 GPU) в течение 21 дня. Для обеспечения стабильности использовалась инфраструктура партнеров Scaleway и ADASTRA во Франции.

Скорость вывода на стандартных серверных GPU:

• 8 × AMD MI300X: 3 000 токенов/сек (FP16, батч 1).
• 8 × NVIDIA H200: 2 100 токенов/сек (FP16, батч 1).

Для повышения точности при генерации кода можно использовать стохастическое декодирование (pass@N). Поскольку один проход занимает менее 0,3 секунды, генерация нескольких вариантов решения становится экономически и технически целесообразной.

Операционные последствия и скрытые риски

• Зависимость от движка: Модель требует специфического исполнения для достижения заявленной скорости. Стандартные инструменты (например, llama.cpp в формате GGUF) не поддерживают кастомную архитектуру, что ограничивает гибкость развертывания на стороннем оборудовании.
• Ограниченный контекст: Длина контекста зафиксирована на уровне 4 096 токенов. Это связано с архитектурными решениями для минимизации задержек, что делает модель непригодной для задач, требующих анализа больших объемов текста.
• Специализация против универсальности: Модель не является универсальной. Она оптимизирована для задач программирования. Попытки использовать её для творческого письма или сложного анализа общих текстов могут дать результаты хуже, чем у более крупных универсальных моделей.
• Лицензирование: Веса модели распространяются под лицензией Apache 2.0, но токенизатор основан на Llama 2, что накладывает дополнительные обязательства по соблюдению условий сообщества Meta⋆.

На фоне этого: Выпуск модели демонстрирует тренд на ко-дизайн (совместное проектирование) алгоритмов и аппаратного обеспечения. Для бизнеса это сигнал, что дальнейший рост производительности ИИ-агентов будет зависеть не только от увеличения параметров, но и от глубокой интеграции с конкретным «железом».