Как ускорить обучение моделей в Hugging Face с помощью packing и Flash Attention 2 в 2026 году

Как ускорить обучение LLM в Hugging Face: короткий ответ

Packing уменьшает количество пустых PAD-токенов в батче. В связке с DataCollatorWithFlattening и поддерживаемым режимом Flash Attention 2 несколько последовательностей можно передать в flattened batch без стандартного выравнивания до длины самого длинного примера. Это сокращает лишние вычисления и иногда повышает полезный throughput, то есть число обработанных токенов за секунду.

Ускорение не гарантировано для каждой модели. Результат зависит от распределения длин примеров, размера батча, максимальной длины последовательности, версии Transformers и TRL, типа GPU, сборки flash-attn и выбранного attention backend. При наивном склеивании последовательностей появляется риск cross-example attention: токены одного примера могут влиять на вычисления для следующего.

Практическая схема состоит из трех частей: токенизатор не добавляет лишний padding, collator формирует плоское представление и локальные position_ids, а модель использует совместимый путь variable-length attention через Flash Attention 2. После этого baseline с обычным padding сравнивают с новым режимом по одинаковому количеству полезных токенов.

Что именно меняется в одном шаге обучения

Представим четыре токенизированных примера длиной 80, 200, 440 и 1024 токена. При обычном padding форма батча равна 4 x 1024, то есть GPU получает 4096 позиций. Полезных токенов в примерах 1744, а 2352 позиции нужны только для выравнивания. Полезная доля такого батча составляет около 42,6 процента. Это арифметический пример, не результат бенчмарка.

При padding-free flattening реальные токены размещаются плотнее. Attention backend получает сведения о границах последовательностей и может обработать переменную длину без полного padded-тензора. В обычном dense attention форма тензора все равно определяет значительную часть вычислений. Flash Attention 2 снижает накладные расходы на хранение промежуточных attention-матриц, а varlen-путь позволяет не тратить вычисления на padding при корректной передаче границ.

Вес модели и количество обучаемых параметров при этом не меняются. LoRA-адаптер не становится меньше из-за packing. Меняется стоимость обработки батча, поэтому пиковая VRAM может снизиться, время шага может сократиться, а tokens per second может вырасти. Конкретный эффект нужно измерять на собственном датасете.

Packing не означает безусловное ускорение

Если четыре примера почти одинаковы по длине, padding overhead мал и плотная упаковка почти ничего не дает. Дополнительный collator может даже добавить небольшую нагрузку на CPU. При узком месте в загрузке данных, коммуникациях между GPU, сохранении чекпоинтов или вычислении градиентов attention перестает быть главным ограничением.

Еще один сценарий связан с совместимостью. Модель может не принимать position_ids в нужном формате, attention backend может перейти на другой путь, а конкретная версия библиотеки может не поддерживать нужное сочетание параметров. Поэтому сравнивают одинаковые precision, sequence length, effective batch size, gradient checkpointing и число полезных токенов.

Почему обычный padding тратит ресурсы на instruction-tuning

При instruction-tuning один пример обычно содержит инструкцию, контекст и ответ. Их общая длина меняется от строки к строке. Стандартный collator собирает примеры в тензор формы [batch_size, max_length_in_batch]. Более короткие последовательности дополняются PAD-токенами, а соответствующие позиции получают нулевое значение в attention mask.

Маска сообщает модели, что padding не нужно учитывать при вычислении результата. Она не всегда избавляет GPU от обработки формы тензора. Для dense attention короткая последовательность все равно занимает место в квадратной матрице размера, заданного самой длинной строкой.

Как длина самого длинного примера определяет стоимость батча

Для батча с длинами l1, l2 и l3 полезный объем равен:

useful_tokens = l1 + l2 + l3

После padding объем входа равен:

padded_tokens = batch_size * max(l1, l2, l3)

Разницу можно оценить формулой:

padding_overhead = 1 - useful_tokens / padded_tokens

