$cat ~/posts/bend-ml

bilíngue · en / pt

← /blog

um pytorch onde erro de shape não compila


No dia 29 de setembro o Taelin, criador do Bend, publicou um longo despejo sobre o futuro da linguagem. A empresa dele estava levantando uma rodada de investimento, ia montar dois times novos, e no meio do texto tinha uma lista de desejos para o ecossistema. Um item da lista estava escrito assim: "AI framework / PyTorch / llama.bend / etc.".

Alguém perguntou, nas respostas, o que uma pessoa deveria fazer para entrar num desses times. A resposta foi curta: construa alguma coisa em Bend, que mostre tração para os investidores e impressione o time.

Eu já tinha escrito sobre a ideia da linguagem e feito as extensões de editor. Faltava usar de verdade, numa coisa grande. Então peguei o item da lista e passei os dias 3 e 4 de outubro nele. O resultado se chama bend-ml: uma biblioteca de machine learning em Bend, cinco pacotes publicados, um classificador de dígitos escritos à mão e o GPT-2 rodando, tudo com o resultado conferido número a número contra o PyTorch.

Este post é a história inteira. A primeira metade não assume que você sabe nada de machine learning, de tipos ou de provas: eu explico cada peça antes de usar. A segunda metade desce até o código, os experimentos e os números, inclusive os números em que o projeto perde feio.

o projeto em 60 segundos. o vídeo é gerado, não gravado: cada saída de terminal e cada número vêm das execuções reais.
  1. 2026-09-19o manifestoescrevo sobre a aposta do bend2 sem ter rodado uma linha.
  2. 2026-09-29a lista de desejoso taelin publica o despejo: "AI framework / PyTorch / llama.bend".
  3. 2026-10-03v0.1 a v1.0lemas, tokenizer, tensores, autograd, MNIST e GPT-2. tudo certo, tudo lento.
  4. 2026-10-04v2 e v2.1os experimentos de performance, a GPU, o CI e as provas que faltavam.
cada barra é o intervalo até o marco seguinte. o projeto inteiro coube nas duas últimas linhas.

Eu não tenho formação em teoria de tipos nem em provas formais. Meu dia a dia é AI engineering, com agentes e NLP, mas por cima dos modelos: a parte de dentro deles, o treino, os gradientes e os tensores, eu venho estudando por fora. O projeto também foi meu curso intensivo nisso tudo. Quando eu digo "provado" neste post, quem atesta não sou eu: é o checador da linguagem, e qualquer pessoa pode rodar ele de novo.

o que é treinar um modelo

Antes de qualquer código, vale entender o que um framework como o PyTorch faz, porque é isso que eu quis reconstruir.

Um modelo de machine learning é uma função com muitos botões. Você dá uma entrada (a foto de um dígito, um pedaço de texto) e ela devolve uma saída (qual dígito é, qual a próxima palavra). Os botões são números, chamados de pesos, e cada combinação de pesos dá uma função diferente. Treinar é girar esses botões, pouco a pouco, até a função acertar.

O jeito de girar é sempre o mesmo. Você mostra um exemplo, mede o quanto a função errou, descobre para que lado cada botão deveria girar para errar menos, e gira um tiquinho. Depois repete com o próximo exemplo, milhões de vezes. Esse "para que lado girar" é o gradiente, e eu volto nele daqui a pouco.

Os pesos não ficam soltos: ficam organizados em tabelas. Uma tabela de números com linhas e colunas é uma matriz; quando tem mais dimensões, o nome genérico é tensor. E a operação que um modelo faz o tempo inteiro, bilhões de vezes, é multiplicar matrizes.

multiplicar matrizes, e o erro de shape

Multiplicar uma matriz A por uma matriz B é assim: para cada linha de A e cada coluna de B, você multiplica os números um a um e soma tudo. O resultado é um número da matriz C.

cada número de C sai de uma linha de A e uma coluna de B. por isso a linha de A e a coluna de B precisam ter o mesmo comprimento.

Repare na regra escondida aí: para combinar uma linha de A com uma coluna de B número a número, as duas precisam ter o mesmo tamanho. Ou seja, o número de colunas de A tem que ser igual ao número de linhas de B. Uma matriz de 2×3 multiplica uma de 3×4, mas não uma de 4×5. Esse par de números, linhas por colunas, é o shape da matriz.

Num modelo de verdade, os shapes passam por dezenas de camadas, cada uma transformando o shape da anterior. E errar um shape é o bug mais comum de quem escreve modelo. No PyTorch, o erro aparece assim, em tempo de execução, depois que o programa já está rodando:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x3 and 4x5)

Às vezes ele aparece em segundos. Às vezes aparece depois de horas de treino, numa camada que só roda no fim de uma época. A pergunta que guiou o projeto foi: e se esse erro nunca pudesse chegar a rodar?

o que é um tipo

Em quase toda linguagem existe a ideia de tipo: a etiqueta que diz que tipo de coisa um valor é. 3 é um inteiro, "olá" é um texto. Em Python os tipos existem, mas ninguém confere antes de rodar: se você soma um número com um texto, o erro só aparece quando aquela linha executa. Em linguagens com checagem estática, um programa chamado checador de tipos lê o código antes de rodar e recusa o que não faz sentido. Somar número com texto nem compila.

Só que o checador de tipos comum não sabe nada de shape. Para ele, uma matriz de 2×3 e uma de 4×5 são do mesmo tipo: "matriz". O erro de shape passa.

O Bend2 tem tipos dependentes. O nome assusta, mas a ideia é simples: um tipo pode depender de um valor. A etiqueta pode carregar números. Então em vez de "matriz", o tipo pode ser "matriz de 2 linhas e 3 colunas", escrito Mat<2, 3>. E aí a regra da multiplicação vira uma regra de tipos: Mat<n, k> vezes Mat<k, m> dá Mat<n, m>, e o k tem que ser o mesmo dos dois lados.

Teste você mesmo. Escolha os shapes de A e B:

