Treinamento sem padding

O treinamento sem padding (também chamado de packing) concatena várias amostras em uma única sequência, em vez de aplicar padding a cada uma até um comprimento fixo. O modelo precisa saber onde cada amostra termina para que a atenção (linear) não misture tokens entre as amostras.

Existem duas maneiras de fornecer essas fronteiras.

• Prepare-os previamente com um data collator.
• Infira-os de position_ids em tempo de execução.

A abordagem recomendada é o data collator. Este guia explica o motivo e aborda as ressalvas do caminho position_ids.

Inferir fronteiras a partir de position_ids não é a abordagem recomendada e só funciona para modelos de atenção padrão. Modelos de atenção linear, como Qwen3-Next e Qwen3.5 (Gated DeltaNet), e modelos baseados em convolução ignoram as fronteiras de position_ids e exigem o data collator. Consulte Modelos de atenção linear e convolução.

Prepare fronteiras com um data collator

Preparar os kwargs de limite antecipadamente elimina os problemas acima e se comporta de forma idêntica, independentemente de você compilar ou não.

Use DataCollatorWithFlattening para achatar cada lote e retornar as informações de fronteira. Defina return_flash_attn_kwargs=True para que o collator pré-compute as fronteiras, em vez de deixá-las para serem inferidas a partir de position_ids em tempo de execução. Passe-o para Trainer e não adicione uma attention_mask, pois o lote achatado já codifica as fronteiras e uma máscara conflita com o layout empacotado.

O Padding-free depende de uma implementação do FlashAttention para modelos de atenção padrão, pois apenas os kernels do FlashAttention expõem o caminho de comprimento variável que um lote achatado precisa.

Instale a biblioteca kernels, que obtém um kernel FlashAttention pré-compilado sem exigir uma compilação local. Ela também funciona como alternativa quando o flash-attn não está instalado localmente. Carregue o modelo com attn_implementation="kernels-community/flash-attn2".

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

model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
dataset = dataset.map(
    lambda example: tokenizer(example["text"], truncation=True, max_length=512),
    remove_columns=dataset.column_names,
)

# return_flash_attn_kwargs=True precomputes the sequence boundaries
data_collator = DataCollatorWithFlattening(return_flash_attn_kwargs=True)

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="padding-free-llama"),
    train_dataset=dataset,
    data_collator=data_collator,
)
trainer.train()

Inferir limites a partir de position_ids

O FlashAttention pode detectar lotes sem padding apenas a partir de position_ids e permanece por compatibilidade reversa, pois frameworks downstream, como o TRL, dependem dele.

Confiar em position_ids traz dois problemas.

• A detecção de sequências empacotadas a partir de position_ids é uma verificação dinâmica e dependente dos dados. Ela funciona sem compilação, mas, sob torch.compile, causa quebras de grafo. A verificação está atualmente restrita a batch_size == 1 para limitar a frequência com que é executada, já que os tamanhos de lote reais costumam ser maiores.
• O FlashAttention compilado força alguns kwargs a serem ints Python puros. Inferi-los a partir de position_ids em tempo de execução força sincronizações de dispositivo para host e, em versões mais antigas do PyTorch, uma quebra de grafo extra devido à conversão de tensor para int.

Modelos de atenção linear e convolução

Gated DeltaNet (GDN), outras camadas de atenção linear e convoluções causais não possuem um caminho apenas com position_ids, por design. Preparar os dados com o collator é a única opção suportada para esses modelos.

Não dependa apenas de position_ids para modelos GDN, linear-attention ou de convolução causal. Prepare os kwargs de fronteira, incluindo seq_idx, com o data collator.

Para esses modelos, defina tanto return_flash_attn_kwargs=True quanto return_seq_idx=True.

from transformers import DataCollatorWithFlattening

data_collator = DataCollatorWithFlattening(
    return_flash_attn_kwargs=True,
    return_seq_idx=True,
)

Os pacotes de kernels exatos dependem da implementação original do modelo. Modelos Gated DeltaNet, como Qwen3-Next e Qwen3.5, usam flash-linear-attention, e modelos baseados em Mamba, como Bamba, usam mamba-ssm. Ambos dependem de causal-conv1d para a convolução. Sem os kernels corretos, o modelo recorre a implementações de referência que ignoram os kwargs de fronteira e misturam tokens entre amostras.

Muitos desses kernels também estão disponíveis por meio da biblioteca kernels, que pode buscar um build compatível para você. O flash-linear-attention geralmente ainda requer uma instalação direta.

Quando os kwargs de fronteira estão ausentes, os kernels tratam silenciosamente o lote inteiro como uma única sequência. Nenhum erro ou aviso é gerado, pois uma verificação em tempo de execução adicionaria um ramo dependente dos dados, o que conflita com torch.compile.

Próximas etapas

• Consulte o guia de data collators para conhecer outros collators.
• Consulte a referência da API de DataCollatorWithFlattening para o conjunto completo de argumentos.
• Leia Improving Hugging Face Training Efficiency Through Packing with Flash Attention para benchmarks e uma explicação mais detalhada.