Для длин 120, 300, 600 и 1024 получаем 2044 полезных токена при 4096 позициях в padded batch. Накладная доля равна примерно 50,1 процента. При другом размере батча, токенизаторе или фильтрации данных цифра изменится.

Attention не единственный потребитель памяти. При обучении хранятся активации, состояния оптимизатора и градиенты. Padding увеличивает объем активаций, а при длинных последовательностях стоимость attention растет особенно заметно. Flash Attention 2 уменьшает память для самого attention, но не превращает любую архитектуру в padding-free модель автоматически.

Почему проблема заметнее на смешанных instruction-датасетах

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

Однородный датасет дает меньший потенциал. Например, если все записи занимают примерно 900-1024 токена при лимите 1024, padding overhead невелик. Если рядом встречаются записи длиной 64 и 1024 токена, разница становится существенной.

Для LoRA и других PEFT-сценариев это особенно практично при ограниченной VRAM. Packing не снижает число параметров адаптера и не заменяет 4-битную загрузку, gradient checkpointing или настройку batch size. Он сокращает пустое место в последовательностях, а остальные параметры нужно держать под контролем отдельно.

Packing в Hugging Face: как несколько примеров превращаются в плоский батч

В обучающем пайплайне встречаются три связанные операции. Concatenation соединяет последовательности в один поток. Packing укладывает несколько примеров в доступные блоки заданного размера, стараясь уменьшить незаполненное место. Padding-free flattening меняет форму уже собранного батча, убирая обычную сетку с PAD-позициями.

Эти операции отвечают за разные задачи. Укладка определяет, какие примеры попадут в один блок. Flattening определяет, как блок передадут модели. Изоляция примеров требует отдельного механизма attention или совместимой интеграции с varlen kernel.

Packing, concatenation и padding-free flattening: не одно и то же

Подход	Представление	Границы примеров	Основной риск
Обычный padding	Все строки выровнены до общей длины	Задаются attention mask	Вычисления на PAD-позициях
Простое concatenation	Примеры соединяются в один поток через EOS или другой разделитель	Могут быть не изолированы	Cross-example attention и изменение контекста для следующего примера
Packing с блоковыми границами	Несколько примеров занимают один блок фиксированного размера	Задаются block-diagonal mask или эквивалентными индексами	Накладные расходы маски и ограниченная поддержка kernel
Flattening с поддерживаемым Flash Attention 2	Токены батча передаются плотным плоским потоком	Передаются через служебные данные и varlen-путь	Несовместимость модели или fallback на другой backend

DataCollatorWithFlattening относится прежде всего к последнему уровню. Он может объединить токены примеров в flattened representation, но сам по себе не обязан выполнять сложный bin packing по всему датасету. В обычном Trainer фактическая укладка в блоки может происходить на этапе подготовки датасета или через специальный batch sampler.

Что происходит с labels и границами примеров

Для causal language modeling collator должен сохранить соответствие между input_ids и labels. Если пример содержит 120 токенов, обе последовательности должны описывать те же 120 позиций после flattening. Для masked loss обычно используют -100 на токенах, по которым градиент считать не нужно.

При instruction-tuning часто маскируют часть prompt и оставляют целевыми только токены ответа. Эта логика должна сохраниться после объединения. Нельзя рассчитывать, что EOS автоматически решит вопрос с labels. EOS может служить разделителем, но решение о его включении в loss зависит от выбранной схемы подготовки данных.

Разделитель не создает изоляцию attention сам по себе. Если backend видит один обычный causal-поток, токены после EOS могут обращаться к предыдущим токенам. Для независимого обучения примеров нужны корректные границы в attention-представлении или специально поддержанный путь variable-length attention.

DataCollatorWithFlattening Hugging Face и границы примеров

DataCollatorWithFlattening подготавливает батч для padding-free обучения. В типичном сценарии он объединяет значения input_ids из нескольких features, объединяет labels и при включенном параметре return_position_ids=True создает позиции для каждого исходного примера.