matriz A
matriz B
Mat<2, 3>
·
Mat<4, 5>
=
?—
SOME PROOFS FAIL
Error:
- expected : Mat<3n, 5n>
- observed : Mat<4n, 5n>
a regra é a mesma do bend-ml: as colunas de A precisam ser iguais às linhas de B. quando não são, a mensagem imita a do compilador de verdade.

Um framework em Bend consegue oferecer isso e o PyTorch não, então essa virou a proposta do projeto: um PyTorch onde erro de shape não compila.

o que é uma prova

O Bend2 vai além dos tipos dependentes. Ele tem uma terceira coisa, além de função (def) e dado (type): a lei (law). Uma lei é uma afirmação sobre o programa, do tipo "para todo número a e todo número b, a + b é igual a b + a". E a linguagem exige uma prova de cada lei.

Para quem nunca viu prova formal, a melhor imagem é a de dominós enfileirados. Se você mostra que o primeiro dominó cai, e mostra que qualquer dominó que cai derruba o seguinte, então todos caem, não importa quantos sejam. Isso se chama indução, e é o jeito de provar uma coisa sobre infinitos números com uma quantidade finita de trabalho: você prova para o zero, e prova que se vale para um número, vale para o próximo.

No Bend, uma prova é uma função comum. O tipo da função é a afirmação, e o corpo é o argumento. O caso do zero é o primeiro dominó; a chamada recursiva da função é o "se vale para o anterior". Se a função passa no checador de tipos, a afirmação está provada.

Quem confere isso é o kernel, o pedaço do checador que decide se uma prova está certa. O Bend tem dois: o normal, rápido, e um segundo, chamado com --verdict, que foi escrito e provado correto em Lean, outra linguagem de provas, essa já usada por matemáticos profissionais. O compilador do Bend foi escrito quase todo por IA e não foi auditado inteiro; o kernel do --verdict é a parte que humanos auditaram. Por isso, no projeto, toda lei tinha que passar nos dois.

E aqui entra a tese do Bend: "uma linguagem evoluída por IA, garantida por matemática". Quem escreve o código pode errar; a lei, verificada pelo kernel, garante que certas propriedades valem sempre. ML é exatamente o tipo de código em que ninguém confia muito no que está escrito, e ninguém prova nada.

o reconhecimento

Com a ideia na mão, a primeira tarefa não foi código. O Bend2 é novo e muda rápido, e a sintaxe dele não tem nada a ver com a do Bend1, que era estilo Python. Então eu li a documentação, os exemplos e a biblioteca padrão, e anotei tudo que ia precisar. Algumas descobertas mudaram o plano:

  • fixar a versão: tudo foi feito na 2.0.35, e o projeto se recusa a rodar em outra. Numa linguagem que muda toda semana, isso é sobrevivência.
  • números: só existem Nat (naturais: 0, 1, 2...), U32 (inteiros de 32 bits) e F32 (números com vírgula de 32 bits). Não existe F64. Para ML, F32 serve, mas ele volta a aparecer na hora das provas.
  • a biblioteca padrão quase não tem lemas: o Base tem números, listas e textos, mas nenhuma prova de que a soma é comutativa, por exemplo. Lema é uma lei pequena, que serve de degrau para provar outras.
  • o BendHub, o repositório de pacotes: login só pelo GitHub, versões de quatro números (0.1.2.0), e publicação pública e permanente. Versão publicada não se apaga.
  • para onde o Bend compila: C, JavaScript e CUDA, que é o que roda em placa de vídeo NVIDIA.

A primeira coisa que eu escrevi foi uma prova de conceito pequena: uma matriz com o shape no tipo, um matmul, e dois programas que não deviam compilar. Quando eles não compilaram, pelo motivo certo, o resto do plano fez sentido.

o shape mora no tipo

Daqui em diante o post fica mais técnico. O tipo central do primeiro pacote, o bend-ml-tensor, é esse:

type Mat<-r: Nat, -c: Nat> is Data:
  Mat{rows: List<&2, List<&2, F32>>}

Uma matriz de r linhas e c colunas guarda uma lista de linhas, cada linha uma lista de F32. O - antes de r e c marca as dimensões como apagadas: elas existem só para o checador e somem do programa compilado. Quer dizer que a garantia não custa nada em tempo de execução. Com isso a multiplicação fica com a assinatura que a matemática pede:

def Mat.matmul(-n: Nat, +k: Nat, +m: Nat, a: Mat<n, k>, b: Mat<k, m>) -> Mat<n, m>:

O k aparece nos dois argumentos. Se a coluna de um não bate com a linha do outro, não existe k que satisfaça, e o programa não compila:

bend recusando o produto de uma matriz 2x3 por uma 4x5: expected Mat<3n, 5n>, observed Mat<4n, 5n>
a saída real do compilador. o programa não roda: ele nem chega a existir.

A mensagem diz exatamente o que aconteceu: o segundo argumento devia ser Mat<3, 5>, para casar com as 3 colunas do primeiro, e chegou um Mat<4, 5>.

O reshape é mais interessante. Ele reorganiza os mesmos números num shape diferente, por exemplo 12 números de 2×6 para 3×4. Isso só faz sentido se o total de números não muda. No bend-ml, o reshape exige uma prova disso como argumento:

def Mat.reshape(-r1: Nat, -c1: Nat, r2: Nat, +c2: Nat,
                -p: {Nat.mul(r1, c1) == Nat.mul(r2, c2) : Nat},
                a: Mat<r1, c1>) -> Mat<r2, c2>:

O parâmetro p tem um tipo estranho: {r1*c1 == r2*c2}. Um valor desse tipo é uma prova de que os dois produtos são iguais. Para chamar a função, você precisa entregar essa prova. De 2×6 para 3×4, a prova é só {==}, que significa "computa os dois lados e vê que dá igual": 12 e 12. De 2×6 para 5×3, não existe prova de que 12 é igual a 15, e o compilador recusa. E como o parâmetro também é apagado (-p), a prova não custa nada em tempo de execução.

Quando as dimensões não são números fixos, a prova vira uma lei de verdade. Transpor r×c em c×r exige saber que r*c == c*r para quaisquer r e c, e isso é a comutatividade da multiplicação, que precisou ser provada antes.

