ruGPT3XL alcanza un contexto de 8k: el modelo restaurado supera el límite de 2k casi sin pérdidas

ruGPT3XL, restaurado a partir de un antiguo checkpoint de Megatron-LM, ha obtenido una ventana de contexto completa de 8.000 tokens en lugar de los 2.000 originales, y casi no ha perdido calidad en textos cortos. De paso, el autor del proyecto corrigió un error clave en el mecanismo de atención por el que la versión temprana del modelo parecía bastante peor que el original, aunque formalmente arrancaba y generaba texto.

El proyecto empezó como una restauración técnica de un antiguo modelo en ruso: los pesos de ruGPT3XL se convirtieron al formato Hugging Face, se preparó una versión GGUF para llama.cpp y se pusieron en marcha las pruebas. Fue en esa fase cuando se descubrió que la conversión no funcionaba del todo bien. En lugar del sparse attention original, el modelo usaba en realidad el dense attention habitual de GPT-2, por lo que la calidad caía en picado en secuencias largas. La métrica de perplexity lo confirmó enseguida: la primera comprobación dio una PPL de 50.1, mientras que para la ruGPT3XL original se declara 12.05.

Tras restaurar el alternating sparse attention, la situación cambió de forma radical. La PPL bajó a 11.68, es decir, el modelo volvió a acercarse a los valores originales y a calcular la atención como preveía la arquitectura de partida. En paralelo hubo que actualizar también el soporte en llama.cpp: el parche anterior trasladaba los pesos a GGUF, pero no implementaba el sparse attention en sí, así que la versión local también calculaba con el esquema denso.

Además, el autor corrigió un error en la máscara para lotes con más de un ejemplo y añadió aceleración mediante SDPA, torch.compile y Triton. En una RTX 4090 esto aportó un aumento de velocidad de entrenamiento de aproximadamente 1.85 veces respecto a la implementación eager base.

El objetivo principal de la siguiente etapa era práctico: eliminar el viejo límite de 2048 tokens, que dificultaba trabajar con chats y documentos largos. Pero en ruGPT3XL no basta con cambiar un número en la configuración. El modelo usa learned absolute positional embeddings, que no extrapolan bien a posiciones nuevas sin reentrenamiento, y el esquema de sparse attention también depende de la longitud máxima del contexto.

Por eso la ampliación se hizo por etapas: primero de 2k a 4k y después de 4k a 8k. Para las posiciones nuevas se aplicó tiling de los embeddings posicionales, con el fin de no romper las secuencias cortas ya aprendidas, y el dataset se mezcló con ejemplos largos y cortos en proporción 60 a 40. El entrenamiento con el dataset Gazeta duró unas 2.6 horas en el primer paso y 3.9 horas en el segundo.

El resultado es pulcro, más que vistoso. En la ventana original de 2k, la versión final de 8k obtuvo una PPL de 11.77 frente a 11.68 del modelo base: una regresión de solo 0.09. A 4k el valor final fue 11.99 y en la ventana completa de 8k, 13.00, algo muy digno para multiplicar por cuatro el contexto.

En memoria el experimento también resultó viable: gracias al sparse attention, el aumento del consumo no se convirtió en una catástrofe, y tanto el entrenamiento como la inferencia cupieron en una RTX 4090 con 48 GB de VRAM. Durante el entrenamiento a 8k surgió otro problema práctico, la fragmentación de la memoria de CUDA, pero se sorteó configurando expandable_segments, tras lo cual el consumo máximo bajó de 46.8 a 38.5 GB.

La velocidad de generación cae, como es lógico, al crecer el prompt, pero con el contexto completo de 8k el modelo sigue manteniendo unos 38 tokens por segundo, así que no se trata solo de un truco de investigación, sino de un escenario local perfectamente utilizable.

Para el segmento open source en ruso es una señal importante: incluso los modelos antiguos pueden no limitarse a archivarse, sino llevarse a un estado de trabajo moderno si se restauran con cuidado los detalles arquitectónicos y no se escatima en validación. En el caso de ruGPT3XL ya no hablamos de una actualización cosmética, sino de un aumento real de utilidad: el modelo se acerca más al original en calidad, ha ganado soporte en herramientas populares y ha aprendido a trabajar con contexto largo sin pérdidas serias en tareas cortas.