Ожидаемые поля зависят от версии Transformers и структуры features. Базовый набор обычно включает input_ids и labels. Поле position_ids появляется при соответствующей настройке collator. attention_mask может отсутствовать в flattened batch, если модель и backend используют другой набор служебных индексов для variable-length attention.

Проверяйте сигнатуру установленной версии перед запуском. Название класса еще не гарантирует, что конкретная модель умеет принять его результат.

Какие данные должен вернуть collator

Минимальная проверка одного батча должна ответить на четыре вопроса:

• Есть ли input_ids и labels, если модель должна считать loss.
• Совпадает ли количество элементов в этих полях после flattening.
• Есть ли position_ids, если модель использует их для выбранного attention backend.
• Не добавился ли повторный padding на этапе preprocessing или внутри другого collator.

У flattened batch форма полей может отличаться от привычной формы [B, L]. Часто получается одномерный поток с общей длиной токенов батча, но точная форма определяется контрактом конкретной версии Transformers и модели.

Почему наивное склеивание может ухудшать обучение

Пусть первый пример содержит инструкцию о настройке базы данных, а следующий содержит код для другой задачи. При простом concatenation второй пример получает в causal-контексте токены первого. Модель может использовать чужой контекст при расчете logits, а loss перестает соответствовать независимому набору примеров.

У такой схемы меняется не только скорость. Меняется задача, которую видит модель во время обучения. На коротком датасете это может быть незаметно, но при большом числе границ эффект способен проявиться в loss и качестве ответов.

Корректная схема передает модели информацию о длинах и начале каждой последовательности. В зависимости от стека это могут быть block-diagonal attention mask, cumulative sequence lengths или другая служебная структура. position_ids помогают восстановить локальные позиции, но не служат универсальной заменой attention mask.

Flash Attention 2 и position_ids: как сохраняются позиции внутри packed-последовательности

После flattening позиции нельзя нумеровать только глобально. Для двух примеров длиной 3 и 2 поток может выглядеть так:

input_ids: [A0, A1, A2, B0, B1]

position_ids: [0, 1, 2, 0, 1]

Вариант [0, 1, 2, 3, 4] продолжает позиции первого примера и не описывает локальную позицию токенов второго. Для моделей с абсолютными позициями или RoPE это может изменить позиционную информацию, которую получает модель.

Зачем position_ids сбрасываются на границах примеров

Сброс позиции сообщает модели, что первый токен B0 находится в начале своей последовательности, даже если физически он идет сразу после A2 в flattened buffer. Такой массив сохраняет локальную нумерацию внутри каждого примера.

Позиции не отвечают на вопрос, может ли B1 читать A0. Этот вопрос решает attention-путь. В настоящем varlen режиме backend знает границы отдельных последовательностей и считает attention внутри них. При обычном causal attention один плоский поток может сохранить доступ через границу.

Что Flash Attention 2 дает этой схеме

Flash Attention 2 использует более эффективный способ вычисления attention, уменьшая обращение к большой промежуточной матрице и лучше используя память GPU. В режиме variable-length attention он может обработать последовательности разной длины через плотный набор токенов и служебные сведения о границах.

Flash Attention 2 не выполняет packing за collator. Он не решает, какие примеры нужно объединить, не создает labels и не определяет, какие prompt-токены исключить из loss. Эти задачи остаются у preprocessing, collator и Trainer.

Для работы нужны совместимые модель, PyTorch, CUDA, GPU, версия пакета flash-attn и attention implementation. Если модель не поддерживает передаваемый формат, загрузка или forward могут завершиться ошибкой. В другой конфигурации стек способен выбрать eager или SDPA-путь, что меняет производительность.

Скорость обучения и скорость генерации измеряются по-разному. Для понимания различий между prefill, decode и поведением конкретного GPU-рантайма полезен разбор сравнения TensorSharp и llama.cpp. Его результаты нельзя переносить на training throughput: там другие kernels, формы батчей и критерии измерения.

