PyTorch: o runtime de referência

01 · Conceito

Conceito

Dois engines respondem ao mesmo prompt com os mesmos pesos, a mesma seed e decoding guloso, e devolvem textos diferentes. Um dos dois está errado. Qual? Você não resolve essa discussão com benchmarks, e não resolve lendo o código-fonte de nenhum dos dois, porque ambos são milhares de linhas de escalonamento e seleção de kernel. Você resolve tendo uma referência — uma definição do que este modelo calcula que seja lenta, sem graça e inspecionável. Essa referência é o PyTorch eager rodando o próprio código de modelagem do modelo.

Eager significa que cada operação executa quando a linha em Python roda. Nenhum graph é capturado, nada é fundido, nenhum kernel é escolhido para um padrão fundido que só existe na imaginação de um compilador. Uma multiplicação de matrizes é uma multiplicação de matrizes, na ordem escrita, sobre tensors que você pode imprimir. A lição 7.11 mandou você para cá exatamente por isso: antes de comparar stacks, você precisa de um oráculo, e o oráculo é o runtime com o menor número de peças móveis entre os parâmetros e os logits.

Então percorra a árvore. Na base está uma tabela de embedding com 248.320 linhas de largura 5120 — a geometria da lição 1.5, cerca de 1,271 bilhão de parâmetros de pura consulta. Acima dela, 64 camadas de decoder. Acima delas, um RMSNorm final e uma lm_head que não é amarrada ao embedding, contribuindo outros 1,271 bilhão de parâmetros próprios. Ao lado de tudo isso, para entrada de imagem e vídeo, uma vision tower de 27 camadas com hidden size 1152 e patch size 16, cuja saída é projetada para 5120 de modo que os patches visuais entrem na mesma sequência em que vivem os tokens de texto (lição 4.17).

Dentro de uma camada de decoder, duas coisas são invariantes. Há um fluxo residual pre-norm com RMSNorm em eps 1e-6 antes do mixer e antes da FFN (lições 2.8 e 4.10), e há uma feed-forward network com gating projetando de 5120 para 17408 e de volta (lição 4.9). É nessa FFN que a maior parte do modelo fisicamente está. O terceiro slot é o que muda. A config expõe full_attention_interval: 4, e o layout é dezesseis repetições de três camadas Gated DeltaNet seguidas de uma camada de gated attention — 48 mixers que carregam um estado recorrente de tamanho fixo (lição 4.15) e 16 que fazem attention quadrática completa com 24 query heads compartilhando 4 KV heads em head dimension 256 (lições 4.2 e 4.7).

Aqui está o retorno concreto de ler uma árvore de módulos em vez de uma ficha técnica. Percorra as 64 camadas, pergunte a cada uma que tipo de mixer ela abriga, e conte: 16 de attention, 48 de DeltaNet. Essa contagem não é curiosidade — é a entrada da derivação de cache da lição 7.2. Dezesseis camadas guardam keys e values por token; quarenta e oito guardam um estado cujo tamanho não depende de quão longa a conversa fica. Os 64 KiB por token de que todo plano de capacidade desta track e da próxima depende é algo que você pode verificar percorrendo esta árvore, em vez de aceitar por fé.

Agora a virada errada clássica, e é uma boa. Vindo do bloco uniforme da lição 4.11, onde hidden size é igual ao número de heads vezes a head dimension, você espera que toda projeção de attention seja quadrada: 5120 de entrada, 5120 de saída. Imprima os shapes e a projeção de query é 5120 para 12288. Seu primeiro instinto será que o checkpoint está corrompido ou que você leu a config errado. Nem um nem outro. Faça a conta: 24 query heads × 256 de head dimension = 6144, que simplesmente não é 5120. Os projetistas da Qwen desacoplaram a largura da attention da largura do residual, então q_proj alarga para 12288, divide por head em um query state de 6144 e uma porta de 6144, e o_proj estreita esses 6144 de volta para 5120, enquanto k_proj e v_proj produzem 4 × 256 = 1024 cada.

Com uma referência em mãos, comparar engines vira experimento em vez de opinião. Alimente os mesmos ids de token no eager e no engine sob teste, decodifique de forma gulosa e encontre a primeira posição em que o token amostrado difere. Depois pergunte se a divergência veio do tokenizer, do chat template, da configuração de sampling ou dos kernels. O que você não deve exigir é igualdade bit a bit: reduções em bf16 dependem da ordem de acumulação, então qualquer mudança de tiling ou de composição do batch perturba os últimos bits. A régua é concordância nos tokens sob decoding guloso e estatísticas de distribuição compatíveis sob sampling — não floats idênticos.

O eager é deliberadamente a ferramenta errada para serving. Cada operação paga um frame Python e um dispatch, nada funde, memória é alocada por requisição, e não há cache paginado nem continuous batching. O decode com batch 1 passa boa parte da vida esperando esse overhead em vez de esperando aritmética, que é exatamente o problema atacado pela lição 8.3. Guarde a imagem: os parâmetros são um dicionário de tensors, a árvore de módulos são as tracks 1 a 4 tornadas executáveis, e todo stack mais rápido desta track é uma alegação de equivalência com o que este runtime faz de forma lenta e honesta.