O mesmo vale dentro do treino. No backprop, que eu explico mais adiante, o gradiente de uma matriz de pesos W de 784×128 é Xᵀ·dY, e o tipo dele é Mat<784, 128>. Pedir esse gradiente como 128×784, o erro clássico de quem transpõe do lado errado, também não compila. O passo de treino inteiro do MNIST é checado assim, camada por camada.

as provas que faltavam

Como o Base não tem lemas, o primeiro pacote publicado nem foi de ML: foi o bend-ml-nat-lemmas, com 15 leis sobre Nat e listas. Soma comutativa e associativa, multiplicação comutativa, associativa e distributiva, elemento neutro, concatenação de listas e o tamanho dela.

A comutatividade da soma fica assim:

law add_comm:
  for  a: Nat
  for +b: Nat
  {Nat.add(a, b) == Nat.add(b, a) : Nat}

def add_comm(a, b):
  match a:
    case 0n:
      %add_zero(b) : {_ == Nat.add(b, 0n) : Nat}
      {==}
    case 1n+p:
      %add_succ(b, p) : {1n+Nat.add(p, b) == _ : Nat}
      %add_comm(p, b) : {1n+Nat.add(p, b) == 1n+_ : Nat}
      {==}

Vale ler linha por linha, porque é a mesma estrutura de toda prova do projeto.

A law é a afirmação: para todo a e todo b, a + b == b + a. O + antes de b eu explico mais adiante; por enquanto, ele diz que b pode ser usado mais de uma vez.

O def com o mesmo nome é a prova. O match a separa os dois casos de um número natural: ou ele é zero (0n), ou ele é "um mais alguma coisa" (1n+p), onde p é o número anterior. São os dois dominós.

No caso zero, o objetivo é 0 + b == b + 0. O lado esquerdo o Bend já sabe simplificar para b, mas o direito não. O %add_zero(b) usa outro lema, b + 0 == b, para reescrever o objetivo, e aí os dois lados viram b e o {==} fecha.

No caso 1+p, o objetivo é (1+p) + b == b + (1+p). O %add_succ reescreve o lado direito para 1 + (b + p). O %add_comm(p, b) é a chamada recursiva: a própria lei, para o número menor p. É o "se vale para o anterior". Ela troca p + b por b + p, os dois lados ficam iguais, e o {==} fecha.

Não tem linguagem de tática, como no Lean ou no Coq: você prova escrevendo código. Para quem vem de Python como eu, ler isso pela primeira vez parece magia; depois da décima prova, parece só recursão com um tipo muito exigente. As provas de soma e multiplicação seguem a demo oficial de provas numéricas do Bend, adaptadas para o Nat.add e o Nat.mul do Base.

O lema que mais importou foi product_append, que diz que o produto de uma lista concatenada é o produto de cada parte. Uma lista de dimensões é o shape de um tensor, e o produto dela é o número de elementos: é esse lema que sustenta o reshape com shapes arbitrários.

Esse também foi o primeiro pacote no BendHub. Publicar é um comando, mas como não tem volta, virou regra: só publico depois de ALL PROOFS CHECK, de --verdict passar, de não ter nenhum @unsafe (o "confia em mim" da linguagem, que desliga a checagem) nem ?TODO (um buraco deixado para depois) no arquivo, e de importar o pacote do próprio BendHub numa pasta limpa e usar uma lei dele.

texto vira número: o tokenizer

Um modelo de linguagem não lê letras: lê números. Então todo texto que entra no GPT-2 passa antes por um tokenizer, que corta o texto em pedaços e troca cada pedaço por um número. Esses pedaços são os tokens.

O jeito mais simples seria um token por letra. Mas aí uma frase vira uma sequência enorme. O outro extremo, um token por palavra, precisaria de um dicionário com todas as palavras de todas as línguas. O GPT-2 usa um meio-termo chamado BPE (byte pair encoding). Começa com os 256 bytes possíveis, que conseguem representar qualquer texto, e vai aprendendo junções: o par de tokens que mais aparece junto vira um token novo. Depois o próximo par mais frequente, e assim por diante. Palavras comuns acabam virando um token só; palavras raras ficam em pedaços.

Digite qualquer coisa e mexa no controle para ver o treino acontecer:

tokens (5)
aaabdaaabac
regras (da mais antiga)
  1. a + a → aa
  2. aa + a → aaa
  3. aaa + b → aaab
decodificado
aaabdaaabac

✓ decode(encode(texto)) == texto

o mesmo algoritmo do pacote: o par mais frequente vira um token novo, o empate vai para o par que apareceu primeiro, e decodificar expande cada token de volta nos bytes. a última linha confere que o texto voltou inteiro.

Repare na última linha: decodificar o que foi codificado devolve o texto original. Parece óbvio, mas é exatamente o tipo de coisa que quebra em produção, num caractere estranho, numa ordem de regras diferente, num emoji. E é a lei principal do segundo pacote, o bend-ml-bpe-tokenizer.

No pacote, um token é um byte cru, B{n}, ou um token criado pela regra de id k, M{k}. A tabela é uma lista de regras da mais antiga para a mais nova, igual ao merges.txt do GPT-2. E a lei é essa:

a lei roundtrip como o bend a enuncia, seguida de ALL PROOFS CHECK
o enunciado que o kernel verificou, exatamente como está no pacote publicado.

decode(encode(s)) == s, para qualquer sequência de bytes, desde que a tabela seja bem formada, ou seja, que nenhum id de regra se repita. A ideia da prova é que, quando a regra k junta o par (a, b) em M{k}, o M{k} expande de volta para a expansão de a seguida da de b, então decodificar não muda nada. Isso só vale se k ainda não existia na tabela, que é justamente o que "bem formada" garante. Os lemas auxiliares provam que colocar a regra nova por cima da tabela não muda a expansão dos tokens que já existiam.

Vieram junto vocab_bound (todo token que o encode emite é um byte ou foi criado por uma regra da tabela, nunca um id desconhecido) e dec_append (decodificar duas listas juntas é decodificar cada uma e concatenar os bytes).