Чем этот подход отличается от block-diagonal mask

Block-diagonal mask явно запрещает внимание между разными примерами. Для трех последовательностей она формирует три независимых блока на диагонали, а позиции вне блоков маскируются. Механизм понятен, но хранение и обработка произвольной маски могут ограничить преимущество fused attention.

Flattened input с varlen-метаданными описывает те же границы компактнее, если конкретный kernel умеет работать с таким форматом. Это может снизить расходы на маску. Цена подхода, зависимость от forward модели и версии attention backend.

Простое concatenation требует меньше служебных данных, но не гарантирует независимость примеров. Поэтому три схемы нужно сравнивать по корректности границ, форме тензоров, пропускной способности и сложности диагностики.

Когда packing действительно ускоряет обучение LLM в Hugging Face

Потенциал метода заранее оценивают по распределению длин. Для каждого батча полезно посчитать сумму токенов и сравнить ее с размером padded batch. Чем ниже полезная доля, тем больше пространства для улучшения, если внимание действительно ограничивает скорость.

Сценарии с наибольшим потенциалом

• Смешанный instruction-датасет, где короткие вопросы соседствуют с длинными ответами.
• Диалоги с большой разницей в количестве реплик.
• Корпус с фрагментами кода, документами и обычными короткими инструкциями в одном потоке.
• Обучение локальной LLM на GPU с жестким ограничением VRAM.
• Сценарий, где dataloader уже загружает GPU стабильно, а основное время шага уходит на forward и backward attention.

Например, для длин 64, 256 и 1024 при лимите 1024 полезная доля составляет 1344 / 3072, то есть около 43,8 процента. Это сигнал проверить packing. Он не предсказывает итоговый прирост tokens per second, потому что в реальном шаге участвуют embedding, MLP, коммуникации и операции оптимизатора.

Токенизатор тоже влияет на результат. Один и тот же набор текстов может давать разные длины при использовании разных vocabulary и правил нормализации. Потенциал оценивают после токенизации, а не по числу символов или слов.

Когда выигрыш будет небольшим или исчезнет

• Примеры имеют близкую длину, поэтому padding overhead мал.
• GPU ждет CPU, сеть, дисковое чтение или синхронизацию между устройствами.
• Модель использует attention backend, который не поддерживает flattened batch и position ids в нужном формате.
• Установка Flash Attention 2 не загрузилась, а запуск перешел на более медленный backend.
• Размер microbatch настолько мал, что расходы на подготовку данных заметны относительно forward.
• Настройки sequence length, gradient accumulation или effective batch size изменились одновременно с packing.

Packing не заменяет подбор длины последовательности и batch size. Слишком большой лимит сохраняет длинные вычисления, а слишком маленький лимит обрезает контекст. Gradient accumulation увеличивает effective batch size, но не убирает padding внутри каждого microbatch.

Минимальный протокол сравнения

1. Зафиксируйте модель, токенизатор, датасет, порядок данных, precision, GPU и число процессов.
2. Запустите baseline с обычным padding и измерьте время шага, tokens per second, loss и пиковую VRAM.
3. Посчитайте полезные токены и padding overhead на том же участке данных.
4. Включите flattening или packing, сохранив sequence length, effective batch size и gradient accumulation.
5. Проверьте, какой attention backend реально выбран, и зафиксируйте предупреждения о fallback.
6. Сравните результаты после нескольких прогретых шагов и отдельно проверьте loss на одинаковом наборе labels.

Для throughput используйте формулу useful_tokens / elapsed_time. Если один запуск считает PAD-токены, а второй показывает общее число элементов flattened buffer, сырые значения нельзя сравнивать напрямую.

Как включить packing через обычный Trainer

В обычном Trainer есть два уровня настройки. Токенизатор готовит независимые примеры с ограничением max_length и без padding. DataCollatorWithFlattening объединяет features текущего батча в плотное представление. Полноценная укладка большого числа примеров в блоки фиксированной длины может потребовать отдельной подготовки датасета или batch sampler.

