$cat ~/posts/bend-ml-v3

bilingual · en / pt

bend-ml after the launch: copying was the cost, and the bug was in the compiler

2026-10-07·en·22 min

Languages and compilers#Bend#compilers#performance

The previous post told how I built bend-ml, a machine learning library in Bend where a shape error doesn't compile. It ended with a sentence about what would come next:

$ quote ~/posts/bend-ml

The obvious next step is splitting the weight array between tasks without copying: since the Array is a tree underneath, its branches can be separated and handed one to each task.

a pytorch where a shape error doesn't compile2026-10-04en

I took that step. Along the way the parallelism vanished without a single error, and the reason was in Bend's compiler, not in my code. I opened an issue, went after the cause, and the fix was accepted into the official repository. This post covers those three days the way the other one covered the first two: explaining each piece before using it, and with every number, including the ones that didn't improve.

2026-10-04v2.11 daywhere the previous post left off: flat array, CI, 24 laws.
2026-10-05launch and v31 daythe questions from the launch, the bands, and the parallelism that vanished.
2026-10-06v3.1 and three issues1 dayloading the weights, the one-thread cost, and three reports for the bend team.
2026-10-07the fixthe cause in the compiler, the PR, and the merge.

each bar is the interval until the next milestone.

the questions from the launch

On the 5th I posted the project on X, as an answer to Taelin's wishlist. Three questions came back and became the plan for v3.

The first: "doesn't dfdx already do this?". dfdx is an ML library in Rust that also checks shapes before running, using a Rust feature called const generics: fixed numbers inside the type. The question was fair, and the answer is in one detail: fixed numbers. When a size is only known while the program runs (the batch size read from a file, for example), it becomes a regular variable in Rust, and from then on dfdx checks the shape the old way, with a run-time check that brings the program down if it fails.

In Bend, a size that only exists at run time is still a variable inside the type. This reshape regroups a matrix of n rows and 6 columns into n·2 rows and 3 columns, with n coming from the command line:

def regroup(+n: Nat, x: T.Mat<n, 6n>) -> T.Mat<Nat.mul(n, 2n), 3n>:
  T.Mat.reshape(n, 6n, Nat.mul(n, 2n), 3n, NL.mul_assoc(n, 2n, 3n), x)

Remember that reshape demands a proof that the total number of elements doesn't change? Here the total is n·6 on one side and (n·2)·3 on the other, and the proof is mul_assoc, the associativity of multiplication, one of the lemmas in the first package. It holds for every n, so the checker accepts the function once, for whatever size arrives. Asking for n×5 instead doesn't compile: there is no proof that n·6 equals n·5.

That catches a bug no fixed-number checker can catch. In GPT-2, the attention's output projection has 768×768 weights. The weights are stored one row per output, so the right product is X·Wᵀ. If someone writes X·W, with 768 on both sides the shapes fit just the same, and the program runs and gives wrong numbers. Written for any input and output size, the same layer doesn't compile:

def proj(+n: Nat, +din: Nat, +dout: Nat, x: TA.Mat<n, din>, w: TA.Mat<dout, din>) -> TA.MMul<n, din, dout>:
  TA.Mat.matmul(n, din, dout, 0n, x, w)

Error:
- expected : TA.Mat<din, dout>
- observed : TA.Mat<dout, din>

To the checker, din and dout are different names, so W has dout rows and not din, even if the program only ever uses the layer with 768 and 768.

The second question was about compile times: with types this detailed, what happens on a deep network? I measured. Checking 8 layers with all-different sizes takes 0.22 seconds; 128 layers, 0.27. The build is dominated by clang, the C compiler: 1.1 to 4.7 seconds. The whole GPT-2 checks in 0.38 seconds.

The third wasn't a question. The author of Bend-Conv, a tensor and convolution library in Bend, pointed out that the real cost of parallelism was copying data. That was exactly the problem from experiment 3 in the previous post, and it took up the rest of v3.

copying was the cost

A reminder of experiment 3: to generate text, GPT-2 multiplies matrices by a vector, a single row of input. Parallelizing that by splitting the matrix into blocks made everything 10 times slower, because each task got a full copy of the weights, and copying cost more than doing the math. The sentence at the end of the post was the way out: instead of copying, hand each task only its own piece.

Bend-Conv already stored matrices in row bands. That's the idea I took from it, and what I added was the type. In the package, the banded matrix is called Bands:

type Bands<-r: Nat, -c: Nat> is Type:
  BLeaf{m: Mat<r, c>}
  BNode{x: Bands<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>}

A Bands of r rows is one of two things. Either it's a leaf (BLeaf) holding a regular matrix of r rows, or it's a node (BNode) that splits the rows into two halves, each one another Bands: the top one with half(r) rows, the bottom one with the rest. Each leaf is a band, with its own array.

splitting the tree copies nothing: each task gets its own branch. the only thing copied is the input vector, which is small.

