ruGPT3XL Gana Contexto de 8k: Modelo Restaurado Supera Límite de 2k con Pérdidas Mínimas

ruGPT3XL, restaurado de un antiguo checkpoint de Megatron-LM, recibió una ventana de contexto completa de 8 mil tokens en lugar de los 2 mil originales — y prácticamente no perdió en calidad en textos cortos. Al mismo tiempo, el autor del proyecto corrigió un error crítico en el mecanismo de atención que hacía que la versión inicial del modelo se viera notablemente peor que el original, aunque formalmente funcionara y generara texto. El proyecto comenzó como una restauración técnica de un antiguo modelo en lengua rusa: los pesos de ruGPT3XL se convirtieron al formato de Hugging Face, se preparó una versión GGUF para llama.

cpp y se ejecutaron pruebas. En esta etapa, se descubrió que la conversión no era del todo correcta. En lugar de usar la atención dispersa original, el modelo estaba utilizando atención densa común de GPT-2, por lo que la calidad disminuyó drásticamente en secuencias largas.

Esto fue confirmado rápidamente por la métrica de perplejidad: la primera verificación mostró PPL 50,1, mientras que el ruGPT3XL original tenía 12,05. Después de restaurar la atención dispersa alternada, la situación cambió dramáticamente. PPL cayó a 11,68, lo que significa que el modelo volvió a los valores originales y comenzó a calcular la atención como se pretendía en la arquitectura original.

En paralelo, el soporte en llama.cpp tuvo que ser actualizado: el parche anterior transfería pesos a GGUF, pero no implementaba la propia atención dispersa, por lo que la versión local también estaba calculando usando esquema denso. El autor corrigió adicionalmente un error en la máscara para lotes más largos que un ejemplo y agregó aceleración a través de SDPA, torch.

compile y Triton. En una RTX 4090, esto dio un aumento de velocidad de entrenamiento de aproximadamente 1,85x en relación con la implementación base. El objetivo principal de la siguiente etapa era práctico: eliminar el antiguo límite de 2048 tokens que dificultaba el trabajo con chats largos y documentos.

Pero para ruGPT3XL esto no es suficiente simplemente cambiar un número en la configuración. El modelo utiliza incrustaciones de posición absoluta aprendidas, que no pueden extrapolar adecuadamente a nuevas posiciones sin entrenamiento adicional, y el esquema de atención dispersa también depende de la longitud máxima del contexto. Por lo tanto, la expansión se realizó en etapas: primero de 2k a 4k, luego de 4k a 8k.

Para nuevas posiciones, se aplicó mosaico de incrustaciones posicionales para evitar romper secuencias cortas ya aprendidas, y el conjunto de datos se mezcló con ejemplos largos y cortos en una proporción de 60 a 40. El entrenamiento en el conjunto de datos Gazeta tomó alrededor de 2,6 horas en la primera etapa y 3,9 horas en la segunda. El resultado resultó cuidadoso, no demostrativo.

En la ventana original de 2k, la versión final de 8k mostró PPL 11,77 versus 11,68 en el modelo base, lo que significa que la regresión fue solo 0,09. En 4k, la cifra final fue 11,99 y en la ventana completa de 8k — 13,00, lo que se ve muy sólido para un aumento cuatro veces mayor de contexto. En términos de memoria, el experimento también resultó viable: gracias a la atención dispersa, el aumento en el consumo no se convirtió en catastrófico, y tanto el entrenamiento como la inferencia se ajustaron en una RTX 4090 con 48 GB VRAM.

Durante el entrenamiento en 8k, surgió otro problema práctico — fragmentación de memoria CUDA — pero se logró resolver configurando expandable_segments, después de lo cual el consumo máximo disminuyó de 46,8 a 38,5 GB. La velocidad de generación ciertamente disminuye con un prompt creciente, pero en el contexto completo de 8k el modelo aún mantiene alrededor de 38 tokens por segundo, así que esto no es solo un truco de investigación, sino un escenario local completamente operacional. Para el segmento de código abierto de habla rusa, esta es una señal importante: incluso los modelos antiguos pueden no solo ser archivados, sino llevados a condiciones de trabajo contemporáneas si se restauran cuidadosamente los detalles arquitectónicos y no se escatima en validación.

En el caso de ruGPT3XL, esto ya no es una actualización cosmética, sino un aumento real en la utilidad: el modelo se volvió más cercano al original en calidad, recibió soporte en herramientas populares y aprendió a trabajar con contexto largo sin pérdida seria en tareas cortas.