Базовая схема конфигурации

Ниже приведен короткий шаблон. Значение 2048 служит примером, его нужно заменить лимитом конкретной модели и задачей. Поле text предполагает заранее подготовленный текст. Для response-only loss вместо копирования input_ids в labels нужно создать маску prompt и выставить -100 на исключаемых позициях.

import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
from transformers import DataCollatorWithFlattening
from transformers import Trainer, TrainingArguments

model_id = 'your-model'
max_length = 2048

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    attn_implementation='flash_attention_2',
)

def tokenize(example):
    encoded = tokenizer(
        example['text'],
        truncation=True,
        max_length=max_length,
        padding=False,
        add_special_tokens=True,
    )
    encoded['labels'] = encoded['input_ids'].copy()
    return encoded

tokenized_dataset = dataset.map(
    tokenize,
    remove_columns=dataset.column_names,
)

collator = DataCollatorWithFlattening(
    return_position_ids=True,
)

training_args = TrainingArguments(
    output_dir='out',
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    bf16=True,
    remove_unused_columns=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=collator,
)

trainer.train()

Параметр attn_implementation='flash_attention_2' передает модели требуемый attention backend, если эта модель и установленный стек его поддерживают. Параметр return_position_ids=True version-sensitive: проверьте его наличие и поведение в своей версии Transformers.

В этом шаблоне padding=False предотвращает padding на этапе токенизации. Если другой preprocessing уже создал padded-массивы, collator не сможет восстановить потерянную экономию. Общее число токенов flattened batch тоже должно укладываться в ограничения конкретного forward и memory budget. В varlen-пути отдельные примеры могут обрабатываться как независимые последовательности, но fallback способен трактовать поток иначе.

Что проверить перед первым запуском

Сначала вызовите collator вручную на нескольких features. Это позволяет поймать ошибку формы до старта обучения.

features = [tokenized_dataset[index] for index in range(4)]
batch = collator(features)

print({
    name: tuple(value.shape)
    for name, value in batch.items()
})
print(batch.get('position_ids'))

Ожидаемая проверка не сводится к наличию ключа position_ids. Убедитесь, что длины input_ids и labels совпадают, позиции перезапускаются на границах примеров, а модель принимает этот ключ в forward. После одного тестового шага проверьте, что loss числовой и не содержит неожиданных пропусков.

Как настроить packing в TRL через SFTTrainer

SFTTrainer скрывает часть preprocessing и удобен для instruction-tuning. В актуальных версиях TRL параметры packing и padding_free описывают разные уровни. packing=True просит укладывать несколько примеров в доступную длину, а padding_free=True включает плотную передачу последовательностей без обычного padding, если такая связка поддерживается текущей версией.

Точные имена аргументов меняются. В одной версии используется processing_class, в старой конфигурации может применяться аргумент tokenizer. Аналогичная ситуация встречается с max_length и устаревшим именем параметра максимальной длины. Перед запуском сверяйте сигнатуру установленного SFTConfig и SFTTrainer.

Packing и padding_free в SFTConfig

packing определяет способ размещения примеров внутри доступных блоков. padding_free влияет на представление батча после такой подготовки. Первый параметр не гарантирует удаление padding во всех конфигурациях, второй не превращает неподдерживаемую модель в совместимую с Flash Attention 2.

Не передавайте одновременно собственный collator, который уже flatten-ит batch, и автоматический режим TRL без проверки контракта. Двойная обработка может привести к неправильной форме, повторному объединению labels или потере служебных полей.

Минимальный пример SFTTrainer

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

model_id = 'your-model'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    attn_implementation='flash_attention_2',
)

training_args = SFTConfig(
    output_dir='out-sft',
    max_length=2048,
    packing=True,
    padding_free=True,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    bf16=True,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)

trainer.train()

