Миссия Reka — разработка ИИ, который понимает физический мир, моделирует его и действует в нем. Мы создаем наши модели для решения самых сложных задач реального мира, обеспечивая работу всего: от интеллектуальных роботов до носимых устройств и медиа следующего поколения.
Наряду с рассуждением и визуальным пониманием, генерация видео является ключевой возможностью для достижения этой цели, создавая основу для систем, которые моделируют возможные варианты будущего и взаимодействуют с динамическими средами. Однако для многих наиболее перспективных приложений генерации высококачественного видео недостаточно. Модель должна генерировать его в реальном времени. Модель реального времени может непрерывно реагировать на новые входные данные от пользователя, агента или среды, поддерживая интерактивные приложения, в которых модель не просто описывает физический мир, но и участвует в нем.
Сегодня мы делаем первый шаг к этой цели: представляем версию нашей существующей модели генерации видео, работающую в реальном времени. Наша команда тесно сотрудничала с NVIDIA для достижения этого рубежа.
Наша система, доступная в закрытой бета-версии, генерирует бесконечный непрерывный видеопоток авторегрессионным способом. Пользователи могут управлять потоком по мере его создания, предоставляя инструкции на естественном языке.
Мы достигли этого без ущерба для визуального качества: генерация плавно работает в разрешении 720p при 24 кадрах в секунду.
Обеспечение генерации в реальном времени для модели такого масштаба было сложной технической задачей. Наш трансформер с 30 миллиардами параметров значительно больше большинства моделей генерации видео. Чтобы заставить модель такого размера работать достаточно быстро, потребовалась оптимизация как самой модели, так и движка вывода, который ее обслуживает. Результат? Мы добились производительности в реальном времени даже на оборудовании предыдущего поколения, таком как графические процессоры NVIDIA H100.
Дистилляция временных шагов с помощью NVIDIA FastGen
Основным узким местом при генерации с помощью диффузионных моделей является итеративный процесс удаления шума. Диффузионная модель не создает образец за один проход: она начинает с шума и уточняет его в течение многих прямых проходов, удаляя немного шума на каждом шаге. Ситуация усугубляется использованием классификаторно-свободного руководства (CFG), которое улучшает соответствие промпту и визуальное качество за счет оценки модели по негативному промпту и удаления его влияния из результата. В совокупности наша конфигурация по умолчанию для исходной модели требовала более 100 прямых проходов для генерации каждого сегмента. При 30 миллиардах параметров это делает генерацию в реальном времени практически недостижимой.
Распространенной стратегией смягчения этой проблемы является кэширование без обучения. Такие методы, как TeaCache и Cache-DiT, используют тот факт, что промежуточные активации очень мало меняются между соседними шагами удаления шума, и повторно используют их вместо повторного вычисления. Однако ускорение, обеспечиваемое этими методами, скромно и не гарантировано, поскольку степень безопасности повторного использования зависит от входных данных и траектории выборки.
Принципиальным решением является дистилляция временных шагов: адаптация модели для выполнения той же работы за гораздо меньшее количество шагов генерации. Вместо того чтобы аппроксимировать траекторию выборки более дешевым способом, дистилляция меняет саму модель: «студент» учится покрывать за один шаг то, что исходная модель могла выполнять за несколько. Кроме того, руководство (guidance) обычно встраивается в веса самой модели, что устраняет необходимость в дополнительных прямых проходах CFG.
Для выполнения дистилляции мы использовали NVIDIA FastGen — гибкую платформу, поддерживающую широкое семейство методов дистилляции, включая модели согласованности, сопоставление распределений, самообучение и другие. Мы рассматривали возможность реализации этих методов непосредственно в нашей кодовой базе для обучения, но решили интегрировать их в FastGen: это дало нам готовый доступ к большому семейству методов внутри надежной и масштабируемой платформы, что позволило сравнивать подходы без необходимости перестраивать каждый из них независимо.
Интеграция и ранние эксперименты проводились в тесном сотрудничестве с инженерами NVIDIA. Несмотря на размер нашей модели, мы смогли разместить необходимые рабочие нагрузки по обучению на нашем кластере NVIDIA H100 и эффективно выполнить процесс дистилляции в масштабе.
Наша финальная дистиллированная модель была обучена в ходе многоэтапного процесса, достигнув 11,8-кратного ускорения работы по сравнению с исходной моделью без существенного снижения качества генерации согласно нашей слепой оценке людьми.
Эффективный движок вывода
Быстрая модель не автоматически означает систему реального времени. В обычных офлайн-условиях генерации накладные расходы фреймворка, как правило, незначительны по сравнению с общим временем генерации. В системе реального времени важна каждая миллисекунда. Обмен данными между процессами, задержки планировщика, передача памяти, предварительная и постобработка — все это может прервать конвейер генерации или помешать системе поддерживать целевую частоту кадров.
Чтобы решить эту проблему, мы разработали проприетарный движок вывода на базе vLLM-Omni, сосредоточившись на следующих расширениях:
- Модульный уровень оркестрации реального времени. Мы разработали новый уровень, который координирует генерацию как постоянный поток, управляя зависимостями между последовательными сегментами и продолжая принимать новые инструкции от пользователя. Клиенты могут взаимодействовать с моделью через новую конечную точку реального времени. Дизайн является модульным, что упрощает интеграцию любой модели, поддерживаемой vLLM-Omni или нашим расширением.
Модульный уровень оркестрации реального времени. Мы разработали новый уровень, который координирует генерацию как постоянный поток, управляя зависимостями между последовательными сегментами и продолжая принимать новые инструкции от пользователя. Клиенты могут взаимодействовать с моделью через новую конечную точку реального времени. Дизайн является модульным, что упрощает интеграцию любой модели, поддерживаемой vLLM-Omni или нашим расширением.
- Оптимизация системы с низкой задержкой. Мы радикально сократили накладные расходы фреймворка, возникающие из-за межпроцессного взаимодействия (IPC), передачи тензоров и логики постобработки/декодирования.
Оптимизация системы с низкой задержкой. Мы радикально сократили накладные расходы фреймворка, возникающие из-за межпроцессного взаимодействия (IPC), передачи тензоров и логики постобработки/декодирования.
- Квантование с низкой точностью. Продолжая наследие Reka Quant, мы применили квантование с низкой точностью как к основной сети трансформера, так и к текстовому энкодеру, сумев уместить всю систему в конфигурации с ограниченной памятью на одном узле (например, установки NVIDIA H100) без использования выгрузки на CPU, при этом также получив преимущество от более быстрых вычислений с низкой точностью на современных ускорителях.
Квантование с низкой точностью. Продолжая наследие Reka Quant, мы применили квантование с низкой точностью как к основной сети трансформера, так и к текстовому энкодеру, сумев уместить всю систему в конфигурации с ограниченной памятью на одном узле (например, установки NVIDIA H100) без использования выгрузки на CPU, при этом также получив преимущество от более быстрых вычислений с низкой точностью на современных ускорителях.
Путь вперед: World Language Action Models (WLAM)
Генерация в реальном времени — это важная веха, но не конечная цель. В настоящее время наша команда обучает первую версию нашей будущей модели World Language Action Model (WLAM): омнимадели, разработанной для объединения рассуждений, визуального восприятия, генерации и действий в рамках единой системы. Сосредоточившись на эгоцентрических приложениях и взаимодействии с физическим миром, WLAM будет воспринимать окружающую среду, моделировать физические результаты в реальном времени и помогать в планировании и выполнении действий в реальном мире.
Мы разработали наш механизм логического вывода с учетом этого направления. Его модульная архитектура будет поддерживать генерацию в реальном времени с помощью WLAM с первого дня, что позволит нам перенести инфраструктуру и накопленный опыт оптимизации в наши модели следующего поколения.
Генерация в реальном времени в настоящее время доступна в рамках закрытого бета-тестирования. Чтобы узнать больше о коммерческом применении или заказать демонстрацию в реальном времени, свяжитесь с нами по адресу contact@reka.ai.
