Ir para o conteúdo

Treinando LLMs do zero com PyTorch

Este artigo descreve o conteúdo do repositório train-llm-from-scratch, que oferece uma abordagem completa para treinar grandes modelos de linguagem (LLMs) do início ao fim, utilizando apenas PyTorch e datasets públicos.

Visão geral

O projeto implementa um arquitetura transformer from scratch, seguindo exatamente as propostas do paper Attention is All You Need. Todo o código está em Python e depende exclusivamente do PyTorch, permitindo treinar modelos com escala de milhões a bilhões de parâmetros em uma única GPU.

Observacao

Embora o treinamento seja possível em uma GPU única, o repositivo também suporta escalonamento multi‑GPU via Distributed Data Parallel (DDP) e precisão bf16.

Arquitetura e treinamento

O script de treinamento carrega os datasets selecionados (Alpaca, Dolly, Anthropic HH‑RLHF, UltraFeedback, GSM8K), tokeniza o texto e alimenta o transformer treinado com a função de perda padrão de modelagem de linguagem. Não há uso de bibliotecas de alto nível como 🤗 Transformers; tudo, desde a camada de atenção até os blocos de feed‑forward, foi escrito manualmente.

Dica

Para iniciar o treinamento rápido, basta clonar o repositório, instalar as dependências (Python ≥ 3.8 e PyTorch) e executar o script de pré‑treinamento indicado no README.

Suite de pós‑treinamento

Além do pré‑treinamento, o repositório inclui uma sequência completa de alinhamento:

  1. Supervised Fine‑Tuning (SFT) – ajuste fino com instruções formatadas.
  2. Reward Model – modelo que classifica respostas conforme preferência humana.
  3. Otimização por política – escolha entre PPO, DPO ou GRPO/RLVR para otimizar a política usando o modelo de recompensa.

Cada estágio é implementado em PyTorch puro, sem recorrer a pacotes externos como trl ou peft, o que facilita a inspeção e modificação do código.

Atencao

O treinamento dos estágios de pós‑treinamento pode exigir recursos significativos de memória; recomenda‑se usar múltiplas GPUs com DDP quando trabalhar com modelos acima de alguns centenas de milhões de parâmetros.

Exemplos de saída

Ao final do treinamento de um modelo de 13 milhões de parâmetros, o repositório exibe um trecho de texto gerado, ilustrando a capacidade do modelo de produzir frases coerentes (embora ainda com algumas inconsistências típicas de modelos pequenos). O bloco abaixo mostra exatamente o que foi apresentado no README:

In ***1978, The park was returned to the factory-plate that  
the public share to the lower of the electronic fence that  
follow from the Station's cities. The Canal of ancient Western  
nations were confined to the city spot. The villages were directly  
linked to cities in China that revolt that the US budget and in Odambinais is uncertain and fortune established in rural areas.  

Pontos-chave

  • Implementação completa de um transformer do zero usando apenas PyTorch.
  • Scripts para pré‑treinamento e pós‑treinamento (SFT → Reward Model → {PPO, DPO} → GRPO/RLVR) em PyTorch puro.
  • Suporte a treinamento em GPU única e multi‑GPU (DDP + bf16) com datasets públicos reais.
  • Código totalmente comentado e acompanhado por guia detalhado em POST_TRAINING.md.
  • Exemplo de saída fornecido para demonstração rápida das capacidades do modelo.

Ferramentas e Tecnologias

  • [[Python]]
  • [[PyTorch]]

Nota pessoal

https://github.com/FareedKhan-dev/train-llm-from-scratch

Tags

llm #pytorch #transformador #post-treinamento #datasets