What makes this work with no tricks is that Bend's Array is, underneath, a plain tree from the standard library. Opening a node with match hands over both halves right away, without copying and without @unsafe. Because of affine types, whoever opens the node becomes the owner of the halves and can give one to each task. The input vector x is still copied for each task, but it has 768 numbers, against the 38 million of GPT-2's largest matrix.

And the type guarantees the part that goes wrong easily: the two halves add up to exactly r rows, none lost and none counted twice. That's a law:

law half_cover:
  for +n: Nat
  {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat}

A matrix built with the wrong number of rows doesn't compile. And every function that walks the tree carries an erased proof, {rr == r}, tying the row count used at run time (rr) to the one in the type (r). The proof costs nothing when it runs, and it's the one that gave me trouble, as you'll see in a moment.

Experiment 11. The same matrix-by-vector product, on GPT-2's shapes, with 16 threads (for bands, the best of 16 and 32 bands):

matrix	sequential	blocks with copies	bands, no copies
2304 × 768	0.80 ms	10.4 ms	0.32 ms
3072 × 768	1.07 ms	23.5 ms	0.40 ms
50257 × 768	25.1 ms	337 ms	4.5 ms

The middle column is experiment 3's idea, with blocks that copy the weights: slower than doing everything on one thread. Without copies, the largest matrix, the logits one (which scores each of the 50257 tokens), got 5.6 times faster than the sequential product. And with the same result bit for bit: each band runs the same math as the sequential product, so GPT-2's 11 test prompts give the same tokens and the same logits as v2.1, digit for digit.

But that table is from the end of the day. The first version of the bands wasn't faster.

the parallelism that vanished

The first version compiled, gave the right result, and took the same time on 1 and on 16 threads. No error, no warning.

To see why, you need a piece of how Bend parallelizes. In Bend, you write two calls side by side in the same let, and it runs both at the same time:

p q = bmv(..., top_half, ...) bmv(..., bottom_half, ...)

That's called a parallel let. The compiler turns the program into C, and in the C, to run the two calls in parallel, it creates a join function: the part that waits for both halves to finish and puts the results together. If there's no join, the two calls become a plain sequence, one after the other.

the result is the same in both cases. only the time changes, which is why no test complains.

What led me to the cause was reading the C the compiler generates (bend X.bend -o x.c). In the slow version, the join function wasn't there. And the difference between the slow version and one that worked was only the order of the parameters:

bmv(+c, -r, b, xp, +rr, -e)    no join: 1 core
bmv(+c, -r, b, +rr, -e, xp)    join: 16 cores

The -e is the erased proof {rr == r}. It disappears from the compiled program, so in theory its position shouldn't matter. But when it was the last parameter, the parallel let became two calls in sequence. I moved the proof, with a real parameter after it, and the join came back.

I shrank the problem down to the smallest program that showed the same effect: a tree of calls that sums leaves, with an erased parameter -z that does nothing. With -z at the end, 1.01 seconds on 1 thread and 1.01 on 16. With -z in the middle, 0.17 on 16. On GPT-2's logits matrix, the same change of position took the product from 26.7 ms to 3.3 ms.

There was a second, smaller trap. The half function, which computes the size of the top half, was recursive in the simplest way (half(n) = 1 + half(n − 2)) and was called at every node of the tree. Recursion like that, which still has work to do after it returns, gets in the way of the parallelism around it. Rewritten as a loop that carries the result in an accumulator, the same batch of 40 products went from 0.69 to 0.36 seconds.

Both traps became rules in the project: when parallel code doesn't scale, compare the time on 1 and on 16 threads and read the generated C. And so the first one doesn't come back, CI compiles the banded product to C and fails if the join function isn't there. I checked that the check catches it: with the proof moved back to the end, it fails.

On GPT-2, generating 8 tokens, v3 got faster than v2.1 on 16 threads. The exact numbers moved with the machine's load, and I come back to them at the end. Something else showed up too: on one thread, v3 was about 12% slower than v2.1. That was the start of v3.1.

v3.1: one thread, loading the weights, and a negative result

On a single thread, bands don't help: the tasks run one after the other, and what's left is the cost of building the tree and joining the results. I found two causes. The typed API converted the result to a Mat and the demo converted it back to a list; a Bands.matvec_l, which takes and returns lists, removed that round trip. And with 2 threads or fewer the demo now uses a single band, which runs exactly v2.1's math. Measuring 36 forward passes of the model:

36 GPT-2 forward passes	1 thread	16 threads
v2.1	4.7 s	4.7 s
v3.0	5.2 s	3.8 to 4.2 s
v3.1	4.8 s	3.6 to 3.8 s

The gain on the whole model is smaller than on the product alone because attention, which still runs on lists, takes a growing share as the text gets longer.

Experiment 12. The other target was loading the weights. A profile showed about 50 nanoseconds per number, spent turning each 1 MB block into a list of bytes, then a list of F32, and only then writing to the array. Decoding straight into the array, 4 bytes become an F32 and go right into place, and loading dropped by about 25%.

