Reimplementando o modelo GPT-2 de forma barata com `llm.c`

• Explica como reproduzir o modelo GPT-2 (124M) com llm.c em 90 minutos por US$ 20
• O GPT-2 (124M) é o menor modelo lançado pela OpenAI em 2019
• Usando um nó Lambda com 8X A100 80GB SXM, o custo é de cerca de US$ 14 por hora, totalizando aproximadamente US$ 20
  • Também é possível treinar com uma única GPU, mas leva mais tempo (4–24 horas)

Comparação de resultados

• No conjunto de validação FineWeb, mostrou desempenho melhor do que o checkpoint público divulgado pela OpenAI
  • No entanto, como o GPT-2 foi treinado com WebText, a comparação não é totalmente justa
• A acurácia no HellaSwag também foi medida, alcançando 29,9, próximo dos 33,7 do GPT-3 Small (124M)
  • Já supera os 29,4 do GPT-2 (124M)
  • Porém, aqui o treinamento foi feito com 10B tokens, enquanto o GPT-3 foi treinado com 300B tokens

Configuração mínima do ambiente

• Requer GPU (recomendado usar Lambda Labs)
• Guia baseado em Linux x86 64-bit Ubuntu 22.04 com CUDA 12
• Instalar o miniconda e depois a versão nightly do PyTorch (opcional)
• Instalar os pacotes necessários para o tokenizer
• Instalar cuDNN para melhorar a velocidade (opcional)
• Instalar MPI ao usar várias GPUs (opcional)
• Pré-processar o dataset FineWeb de 10B tokens (~1 hora)
• Compilar o llm.c (precisão mista, usando cuDNN FlashAttention)

Execução do treinamento

• Exemplo de comando para uso com uma única GPU
• Em múltiplas GPUs (8), executar com mpirun
• Explicação dos principais parâmetros
  • -i, -j : caminhos dos dados de treino/validação
  • -o : caminho para salvar logs e checkpoints
  • -e : inicialização do modelo (depth 12 GPT-2)
  • -b : tamanho do microbatch (reduzir se faltar memória)
  • -t : comprimento máximo da sequência
  • -d : tamanho total do batch (ver artigo do GPT-3)
  • -r : configuração de Recompute (economia de memória)
  • -z : ZeRO-1 (sharding do estado do otimizador)
  • Outras configurações como weight decay, taxa de aprendizado e intervalo de checkpoints

Processo de treinamento

• Com base em 10B tokens de treino e batch size de 0,5M, são esperados cerca de 20K steps
• Em uma A100 40GB PCIe, são exibidos o tempo por step, MFU e vazão de tokens
• No início do treinamento há gradient exploding, mas isso é resolvido com clipping

Visualização

• É fornecido um notebook Jupyter para fazer parsing do arquivo de log e visualizar as curvas de treinamento

Tokenizer

• Necessário para converter tokens inteiros em strings
• Pode ser gerado com um script em PyTorch

Sampling

• Atualmente não está otimizado para inferência
• Com pequenas alterações no código, é possível fazer sampling unconditional/conditional

Estrutura do código

• A maior parte da implementação está no arquivo train_gpt2.cu
• As primeiras 500 linhas tratam da configuração de MPI, NCCL, cuDNN, cuBLAS etc.
• As 1500 linhas seguintes cobrem o forward/backward do Transformer
• As próximas 1000 linhas implementam o modelo GPT-2
• As últimas 1000 linhas incluem o loop de treinamento, parsing de argumentos etc.

Modelo 350M

• 10B tokens não são suficientes, então são usados 30B tokens
• Em 8X A100 80GB, leva 14 horas e custa cerca de US$ 200

FAQ

• É possível fazer sampling?: Sim, mas é ineficiente.
• É possível conversar com ele?: No momento, só é possível fazer pré-treinamento; fine-tuning para chat não é suportado.
• Treinamento distribuído em múltiplos nós: É possível, mas ainda não foi testado.
• É determinístico em nível de bits?: Quase determinístico, mas alguns patches de kernel são necessários.
• É possível treinar com FP8?: Atualmente o treinamento é feito em BF16; suporte a FP8 deve chegar em breve.
• Há suporte para GPUs que não sejam NVIDIA?: No momento, só há suporte para C/CUDA.

Opinião do GN⁺

• O GPT-2 é um modelo muito importante como ponto de partida dos LLMs modernos. Depois dele, o GPT-3 e outros LLMs também não diferem tanto em essência.
• Este projeto permite que qualquer pessoa treine diretamente um modelo no nível do GPT-2 com custo razoável. Parece ser uma grande ajuda para aumentar a compreensão sobre LLMs.
• No entanto, como ainda não está otimizado para inferência, seu uso em serviços reais é limitado. Também não há suporte para fine-tuning em modelo conversacional.
• Atualmente só há suporte para GPUs NVIDIA, mas espera-se que no futuro haja suporte a várias plataformas, como AMD e Apple Silicon.
• Projetos open source com objetivo semelhante incluem Megatron-LM, DeepSpeed e FairSeq. Cada um tem seus prós e contras, então vale escolher de acordo com o uso.
• É um projeto muito animador para fortalecer o ecossistema de desenvolvimento de LLMs. As próximas evoluções prometem bastante.