Упаковка последовательностей ускорит обучение ИИ, но требует кастомной маски внимания

Статья сообщества HuggingFace от 7 октября 2024 года описывает метод оптимизации обучения больших языковых моделей (LLM), который позволяет существенно сократить время вычислений. Авторы предлагают использовать технику «упаковки последовательностей» (packed sequences), чтобы исключить пустые токены-заполнители, которые обычно занимают память видеокарт, но не несут смысла. Вместо того чтобы доводить короткие предложения до фиксированной длины, их объединяют в одну длинную цепочку. Это позволяет обрабатывать больше реальных данных за один шаг обучения, но требует специальной настройки механизма внимания, чтобы модель не путала границы между разными предложениями.

Важный нюанс: Без правильной маски внимания модель может начать «подглядывать» в конец предыдущего предложения при генерации следующего, что приведет к искажению логики и снижению качества обучения.

Суть метода упаковки последовательностей

Стандартный подход к обучению трансформеров требует, чтобы все предложения в одном пакете (batch) имели одинаковую длину. Если одно предложение короче другого, к нему добавляются специальные символы-заполнители (padding). Видеокарты вынуждены тратить ресурсы на обработку этих пустых токенов, что замедляет процесс.

Метод упаковки решает эту проблему иначе:

• Конкатенация: Несколько коротких предложений соединяются в одну длинную последовательность, разделенных специальными маркерами конца предложения (EOS).
• Экономия ресурсов: Исключаются пустые токены, что снижает нагрузку на память GPU.
• Увеличение пропускной способности: За один шаг обучения модель обрабатывает больше полезных токенов, что ускоряет обучение.

Однако простое склеивание создает риск: модель может ошибочно считать, что конец первого предложения связан с началом второго. Чтобы этого избежать, необходимо модифицировать маску внимания.

Стоит учесть: При использовании упаковки критически важно сбрасывать позиционные ID (position IDs) для каждого нового предложения внутри пакета, чтобы модель воспринимала их как независимые начала, а не как продолжение предыдущего текста.

Техническая реализация и совместимость

Для корректной работы метода требуется изменить стандартную причинно-следственную маску (causal mask). Она должна не только скрывать будущие токены, но и блокировать доступ к токенам из предыдущих предложений внутри одной упакованной последовательности.

В статье приведен алгоритм на PyTorch для создания такой маски:

1. Поиск границ: Скрипт находит индексы всех токенов EOS (End of Sequence).
2. Построение маски: На основе этих индексов создается матрица, где доступ к токенам из предыдущих предложений блокируется (становится False).
3. Корректировка позиций: Позиционные идентификаторы для каждого нового предложения внутри пакета обнуляются или начинаются с единицы заново.

Для пакетной обработки (batch processing) алгоритм усложняется, так как нужно учитывать глобальные индексы в рамках всего пакета, а не одного примера. Код в статье демонстрирует, как реализовать это без использования медленных циклов, используя векторизованные операции PyTorch.

В комментариях к статье обсуждается совместимость метода с современными библиотеками:

• FlashAttention и SDPA: Автор статьи предполагает, что метод будет работать с этими ускоренными реализациями внимания, особенно через API flexattention в PyTorch, которое поддерживает маски для «рваных» последовательностей.
• SFTTrainer (TRL): Пользователи выяснили, что текущая реализация параметра packing=True в библиотеке trl (SFTTrainer) объединяет последовательности, но не создает специальную маску внимания и не корректирует позиционные ID. Она полагается на то, что модель сама научится игнорировать токены после EOS, следуя подходу из статьи GPT-3. Это может быть менее эффективно, чем явная блокировка внимания.

На фоне этого: Использование готовых инструментов без проверки их внутренней логики может привести к потере преимуществ упаковки, если они не реализуют явную блокировку внимания между предложениями.

Практические выводы для разработчиков

Внедрение упаковки последовательностей меняет подход к подготовке данных и настройке обучения.

• Снижение затрат на обучение: Устранение пустых токенов позволяет ускорить обучение на том же оборудовании или сократить количество необходимых GPU для достижения того же результата.
• Необходимость кастомизации: Стандартные реализации в популярных библиотеках (например, transformers или trl) могут не поддерживать полную логику маскировки для упакованных последовательностей. Разработчикам, вероятно, придется писать кастомные функции для создания масок и корректировки позиционных ID.
• Зависимость от реализации внимания: Эффективность метода зависит от того, как библиотека обработки внимания обрабатывает кастомные маски. Использование flexattention в PyTorch может стать стандартом для таких задач.
• Риск «загрязнения» пакетов: Если маска настроена неверно, токены из одного примера могут влиять на обучение другого, что приведет к ухудшению качества модели.

Важный нюанс: Метод упаковки не является «волшебной таблеткой», если модель сама не умеет корректно обрабатывать границы предложений; явная блокировка внимания через маску часто дает более предсказуемый результат, чем надежда на самообучение модели.