ruGPT3XL con contexto de 8k: atención dispersa

Extendiendo el contexto de ruGPT3XL a 8k tokens: Atención dispersa y ajuste fino

Los desarrolladores adaptaron exitosamente ruGPT3XL de Megatron-LM al formato HuggingFace con soporte GGUF en llama.cpp. Las pruebas revelaron una degradación en la Perplejidad (PPL) a 50.1 en el conjunto de datos gazeta, en lugar del 12.05 esperado. La causa fue el reemplazo de la atención dispersa por la atención densa nn.MultiheadAttention de GPT-2.

Restaurar el mecanismo original de Megatron-LM redujo la PPL a 11.68. La correlación con las métricas declaradas del modelo original alcanzó R=0.93.

Comparación de PPL en gazeta

| Modelo | PPL (gazeta) |

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

| ruGPT3Small | — |

| ruGPT3Medium | 25.2 |

| ruGPT3Large | 16.8 |

| ruGPT3XL (denso)| 50.1 |

| ruGPT3XL (disperso)| 11.68 |

El mecanismo de atención dispersa en ruGPT3XL

La atención dispersa utiliza un patrón alternante del artículo "Generating Long Sequences with Sparse Transformers" (arXiv:1904.10509):

• Capas pares (0,2,4…): atención dispersa por bloques. Cada token ve una ventana local de 128 tokens + bloques globales a intervalos regulares. Diferentes cabezales usan diferentes posiciones globales.
• Capas impares (1,3,5…): atención causal densa estándar.

Esto asegura un crecimiento de memoria casi lineal en lugar de cuadrático. Con un aumento de contexto de 4x, la memoria para la caché KV y las activaciones crece 3-4 veces.

En llama.cpp, se añadió la arquitectura LLM_ARCH_RUGPT3XL (PR #21161) con soporte completo para atención dispersa. Se corrigió un error en la máscara para lotes >1 durante el entrenamiento.

Optimización mediante Triton y SDPA

La implementación mediante matmul + softmax + matmul (modo eager) produce ~6280 tok/s en una RTX 4090.

Aceleración:

• F.scaled_dot_product_attention (SDPA): +40% (~8800 tok/s).
• SDPA con tamaño de lote aumentado (5×2048): otro +25%.
• torch.compile con Inductor: ×1.85 sobre la línea base (~11600 tok/s), genera kernels Triton.

Estrategia para extender el contexto a 8k

Limitaciones arquitectónicas:

• Los Embeddings de Posición Absoluta Aprendidos (nn.Embedding(2048,2048)) no extrapolan.
• La cuadrícula de atención dispersa depende de max_position_embeddings // sparse_block_size.

Principios del ajuste fino

• Mosaico de embeddings posicionales: Duplicación cíclica de 0-2047 para nuevas posiciones. Mantiene estabilidad en contexto corto.
• Conjunto de datos mixto: 60% ejemplos largos (gazeta vía EOS), 40% fragmentos cortos.
• Extensión por pasos: 2k→4k (2.6h), luego 4k→8k (3.9h) en RTX 4090.

Parámetros: entrenamiento en gazeta, 3 épocas, lr=5e-6 decaimiento coseno, bfloat16, verificación de gradiente. OOM resuelto mediante PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (38.5 GB pico).

Modelo: evilfreelancer/ruGPT3XL-8k.

Resultados de PPL en prueba gazeta

| Modelo | 2048 | 4096 | 8192 |

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

| ruGPT3XL base | 11.68| — | — |

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

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

Regresión mínima en 2k, el modelo conserva calidad.

Conclusiones clave

• La atención dispersa es crítica: la implementación densa produce PPL 50.1 en lugar de 11.68.
• El mosaico APE es efectivo: permite ajuste fino sin degradar el contexto base.
• 8k es factible en RTX 4090: 38.5 GB con expandable_segments.
• El enfoque por pasos minimiza el sobreajuste: 2k→4k→8k en 6.5 horas.
• Triton/SDPA acelera 1.85x: torch.compile genera kernels optimizados.

— Editorial Team