Depois veio a parte que eu mais gosto: provar que o train sempre produz uma tabela bem formada (train_wf). A prova é por indução sobre o laço de treino, com o invariante "todo id já usado é menor que o próximo". Invariante é uma coisa que vale antes e depois de cada passo do laço; achar o invariante certo é metade de qualquer prova sobre laços. Juntando as duas leis, sai roundtrip_trained: treine uma tabela em qualquer corpus, com quantas regras quiser, e o roundtrip vale, sem hipótese nenhuma.

O que não é provado é qual par o train escolhe juntar. E não precisa ser: a correção do roundtrip não depende disso. Essa parte é testada contra uma implementação de referência em Python, com texto em inglês, texto com acento (coração, ação), emoji e caracteres chineses. train, encode e decode dão exatamente o mesmo resultado nos dois lados.

como uma rede aprende: o gradiente

Voltando aos botões. Para saber para que lado girar cada peso, o treino calcula o gradiente: para cada peso, quanto o erro muda se aquele peso aumentar um pouquinho. Se aumentar o peso aumenta o erro, você gira para baixo; se diminui, gira para cima. É a inclinação do terreno, e o treino é descer a ladeira do erro.

Calcular o gradiente de uma função com milhões de pesos parece impossível, mas tem um truque. Toda função de um modelo é feita de operações pequenas (somas, multiplicações), encadeadas. E existe uma regra do cálculo, a regra da cadeia, que diz como combinar as inclinações das partes para ter a inclinação do todo. Um framework faz isso sozinho, e o nome disso é diferenciação automática, ou autograd.

Tem dois jeitos de aplicar a regra da cadeia. O modo direto acompanha a inclinação junto com o valor, da entrada para a saída. O modo reverso calcula todos os valores primeiro e depois volta da saída para a entrada, distribuindo o gradiente para trás. O reverso é o que todo framework de ML usa, porque com uma volta só ele calcula o gradiente de todos os pesos de uma vez. É isso que se chama backpropagation.

na volta, cada multiplicação passa adiante o gradiente vezes o valor do outro fator, e cada soma passa o gradiente igual. o que chega em x por caminhos diferentes se soma.

O terceiro pacote, o bend-ml-autograd, implementa os dois modos e prova que eles concordam. A expressão é uma árvore de constantes, da variável, de somas e de produtos:

type NE is Data:
  NCst{n: Nat}
  NX{}
  NAdd{a: NE, b: NE}
  NMul{a: NE, b: NE}

O modo reverso recebe o gradiente g que chega de cima e o distribui para os filhos. Numa multiplicação, cada lado recebe g vezes o valor do outro lado, exatamente como no diagrama:

def nbwd(e: NE, +x: Nat, +g: Nat) -> Nat:
  match e:
    case NCst{n}:
      0n
    case NX{}:
      g
    case NAdd{a, b}:
      Nat.add(nbwd(a, x, g), nbwd(b, x, g))
    case NMul{+a, +b}:
      Nat.add(nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))))

E a lei:

law reverse_eq_forward:
  for +e: NE
  for +x: Nat
  {nbwd(e, x, 1n) == nfwd(e, x) : Nat}

A prova passa por um lema mais forte, nbwd(e, x, g) == g * nfwd(e, x) para qualquer g, porque a indução só fecha se o gradiente que chega de cima for genérico. Com g = 1, sai a lei. É um padrão comum: às vezes a afirmação que você quer não é provável direto, mas uma mais geral é.

Aqui aparece a regra que eu escrevi no começo do projeto e mantive até o fim: float não é número real. F32 arredonda: 0.1 + 0.2 não é exatamente 0.3, e a soma nem sempre é associativa. Provar propriedade numérica sobre F32 é pedir para provar coisa falsa. Então a lei é provada sobre Nat, onde a aritmética é exata, e os números de verdade são testados contra o PyTorch, com gradient checking: você calcula o gradiente pelo autograd e confere com uma aproximação numérica, mexendo de leve na entrada e medindo a mudança na saída. Prova para a estrutura, teste para os números. Essa divisão atravessa o projeto inteiro.

MNIST: ensinando a ler dígitos

Com os pacotes no lugar, vieram as duas demos. A primeira é o "hello world" de machine learning: o MNIST, um conjunto de 70 mil fotos de dígitos escritos à mão, de 0 a 9, cada uma com 28×28 pixels em tons de cinza. A tarefa é olhar a foto e dizer qual dígito é.

O modelo é uma MLP (multilayer perceptron), a rede neural mais simples que existe. Os 784 pixels (28×28) entram, passam por uma camada de 128 neurônios, e saem 10 números, um para cada dígito; o maior é o palpite. Cada camada é um produto de matrizes mais um vetor de ajuste (o bias), seguido de uma função que dobra o espaço, a ReLU (que zera os negativos).

O treino olha 100 fotos de cada vez, um batch, e os shapes de um passo ficam assim:

X      : Mat<100, 784>     100 fotos, 784 pixels cada
W1     : Mat<784, 128>     pesos da primeira camada
H = X·W1 + b1 : Mat<100, 128>
W2     : Mat<128, 10>      pesos da segunda camada
Z = H·W2 + b2 : Mat<100, 10>   10 notas por foto

dW2 = Hᵀ·dZ : Mat<128, 10>     o gradiente tem o shape do peso
dW1 = Xᵀ·dH : Mat<784, 128>

No fim de cada passo, cada peso anda um pouquinho contra o gradiente (taxa de aprendizado 0,1). Uma época é passar pelas 60 mil fotos de treino uma vez: 600 batches.

Para conferir o resultado, eu fiz os pesos iniciais no script de referência em PyTorch e exportei para o Bend, e os dois lados leem os mesmos batches na mesma ordem. Assim dá para comparar número com número, não só "acurácia parecida". E bateu: a perda da primeira época é 0.5204771 nos dois, com 9129 acertos em 10000 fotos de teste. Depois 0.27043572 e 9298, depois 0.2156194 e 9418. Igual até a sétima casa.

GPT-2: prevendo a próxima palavra

