Qual é o modelo de IA mais poderoso que pode ser treinado em um notebook em 5 minutos?
(seangoedecke.com)- Em um MacBook Pro, foi treinado em 5 minutos um modelo transformer no estilo GPT com cerca de 1,8M de parâmetros usando aproximadamente 20M de tokens do TinyStories, alcançando cerca de 9,6 de perplexidade
- As principais limitações do treinamento em 5 minutos são o tamanho do modelo e a quantidade de tokens processáveis; modelos maiores convergem mais lentamente e ficam menos eficazes com menos dados
- Na otimização de desempenho, usar MPS, compilar/quantizar/acumular gradientes ou substituir o PyTorch é menos eficaz do que escolher um modelo pequeno
- Um conjunto de dados simples e consistente como o TinyStories tem impacto mais positivo no desempenho de modelos pequenos do que dados enciclopédicos
- A arquitetura transformer apresentou resultados superiores a abordagens de LSTM ou diffusion em condições de tamanho reduzido e tempo curto de treinamento
Visão geral
Este texto apresenta os resultados de um experimento para descobrir o modelo de linguagem de IA com melhor desempenho que pode ser treinado em 5 minutos em um notebook (MacBook Pro), além de trazer insights sobre a estratégia ideal de treinamento, escolha de dataset e arquitetura de modelo.
Resumo dos resultados do experimento
- Foi treinado um modelo transformer no estilo GPT com cerca de 1,8M de parâmetros em aproximadamente 20M de dados do TinyStories, registrando 9,6 de perplexidade
- Os exemplos gerados são curtos, mas no formato de histórias consistentes, com gramática inglesa em geral corretamente mantida
- Destaca-se que os resultados do modelo, em um nível prático dentro de 5 minutos, foram melhores do que o esperado
Contexto e limitações do experimento
- O experimento começou a partir da curiosidade pouco realista de treinar rapidamente um modelo poderoso em um ambiente de notebook
- Na prática, seria possível treinar modelos mais poderosos na nuvem com GPUs de alto desempenho (como H100), mas a condição limitante do experimento é o tempo: 5 minutos
- Quanto maior o modelo, mais lenta fica a velocidade de processamento de tokens, dificultando obter bons resultados em 5 minutos
- Modelos pequenos demais (ex.: 10K parâmetros) não conseguem aprender complexidade suficiente
- A faixa prática é de aproximadamente 1M a 2M de parâmetros
Otimização de throughput
- O uso de MPS (Metal Performance Shaders da Apple) é o mais eficaz
- Diversas otimizações matemáticas como
torch.compile, float16 e MLX tiveram ganho de desempenho menor do que o esperado ou até pioraram o resultado - O acúmulo de gradientes ajuda na gestão de memória, mas na prática provoca forte queda de velocidade
- Para ser eficiente, o modelo precisa conseguir fazer updates de pesos rapidamente na memória interna
Escolha do dataset
- Com uma quantidade limitada de tokens (cerca de 10~20M), ao usar primeiro dados simples em inglês da wiki, como Simple English Wikipedia, houve consistência gramatical, mas faltou consistência semântica
- Como o conteúdo é centrado em nomes próprios e em enumerações de fatos que parecem forçadas, há limite para gerar conteúdo realmente significativo
- Ao usar o dataset TinyStories, os resultados foram muito mais consistentes e significativos, porque a estrutura narrativa é clara e a linguagem é simples
- Como são histórias em nível de uma criança de 4 anos, até modelos pequenos conseguem aprender bem
Tokenizer e tokenização
- O treinamento do tokenizer não está incluído nos 5 minutos e, como a escala de dados é pequena, a necessidade de otimização é baixa
- Aprender tokens multibyte é mais fácil para o treinamento do modelo
Experimentos com arquitetura de modelo
-
Foi usada a arquitetura transformer (estilo GPT-2)
- Ajuste de hiperparâmetros como 2~3 camadas, funções de ativação como SwiGLU e positional embedding
- O LSTM teve desempenho próximo, mas o transformer foi superior em termos de perplexidade
- Dropout e mixture-of-experts são ineficientes nesse porte pequeno
- Curriculum learning teve pouco efeito porque o tempo de treinamento era curto demais
-
Tentativa com modelo diffusion (D3PM)
- Como linguagem natural é composta por tokens discretos, o processo de difusão acabou gerando apenas tokens aleatórios sem sentido e fracassou
- É mais difícil formar rapidamente estruturas de frase do que com transformer ou LSTM
Relação entre tamanho do modelo e throughput em tokens/segundo
- Modelos com 1M~2M de parâmetros são o sweet spot mais ideal
- Se forem grandes demais, não convergem em 5 minutos; se forem pequenos demais, atingem o limite de desempenho logo no começo do treinamento
- A Chinchilla scaling law foi, em geral, consistente com os resultados do experimento
- O tamanho ideal do modelo é o total de tokens de treinamento dividido por 20, e isso também foi confirmado neste experimento
Conclusão e implicações
- Mesmo com tempo muito curto e hardware pequeno, é possível treinar um modelo de storytelling consistente
- Treinar por 5 minutos não é adequado para desenvolver modelos poderosos, mas tem valor para experimentos de design de modelos pequenos e ultraleves, além de otimização de hardware e arquitetura
- Com a evolução futura das GPUs para notebook e das estruturas de modelo, existe potencial para avanços no desempenho de modelos treináveis em apenas alguns minutos
1 comentários
Comentários do Hacker News
Gostaria de ver a palestra dele sobre criptografia quântica; perguntou se alguém sabe qual é o link da palestra mencionada no começo do vídeo.
O treinamento otimizado de modelos pequenos é importante não só pela acessibilidade, mas também para a pesquisa científica sobre LLMs; assim como a biologia usa organismos simples como leveduras, é preciso estudar o transformador mais simples que ainda exiba os comportamentos interessantes dos modelos grandes para fins de compreensão e controle.
Um dos podcasts mais interessantes que ouvi recentemente foi sobre o artigo e o dataset TinyStories; esse dataset contém apenas palavras e conceitos simples, como os de histórias infantis, e mesmo assim permite que modelos pequenos gerem inglês com gramática, variedade e raciocínio; o podcast com os autores também explica muito bem as capacidades dos LLMs usando um exemplo de pesquisa pequeno e controlado; mantendo a analogia biológica, o dataset talvez corresponda a uma placa de ágar, isto é, um ambiente muito simples e controlado; os links relacionados são o episódio do podcast e o artigo TinyStories.
Muitas empresas podem resolver problemas reais de negócio com modelos pequenos usando datasets privados, como o histórico de compras dos usuários; avanços vindos dos grandes modelos de linguagem também podem ser aplicados diretamente a problemas menores se a sequência de entrada puder ser representada em uma linguagem especializada.
É amplamente conhecido que comportamentos e otimizações observados em modelos pequenos muitas vezes não se reproduzem bem em modelos grandes.
O que o autor está fazendo aqui é pré-treinamento, algo que normalmente é feito por criadores de modelos como Google ou Meta; para negócios, fine-tuning ou, em menor medida, pré-treinamento adicional é bem mais prático; enfatiza que o autor está tentando isso por razões acadêmicas.
Tenho interesse em modelos que rodem rápido em notebook, mas só o processo de treinamento pode levar dias ou até mais.
Acho que seria melhor usar energia em vez de tempo como critério; ou seja, treinar o melhor modelo possível dentro de um orçamento energético dado em joules seria a comparação de verdade, e isso tornaria mais justa a comparação entre um MBP e um H100.
O ponto central aqui não é eficiência, e sim “acessibilidade”, porque um H100 não é um produto do dia a dia, enquanto um notebook é.
Pelo que entendo, o Mac é mais competitivo em consumo de energia; ele não puxa tanta potência quanto uma GPU da Nvidia; dito isso, dá para alugar um H100 por menos de 10 dólares por hora, então também seria interessante comparar o desempenho de modelos que possam ser treinados em menos de 1 hora.
Acho que qualquer critério serve; mesmo que seja um pouco arbitrário, não vejo problema.
Se a ideia é expressar um indicador prático baseado em notebook, ou mais especificamente em MacBook Pro, seria bom deixar esse objetivo mais claro.
Já dá uma vibe de Olimpíada de eficiência em IA; notebook, desktop, celular, 5 minutos, 1 hora, 1 dia, 1 semana, no barco ou com uma cabra, tanto faz, dito em tom bem-humorado.
Com uma cabra? Acho que queria dizer Llama; a rima é limitada, mas com sotaque de Boston talvez fique mais engraçado.
Há um romance do Vernor Vinge em que humanos constroem computadores portáteis de xadrez e os usam como auxiliares durante as partidas; seria curioso se, em um torneio, além do relógio de xadrez, também fornecessem energia, para que os participantes pudessem pensar em jogadas favoráveis para sua própria IA.
Frase em tom exagerado sobre extrair desempenho extremo de um Mac Studio M3 Ultra 512GB, com a piada de que, com esse barco, até a cabra dá para levar.
Piada de que a cabra tem parâmetros demais, quase nível GPT-4.
Pagaria por um GoatLM.
No exemplo nonsense "Paris, France is a city in North Carolina...", gostou da expressão "officially major people" e comentou que tem curiosidade sobre como usá-la em conversas do dia a dia.
Comentário de que isso lembra o artigo sobre “cramming” de alguns anos atrás, compartilhando o link do artigo, que trata da tentativa do autor de treinar o melhor modelo possível em um notebook moderno ao longo de um dia.
Acha que só de aplicar alguns truques trazidos das tentativas de speedrun do GPT-2 — Muon, melhor inicialização de pesos e ajuste cuidadoso da taxa de aprendizado — já daria para melhorar bastante; o material relacionado está aqui.
Acha que falta uma demoscene na IA, no sentido da cultura de mostrar proeza técnica com poucos recursos, como nas demos gráficas dos anos 90.
Considera valioso querer criar modelos menores e mais especializados, e até produzi-los sob demanda quando necessário; não é preciso um modelo grande que saiba tudo, e sim um modelo muito rápido e focado a laser apenas no domínio em que trabalho; gostaria de poder pedir a um LLM grande “escreva um script para treinar um modelo otimizado para <tarefa necessária>” e então rodar esse modelo; mas, enquanto escrevia o comentário, o Google lançou o Gemma 3 270M.
Usando o exemplo de nonsense como "Paris, France is a city in North Carolina...", comenta que, se houvesse apenas a capacidade de dizer “não sei”, os modelos pequenos já seriam muito mais úteis; o motivo de precisar de LLMs enormes é a amplitude do escopo, para evitar que inventem respostas; seria ótimo se desse para treinar um chatbot de atendimento ao cliente em um notebook em um tempo razoável, mas fora daquele domínio a chance de ele responder algo seriamente sem sentido seria alta.