ruGPT3XL 带 8k 上下文：稀疏注意力

将ruGPT3XL上下文扩展至8k令牌：稀疏注意力机制与微调策略

开发者成功将ruGPT3XL从Megatron-LM格式适配至HuggingFace格式，并支持llama.cpp中的GGUF。测试显示，在gazeta数据集上，困惑度（PPL）下降至50.1，而非预期的12.05。原因是稀疏注意力被GPT-2的密集nn.MultiheadAttention所替代。

恢复Megatron-LM的原始机制后，PPL降至11.68。与原始模型声称指标的相关系数达到R=0.93。

gazeta数据集上的PPL对比

| 模型 | PPL (gazeta) |

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

| ruGPT3Small | — |

| ruGPT3Medium | 25.2 |

| ruGPT3Large | 16.8 |

| ruGPT3XL (密集) | 50.1 |

| ruGPT3XL (稀疏) | 11.68 |

ruGPT3XL中的稀疏注意力机制

稀疏注意力采用论文《Generating Long Sequences with Sparse Transformers》（arXiv:1904.10509）中的交替模式：

• 偶数层（0,2,4…）：块稀疏注意力。每个令牌关注128个令牌的局部窗口加上规则间隔的全局块。不同头使用不同的全局位置。
• 奇数层（1,3,5…）：标准因果密集注意力。

这确保了近乎线性的内存增长，而非二次增长。上下文长度增加4倍时，KV缓存和激活的内存增长3-4倍。

在llama.cpp中，添加了LLM_ARCH_RUGPT3XL架构（PR #21161），支持完整的稀疏注意力。修复了训练时批次大小>1的掩码错误。

通过Triton和SDPA优化

通过matmul + softmax + matmul（eager模式）实现，在RTX 4090上达到约6280令牌/秒。

加速效果：

• F.scaled_dot_product_attention（SDPA）：+40%（约8800令牌/秒）。
• SDPA增加批次大小（5×2048）：再+25%。
• torch.compile与Inductor：比基线快1.85倍（约11600令牌/秒），生成Triton内核。

扩展上下文至8k的策略

架构限制：

• 学习的绝对位置嵌入（nn.Embedding(2048,2048)）无法外推。
• 稀疏注意力网格依赖于max_position_embeddings // sparse_block_size。

微调原则

• 平铺位置嵌入：将0-2047循环复制到新位置。保持短上下文稳定性。
• 混合数据集：60%长示例（通过EOS的gazeta），40%短片段。
• 逐步扩展：2k→4k（2.6小时），然后4k→8k（3.9小时），在RTX 4090上完成。

参数：gazeta训练集，3个周期，lr=5e-6余弦衰减，bfloat16，梯度检查点。通过PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True解决OOM（峰值38.5 GB）。

模型：evilfreelancer/ruGPT3XL-8k。

gazeta测试集上的PPL结果

| 模型 | 2048 | 4096 | 8192 |

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

| ruGPT3XL基础版 | 11.68| — | — |

| ruGPT3XL-4k | 11.75| 12.04| — |

| ruGPT3XL-8k | 11.77| 11.99| 13.00 |

在2k处回归最小，模型保持质量。

关键要点

• 稀疏注意力至关重要：密集实现导致PPL为50.1，而非11.68。
• 平铺绝对位置嵌入有效：支持微调而不降低基础上下文质量。
• RTX 4090上实现8k可行：使用expandable_segments时内存为38.5 GB。
• 逐步方法最小化过拟合：2k→4k→8k耗时6.5小时。
• Triton/SDPA加速1.85倍：torch.compile生成优化内核。

— Editorial Team