A segunda demo é o GPT-2 small, o modelo de linguagem que a OpenAI publicou em 2019, com 124 milhões de pesos. Ele faz uma coisa só: dado um texto, prevê o próximo token. Para gerar uma frase, você pede o próximo token, coloca no fim do texto e pede de novo.

Por dentro, ele é uma pilha de 12 camadas iguais, cada uma com duas partes. A atenção deixa cada posição do texto olhar para as posições anteriores e decidir quais importam (em "a capital da França é", a palavra "França" importa muito para o que vem depois). A MLP transforma o resultado, como no MNIST. Entre as partes vêm normalizações (LayerNorm) e uma função suave parecida com a ReLU (GELU).

Fazer isso rodar em Bend exigiu mais peças do que eu imaginava:

  • os pesos vêm do arquivo oficial, o model.safetensors, que um script em Python separa em um arquivo binário por tensor. O Bend lê cada arquivo em blocos com File.read_at e converte cada 4 bytes em um F32.
  • o tokenizer usa o meu pacote de BPE com as 50 mil regras do GPT-2. Só que o GPT-2 não aplica as regras em ordem: ele procura sempre o par de menor posição na tabela. Eu apliquei em ordem e conferi contra o tiktoken, o tokenizer oficial, que para tabelas treinadas as duas coisas dão o mesmo resultado.
  • o pré-tokenizador: antes do BPE, o GPT-2 quebra o texto com uma expressão regular que separa contrações, letras, números, pontuação e espaços. O Bend não tem expressão regular, então eu escrevi o pré-tokenizador à mão, byte a byte, imitando cada alternativa da regex.
  • a permutação dos bytes: o GPT-2 renumera os 256 bytes com uma tabela própria, que vira mais um arquivo.
  • o cache: para não recalcular o texto inteiro a cada token novo, cada camada guarda as chaves e valores da atenção das posições anteriores, o KV cache.
  • a geração: gulosa, sempre o token de nota mais alta, que é a mais fácil de comparar.

O resultado: os mesmos tokens que o PyTorch, com as notas de cada token (os logits) iguais até a quarta casa.

$ ./gpt2_fast "The capital of France is" 8
text: The capital of France is the capital of the French Republic, and

(O GPT-2 small não sabe geografia muito bem. Mas ele erra igualzinho nos dois lados, que é o que importa aqui.)

a v1 funcionava. e era 1800 vezes mais lenta

No fim do primeiro dia, tudo estava certo e tudo estava lento.

Uma época de MNIST levava 544 segundos. O PyTorch leva 0,3. Umas 1800 vezes mais lento. O GPT-2 levava 3 segundos por token, e uns 10 segundos só para carregar os pesos.

A tentação era concluir que o Bend é lento e parar por ali. Mas eu não tinha medido nada, só cronometrado o total. Então a v2 começou com um plano diferente: não otimizar nada sem antes ter um número que dissesse onde estava o problema. O segundo dia virou uma série de experimentos, cada um com uma hipótese e uma medição antes de qualquer mudança no código. Todos estão nas notas do projeto, com o comando para reproduzir.

por que lista é lenta

O primeiro suspeito era a estrutura de dados. A v1 guardava a matriz como lista de linhas, e cada linha era uma lista encadeada: cada número fica num lugar qualquer da memória, junto com o endereço do próximo. Para chegar no décimo número, você passa pelos nove anteriores.

A alternativa é um array: os números lado a lado, e você chega em qualquer um direto pela posição, o índice.

na lista, ler um número é seguir setas. no array, é fazer uma conta com o índice.

Para o processador, a diferença é enorme. Ele busca a memória em blocos e tenta adivinhar o que você vai ler depois. Com um array, a próxima leitura está do lado da anterior e já veio no mesmo bloco. Com uma lista, cada número pode estar em qualquer lugar, e cada leitura espera a anterior terminar para saber o endereço da próxima.

Experimento 1. Reescrevi o mesmo produto de matrizes (100 milhões de multiplica-soma, uma thread) com um Array<F32> plano, acessado por índice:

tempomultiplica-somas/s
listas2,274 s44 M
Array<F32> por índice0,046 s2.200 M

Quarenta e nove vezes, só trocando a estrutura de dados. A conclusão da v1 ("o Bend é 1800 vezes mais lento") era efeito da estrutura de dados, não da linguagem.

O Array do Bend tem um detalhe: por baixo, ele não é um bloco contínuo de memória, é uma árvore, e o índice escolhe o caminho dos galhos. Isso apareceu em mais dois números. Percorrer a árvore inteira em vez de indexar foi 70 vezes pior. E um array de 2^20 posições em vez de 2^17 deixou o mesmo produto 60% mais lento, porque a profundidade da árvore entra em cada leitura. Regra que saiu disso: alocar o menor array possível.

tipos afins, na prática

O array trouxe um conceito do Bend que eu só tinha lido sobre: valores afins. Um valor afim pode ser usado no máximo uma vez. É isso que deixa o Bend gerenciar memória sem coletor de lixo: se o compilador sabe que ninguém mais aponta para um valor, pode liberar na hora ou reaproveitar o espaço.

Para um array, isso tem uma consequência curiosa: ler um elemento "usa" o array. Então Array.get devolve o elemento e o array de volta, para você continuar usando. Isso vazou para a API do pacote: toda operação devolve o que leu, junto com o resultado, num registro com tipo próprio. O produto de matrizes devolve isso:

type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
  MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}

Os dois fatores de volta, e o produto. Estranhei na primeira hora; depois ficou natural, porque deixa explícito quem é dono de cada matriz em cada momento. E os marcadores dos parâmetros completam a história: + permite reusar um valor Data (números, listas comuns), - apaga o parâmetro do programa compilado, e um parâmetro sem marcador pode ser usado uma vez só.

paralelismo: quando dividir ajuda e quando atrapalha

O processador do meu notebook tem 22 threads, ou seja, consegue fazer 22 coisas ao mesmo tempo. O Bend é feito para isso: ele paraleliza automaticamente partes do programa que não dependem umas das outras. Um produto de matrizes é perfeito para isso, porque cada linha do resultado é independente das outras.

