Como fazer fine-tuning de um modelo de linguagem 70B em casa

• Open source que combina FSDP(Fully Sharded Data Parallel) e QLoRA(Quantization + Low-Rank Adaptation")
• Permite fazer o ajuste fino de um modelo de linguagem de 70B de forma eficiente com um computador desktop comum e duas GPUs gamer padrão

Contexto

• Apesar de não haver uma diferença tão grande de desempenho entre hardware de nível de data center e GPUs gamer instaladas em desktops, a diferença de preço é enorme
• GPUs gamer têm pouca memória, o que dificulta treinar modelos de linguagem de grande escala
• A Answer.AI quis encontrar uma forma de treinar os melhores modelos open source usando GPUs gamer baratas para resolver esse problema

QLoRA: treinando modelos maiores em uma única GPU

• QLoRA é a combinação de quantização com LoRA, reduzindo os pesos do modelo para 4 bits e diminuindo o uso de memória da GPU.
• LoRA evita as limitações de memória ao adicionar pequenas matrizes adaptadoras para treinamento, em vez de treinar todo o modelo de linguagem de grande escala.
• O QLoRA permite treinar um modelo 65b em uma placa de 48GB, mas ainda não é suficiente para treinar um modelo 70b em uma GPU gamer de 24GB.

FSDP: ampliando a escala do treinamento com várias GPUs

• FSDP(Fully Sharded Data Parallel) distribui os parâmetros do modelo entre várias GPUs para que possam ser usados ao mesmo tempo.
• O FSDP permite aplicar o desempenho do DDP(Distributed Data Parallel) também a modelos grandes demais para caber em uma única GPU.
• Por exemplo, um modelo 70b (70 bilhões de parâmetros) não quantizado usa 140GB de RAM (porque cada parâmetro é armazenado com 16 bits, ou seja, 2 bytes), então até uma placa H100 da NVIDIA (custando cerca de $40,000 por placa!) fica aquém. 80GB de RAM são necessários
  • Porém, com FSDP é possível combinar 4 GPUs H100 para formar um total de 320GB de RAM. E economizar $150000

Combinando FSDP e QLoRA

• A Answer.AI reduz o tamanho do modelo com QLoRA e o distribui entre várias placas gamer de 24GB com FSDP, possibilitando um treinamento eficiente.
• Com isso, tornou-se possível treinar um modelo 70b com duas GPUs gamer comuns.
• No fim de 2023, a equipe discutiu a ideia de combinar FSDP e QLoRA e depois resolveu os problemas necessários para isso.
• Nesse processo, várias tecnologias e bibliotecas foram combinadas e usadas
  • Armazenar parâmetros quantizados com o mesmo tipo de dado do "Computation Type" do modelo
  • Gradient checkpointing, que em vez de armazenar o gradiente completo, salva ativações em vários "checkpoints" de todo o modelo e depois reexecuta etapas de forward pass conforme necessário para recalcular os gradientes
  • CPU offloading, que armazena pesos na RAM da CPU em vez da GPU quando não estão em uso, reduzindo drasticamente a memória de GPU necessária
  • Flash Attention 2: calcula Attention com eficiência usando kernels Cuda otimizados para memória
  • Descoberta do HQQ(Half-Quadratic Quantization), integrado ao FSDP (50 vezes mais rápido e preciso em comparação ao GPTQ)

Como usar FSDP/QLoRA

• Para usar FSDP, são necessárias duas ou mais GPUs. (No Runpod, é possível usar uma máquina dual 3090 por $0.6/hora)
• Instale as versões mais recentes de Transformers, PEFT, Bitsandbytes (e HQQ, se for usar), clone o repositório da Answer.AI e execute seguindo o README.
• Para treinar Alpaca em 2 placas de 24GB:
  python train.py --train_type qlora --dataset alpaca --batch_size 8 --gradient_accumulation_steps 2 --output_dir qlora_output --log_to wandb
• Ainda está em estágio inicial (alpha/preview)
• Resultados de benchmark também devem ser divulgados nas próximas semanas

Opinião do GN⁺

• Este artigo deve ser muito interessante para quem trabalha com IA e machine learning. Em especial, chama atenção por apresentar um novo método para treinar modelos de linguagem de grande escala com eficiência.
• Se essa tecnologia passar a ser usada com sucesso em larga escala, poderá reduzir bastante a barreira de custo em pesquisa e desenvolvimento de IA. Isso também pode abrir grandes oportunidades para pequenos laboratórios e desenvolvedores individuais.
• No entanto, para que essa abordagem seja realmente amplamente adotada, estabilidade e facilidade de uso serão fatores importantes. Como ainda está em estágio inicial, o sistema precisa ser melhorado com testes e feedback da comunidade.
• Ao adotar uma nova tecnologia, sempre é preciso considerar compatibilidade com sistemas existentes, desempenho e custo. O benefício desta tecnologia é melhorar a acessibilidade ao treinamento de modelos de grande escala, mas como ainda está em estágio inicial, pode haver desvantagens em estabilidade e suporte.

[GN⁺]

Comentários do Hacker News

• O texto é excelente, mas é uma pena que houve tantas chances de usar a palavra 'nibble' e todas foram perdidas.

  • O comentário elogia o texto, mas lamenta que a palavra 'nibble' não tenha sido aproveitada.

• Como consumidor, eu teria interesse em uma tabela mostrando quanto tempo levaria para ajustar um modelo 70b em um sistema com pesos de 4 bits, por exemplo com tamanho do modelo, tamanho do treinamento, custo do treinamento e a perda de qualidade esperada com essa técnica.

  • O comentário demonstra interesse no custo final do fine-tuning e pede informações detalhadas, incluindo tamanho do modelo e custo de treinamento.

• Fico me perguntando se seria possível aplicar codificação (-1,0,1). Quando esse modelo sair, gostaria de saber se a answer.ai conseguiria usar a mesma técnica para ajustar um modelo grande com duas GPUs pequenas, além de uma análise de custo/benefício.

  • O comentário levanta a possibilidade de aplicar uma nova técnica de codificação e pede uma análise de custo versus benefício.

• Esse tipo de pesquisa é muito útil para permitir que usuários individuais concorram no mercado de LLMs. É um ótimo trabalho e muito alinhado com a ideia de 'zero a um'.

  • O comentário avalia positivamente a utilidade da pesquisa para usuários individuais e sua contribuição para resolver um problema específico.

• O texto não mencionou que todos os trabalhos mais bem colocados no desafio de eficiência da NeurIPS usaram QLoRA, e isso é um ótimo exemplo da importância e utilidade do QLoRA.

  • O comentário destaca a importância e utilidade da técnica QLoRA, mencionando seu reconhecimento em uma competição.

• É a abordagem mais interessante e construtiva que já vi em LLMs 'open source'. Espero que a academia também a adote como padrão.

  • O comentário elogia bastante a abordagem para LLMs open source e deseja que ela vire padrão na academia.

• Tomara que o treinamento rápido no Metal chegue mais cedo. No momento, o Mac Studio é muito mais barato e rápido o suficiente para inferência em modelos grandes, mas quase não existem soluções de treinamento.

  • O comentário expressa expectativa por treinamento rápido na plataforma Metal e opina sobre a eficiência atual do Mac Studio.

• A qualidade e a clareza da escrita técnica são excelentes.

  • O comentário elogia a qualidade e a clareza da documentação técnica.

• Fico me perguntando se é possível retreinar os pesos de um LLM (70B) com meus próprios dados.

  • O comentário pergunta sobre a possibilidade de retreinar um LLM com dados próprios.

• Para quem faz fine-tuning de LLMs com orçamento limitado de hardware, isso é um avanço revolucionário.

  • O comentário avalia a técnica como um grande avanço para usuários com orçamento de hardware limitado.

• Fico curioso para saber por que escolheram DeepSpeed em vez de FSDP.

  • O comentário expressa curiosidade sobre a escolha entre FSDP e DeepSpeed.

• Queria saber se houve alguma ferramenta usada para tuning de desempenho quando vocês conseguiram fazer o fine-tuning do modelo.

  • O comentário pergunta sobre ferramentas usadas para ajuste de desempenho.

• Vocês já tentaram usar esparsificação? Me parece que isso poderia permitir fine-tuning com menos RAM.

  • O comentário sugere o uso de esparsificação e a possível redução de RAM com essa técnica.