Шаблон рассчитан на dataset, который текущий TRL умеет прочитать как текстовый или conversational dataset. Для диалогов проверьте структуру сообщений, роли и поле содержимого. Если версия не знает padding_free, нельзя считать, что режим включился молча: уберите неподдерживаемый аргумент только после проверки альтернативного collator и закрепите рабочую конфигурацию.

Частые ошибки в TRL

• Документация относится к другой версии TRL, поэтому SFTConfig отклоняет параметр или игнорирует ожидаемое поведение.
• Тексты заранее дополнены до общей длины, и packing=True работает поверх уже заполненных последовательностей.
• Модель не поддерживает нужную комбинацию position_ids и Flash Attention 2.
• Conversational dataset содержит нестандартные имена полей, из-за чего шаблон диалога применяется неправильно.
• Собственный data_collator конфликтует с collator, который создает SFTTrainer.
• В логах появился fallback на eager или SDPA, но время шага сравнивают с baseline без учета этого изменения.

Диагностику начинайте с одного батча и одного forward. После этого проверьте несколько шагов с фиксированным набором данных. Такой порядок помогает отличить ошибку API от отсутствия реального выигрыша на конкретном распределении длин.

Ограничения, совместимость и финальный чек-лист

Совместимость определяется не названием семейства модели, а деталями ее forward, позиционного кодирования и attention integration. Две модели с одинаковым типом causal language model могут по-разному обрабатывать position_ids, формы масок и flattened input.

Какие модели и окружения нужно проверять отдельно

• Модель принимает position_ids и использует их ожидаемым способом.
• Выбранный attention backend действительно поддерживает Flash Attention 2 для этой архитектуры.
• Версии PyTorch, CUDA и пакета flash-attn совместимы между собой.
• GPU поддерживает выбранную precision и нужные CUDA kernels.
• Особенности RoPE, sliding window attention, GQA, мультимодальных входов или специальных масок не конфликтуют с flattening.
• Gradient checkpointing включен одинаково в baseline и packed-варианте.

Универсального списка поддерживаемых моделей для этого режима недостаточно. Проверяйте документацию конкретной модели, сигнатуру ее forward и фактический backend во время запуска. Минимальный smoke test должен включать collate, forward, backward и проверку одного значения loss.

Как отличить реальное ускорение от неудачной настройки

Если tokens per second не вырос, сначала проверьте путь выполнения. Сообщение о загрузке Flash Attention 2 еще не доказывает, что каждый forward использует нужный varlen режим. Профилирование и логи backend полезнее, чем вывод по одному короткому запуску.

Сравнение должно учитывать полезные токены, а не только количество элементов в тензоре. Зафиксируйте:

• время шага после прогрева;
• полезные токены за шаг;
• tokens per second;
• пиковую VRAM;
• долю padding в baseline;
• loss и способ усреднения по labels;
• наличие fallback или предупреждений о несовместимости.

Если loss изменился сильнее ожидаемого, проверьте границы примеров, EOS, маску prompt и cross-example attention. Если VRAM не уменьшилась, узким местом может быть MLP, optimizer states или сохранение активаций. Если время шага осталось прежним, GPU может ждать dataloader или коммуникации.

Итог: кому стоит попробовать packing

Packing стоит проверить при instruction-tuning с большим разбросом длин, особенно когда короткие инструкции регулярно попадают в один батч с длинными ответами, кодом или документами. Для локального обучения это может освободить часть VRAM и увеличить полезный throughput, но результат зависит от всей конфигурации.

Начните с подсчета padding overhead, затем протестируйте baseline и packed-вариант на одинаковом участке данных. В обычном Trainer проверьте DataCollatorWithFlattening и передачу position_ids. В TRL проверьте совместимость packing, padding_free, SFTConfig и SFTTrainer. Длительный запуск имеет смысл только после проверки формы батча, backend, loss, VRAM и полезных токенов за секунду.

Подписаться на канал