Experimento 2. Dividi o produto em blocos de linhas e mandei cada bloco para uma tarefa. Como os valores são afins, cada tarefa precisa da sua própria cópia das matrizes (Array.clone). Minha hipótese era que as cópias compartilhariam memória por baixo e as tarefas iam disputar o acesso. Medi: 16 tarefas lendo cópias escalam igual a 16 tarefas com arrays próprios. Hipótese refutada, e eu fiquei feliz de ter medido em vez de otimizar para um problema que não existia. O ganho no MNIST foi de uns 2 vezes.

Experimento 3. No GPT-2, gerando um token por vez, as contas são matriz vezes vetor: uma linha só de entrada. Paralelizei do mesmo jeito e ficou 10 vezes mais lento, com 1, 8 ou 16 threads.

o que decide se a cópia compensa é quantas vezes cada peso é usado depois de copiado.

A causa é aritmética simples. Clonar custa uns 0,1 nanossegundo por número, por tarefa. Em matriz vezes vetor, cada peso é lido uma vez só, e a conta com ele custa uns 0,4 nanossegundo. Copiar a matriz para várias tarefas custa mais que fazer a conta inteira sozinho. No MNIST, matriz vezes matriz com 100 fotos, cada peso é lido 100 vezes, e a cópia some no ruído. Eu só descobri isso medindo.

o que mexeu o ponteiro: listas para array 49x mais rápido, blocos paralelos 1,8x, GPU em loop plano 3,8x mais rápida, GPU nos nossos kernels 3 a 18x mais lenta, copiar a matriz por tarefa 10x mais lento
cada linha é um experimento nas notas do projeto, com o número medido antes da mudança.

Com os números na mão, a v2 virou um pacote novo, o bend-ml-tensor-array: o mesmo Mat<r,c> com o shape no tipo, só que sobre um Array<F32> plano, linha por linha (o número da linha i, coluna j fica na posição i*c + j).

Um detalhe que eu gosto nele: o treino precisa de três produtos diferentes, X·W (ida), H·Wᵀ e Xᵀ·dZ (volta). Transpor uma matriz na memória custa caro. Em vez disso, o pacote tem um único produto genérico, o gemm, que recebe os passos com que os índices andam em cada matriz:

def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32,
         a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:

O elemento da linha i, posição p de A fica em i*arow + p*sa. Para A normal, arow = k e sa = 1. Para ler A transposta sem transpor nada, basta trocar: arow = 1 e sa = n. Os três produtos públicos (matmul, matmul_nt, matmul_tn) são o mesmo gemm com passos diferentes, e os três com o shape certo no tipo. Todos têm blocos paralelos opcionais, e o GPT-2 roda os dele em sequência, porque o experimento 3 mandou.

No fim, a época de MNIST foi de 544 segundos para uns 7 segundos, e o GPT-2 de 3 segundos para 0,1 segundo por token.

a GPU parecia inútil

Placa de vídeo é outro bicho. Uma CPU tem poucos núcleos muito espertos; uma GPU tem milhares de núcleos simples, que fazem a mesma conta ao mesmo tempo em dados diferentes. É perfeita para ML, e o Bend compila para CUDA, a plataforma das placas NVIDIA, marcando uma chamada com !.

Eu tenho uma RTX 4050. O primeiro problema foi instalar: o Bend pede CUDA 12, e o pacote do meu Arch é o 13, que precisaria de sudo e quebraria outras coisas. Investigando, descobri que o Bend respeita a variável CUDA_HOME e só precisa de dois pedaços do kit da NVIDIA (cudart e nvrtc). Baixei os arquivos avulsos, descompactei em ~/.local/cuda12, e funcionou sem sudo.

Aí veio o primeiro resultado: a GPU era mais lenta que a CPU em tudo. De 3 a 18 vezes, sem nenhum ponto em que a curva virasse. Minha primeira anotação foi "a GPU não serve para isso". Relendo, achei a conclusão rápida demais: se uma GPU perde para a CPU em tudo, o mais provável é que eu esteja usando errado.

E estava. O guia do Bend diz que uma chamada ! precisa se abrir em umas 16 mil folhas, uma por linha de execução da GPU, e terminar em laços simples, só com números. Meus testes usavam de 32 a 1024 tarefas: a placa ficava quase toda parada. Com 16.384 folhas e um laço numérico simples, a GPU ganhou, e a vantagem crescia com o trabalho de cada folha:

16.384 folhas, laço simples de F32GPUCPU (22 threads)
20 mil iterações cada0,217 s0,054 s
200 mil iterações cada0,254 s0,300 s
2 milhões de iterações cada0,479 s1,808 s

3,8 vezes mais rápida no último caso. O custo fixo de uns 0,2 segundo é compilar o código da GPU na hora.

Mas os produtos de matriz continuaram perdendo, e o motivo é bonito de entender. Lembra da lista encadeada e da árvore do array? Cada leitura depende de um endereço que vem da leitura anterior. A GPU é ótima para fazer a mesma conta em muitos dados que já estão lado a lado; é péssima para seguir endereços. E a memória que o Bend usa na GPU é "gerenciada": as páginas atravessam o barramento entre a CPU e a placa no primeiro acesso. Nossos produtos são correntes de leituras dependentes, o oposto do que uma GPU quer, e é esse formato dos dados que faz eles perderem.

Os benchmarks ficaram na CPU, com o motivo escrito. E a primeira conclusão continua nas notas, corrigida em vez de apagada.

deixando estável

Até ali, tudo funcionava no meu computador e em nenhum outro. Quando eu me perguntei "isso é uma v1 de verdade?", a resposta honesta foi não: um estranho que clonasse o repositório não conseguiria rodar os testes, nada rodava sozinho, e quem tinha escrito os testes era eu, o mesmo que escreveu o código. A última parte do trabalho foi fechar cada uma dessas lacunas.

