Distilação de LLMs mais barata: top-K offline e KL fundida cortam VRAM em até 15x

Distilação de conhecimento, treinar um modelo "aluno" menor para reproduzir o comportamento de um "professor" gigante, virou prática padrão na era dos LLMs abertos. O problema é o custo. Segundo o artigo publicado no blog da Hugging Face pela equipe da Multiverse Computing (Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss), a etapa de distilação é justamente a que mais define a qualidade final e a mais cara de todas. Para quem constrói no Brasil, onde acesso a clusters de GPU é escasso e caro, esse gargalo importa: ele decide se dá ou não para experimentar com modelos comprimidos localmente.

Por que a conta explode

O setup clássico é a distilação online com perda de divergência KL: professor e aluno ficam carregados ao mesmo tempo, e a cada passo o professor faz um forward completo para produzir sua distribuição de saída. O gargalo é a matriz vocabulário × comprimento da sequência, materializada duas vezes por posição de token.

O exemplo da fonte é ilustrativo. O gpt-oss-120b tem vocabulário de 201.088 tokens. Com sequência de 32K e batch 4, só o tensor de probabilidades do professor tem formato 4 × 201.088 × 32.768, o que em bfloat16 já são cerca de 50GB. Somando gradientes, ativações, pesos e estados do otimizador, uma iteração pode picar em torno de 250GB de VRAM, acima do que uma H200 (141GB) ou uma B200 entregam sozinhas.

As duas mudanças

A proposta ataca o problema no nível de sistemas, não de matemática (as perdas continuam equivalentes):

• Distilação offline com top-K. Em vez de recomputar o professor a cada passo, roda-se o professor uma única vez, cacheiam-se os 100 tokens mais prováveis por posição, e o aluno treina contra esse cache. O professor nunca precisa dividir memória com o aluno, e o mesmo cache serve para várias ablações.
• Perda KL fundida e em chunks. A contribuição central. Em vez de construir a grade completa de logits do aluno antes de calcular a perda, o método funde a projeção de saída do modelo dentro do cálculo da perda e processa uma fatia da sequência por vez: projeta os estados ocultos daquele chunk, incorpora o resultado na perda acumulada e descarta antes de seguir. O backward recomputa cada chunk em vez de guardá-lo.

O trade-off é honesto: essa projeção é feita duas vezes (forward e backward), o que custa um pouco de velocidade. Em compensação, a memória de pico cresce linearmente com o comprimento da sequência, em vez de estourar com o produto vocabulário × sequência.

Os números

Em sequências curtas o ganho de memória é modesto e a versão fundida chega a ser um pouco mais lenta. Num benchmark em uma única H200 com Llama 3.1 8B Instruct como professor e um aluno Llama de 3,2B a 8K de contexto, o pico caiu de 102,8GB (online) para 58,3GB (offline fundida), com curvas de perda praticamente idênticas, ou seja, o cache top-100 é lossless frente à distilação online.

A vantagem real aparece no contexto longo. Num benchmark isolado da própria loss, a 32K tokens o pico de memória caiu de 85,2 GiB (dense) para 5,45 GiB (fundida), uma redução de 15,6x, e a versão dense simplesmente falha a partir de 64K. A 256K, a variante fundida usa 11,6 GiB contra 134,2 GiB da segunda melhor.

O impacto prático mais concreto: distilando um GPT-OSS 20B a 32.768 tokens, a memória liberada permitiu sair de quatro nós de GPU para um só. O tempo por passo caiu de 57,0s para 12,23s (~5x mais rápido) e a vazão por GPU subiu de 74,2 para 345,7 TFLOP/s.

O que fica para quem constrói

O aluno resultante, comprimido de 8B para ~3,2B parâmetros, mantém boa parte da acurácia do professor em BoolQ e HellaSwag e fica a cerca de nove pontos em MMLU, com menos da metade dos parâmetros. Nove pontos em MMLU não é pouco e vale testar no seu caso de uso antes de assumir paridade.

A implementação da loss em chunks foi aberta em github.com/CompactifAI/Full-Chunked-KL-Loss. Vale a ressalva de que o texto original é de uma empresa apresentando o próprio método, o que pede leitura crítica. Ainda assim, a ideia de fundir a projeção na perda para evitar materializar a grade completa é um padrão de engenharia reaproveitável em qualquer pipeline de treino sensível a memória, não só distilação, e é isso que a torna interessante além do case de marketing.

Fonte: Hugging Face Blog