ruGPT3XL Gagne un Contexte de 8k : le Modèle Restauré Dépasse la Limite de 2k avec des Pertes Minimales

ruGPT3XL, restauré à partir d'un ancien checkpoint Megatron-LM, a reçu une fenêtre de contexte complète de 8 mille tokens au lieu des 2 mille originaux — et a pratiquement rien perdu en qualité sur les textes courts. En même temps, l'auteur du projet a corrigé une erreur critique dans le mécanisme d'attention qui rendait la version initiale du modèle nettement pire que l'original, bien qu'elle fonctionnait formellement et générât du texte. Le projet a commencé comme une restauration technique d'un ancien modèle en langue russe : les poids de ruGPT3XL ont été convertis au format Hugging Face, une version GGUF a été préparée pour llama.

cpp et des tests ont été exécutés. À ce stade, il s'est avéré que la conversion n'était pas tout à fait correcte. Au lieu d'utiliser l'attention dispersée originale, le modèle utilisait en fait l'attention dense ordinaire de GPT-2, de sorte que la qualité s'effondrait sur les longues séquences.

Cela a été rapidement confirmé par la métrique de perplexité : le premier contrôle a montré PPL 50,1, tandis que le ruGPT3XL original était rapporté à 12,05. Après la restauration de l'attention dispersée alternée, la situation a changé radicalement. PPL est tombé à 11,68, ce qui signifie que le modèle est revenu aux valeurs originales et a commencé à calculer l'attention comme prévu dans l'architecture originale.

En parallèle, le support dans llama.cpp a dû être mis à jour : le patch précédent transférait les poids vers GGUF, mais n'implémentait pas l'attention dispersée elle-même, donc la version locale calculait également selon un schéma dense. L'auteur a en outre corrigé une erreur dans le masque pour les lots plus longs qu'un exemple et a ajouté l'accélération via SDPA, torch.

compile et Triton. Sur une RTX 4090, cela a donné une augmentation de la vitesse d'entraînement d'environ 1,85x par rapport à l'implémentation de base. L'objectif principal de l'étape suivante était pratique : supprimer l'ancienne limite de 2048 tokens qui entravait le travail avec les longs chats et documents.

Mais pour ruGPT3XL, il ne suffit pas de changer simplement un nombre dans la configuration. Le modèle utilise des embeddings de position absolue appris, qui ne peuvent pas extrapoler correctement vers de nouvelles positions sans entraînement supplémentaire, et le schéma d'attention dispersée dépend également de la longueur maximale du contexte. Par conséquent, l'expansion a été effectuée par étapes : d'abord de 2k à 4k, puis de 4k à 8k.

Pour les nouvelles positions, le tiling des embeddings de position a été appliqué pour éviter de casser les courtes séquences déjà apprises, et l'ensemble de données a été mélangé à partir d'exemples longs et courts dans un rapport de 60 à 40. L'entraînement sur l'ensemble de données Gazeta a pris environ 2,6 heures à la première étape et 3,9 heures à la deuxième. Le résultat s'est avéré soigné, pas démonstratif.

Sur la fenêtre originale de 2k, la version finale de 8k a montré PPL 11,77 contre 11,68 dans le modèle de base, ce qui signifie que la régression n'était que de 0,09. À 4k, le chiffre final était 11,99, et à la fenêtre complète de 8k — 13,00, ce qui semble très solide pour une augmentation quatre fois plus grande du contexte. En termes de mémoire, l'expérience s'est également avérée viable : grâce à l'attention dispersée, l'augmentation de la consommation n'a pas été catastrophique, et l'entraînement et l'inférence tenaient sur une RTX 4090 avec 48 GB VRAM.

Lors de l'entraînement sur 8k, un autre problème pratique s'est posé — la fragmentation de la mémoire CUDA — mais il a pu être contourné en configurant expandable_segments, après quoi la consommation maximale a chuté de 46,8 à 38,5 GB. La vitesse de génération diminue bien sûr avec une invite croissante, mais au contexte complet de 8k, le modèle maintient toujours environ 38 tokens par seconde, donc ceci n'est pas seulement un tour de recherche, mais un scénario local entièrement opérationnel. Pour le segment open-source de langue russe, c'est un signal important : même les anciens modèles ne peuvent pas seulement être archivés, mais être amenés à un état de fonctionnement contemporain si on restaure soigneusement les détails architecturaux et qu'on n'épargne pas sur la validation.

Dans le cas de ruGPT3XL, il ne s'agit plus d'une mise à jour cosmétique, mais d'une augmentation réelle d'utilité : le modèle est devenu plus proche de l'original en qualité, a reçu le support des outils populaires et a appris à fonctionner avec un contexte long sans perte grave sur les tâches courtes.