Otimizando Memória para Inferência e Ajuste Fino de Modelos de Linguagem Grande

Modelos de linguagem grande (LLMs) como GPT-4, Bloom e LLaMA alcançaram capacidades notáveis ao escalar para bilhões de parâmetros. No entanto, implantar esses modelos maciços para inferência ou ajuste fino é desafiador devido às suas imensas necessidades de memória. Neste blog técnico, exploraremos técnicas para estimar e otimizar o consumo de memória durante a inferência e o ajuste fino de LLMs em vários conjuntos de hardware.

Entendendo os Requisitos de Memória

A memória necessária para carregar um LLM é determinada principalmente pelo número de parâmetros e pela precisão numérica usada para armazenar os parâmetros. Uma regra simples é:

• Carregar um modelo com X bilhões de parâmetros requer aproximadamente 4X GB de VRAM em precisão de ponto flutuante de 32 bits
• Carregar um modelo com X bilhões de parâmetros requer aproximadamente 2X GB de VRAM em precisão de 16 bits bfloat16/float16

Por exemplo, carregar o modelo GPT-3 com 175 bilhões de parâmetros exigiria aproximadamente 350 GB de VRAM em precisão bfloat16. Até o momento, os maiores GPUs comerciais disponíveis, como o NVIDIA A100 e H100, oferecem apenas 80 GB de VRAM, necessitando técnicas de paralelismo de tensor e paralelismo de modelo.

Durante a inferência, a pegada de memória é dominada pelos parâmetros do modelo e pelos tensores de ativação temporários produzidos. Uma estimativa de alto nível para o uso de memória de pico durante a inferência é a soma da memória necessária para carregar os parâmetros do modelo e a memória para ativações.

Quantificando a Memória de Inferência

Vamos quantificar os requisitos de memória para a inferência usando o modelo OctoCode, que tem cerca de 15 bilhões de parâmetros no formato bfloat16 (~ 31 GB). Usaremos a biblioteca Transformers para carregar o modelo e gerar texto:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Pergunta: Por favor, escreva uma função em Python para converter bytes para gigabytes.\n\nResposta:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Saída:

O uso de memória de pico na GPU é de aproximadamente 29 GB, o que está de acordo com nossa estimativa de 31 GB para carregar os parâmetros do modelo em formato bfloat16.

Otimizando a Memória de Inferência com Quantização

Embora a precisão bfloat16 seja comum para treinamento de LLMs, pesquisadores descobriram que a quantização dos pesos do modelo para tipos de dados de precisão mais baixa, como inteiros de 8 bits (int8) ou inteiros de 4 bits, pode reduzir significativamente o uso de memória com perda mínima de precisão para tarefas de inferência como geração de texto.

Vamos ver as economias de memória com a quantização de 8 bits e 4 bits do modelo OctoCode:

</div>
# Quantização de 8 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Saída:

# Quantização de 4 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Saída:

Com a quantização de 8 bits, a necessidade de memória cai de 31 GB para 15 GB, enquanto a quantização de 4 bits reduz ainda mais para apenas 9,5 GB! Isso permite executar o modelo OctoCode de 15 bilhões de parâmetros em GPUs de consumidor como o RTX 3090 (24 GB de VRAM).

No entanto, note que a quantização mais agressiva, como a de 4 bits, às vezes pode levar a uma degradação da precisão em comparação com a precisão de 8 bits ou bfloat16. Há um trade-off entre economia de memória e precisão que os usuários devem avaliar para o seu caso de uso.

A quantização é uma técnica poderosa que pode permitir a implantação de LLMs em ambientes com recursos limitados, como instâncias de nuvem, dispositivos de borda ou até mesmo telefones celulares, reduzindo drasticamente a pegada de memória.

Estimando Memória para Ajuste Fino

Enquanto a quantização é usada principalmente para inferência eficiente, técnicas como paralelismo de tensor e paralelismo de modelo são cruciais para gerenciar as necessidades de memória durante o treinamento ou ajuste fino de modelos de linguagem grande.

O consumo de memória de pico durante o ajuste fino é tipicamente 3-4 vezes maior do que a inferência devido a necessidades adicionais de memória para:

• Gradientes
• Estados do otimizador
• Ativações da passagem direta armazenadas para retropropagação

Uma estimativa conservadora é que o ajuste fino de um LLM com X bilhões de parâmetros requer cerca de 4 * (2X) = 8X GB de VRAM em precisão bfloat16.

Por exemplo, ajustar o modelo LLaMA de 7 bilhões de parâmetros exigiria aproximadamente 7 * 8 = 56 GB de VRAM por GPU em precisão bfloat16. Isso excede a capacidade de memória dos GPUs atuais, necessitando técnicas de ajuste fino distribuído.

Técnicas de Ajuste Fino Distribuído

Vários métodos de ajuste fino distribuído foram propostos para superar as limitações de memória de GPU para modelos grandes:

1. Paralelismo de Dados: A abordagem clássica de paralelismo de dados replica o modelo inteiro em várias GPUs enquanto divide e distribui os lotes de dados de treinamento. Isso reduz o tempo de treinamento linearmente com o número de GPUs, mas não reduz a necessidade de memória de pico em cada GPU.
2. ZeRO Stage 3: Uma forma avançada de paralelismo de dados que divide os parâmetros do modelo, gradientes e estados do otimizador em GPUs. Isso reduz a memória em comparação com o paralelismo de dados clássico, mantendo apenas os dados partitionados necessários em cada GPU durante diferentes fases do treinamento.
3. Paralelismo de Tensor: Em vez de replicar o modelo, o paralelismo de tensor divide os parâmetros do modelo em linhas ou colunas e os distribui em GPUs. Cada GPU opera em um conjunto partitionado de parâmetros, gradientes e estados do otimizador, levando a economias significativas de memória.
4. Paralelismo de Pipeline: Essa técnica divide as camadas do modelo em diferentes GPUs/trabalhadores, com cada dispositivo executando um subconjunto das camadas. As ativações são passadas entre os trabalhadores, reduzindo a memória de pico, mas aumentando a sobrecarga de comunicação.

Estimar o uso de memória para esses métodos distribuídos não é trivial, pois a distribuição de parâmetros, gradientes, ativações e estados do otimizador varia entre as técnicas. Além disso, componentes diferentes, como o corpo do transformador e a cabeça de modelagem de linguagem, podem exibir comportamentos de alocação de memória diferentes.

A Solução LLMem

Pesquisadores recentemente propuseram LLMem, uma solução que estima com precisão o consumo de memória de GPU quando se aplicam métodos de ajuste fino distribuído a LLMs em várias GPUs.

Estimando o Uso de Memória de GPU para Ajuste Fino de LLMs Pré-Treinados

LLMem considera fatores como a recombinação de parâmetros antes do cálculo (ZeRO Stage 3), a coleta de saída na passagem de volta (paralelismo de tensor) e as diferentes estratégias de alocação de memória para o corpo do transformador e a cabeça de modelagem de linguagem.

Os resultados experimentais mostram que LLMem pode estimar o uso de memória de GPU de pico para ajuste fino de LLMs em uma única GPU com taxas de erro de até 1,6%, superando a taxa de erro média do estado da arte DNNMem de 42,6%. Quando se aplicam métodos de ajuste fino distribuído a LLMs com mais de um bilhão de parâmetros em várias GPUs, LLMem alcança uma taxa de erro média impressionante de 3,0%.