Para começar, make setup && make check num clone limpo instala o Bend fixado na 2.0.35 (conferindo o SHA256 do binário, a impressão digital do arquivo), o Lean 4.34.0 para o kernel auditado, o ambiente Python das referências e os dados do MNIST e do GPT-2, tudo sem sudo. Testei num clone novo, com um HOME vazio, como se fosse outra pessoa.

Depois, o CI, a máquina do GitHub que roda os testes a cada mudança, roda 36 checagens a cada push, GPT-2 incluído, num runner gratuito, em uns 11 minutos. Rodar o GPT-2 inteiro numa máquina de graça era a parte que eu achava que não ia caber, e coube.

Testes escritos pela mesma pessoa que escreveu o código tendem a testar o que ela já sabe que funciona. Então entraram testes gerados aleatoriamente: árvores de expressão aleatórias para o autograd (comparadas com o PyTorch), 80 corpora aleatórios para o BPE (comparados com a referência em Python), shapes aleatórios para os tensores, e 11 prompts para o GPT-2. E um teste que importa cada pacote do BendHub, na versão que o README diz, e aplica uma lei dele. Esse pega publicação quebrada ou README desatualizado.

Para quem quiser auditar, um script imprime o enunciado de cada lei, sem as provas, para quem quiser revisar a especificação sem ler a implementação. Um documento lista o que é provado, o que é testado e o que é confiado.

Faltava também uma prova. O Mat sobre Array assumia que o array tinha espaço para r*c números. O Bend aloca arrays com 2^d posições, e o cálculo de d usava U32.log2, sobre o qual não dá para provar nada. Reescrevi o cálculo com recursão em Nat:

def cap_pair(n: Nat) -> Nat & Nat:
  match n:
    case 0n:
      (0n, 1n)
    case 1n++m:
      cap_step(cap_pair(m), 1n+m)

cap_pair(n) devolve o par (d, 2^d), e cada passo dobra a potência quando n não cabe mais. Com isso deu para provar a lei:

law cap_ok:
  for +n: Nat
  {True{} == leb(n, pow2(cap_depth(n))) : Bool}

n <= 2^cap_depth(n): o array sempre tem espaço. O que sobra de confiança é só o runtime alocar 2^d posições quando pedido.

Entradas ruins passaram a virar erro claro. A perda e a contagem de acertos ganharam variantes que devolvem None quando o número de rótulos não bate com o número de linhas, em vez de calcular lixo. E o tokenizer agora valida o arquivo de regras antes de usar: quantidade par de ids, e nenhuma regra citando um token que ainda não existe:

def ids_ok(ns: List<&2, Nat>, +k: Nat) -> Bool:
  match ns:
    case Nil{}:
      True{}
    case Con{+a, t}:
      match t:
        case Nil{}:
          False{}
        case Con{+b, u}:
          Bool.and(Bool.and(Nat.is_lt(a, Nat.add(256n, k)), Nat.is_lt(b, Nat.add(256n, k))), ids_ok(u, 1n+k))

A regra k só pode citar bytes (abaixo de 256) ou regras anteriores (abaixo de 256 + k).

Por fim, o Unicode. O pré-tokenizador do GPT-2 tratava todo byte acima de 127 como letra. Funcionava para acento, mas não para travessão, emoji ou pontuação chinesa, que a regex do GPT-2 trata como "outro". O conserto: antes de pré-tokenizar, cada byte ganha uma etiqueta com a classe do caractere a que ele pertence, codificada no próprio número:

def tagged(+b: U32, +cls: Nat, +cont: Bool) -> U32:
  U32.add(U32.add(b, U32.mul(U32.from_nat(cls), 256)), U32.mul(U32.from_nat(sel(cont, 1n, 0n)), 2048))

O byte fica nos 8 bits de baixo, a classe (letra, número, outro, espaço) nos bits seguintes, e um bit marca os bytes de continuação de um caractere de vários bytes. A classe vem de uma tabela de 939 faixas de Unicode, gerada a partir do módulo regex do Python, com as mesmas categorias que a regex do GPT-2 usa. O resto do pré-tokenizador trabalha com as classes e tira as etiquetas no fim. Resultado: bate com o tiktoken em 79 textos, incluindo 30 strings de Unicode aleatório.

Duas histórias desse trecho me ensinaram bastante.

A primeira: o README dizia que carregar o GPT-2 usava uns 10 GB de memória. Medindo direito, 10 GB era o espaço de endereçamento que o processo reservava, não a memória usada de fato; a memória residente era 2,3 GB. Depois de ler os pesos em blocos de 1 MB direto para o array, sem montar uma lista no meio, caiu para 1,5 GB, e o carregamento ficou 5 segundos mais rápido.

A segunda: depois de trocar o cálculo de capacidade pela versão provada, a época de MNIST foi de 7 para 16 segundos. Parecia que a prova tinha custado caro. Rodei a versão antiga e a nova lado a lado: as duas deram 16 segundos. O culpado era o protetor de tela do meu desktop, usando 134% de CPU. Com a máquina ociosa, voltou para 7. Se eu não tivesse comparado as duas versões nas mesmas condições, teria jogado fora uma prova boa por causa de um screensaver.

a cara do projeto

A última coisa foi o README e o vídeo. A primeira versão que eu fiz tinha cara de tema de terminal genérico: fundo escuro, azul, ciano e vermelho saturados, janelas de terminal estilo macOS. Ficou feio, e eu refiz do zero.