The next idea was loading the 12 layers at the same time. Bend has IO.fork, which starts several input and output tasks, and the language guide says each one runs its pure code "in parallel, on every core". Loading each layer in an IO.fork made it slower: from 5.1–5.6 to 6.6–7.1 seconds, the same on 1 and on 16 threads. In a minimal test, four forks of a long loop took 1.86 seconds on 1 thread and 1.73 on 16, with CPU time equal to wall-clock time: a single core. The same loops in a parallel let scale. I reverted the change, wrote down the negative result, and wrote a report.

There's also a check at the boundary now: each weight file must hold exactly r·c·4 bytes, otherwise the program stops with a message that says which file, how many bytes it expected and which script to run again, instead of quietly loading half a matrix.

the proof that didn't close

v3.1 had a law I really wanted: "the banded product returns exactly r numbers". I proved half of it. band_len says a band of rows rows returns rows numbers, and join_len says joining two halves adds their lengths. With both, the whole tree looked like an easy induction: each node joins its two halves, and each half, by induction, has the right length.

Except the induction needs to talk about the two recursive calls twice: once to use the result, once to state its length. Each call carries its band's array, and an array is affine, it can only be used once. The checker refused it (consumed more than once), even inside the proof's rewrite. Erased copies didn't solve it either, because of the language's rules on where an erased value can be used.

That law stayed out of the README and is in the notes as debt. The output length is covered by tests: 20 shapes, every depth of the tree, checking every number. The project's rule is that a law that isn't proved yet doesn't show up as proved.

A trick that worked for the smaller laws, and that I wouldn't have found reading the guide: pass the expression that mentions the arrays only in an erased argument, which doesn't count as a use, and the arrays themselves just once.

three reports for the bend team

By the end of the 6th I had three things that weren't in my code, they were in the language. I wrote each one up as a report, with a minimal program that reproduces it, the numbers, and what I had ruled out:

#1374: an erased parameter at the end of the list turns the parallel let into sequential calls.
#1375: IO.fork runs pure work on a single core, and the guide says it runs on all of them.
#1376: a way for several tasks to read the same array without copying and without @unsafe. Today the only way, Array.fork, is @unsafe, and a package that uses it doesn't pass as proved.

Before posting, I checked each one against Bend's latest code, not only the version the project uses. All three were still there.

going after the cause

Working around it in my code fixed my case. #1374 was still there for anyone else who wrote a proof at the end of the list, and worse: with no error at all, just slow. So, on the 7th, I cloned the compiler.

Bend 2's compiler is written in TypeScript, and the part that generates code lives in a single file, comp.ts, with a size limit the repository itself checks. I didn't know the code. I added temporary logs to the step that reorganizes calls before generating C, ran the minimal program in its three versions (proof at the end, proof in the middle, no proof) and compared what each one printed.

The cause, in plain words: before generating C, the compiler breaks a call with several arguments into pieces, applying one argument at a time. When it finds a piece that already has every argument that exists at run time, it treats that piece as a complete call and splits it off into its own sequential step. With the proof in the middle, the piece before it is still incomplete, and nothing gets split off. With the proof at the end, the piece before it already has every real argument. The compiler split the call off as a plain step, and the parallel let was left with two ready values instead of two calls: nothing to parallelize, no join.

The fix: only an argument that exists at run time can make the compiler split the call off. An erased argument stays where it is, as it already did when it sat in the middle. In comp.ts, it was five lines:

