Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Для блочных диффузионных языковых моделей нашли способ быстрее обучаться на длинных контекстах, не меняя функцию потерь и градиенты. Группа Stanford University получила ускорение полной модели до 61%, хотя препринт не прошёл рецензирование и числа получили сами авторы. Метод переносит границу, после которой обмен данными между GPU начинает сдерживать обучение сильнее, чем сами вычисления.
Почему обычный контекстный параллелизм передаёт лишние данные
Блочная диффузионная модель генерирует последовательность блоками слева направо, но токены внутри очередного блока восстанавливает параллельно. Во время обучения каждый искажённый целевой блок видит чистый префикс — правильные предыдущие блоки — и собственную зашумлённую версию.
Обычный контекстный параллелизм делит последовательность по позициям между GPU. Каждая GPU хранит часть контекста, а на слоях внимания процессы обмениваются тензорами ключей и значений K/V, чтобы локальные запросы могли учитывать всю доступную последовательность.
Для блочной диффузионной модели такой способ делит сразу чистую и искажённую копии текста. Чистый префикс действительно нужен нескольким последующим блокам. Искажённые K/V использует только тот блок, которому они принадлежат, однако обычная схема всё равно передаёт их между GPU при прямом проходе, а затем возвращает их градиенты при обратном.
Block parallelism меняет единицу распределения: одна GPU получает целый целевой блок вместе с прямым проходом, функцией потерь и обратным проходом. Потери отдельных блоков складываются, поэтому их можно считать параллельно, а затем суммировать градиенты общих параметров. Искажённые активации и K/V при этом не покидают владельца блока.
Сам по себе этот подход создаёт другое узкое место. GPU, которым достались поздние блоки, приходится хранить длинные чистые префиксы, причём одни и те же участки контекста повторяются на нескольких устройствах.
Context-sharded block parallelism, или CSBP, совмещает распределение по блокам с разделением чистого контекста. Каждая GPU по-прежнему целиком обрабатывает назначенные искажённые блоки, но хранит только свою часть общей чистой последовательности. Между устройствами передаются лишь чистые K/V и их градиенты.
Объём вычислений внимания асимптотически не меняется: модель обрабатывает те же допустимые пары токенов. Ускорение возникает потому, что сеть не перевозит данные, которые всё равно использует одна GPU, а память не расходуется на повторяющиеся префиксы. Ранние и поздние блоки назначают вместе, чтобы выровнять нагрузку между устройствами.
Чем меньше обмена между GPU, тем заметнее выигрыш
Метод проверяли при дообучении NemotronDiffusion и DiffusionGemma, преобразовании авторегрессионных Qwen в блочные диффузионные модели и обучении черновых моделей DFlash2 для спекулятивного декодирования. Эксперименты охватывают полные и разреженные модели, кластеры H200 и H100 и несколько схем параллелизма; CSBP сравнивали с самой быстрой подходящей базовой конфигурацией для каждой нагрузки.
На кластере из 16 H200 при контексте 256K пропускная способность дообучения выросла на 18–45% в зависимости от модели. Пиковое потребление памяти осталось тем же или снизилось. Когда контекст увеличили вдвое, DiffusionGemma 26B-A4B обучалась на 61% быстрее базовой схемы.
Результат согласуется с устройством метода: чем длиннее последовательность и чем больше процессов делят контекст, тем дороже становится обмен искажёнными K/V. Чистый блоковый параллелизм без разделения контекста в отдельных конфигурациях не помещался в память, тогда как CSBP сохранял масштабирование памяти обычного контекстного параллелизма.
Планы меняются только у команд с блочным обучением
CSBP стоит учитывать командам, которые обучают блочные диффузионные модели, преобразуют авторегрессионные контрольные точки в такой формат или строят диффузионную черновую модель для спекулятивного декодирования. Метод дополняет параллелизм по данным, тензорам и экспертам: он не требует менять обучающую цель, но добавляет отдельное распределение целевых блоков.
Самый крупный выигрыш получили не на полной модели, а при обучении DFlash2: на восьми H100 с контекстом 1M пропускная способность выросла в 7,59 раза. Здесь CSBP не только сократил обмен на слоях внимания, но и позволил запустить больше параллельных копий обучения. На самой короткой проверенной конфигурации DFlash2 обычный параллелизм по данным оставался быстрее, поэтому подключать CSBP ко всем запускам без учёта длины контекста невыгодно.
За фиксированные 12 часов DiffusionGemma с CSBP прошла на один процентный пункт больше задач в SWE-bench Verified и Terminal-Bench Lite. Это показывает более быстрый прогресс при одинаковом бюджете времени, а не новый предел качества: вычисления и обучающая цель остаются прежними.
Для обычного авторегрессионного обучения работа планов не меняет. CSBP опирается на то, что известные целевые блоки можно обучать независимо, поэтому он ускоряет обучение, но не генерацию ещё неизвестных токенов. Практический повод пересматривать распределённую архитектуру появляется там, где длинный контекст уже требует нескольких GPU и обмен на внимании занимает заметную часть шага.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
