Erweiterung des ruGPT3XL-Kontexts auf 8k Tokens: Sparse Attention und Fine-Tuning
Entwickler haben ruGPT3XL von Megatron-LM erfolgreich in das HuggingFace-Format mit GGUF-Unterstützung in llama.cpp übertragen. Tests zeigten eine Verschlechterung der Perplexity (PPL) auf 50,1 im gazeta-Datensatz statt der erwarteten 12,05. Die Ursache war der Ersatz der sparse attention durch die dichte nn.MultiheadAttention aus GPT-2.
Die Wiederherstellung des ursprünglichen Mechanismus aus Megatron-LM reduzierte die PPL auf 11,68. Die Korrelation mit den angegebenen Metriken des Originalmodells erreichte R=0,93.
PPL-Vergleich auf gazeta
| Modell | PPL (gazeta) |
|----------------|--------------|
| ruGPT3Small | — |
| ruGPT3Medium | 25,2 |
| ruGPT3Large | 16,8 |
| ruGPT3XL (dicht)| 50,1 |
| ruGPT3XL (sparse)| 11,68 |
Der Sparse-Attention-Mechanismus in ruGPT3XL
Sparse attention verwendet ein abwechselndes Muster aus dem Paper "Generating Long Sequences with Sparse Transformers" (arXiv:1904.10509):
- Gerade Schichten (0,2,4…): Block-sparse attention. Jeder Token sieht ein lokales Fenster von 128 Tokens + globale Blöcke in regelmäßigen Abständen. Verschiedene Heads nutzen unterschiedliche globale Positionen.
- Ungerade Schichten (1,3,5…): Standard kausale dichte Attention.
Dies gewährleistet nahezu lineares Speicherwachstum statt quadratischem. Bei einer 4-fachen Kontexterhöhung wächst der Speicher für KV-Cache und Aktivierungen um das 3-4-fache.
In llama.cpp wurde die Architektur LLM_ARCH_RUGPT3XL (PR #21161) mit voller sparse-attention-Unterstützung hinzugefügt. Ein Maskenfehler für Batches >1 während des Trainings wurde behoben.
Optimierung via Triton und SDPA
Implementierung via matmul + softmax + matmul (eager mode) liefert ~6280 Tok/s auf einer RTX 4090.
Beschleunigung:
F.scaled_dot_product_attention(SDPA): +40% (~8800 Tok/s).- SDPA mit erhöhter Batch-Größe (5×2048): weitere +25%.
torch.compilemit Inductor: ×1,85 über Baseline (~11600 Tok/s), generiert Triton-Kernel.
Strategie zur Kontexterweiterung auf 8k
Architektonische Einschränkungen:
- Gelernte Absolute Positional Embeddings (nn.Embedding(2048,2048)) extrapolieren nicht.
- Das sparse-attention-Gitter hängt von
max_position_embeddings // sparse_block_sizeab.
Fine-Tuning-Prinzipien
- Kachelung der Positions-Embeddings: Zyklische Duplizierung von 0-2047 für neue Positionen. Erhält Stabilität bei kurzem Kontext.
- Gemischter Datensatz: 60% lange Beispiele (gazeta via EOS), 40% kurze Chunks.
- Schrittweise Erweiterung: 2k→4k (2,6h), dann 4k→8k (3,9h) auf RTX 4090.
Parameter: gazeta train, 3 Epochen, lr=5e-6 Kosinusabfall, bfloat16, Gradient Checkpointing. OOM gelöst via PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (38,5 GB Spitze).
Modell: evilfreelancer/ruGPT3XL-8k.
PPL-Ergebnisse auf gazeta test
| Modell | 2048 | 4096 | 8192 |
|----------------|------|------|------|
| ruGPT3XL Basis | 11,68| — | — |
| ruGPT3XL-4k | 11,75| 12,04| — |
| ruGPT3XL-8k | 11,77| 11,99| 13,00 |
Minimale Regression bei 2k, Modell behält Qualität.
Wichtige Erkenntnisse
- Sparse attention ist entscheidend: Dichte Implementierung liefert PPL 50,1 statt 11,68.
- Kachelung von APE ist effektiv: Ermöglicht Fine-Tuning ohne Qualitätseinbußen im Basis-Kontext.
- 8k ist auf RTX 4090 machbar: 38,5 GB mit expandable_segments.
- Schrittweiser Ansatz minimiert Overfitting: 2k→4k→8k in 6,5 Stunden.
- Triton/SDPA beschleunigen um 1,85x: torch.compile generiert optimierte Kernel.
— Editorial Team
Noch keine Kommentare.