- if (m.t.$ === "Var" && !m.args.includes(f.x)) {
+ const on = m.args.includes(f.x);
+ if (m.t.$ === "Var" && !on) {
    return spine(f.f);
  }
- const g = cut(spine(f.f), ty_ann(f.f));
- const x = m.args.includes(f.x) ? go(f.x, false, null) : f.x;
+ const g = on ? cut(spine(f.f), ty_ann(f.f)) : spine(f.f);
+ const x = on ? go(f.x, false, null) : f.x;

on says whether the argument f.x is among the ones that exist at run time. cut is what split the call off; now it only happens when on is true. With the fix, the minimal program went from 1.01 to 0.17 seconds on 16 threads, and bend-ml's bmv got its join even with the proof at the end.

The part that took longest wasn't the fix, it was showing that it didn't break anything. Touching the step every program goes through can change anything. So I generated the C and the JavaScript for all 1,659 programs in Bend's repository (tests, benchmarks and demos) with the compiler before and after, and compared them byte for byte. The JavaScript was all identical. In the C, two files changed, and both have an explanation: in one, eight calls the bug split off became tail calls (the last thing the function does, with nothing kept for later), for the same reason as the bug; in the other, only the name of an internal variable changed, and the binary came out identical. I ran every test in the three ways Bend runs code (interpreter, C and JavaScript), before and after: the same output. The 17 benchmarks built identical binaries. And I said in the PR what I had no way to run, like the tests that need the team's cluster of machines.

I also wrote a test in the repository's format. It has a limitation I left written in it: a program's output doesn't show whether it ran in parallel, so the test checks the result on every path, but wouldn't catch the fix being reverted. The evidence that the join came back is the generated C and the timings, which are in the PR's description.

the merge

The PR was approved and went into the official repository the same day, with no requested changes, exactly as it was sent. #1374 closed with it. It was my first contribution to Bend.

It isn't in a release yet: 2.0.36 came out hours before the merge. Until the next one, bend-ml stays pinned to 2.0.35, with the proof kept off the end in bmv. When the release comes out, the workaround can go, but the join check in CI stays. #1375 and #1376 are still open, with the Bend team.

the numbers, no varnish

All on the same laptop as the previous post (Core Ultra 7 155H, 22 threads, 32 GB). This time the machine was not idle: an Android emulator and a Gradle daemon ran the whole session. So the versions were measured alternately, under the same conditions, and the numbers show up as ranges.

GPT-2 small	v2.1	v3.0	v3.1
36 passes, 16 threads	4.7 s	3.8 to 4.2 s	3.6 to 3.8 s
36 passes, 1 thread	4.7 s	5.2 s	4.8 s
loading the weights	6 to 7 s	5.1 to 6.1 s	3.9 to 5.1 s
peak memory	1462 MB	895 MB	903 MB

matrix · vector product, 16 threads	v2.1	v3
2304 × 768	0.80 ms	0.32 ms
50257 × 768	25.1 ms	4.5 ms

What didn't improve or is still unexplained:

with a single band, on the logits matrix, the product with lists is about 25% slower than v2.1's (29.6 against 23.5 ms). It's what's left of the one-thread cost. Likely cause: the 50 thousand-number result built as a list, against writing into an array. I didn't look into it.
the process's peak virtual size went from 10.6 GB to 41 GB. It's reserved address space, not memory in use (the real peak went down), but I didn't look into the cause.
the README chart still shows v2's numbers, until I can measure on an idle machine.

what's left

27 laws in 5 packages, all passing --verdict, with no @unsafe and no ?TODO. bend-ml-tensor-array went to version 0.1.5.0 on BendHub, with the bands. CI runs 43 checks, GPT-2 and MNIST included, now on the development branch too, and the stable version only reaches the main branch through a pull request.

And what I take from these three days:

the question "doesn't dfdx already do this?" led me to the project's best example: the bug that only shows up when the layer is written for any size.
the person who pointed out the copying problem had already faced it in another library. Publishing early brought in the right idea.
when parallel code doesn't scale, the time on 1 and on 16 threads and the generated C say more than rereading your own code.
a measured negative result (IO.fork) is worth a report, and an unproved law is worth a line in the notes, not a line in the README.
when the problem is in the tool, you can go to it. Bend's compiler was a file I had never opened, and the fix fit in a few lines. What took the most work was the boring part: showing that nothing else changed.

The code and all the notes, with the command to reproduce every number, are in the repository.

O post anterior contou como eu construí o bend-ml, uma biblioteca de machine learning em Bend onde erro de shape não compila. Ele terminou com uma frase sobre o que viria depois:

$ quote ~/posts/bend-ml

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.

um pytorch onde erro de shape não compila2026-10-04pt

Eu dei esse passo. No caminho, o paralelismo sumiu sem nenhum erro, e o motivo estava no compilador do Bend, não no meu código. Abri uma issue, fui atrás da causa, e a correção foi aceita no repositório oficial. Este post conta esses três dias, do jeito que o outro contou os dois primeiros: explicando cada peça antes de usar, e com todos os números, inclusive os que não melhoraram.

2026-10-04v2.11 diao estado do post anterior: array plano, CI, 24 leis.
2026-10-05lançamento e v31 diaas perguntas do lançamento, as bandas, e o paralelismo que sumia.
2026-10-06v3.1 e três issues1 diacarregar os pesos, o custo com uma thread, e três relatos para o time do bend.
2026-10-07a correçãoa causa no compilador, a PR, e o merge.

cada barra é o intervalo até o marco seguinte.

as perguntas do lançamento

No dia 5 eu publiquei o projeto no X, respondendo à lista de desejos do Taelin. Vieram três perguntas que viraram o plano da v3.

A primeira: "o dfdx já não faz isso?". O dfdx é uma biblioteca de ML em Rust que também confere shapes antes de rodar, usando um recurso do Rust chamado const generics: números fixos dentro do tipo. A pergunta era justa, e a resposta está num detalhe: números fixos. Quando o tamanho só é conhecido com o programa rodando (o tamanho do batch lido de um arquivo, por exemplo), ele vira uma variável comum no Rust, e a partir daí o dfdx confere o shape do jeito antigo, com uma verificação em tempo de execução que derruba o programa se falhar.

No Bend, um tamanho que só existe em tempo de execução continua sendo uma variável dentro do tipo. Este reshape reorganiza uma matriz de n linhas e 6 colunas em n·2 linhas e 3 colunas, com n vindo da linha de comando:

def regroup(+n: Nat, x: T.Mat<n, 6n>) -> T.Mat<Nat.mul(n, 2n), 3n>:
  T.Mat.reshape(n, 6n, Nat.mul(n, 2n), 3n, NL.mul_assoc(n, 2n, 3n), x)

Lembra que o reshape exige uma prova de que o total de números não muda? Aqui o total é n·6 de um lado e (n·2)·3 do outro, e a prova é mul_assoc, a associatividade da multiplicação, um dos lemas do primeiro pacote. Ela vale para todo n, então o checador aceita a função uma vez, para qualquer tamanho que chegar. Pedir n×5 em vez disso não compila: não existe prova de que n·6 é igual a n·5.

Isso pega um bug que nenhum checador de números fixos consegue pegar. No GPT-2, a projeção de saída da atenção tem pesos de 768×768. Os pesos ficam guardados uma linha por saída, então o produto certo é X·Wᵀ. Se alguém escreve X·W, com 768 dos dois lados os shapes batem do mesmo jeito, e o programa roda e dá números errados. Escrita para qualquer tamanho de entrada e de saída, a mesma camada não compila:

def proj(+n: Nat, +din: Nat, +dout: Nat, x: TA.Mat<n, din>, w: TA.Mat<dout, din>) -> TA.MMul<n, din, dout>:
  TA.Mat.matmul(n, din, dout, 0n, x, w)

Error:
- expected : TA.Mat<din, dout>
- observed : TA.Mat<dout, din>

Para o checador, din e dout são nomes diferentes, então W tem dout linhas e não din, mesmo que o programa só use a camada com 768 e 768.

A segunda pergunta foi sobre tempo de compilação: com tipos tão detalhados, o que acontece numa rede funda? Medi. Checar 8 camadas com tamanhos todos diferentes leva 0,22 segundo; 128 camadas, 0,27. O build é dominado pelo clang, o compilador de C: de 1,1 a 4,7 segundos. O GPT-2 inteiro checa em 0,38 segundo.

A terceira não foi uma pergunta. O autor do Bend-Conv, uma biblioteca de tensores e convolução em Bend, comentou que o custo real do paralelismo estava em copiar dados. Era exatamente o problema do experimento 3 do post anterior, e foi o que ocupou o resto da v3.

copiar era o custo

Relembrando o experimento 3: para gerar texto, o GPT-2 faz produtos de matriz por vetor, uma linha só de entrada. Paralelizar isso dividindo a matriz em blocos deixou tudo 10 vezes mais lento, porque cada tarefa recebia uma cópia inteira dos pesos, e copiar custava mais que fazer a conta. A frase do fim do post era a saída: em vez de copiar, entregar a cada tarefa só o seu pedaço.

O Bend-Conv já guardava matrizes em faixas de linhas. A ideia que eu peguei dele foi essa, e o que eu acrescentei foi o tipo. No pacote, a matriz em faixas se chama Bands:

type Bands<-r: Nat, -c: Nat> is Type:
  BLeaf{m: Mat<r, c>}
  BNode{x: Bands<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>}

Uma Bands de r linhas é uma de duas coisas. Ou é uma folha (BLeaf) com uma matriz comum de r linhas, ou é um nó (BNode) que divide as linhas em duas metades, cada uma sendo outra Bands: a de cima com half(r) linhas, a de baixo com o resto. Cada folha é uma banda, com o seu próprio array.

dividir a árvore não copia nada: cada tarefa recebe o seu galho. o que se copia é só o vetor de entrada, que é pequeno.

O detalhe que faz isso funcionar sem truque nenhum é que o Array do Bend é, por baixo, uma árvore comum da biblioteca padrão. Abrir um nó com match entrega as duas metades na hora, sem copiar, e sem precisar de @unsafe. Por causa dos tipos afins, quem abre o nó passa a ser dono das metades, e pode entregar uma para cada tarefa. O vetor de entrada x ainda é copiado para cada tarefa, mas ele tem 768 números, contra os 38 milhões da maior matriz do GPT-2.

E o tipo garante a parte que dá errado fácil: as duas metades somam exatamente r linhas, nenhuma se perde e nenhuma conta duas vezes. Isso é uma lei:

law half_cover:
  for +n: Nat
  {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat}

Uma matriz montada com o número de linhas errado não compila. E cada função que percorre a árvore carrega uma prova apagada, {rr == r}, amarrando o número de linhas usado em tempo de execução (rr) ao do tipo (r). A prova não custa nada rodando, e foi ela que me deu trabalho, como você vai ver daqui a pouco.

Experimento 11. O mesmo produto de matriz por vetor, nos shapes do GPT-2, com 16 threads (nas bandas, o melhor entre 16 e 32 bandas):

matriz	sequencial	blocos com cópia	bandas, sem cópia
2304 × 768	0,80 ms	10,4 ms	0,32 ms
3072 × 768	1,07 ms	23,5 ms	0,40 ms
50257 × 768	25,1 ms	337 ms	4,5 ms

A coluna do meio é a ideia do experimento 3, com blocos que copiam os pesos: mais lenta que fazer tudo numa thread só. Sem cópia, a maior matriz, a dos logits (a que dá a nota de cada um dos 50257 tokens), ficou 5,6 vezes mais rápida que a sequencial. E com o mesmo resultado bit a bit: cada banda roda a mesma conta do produto sequencial, então os 11 textos de teste do GPT-2 dão os mesmos tokens e os mesmos logits da v2.1, dígito por dígito.

Mas essa tabela é do fim do dia. A primeira versão das bandas não era mais rápida.

o paralelismo que sumia

A primeira versão compilava, dava o resultado certo e levava o mesmo tempo com 1 e com 16 threads. Nenhum erro, nenhum aviso.

Para entender por quê, precisa de um pedaço de como o Bend paraleliza. No Bend, você escreve duas chamadas lado a lado num mesmo let, e ele roda as duas ao mesmo tempo:

p q = bmv(..., metade_de_cima, ...) bmv(..., metade_de_baixo, ...)

Isso se chama um let paralelo. O compilador transforma o programa em C, e no C, para rodar as duas chamadas em paralelo, ele cria uma função de join: a parte que espera as duas metades terminarem e junta os resultados. Se o join não existe, as duas chamadas viram uma sequência comum, uma depois da outra.

o resultado é o mesmo nos dois casos. só o tempo muda, por isso nenhum teste reclama.

O que me levou à causa foi ler o C que o compilador gera (bend X.bend -o x.c). Na versão lenta, a função de join não estava lá. E a diferença entre a versão lenta e uma que funcionava era só a ordem dos parâmetros:

bmv(+c, -r, b, xp, +rr, -e)    sem join: 1 núcleo
bmv(+c, -r, b, +rr, -e, xp)    com join: 16 núcleos

O -e é a prova apagada {rr == r}. Ela some do programa compilado, então em tese a posição dela não deveria importar. Mas quando ela era o último parâmetro, o let paralelo virava duas chamadas em sequência. Mudei a prova de lugar, com um parâmetro de verdade depois dela, e o join voltou.

Diminuí o problema até o menor programa que mostrava o mesmo efeito: uma árvore de chamadas que soma folhas, com um parâmetro apagado -z que não faz nada. Com -z no fim, 1,01 segundo com 1 thread e 1,01 com 16. Com -z no meio, 0,17 com 16. Na matriz de logits do GPT-2, a mesma troca de posição levou o produto de 26,7 ms para 3,3 ms.

Teve uma segunda armadilha, menor. A função half, que calcula o tamanho da metade, era recursiva do jeito mais simples (half(n) = 1 + half(n − 2)) e era chamada em cada nó da árvore. Uma recursão assim, que ainda tem conta para fazer depois de voltar, atrapalha o paralelismo em volta dela. Reescrita como um laço que carrega o resultado num acumulador, a mesma bateria de 40 produtos caiu de 0,69 para 0,36 segundo.

As duas armadilhas viraram regra no projeto: quando código paralelo não escala, compare o tempo com 1 e com 16 threads e leia o C gerado. E para a primeira não voltar, o CI compila o produto em bandas para C e falha se a função de join não estiver lá. Conferi que a checagem pega o erro: com a prova de volta no fim, ela falha.

No GPT-2, gerando 8 tokens, a v3 ficou mais rápida que a v2.1 com 16 threads. Os números exatos variaram com a carga da máquina, e eu volto neles no fim. O que também apareceu: com uma thread, a v3 era uns 12% mais lenta que a v2.1. Isso foi o começo da v3.1.

v3.1: uma thread, carregar os pesos, e um resultado negativo

Com uma thread só, as bandas não ajudam: as tarefas rodam uma depois da outra, e o que sobra é o custo de montar a árvore e juntar os resultados. Encontrei duas causas. A API com tipos convertia o resultado para Mat e o demo convertia de volta para lista; um Bands.matvec_l, que recebe e devolve listas, tirou essa ida e volta. E com 2 threads ou menos o demo agora usa uma banda só, que roda exatamente a conta da v2.1. Medindo 36 passadas do modelo:

36 passadas do GPT-2	1 thread	16 threads
v2.1	4,7 s	4,7 s
v3.0	5,2 s	3,8 a 4,2 s
v3.1	4,8 s	3,6 a 3,8 s

O ganho no modelo inteiro é menor que o do produto sozinho porque a atenção, que ainda roda sobre listas, ocupa uma fatia cada vez maior conforme o texto cresce.

Experimento 12. O outro alvo foi carregar os pesos. Um perfil mostrou uns 50 nanossegundos por número, gastos em transformar cada bloco de 1 MB numa lista de bytes, depois numa lista de F32, e só então escrever no array. Decodificando direto para o array, 4 bytes viram um F32 e já vão para o lugar, o carregamento caiu uns 25%.

A ideia seguinte era carregar as 12 camadas ao mesmo tempo. O Bend tem o IO.fork, que dispara várias tarefas de entrada e saída, e o guia da linguagem diz que cada uma roda o seu código puro "em paralelo, em todos os núcleos". Carregando cada camada num IO.fork, ficou mais lento: de 5,1–5,6 para 6,6–7,1 segundos, igual com 1 e com 16 threads. Num teste mínimo, quatro forks de um laço longo levaram 1,86 segundo com 1 thread e 1,73 com 16, com o tempo de CPU igual ao tempo de relógio: um núcleo só. Os mesmos laços num let paralelo escalam. Desfiz a mudança, anotei o resultado negativo e escrevi um relato.

Ficou também uma checagem na fronteira: cada arquivo de pesos precisa ter exatamente r·c·4 bytes, senão o programa para com uma mensagem que diz qual arquivo, quantos bytes esperava e qual script rodar de novo, em vez de carregar meia matriz em silêncio.

a prova que não fechou

A v3.1 tinha uma lei que eu queria muito: "o produto em bandas devolve exatamente r números". Provei metade. band_len diz que uma banda de rows linhas devolve rows números, e join_len diz que juntar duas metades soma os tamanhos. Com as duas, a árvore inteira parecia uma indução fácil: cada nó junta as duas metades, e cada metade, por indução, tem o tamanho certo.

Só que a indução precisa falar das duas chamadas recursivas duas vezes: uma para usar o resultado, outra para dizer o tamanho dele. Cada chamada carrega o array da sua banda, e array é afim, só pode ser usado uma vez. O checador recusou (consumed more than once), inclusive dentro da reescrita da prova. Cópias apagadas também não resolveram, por regras da linguagem sobre onde um valor apagado pode ser usado.

Essa lei ficou fora do README e está nas notas como dívida. O tamanho da saída é coberto por testes: 20 shapes, todas as profundidades da árvore, conferindo cada número. A regra do projeto é que uma lei que ainda não foi provada não aparece como provada.

Um truque que funcionou nas leis menores, e que eu não teria achado lendo o guia: passar a expressão que menciona os arrays só num argumento apagado, que não conta como uso, e os arrays em si uma vez só.

três relatos para o time do bend

No fim do dia 6 eu tinha três coisas que não eram do meu código, eram da linguagem. Escrevi cada uma como relato, com um programa mínimo que reproduz, os números, e o que eu tinha descartado:

#1374: um parâmetro apagado no fim da lista transforma o let paralelo em chamadas sequenciais.
#1375: o IO.fork roda o trabalho puro num núcleo só, e o guia diz que roda em todos.
#1376: um jeito de várias tarefas lerem o mesmo array sem copiar e sem @unsafe. Hoje a única forma, Array.fork, é @unsafe, e um pacote com isso não passa como provado.

Antes de postar, conferi cada um no código mais recente do Bend, não só na versão que o projeto usa. Os três continuavam lá.

indo atrás da causa

Contornar no meu código resolvia o meu caso. A #1374 continuava lá para qualquer outra pessoa que escrevesse uma prova no fim da lista, e pior: sem erro nenhum, só lento. Então, no dia 7, eu clonei o compilador.

O compilador do Bend 2 é escrito em TypeScript, e a parte que gera o código fica num arquivo só, o comp.ts, com um limite de tamanho que o próprio repositório confere. Eu não conhecia o código. Coloquei logs temporários no passo que reorganiza as chamadas antes de gerar o C, rodei o programa mínimo nas três versões (prova no fim, prova no meio, sem prova) e comparei o que cada uma imprimia.

A causa, em português: antes de gerar o C, o compilador quebra uma chamada com vários argumentos em pedaços, aplicando um argumento de cada vez. Quando ele encontra um pedaço que já tem todos os argumentos que existem em tempo de execução, trata esse pedaço como uma chamada completa e separa ela num passo próprio, sequencial. Com a prova no meio, o pedaço antes dela ainda está incompleto, e nada é separado. Com a prova no fim, o pedaço antes dela já tem todos os argumentos de verdade. O compilador separava a chamada como um passo comum, e o let paralelo ficava com duas variáveis prontas em vez de duas chamadas: nada para paralelizar, nenhum join.

A correção: só um argumento que existe em tempo de execução pode fazer o compilador separar a chamada. Um argumento apagado fica onde está, como já ficava quando estava no meio. No comp.ts, foram cinco linhas:

- if (m.t.$ === "Var" && !m.args.includes(f.x)) {
+ const on = m.args.includes(f.x);
+ if (m.t.$ === "Var" && !on) {
    return spine(f.f);
  }
- const g = cut(spine(f.f), ty_ann(f.f));
- const x = m.args.includes(f.x) ? go(f.x, false, null) : f.x;
+ const g = on ? cut(spine(f.f), ty_ann(f.f)) : spine(f.f);
+ const x = on ? go(f.x, false, null) : f.x;

O on diz se o argumento f.x está entre os que existem em tempo de execução. O cut é o que separava a chamada; agora ele só acontece quando on é verdadeiro. Com a correção, o programa mínimo foi de 1,01 para 0,17 segundo com 16 threads, e o bmv do bend-ml ganhou o seu join mesmo com a prova no fim.

A parte que levou mais tempo não foi a correção, foi provar que ela não quebrava nada. Mexer no passo que todo programa atravessa pode mudar qualquer coisa. Então eu gerei o C e o JavaScript de todos os 1.659 programas do repositório do Bend (testes, benchmarks e demos) com o compilador antes e depois, e comparei byte a byte. O JavaScript ficou todo idêntico. No C, dois arquivos mudaram, e os dois têm explicação: num deles, oito chamadas que o bug separava viraram chamadas de cauda (a última coisa que a função faz, sem nada guardado para depois), pelo mesmo motivo do bug; no outro, só o nome de uma variável interna mudou, e o binário saiu idêntico. Rodei cada teste nos três jeitos que o Bend executa (interpretador, C e JavaScript), antes e depois: mesma saída. Os 17 benchmarks geraram binários idênticos. E disse na PR o que eu não tinha como rodar, como os testes que precisam do cluster de máquinas do time.

Escrevi também um teste no formato do repositório. Ele tem uma limitação que eu deixei escrita nele: a saída de um programa não mostra se ele rodou em paralelo, então o teste confere o resultado em todos os caminhos, mas não pegaria a correção sendo desfeita. A prova de que o join voltou é o C gerado e os tempos, que estão na descrição da PR.

o merge

A PR foi aprovada e entrou no repositório oficial no mesmo dia, sem pedido de mudança, do jeito que foi enviada. A #1374 fechou junto. Foi a minha primeira contribuição para o Bend.

Ela ainda não está numa versão lançada: a 2.0.36 saiu horas antes do merge. Até a próxima, o bend-ml continua fixado na 2.0.35, com a prova fora do fim no bmv. Quando a versão sair, o contorno pode ir embora, mas a checagem do join no CI fica. A #1375 e a #1376 continuam abertas, com o time do Bend.

os números, sem enfeite

Tudo no mesmo notebook do post anterior (Core Ultra 7 155H, 22 threads, 32 GB). Desta vez a máquina não estava ociosa: um emulador de Android e um daemon do Gradle rodaram a sessão inteira. Por isso as versões foram medidas alternadas, nas mesmas condições, e os números aparecem como faixas.

GPT-2 small	v2.1	v3.0	v3.1
36 passadas, 16 threads	4,7 s	3,8 a 4,2 s	3,6 a 3,8 s
36 passadas, 1 thread	4,7 s	5,2 s	4,8 s
carregar os pesos	6 a 7 s	5,1 a 6,1 s	3,9 a 5,1 s
pico de memória	1462 MB	895 MB	903 MB

produto matriz · vetor, 16 threads	v2.1	v3
2304 × 768	0,80 ms	0,32 ms
50257 × 768	25,1 ms	4,5 ms

O que não melhorou ou ficou sem explicação:

com uma banda só, na matriz de logits, o produto com listas é uns 25% mais lento que o da v2.1 (29,6 contra 23,5 ms). É o resto do custo de uma thread. Causa provável: o resultado de 50 mil números montado como lista, contra escrever num array. Não investiguei.
o tamanho virtual máximo do processo foi de 10,6 GB para 41 GB. É espaço reservado, não memória usada (o pico real caiu), mas a causa eu não investiguei.
o gráfico do README ainda mostra os números da v2, até eu conseguir medir com a máquina ociosa.

o que ficou

27 leis em 5 pacotes, todas passando no --verdict, sem @unsafe e sem ?TODO. O bend-ml-tensor-array foi para a versão 0.1.5.0 no BendHub, com as bandas. O CI roda 43 checagens, GPT-2 e MNIST incluídos, agora também no branch de desenvolvimento, e a versão estável só chega no principal por pull request.

E o que eu levo destes três dias:

a pergunta "o dfdx já não faz isso?" me fez achar o melhor exemplo do projeto: o bug que só aparece quando a camada é escrita para qualquer tamanho.
quem apontou o problema da cópia foi alguém que já tinha enfrentado ele em outra biblioteca. Publicar cedo trouxe a ideia certa.
quando código paralelo não escala, o tempo com 1 e com 16 threads e o C gerado dizem mais que reler o próprio código.
um resultado negativo medido (o IO.fork) vale um relato, e uma lei não provada vale uma linha nas notas, não uma linha no README.
quando o problema está na ferramenta, dá para ir até ele. O compilador do Bend era um arquivo que eu nunca tinha aberto, e a correção coube em poucas linhas. O que deu mais trabalho foi a parte chata: mostrar que nada mais mudou.

O código e todas as notas, com o comando para reproduzir cada número, estão no repositório.

follow the next posts

keep reading

a pytorch where a shape error doesn't compile

llama.bend: language models with typed weights and cache