ruGPT3XL Ganha Contexto de 8k: Modelo Restaurado Ultrapassa Limite de 2k com Mínimas Perdas

ruGPT3XL, restaurado de um antigo checkpoint do Megatron-LM, recebeu uma janela de contexto completa de 8 mil tokens em vez dos 2 mil originais — e praticamente não perdeu em qualidade em textos curtos. Ao mesmo tempo, o autor do projeto corrigiu um erro crítico no mecanismo de atenção que fazia a versão inicial do modelo parecer notavelmente pior que o original, embora formalmente rodasse e gerasse texto. O projeto começou como uma restauração técnica de um antigo modelo em língua russa: os pesos do ruGPT3XL foram convertidos para o formato Hugging Face, uma versão GGUF foi preparada para llama.

cpp e testes foram executados. Neste estágio, descobriu-se que a conversão não era completamente correta. Em vez de usar a atenção esparsa original, o modelo estava na verdade usando atenção densa comum do GPT-2, então a qualidade caiu drasticamente em sequências longas.

Isso foi rapidamente confirmado pela métrica de perplexidade: a primeira verificação mostrou PPL 50,1, enquanto o ruGPT3XL original tinha 12,05. Após restaurar a atenção esparsa alternada, a situação mudou dramaticamente. PPL caiu para 11,68, o que significa que o modelo retornou aos valores originais e começou a calcular atenção conforme pretendido na arquitetura original.

Em paralelo, o suporte no llama.cpp teve que ser atualizado: o patch anterior transferia pesos para GGUF, mas não implementava a própria atenção esparsa, então a versão local também estava calculando usando esquema denso. O autor corrigiu adicionalmente um erro na máscara para lotes maiores que um exemplo e adicionou aceleração através de SDPA, torch.

compile e Triton. Em uma RTX 4090, isso deu um aumento de velocidade de treinamento de aproximadamente 1,85x em relação à implementação base. O objetivo principal da próxima etapa era prático: remover o antigo limite de 2048 tokens que dificultava o trabalho com chats longos e documentos.

Mas para ruGPT3XL isso não é suficiente simplesmente mudar um número na configuração. O modelo usa embeddings de posição absoluta aprendidos, que não conseguem extrapolar adequadamente para novas posições sem treinamento adicional, e o esquema de atenção esparsa também depende do comprimento máximo do contexto. Portanto, a expansão foi feita em etapas: primeiro de 2k para 4k, depois de 4k para 8k.

Para novas posições, o tiling de embeddings posicionais foi aplicado para evitar quebrar sequências curtas já aprendidas, e o conjunto de dados foi misturado de exemplos longos e curtos na proporção de 60 para 40. O treinamento no conjunto de dados Gazeta levou cerca de 2,6 horas na primeira etapa e 3,9 horas na segunda. O resultado se mostrou cuidadoso, não demonstrativo.

Na janela original de 2k, a versão final de 8k mostrou PPL 11,77 versus 11,68 no modelo base, o que significa que a regressão foi apenas 0,09. Em 4k, o valor final foi 11,99 e na janela completa de 8k — 13,00, o que parece muito sólido para um aumento quatro vezes maior de contexto. Em termos de memória, o experimento também se mostrou viável: graças à atenção esparsa, o aumento no consumo não se tornou catastrófico, e tanto treinamento quanto inferência se ajustaram em uma RTX 4090 com 48 GB VRAM.

Durante o treinamento em 8k, surgiu outro problema prático — fragmentação de memória CUDA — mas foi contornado definindo expandable_segments, após o qual o consumo de pico caiu de 46,8 para 38,5 GB. A velocidade de geração certamente diminui com um prompt crescente, mas no contexto completo de 8k o modelo ainda mantém cerca de 38 tokens por segundo, então isso não é apenas um truque de pesquisa, mas um cenário local completamente operacional. Para o segmento de código aberto em língua russa, este é um sinal importante: até mesmo modelos antigos podem não apenas ser arquivados, mas trazidos a condições de trabalho contemporâneas se você restaurar cuidadosamente detalhes arquitetônicos e não economizar na validação.

No caso do ruGPT3XL, este não é mais uma atualização cosmética, mas um aumento real na utilidade: o modelo ficou mais próximo do original em qualidade, recebeu suporte em ferramentas populares e aprendeu a trabalhar com contexto longo sem perda séria em tarefas curtas.