Xiaohongshu представляет BigMac: преодоление компромисса между объемом памяти и скоростью при мультимодальном обучении

Мультимодальные крупномасштабные языковые модели становятся краеугольным камнем искусственного интеллекта следующего поколения, однако развитие инфраструктуры для их обучения долгое время сдерживалось неизбежным компромиссом: системы, оптимизированные под скорость, часто сталкиваются с ограничениями по объему памяти, в то время как системы, в которых приоритет отдается эффективному использованию памяти, как правило, работают медленно. Команда Dots Infra в компании Xiaohongshu определила это узкое место как границу Парето в области обучения мультимодальных конвейеров и теперь преодолела его. 22 июля команда открыла исходный код BigMac — новой парадигмы параллельного обучения с использованием конвейера, разработанной специально для нативных мультимодальных сценариев. Проект теперь доступен на GitHub в репозитории Dots-Infra.

В отличие от стандартных трансформеров, мультимодальные модели по своей сути сложны. Типичная мультимодальная большая языковая модель (MLLM) состоит из трёх отдельных компонентов: модальных кодировщиков, преобразующих изображения и аудио в вложения; базовой структуры LLM для инференса; и генератора, который преобразует вывод LLM обратно в целевые модальности, такие как изображения или речь. Существенные структурные различия между этими компонентами создают серьезные проблемы при их интеграции в единый конвейер обучения.

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

BigMac решает эту проблему с помощью простого, но ключевого принципа: основное внимание по-прежнему уделяется основному конвейеру LLM. Обучение крупномасштабных LLM уже опирается на отработанные стратегии планирования, такие как 1F1B или чередующийся 1F1B, которые глубоко интегрированы в производственные стеки обучения. Вместо того чтобы заменять эти устоявшиеся методы, BigMac использует график LLM в качестве базовой временной шкалы, стратегически вставляя вычисления кодировщика и генератора, когда входные данные готовы, и не нарушая при этом порядок выполнения LLM. Команда называет эту архитектуру «вложенным конвейером с квазизащитой от зависимостей».

Такая конструкция обеспечивает два ключевых преимущества. Во-первых, колебания продолжительности работы кодировщика и генератора больше не распространяются по конвейеру LLM, что позволяет LLM поддерживать оптимальный темп работы. Во-вторых, требования к памяти для модальных активаций сокращаются до O(1) на алгоритмическом уровне. Кодировщикам больше не нужно сохранять активации для всех микропачек до завершения работы конвейера, а генераторы избегают удлинения длинных «хвостов» активаций. По сути, BigMac не жертвует памятью ради кратковременных всплесков производительности; вместо этого он перестраивает логику планирования, чтобы сделать эти две цели совместимыми, а не взаимоисключающими.

Преодоление разрыва между проектированием расписания и фактическим выполнением сопряжено со значительными инженерными трудностями, включая системную интеграцию, определение интерфейсов и отладку производительности. BigMac разработан для решения этих конкретных задач за счёт предоставления трёх ключевых возможностей. Во-первых, он делает глобальное расписание прозрачным: планировщик генерирует во время выполнения исчерпывающую таблицу операторов, охватывающую все ранги конвейера, микропартии и типы модулей. Затем исполнитель распределяет их по локальным последовательностям для каждого ранга, плавно интегрируясь с бэкендами LLM, такими как Megatron Core, модальными средами выполнения и уровнями взаимодействия. Такая прозрачность упрощает проверку и оптимизацию бэкенда.

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

Эффективность BigMac была подтверждена на двух типичных рабочих нагрузках. В задачах MLLM-Understanding с использованием Qwen3-30B-A3B в качестве базовой модели и кодировщика ViT объёмом 1,3 млрд, BigMac обеспечивает ускорение в 1,08–1,1 раза по сравнению с вычислительно эффективным базовым вариантом Optimus и в 1,6–1,9 раза — по сравнению с память-эффективным базовым вариантом Megatron-DistTrain. Важно отметить, что использование памяти остаётся стабильным при увеличении размера партии на один графический процессор, тогда как у Optimus наблюдаются пиковые скачки потребления памяти, которые в конечном итоге приводят к ошибкам «Out-of-Memory» (OOM) при работе с большими партиями. Задачи MLLM-Generation являются более сложными и включают генератор 20B MMDiT, где разрыв в производительности увеличивается ещё больше: Optimus терпит неудачу при всех тестируемых размерах партий из-за OOM, в то время как BigMac избегает этого за счёт эффективного выполнения генератора в обратном порядке и быстрого освобождения активаций. По сравнению с Megatron-DistTrain BigMac по-прежнему обеспечивает ускорение в 1,5–1,9 раза при стабильном использовании памяти. Это подчеркивает основное преимущество вложенных конвейеров: они обслуживают зависимости как кодировщика, так и генератора без чрезмерного удержания активаций или значительного времени простоя.

В настоящее время BigMac является ключевым компонентом инфраструктуры обучения мультимодальных моделей компании Dots и работает в производственной среде Xiaohongshu. Сопутствующая статья «BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training» доступна на arXiv вместе с интерактивными примерами трассировок PP Profiler. Для исследователей и инженеров, сталкивающихся с проблемой компромисса между объемом памяти и скоростью при мультимодальном обучении, этот проверенный в производственной среде проект с открытым исходным кодом предлагает практическое решение, позволяющее значительно сэкономить время и усилия.