Treinamento Eficiente em Memória

🤗 O PEFT torna o fine-tuning eficiente em termos de parâmetros, mas não automaticamente eficiente em termos de memória. Esta visão geral reúne dicas para reduzir a memória de treinamento e links para os guias detalhados.

Considere sempre as bases para reduzir o uso de memória: escolher um modelo base menor, um tamanho de lote menor ou um comprimento de sequência mais curto.

Visão geral da memória de treinamento

Vamos analisar a distribuição da memória de treinamento para que possamos raciocinar sobre possíveis contramedidas. Usaremos um grande modelo de linguagem treinado com o otimizador Adam como exemplo. Ao realizar o fine-tuning completo, teremos as seguintes posições ocupando memória:

1. parâmetros do modelo base: o consumo de memória depende fortemente do dtype escolhido. Quanto menos bits por parâmetro (16 para float16), menos memória será ocupada. Um modelo de 1B em float16 (16 bits/2 bytes por parâmetro) ocupará aproximadamente 1e9 × 2 byte = 1.863 GiB de memória.
2. todos os parâmetros treináveis do modelo base ×3 para gradientes (1×) e estados do otimizador Adam (2×), portanto 5,59 GB de memória
3. memória para as ativações intermediárias entre as camadas, que são difíceis de prever, mas dependem principalmente do dtype de computação utilizado e do comprimento da sequência / tamanho do lote.

Um modelo base menor ou o uso de um dtype de computação menor reduzirá todos os pontos, enquanto o uso de sequências mais curtas ou lotes menores afeta principalmente a memória de gradientes e ativações. O emprego de métodos PEFT reduzirá o número de parâmetros treináveis e, portanto, reduzirá significativamente tanto os gradientes quanto o estado do otimizador, economizando bastante memória.

Escolhendo o método adequado

Nem todo método de PEFT é construído da mesma forma e algumas formulações são mais fáceis de implementar de maneira eficiente em memória. Se você está com um orçamento de memória limitado, faz sentido consultar a suite de comparação de métodos PEFT e filtrar pelo uso máximo de memória do acelerador. O uso médio de memória do acelerador pode ser bastante semelhante entre os métodos, mas nem todo método escala igualmente com ativações e comprimento da sequência; alguns métodos são mais propensos a picos de memória do que outros.

Considere usar tokens treináveis ao direcionar camadas grandes, como cabeças de modelagem de linguagem ou camadas de incorporação, para ajustar finamente tokens específicos.

Quantização

A quantização é uma das melhores formas de reduzir o consumo de memória do modelo base e, dependendo da quantização empregada, também reduzirá a memória de ativação. Como os métodos PEFT ocuparão apenas uma pequena parte do número total de parâmetros, o PEFT usa por padrão uma precisão maior do que a do modelo base. Isso também pode ter o efeito de os adaptadores mitigarem parte da perda de qualidade causada pelos métodos de quantização. Leia o guia de quantização do PEFT.

Compilação

Os modelos que treinamos são compostos por operações como multiplicações de matrizes, somas e atribuições, nas quais cada operação produz um novo resultado e, em seguida, precisa ocupar memória. Se esses resultados intermediários não forem necessários, podemos fundir essas operações e economizar memória. Esta é apenas uma das muitas otimizações que o torch.compile pode fazer para você, então confira o guia do torch.compile do PEFT.

Checkpointing de Gradiente

Você pode trocar memória por computação salvando apenas a cada n-ésimo gradiente entre as camadas e calculando o restante em tempo real. Consulte a documentação de gradient checkpointing do Transformers para saber mais.

Quando não estiver usando Diffusers ou Transformers, você pode precisar implementar sua própria lógica de checkpointing de gradiente, dependendo do framework de treinamento que estiver utilizando.

Perda NLL em blocos

O uso de NLLLoss é muito comum ao treinar modelos de linguagem (ou tarefas de classificação). Você aloca uma matriz de tamanho batch × sequence × vocabulary. Com sequências ou vocabulários particularmente longos, isso pode ficar caro rapidamente.

Ao usar o TRL, você pode usar a integração do kernel Liger ou usar o Chunked NLLLoss. Este último divide a sequência em blocos de tamanho 256 para manter o consumo máximo de memória constante.

Caso o tamanho padrão do chunk não seja ideal para sua configuração, consulte o PR original do TRL para mais informações sobre como ajustar o tamanho do chunk.