ruGPT3XL Gains 8k Context: Restored Model Transcends 2k Limit with Minimal Losses

ruGPT3XL, восстановленная из старого Megatron-LM чекпоинта, получила полноценное контекстное окно 8 тысяч токенов вместо исходных 2 тысяч — и почти не потеряла в качестве на коротких текстах. Заодно автор проекта исправил ключевую ошибку в механизме внимания, из-за которой ранняя версия модели выглядела заметно хуже оригинала, хотя формально запускалась и генерировала текст. Проект начался как техническая реставрация старой русскоязычной модели: веса ruGPT3XL перевели в формат Hugging Face, подготовили GGUF-версию для llama.

cpp и стали прогонять тесты. Именно на этой стадии выяснилось, что конвертация работает не совсем корректно. Вместо оригинального sparse attention в модели фактически использовался обычный dense attention из GPT-2, поэтому на длинных последовательностях качество резко проседало.

Это быстро подтвердилось метрикой perplexity: первая проверка показала PPL 50.1, тогда как у оригинальной ruGPT3XL заявлено 12.05.

После восстановления alternating sparse attention ситуация изменилась радикально. PPL упал до 11.68, то есть модель снова приблизилась к оригинальным значениям и стала считать внимание так, как задумывалось в исходной архитектуре.

Параллельно пришлось обновить и поддержку в llama.cpp: прежний патч переносил веса в GGUF, но не реализовывал сам sparse attention, поэтому локальная версия тоже считала по плотной схеме. Дополнительно автор исправил ошибку в маске для батчей длиннее одного примера и добавил ускорение через SDPA, torch.

compile и Triton. На RTX 4090 это дало прирост скорости обучения примерно в 1.85 раза относительно базовой eager-реализации.

Главная цель следующего этапа была практической: убрать старое ограничение в 2048 токенов, которое мешало работе с длинными чатами и документами. Но у ruGPT3XL для этого недостаточно просто поменять число в конфиге. Модель использует learned absolute positional embeddings, которые не умеют нормально экстраполироваться на новые позиции без дообучения, а схема sparse attention тоже зависит от максимальной длины контекста.

Поэтому расширение делали поэтапно: сначала с 2k до 4k, затем с 4k до 8k. Для новых позиций применили тайлинг позиционных эмбеддингов, чтобы не ломать уже выученные короткие последовательности, а датасет смешали из длинных и коротких примеров в пропорции 60 на 40. Обучение на датасете Gazeta заняло около 2.

6 часа на первом шаге и 3.9 часа на втором. Результат получился аккуратным, а не демонстрационным.

На исходном окне 2k финальная 8k-версия показала PPL 11.77 против 11.68 у базовой модели, то есть регрессия составила всего 0.

09. На 4k итоговый показатель оказался 11.99, а на полном окне 8k — 13.

00, что выглядит очень достойно для четырёхкратного увеличения контекста. По памяти эксперимент тоже оказался жизнеспособным: благодаря sparse attention рост потребления не превратился в катастрофу, а обучение и инференс уместились на RTX 4090 с 48 ГБ VRAM. Во время обучения на 8k возникла ещё одна практическая проблема — фрагментация памяти CUDA, но её удалось обойти настройкой expandable_segments, после чего пиковое потребление снизилось с 46.

8 до 38.5 ГБ. Скорость генерации, конечно, падает с ростом промпта, но на полном 8k-контексте модель всё равно держит около 38 токенов в секунду, так что речь идёт не только об исследовательском трюке, а о вполне рабочем локальном сценарии.

Для русскоязычного open-source сегмента это важный сигнал: даже старые модели можно не просто архивировать, а доводить до современного рабочего состояния, если аккуратно восстановить архитектурные детали и не экономить на валидации. В случае ruGPT3XL речь уже не о косметическом апдейте, а о реальном повышении полезности: модель стала ближе к оригиналу по качеству, получила поддержку в популярных инструментах и научилась работать с длинным контекстом без серьёзной потери на коротких задачах.