ruGPT3XL avec contexte 8k : sparse attention

Extension du contexte de ruGPT3XL à 8 000 tokens : Attention parcimonieuse et affinage

Les développeurs ont adapté avec succès ruGPT3XL de Megatron-LM au format HuggingFace avec prise en charge GGUF dans llama.cpp. Les tests ont révélé une dégradation de la Perplexité (PPL) à 50,1 sur le jeu de données gazeta au lieu des 12,05 attendus. La cause était le remplacement de l'attention parcimonieuse par nn.MultiheadAttention dense de GPT-2.

La restauration du mécanisme original de Megatron-LM a réduit la PPL à 11,68. La corrélation avec les métriques annoncées du modèle original a atteint R=0,93.

Comparaison de la PPL sur gazeta

| Modèle | PPL (gazeta) |

|----------------|--------------|

| ruGPT3Small | — |

| ruGPT3Medium | 25,2 |

| ruGPT3Large | 16,8 |

| ruGPT3XL (dense)| 50,1 |

| ruGPT3XL (parcimonieux)| 11,68 |

Le mécanisme d'attention parcimonieuse dans ruGPT3XL

L'attention parcimonieuse utilise un motif alterné issu de l'article "Generating Long Sequences with Sparse Transformers" (arXiv:1904.10509) :

• Couches paires (0,2,4…) : attention parcimonieuse par blocs. Chaque token voit une fenêtre locale de 128 tokens + des blocs globaux à intervalles réguliers. Différentes têtes utilisent différentes positions globales.
• Couches impaires (1,3,5…) : attention causale dense standard.

Cela garantit une croissance de la mémoire quasi linéaire au lieu de quadratique. Avec une augmentation de contexte de 4x, la mémoire pour le cache KV et les activations croît de 3 à 4 fois.

Dans llama.cpp, l'architecture LLM_ARCH_RUGPT3XL (PR #21161) a été ajoutée avec prise en charge complète de l'attention parcimonieuse. Une erreur de masque pour les lots >1 pendant l'entraînement a été corrigée.

Optimisation via Triton et SDPA

L'implémentation via matmul + softmax + matmul (mode eager) donne ~6280 tok/s sur une RTX 4090.

Accélération :

• F.scaled_dot_product_attention (SDPA) : +40 % (~8800 tok/s).
• SDPA avec taille de lot augmentée (5×2048) : encore +25 %.
• torch.compile avec Inductor : ×1,85 par rapport à la base (~11600 tok/s), génère des noyaux Triton.

Stratégie pour étendre le contexte à 8k

Limitations architecturales :

• Les plongements positionnels absolus appris (nn.Embedding(2048,2048)) n'extrapolent pas.
• La grille d'attention parcimonieuse dépend de max_position_embeddings // sparse_block_size.

Principes d'affinage

• Carrelage des plongements positionnels : Duplication cyclique de 0-2047 pour les nouvelles positions. Maintient la stabilité sur un contexte court.
• Jeu de données mixte : 60 % d'exemples longs (gazeta via EOS), 40 % de segments courts.
• Extension progressive : 2k→4k (2,6 h), puis 4k→8k (3,9 h) sur RTX 4090.

Paramètres : entraînement sur gazeta, 3 époques, lr=5e-6 décroissance cosinus, bfloat16, vérification de gradient. OOM résolu via PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (38,5 Go de pic).

Modèle : evilfreelancer/ruGPT3XL-8k.

Résultats de PPL sur le test gazeta

| Modèle | 2048 | 4096 | 8192 |

|----------------|------|------|------|

| ruGPT3XL de base | 11,68| — | — |

| ruGPT3XL-4k | 11,75| 12,04| — |

| ruGPT3XL-8k | 11,77| 11,99| 13,00 |

Régression minimale à 2k, le modèle conserve sa qualité.

Points clés à retenir

• L'attention parcimonieuse est cruciale : l'implémentation dense donne une PPL de 50,1 au lieu de 11,68.
• Le carrelage APE est efficace : permet l'affinage sans dégrader le contexte de base.
• 8k est réalisable sur RTX 4090 : 38,5 Go avec expandable_segments.
• L'approche progressive minimise le surapprentissage : 2k→4k→8k en 6,5 heures.
• Triton/SDPA accélère de 1,85x : torch.compile génère des noyaux optimisés.

— Editorial Team