A versão final usa a identidade do próprio Bend, tirada do CSS do bend-lang.com: papel quente (#f2eee7), tinta escura, violeta como acento, verde para prova, vermelho para erro. Uma fonte monoespaçada só, a iA Writer Mono, e o cursor violeta piscando depois do nome, que é a assinatura do site deles. Os gráficos seguem o estilo de barras do site: verticais, escala linear, o Bend em violeta e o resto em cinza, com as barras que passam do topo hachuradas.

O vídeo do começo deste post também é gerado, não gravado. É uma página HTML com uma função que desenha o quadro para um tempo t; um navegador sem interface tira 3600 fotos dela e o ffmpeg junta tudo, mesclando os quadros dois a dois para dar o borrão de movimento. Toda saída de terminal e todo número no vídeo são lidos dos arquivos das execuções reais. As 26 células das matrizes do começo atravessam o filme inteiro: viram o logo, se reorganizam no reshape, viram as caixas dos pacotes, os tokens do GPT-2, as barras do benchmark, e voltam para o logo no fim. make media refaz tudo.

onde o bend me pegou

Algumas coisas que eu gostaria de ter sabido no primeiro dia:

  • variáveis afins: usou duas vezes, erro. A saída é marcar o parâmetro com + (reuso permitido em valores Data), ou reorganizar o código para não precisar.
  • recursão só estrutural: o checador de terminação exige que a recursão diminua algo visivelmente, para garantir que o programa termina. Quando o algoritmo não tem essa forma, a saída é um parâmetro de "combustível" que diminui a cada passo. O pré-tokenizador pula de 1 a 4 bytes de cada vez (o tamanho de um caractere UTF-8), então ele recebe como combustível o número de bytes que faltam, e cada passo gasta um:
def tag(fuel: Nat, +tbl: List<&2, U32>, +xs: List<&2, U32>, +n: Nat, +cls: Nat) -> List<&2, U32>:
  match fuel xs:
    case 1n+p Con{b, t}:
      emit(n, cls, False{}, xs, tag(p, tbl, List.drop(&2, U32, xs, n), ...))
    case _ _:
      Nil{}
  • match só em parâmetro: não dá para fazer match numa expressão; vira uma função auxiliar que recebe o valor já calculado como argumento.
  • ordem de declaração: uma função só pode chamar outra declarada antes dela, e recursão mútua não passa. Mais de uma vez a correção foi só trocar duas funções de lugar.
  • padrões aninhados de tupla em match não inferem tipo. A saída foi trocar tuplas por registros com tipo próprio, como o MMul.
  • imports não são reexportados: se o pacote importa uma lei de outro arquivo, quem importa o pacote não vê. Por isso cada pacote publicado é um arquivo só, com a lei e a prova lado a lado.
  • publicar é para sempre: versão publicada não se apaga, então todo pacote só subia depois de todas as provas e de toda a bateria de testes passarem. As versões antigas continuam lá, e o README sempre aponta para a atual.

Essas regras existem por um motivo: são elas que deixam o Bend gerenciar memória sem coletor de lixo, garantir que todo programa termina e checar provas rápido. Mesmo assim, cada uma me custou um erro de compilação para entender.

os números, sem enfeite

benchmarks: época de MNIST, v1 544 s, v2 7 s, PyTorch 0,2 s; GPT-2 por token, v1 3 s, v2 0,1 s, PyTorch 55 ms com 1 thread e 21 ms com 16
escala linear. as barras hachuradas da v1 passam do topo do gráfico; menor é melhor.
v1 (listas)v2 (array)PyTorch
MNIST, uma época544 s~7 s0,2 s
GPT-2 small, por token~3 s~0,1 s21 ms (16 threads), 55 ms (1 thread)
GPT-2, carregar os pesos~10 s~7 s, 1,5 GB~1 s

Tudo medido no mesmo notebook (Core Ultra 7 155H, 32 GB), com a máquina ociosa; o MNIST é a mediana de três execuções. O PyTorch ainda é umas 35 vezes mais rápido no MNIST e de 2 a 5 vezes no GPT-2.

O que sobra da distância está no código gerado. O Bend 2.0.35 gera código escalar: uma conta por vez. Os processadores modernos têm instruções SIMD, que fazem 8 ou 16 contas numa instrução só, e o PyTorch usa bibliotecas como a BLAS, afinadas há décadas para usar cada truque do processador. No mesmo produto de matrizes, o PyTorch faz 62 bilhões de multiplica-soma por segundo numa thread; o Bend, com array, 2,3 bilhões. Fechar essa distância é trabalho do compilador.

o que ficou

Cinco pacotes no BendHub, todos MIT, todos passando no --verdict, sem nenhum @unsafe e nenhum ?TODO:

os cinco pacotes, quem importa quem, e o que é provado, testado ou confiado
verde é provado pelo kernel, cinza é testado contra uma referência, laranja é confiado.
  • bend-ml-nat-lemmas: 15 leis sobre Nat e listas
  • bend-ml-bpe-tokenizer: o tokenizer, com roundtrip, vocab_bound, dec_append, train_wf e roundtrip_trained
  • bend-ml-tensor: Mat<r,c> sobre listas, com reshape provado
  • bend-ml-tensor-array: o mesmo sobre array plano, umas 50 vezes mais rápido, com cap_ok
  • bend-ml-autograd: diferenciação automática, com reverse_eq_forward

São 24 leis no total. E o que eu levo do projeto, mais que os pacotes:

  • tipos carregam os shapes de um passo de treino inteiro, inclusive o shape de cada gradiente, e não custam nada em tempo de execução.
  • prova cobre estrutura, teste cobre número, e saber onde passa a linha entre os dois é metade do trabalho.
  • a base de confiança cabe numa lista: o kernel, as primitivas de F32 e o Array.new. Está escrita no guia de auditoria.
  • a estrutura de dados importa mais que a linguagem: o "1800 vezes mais lento" era lista encadeada, não Bend.
  • medir antes de afirmar. Três conclusões minhas estavam erradas (o tempo do PyTorch por token, a GPU "lenta em tudo", os 10 GB de memória), e as três foram corrigidas por uma medição, não por um argumento. As correções ficaram nas notas, do lado das conclusões erradas.

O próximo passo óbvio é dividir o array de pesos entre as tarefas sem copiar: como o Array é uma árvore por baixo, dá para separar os galhos e entregar um para cada tarefa. Isso destravaria o paralelismo em matriz vezes vetor, que hoje perde para a versão sequencial. Daí para um llama.bend de verdade, o caminho existe.

O código, o vídeo e todas as notas estão no repositório. Clonar, rodar make setup && make check, e conferir cada número deste post é a melhor resposta que eu poderia receber.

Inclusive, escrever este post me deu vontade de falar mais sobre IA por aqui, do mesmo jeito que fiz com criptografia: começando do zero e indo fundo. Em breve sai algo nessa linha.