ruGPT3XL gets an 8k context window: restored model breaks past 2k with almost no loss

ruGPT3XL, restored from an old Megatron-LM checkpoint, now has a full 8,000-token context window instead of the original 2,000 tokens, and it lost almost nothing in quality on short texts. Along the way, the project's author fixed a key bug in the attention mechanism that made the early version of the model look noticeably worse than the original, even though it technically ran and generated text.

The project began as a technical restoration of an old Russian-language model: the ruGPT3XL weights were converted to Hugging Face format, a GGUF version was prepared for llama.cpp, and testing began. It was at this stage that the conversion turned out not to work quite correctly. Instead of the original sparse attention, the model was in fact using ordinary dense attention from GPT-2, so quality dropped sharply on long sequences. The perplexity metric quickly confirmed it: the first check showed a PPL of 50.1, whereas the original ruGPT3XL is stated to have 12.05.

After alternating sparse attention was restored, the situation changed radically. PPL fell to 11.68, meaning the model was once again close to the original values and computed attention the way the source architecture intended. In parallel, llama.cpp support had to be updated too: the previous patch carried the weights over to GGUF but did not implement sparse attention itself, so the local version also computed attention densely.

The author also fixed a bug in the mask for batches longer than one example and added acceleration via SDPA, torch.compile and Triton. On an RTX 4090 this gave a training speedup of roughly 1.85x over the baseline eager implementation.

The main goal of the next stage was practical: remove the old 2048-token limit, which got in the way of long chats and documents. But for ruGPT3XL, simply changing a number in the config is not enough. The model uses learned absolute positional embeddings, which cannot extrapolate properly to new positions without further training, and the sparse attention scheme also depends on the maximum context length.

The extension was therefore done in stages: first from 2k to 4k, then from 4k to 8k. For the new positions, the author tiled the positional embeddings so as not to break the short sequences the model had already learned, and mixed the dataset from long and short examples in a 60-to-40 ratio. Training on the Gazeta dataset took about 2.6 hours for the first step and 3.9 hours for the second.

The result is tidy rather than showy. On the original 2k window, the final 8k version scored a PPL of 11.77 against 11.68 for the base model, a regression of just 0.09. At 4k the final figure was 11.99, and at the full 8k window it was 13.00, which looks very respectable for a fourfold increase in context.

Memory-wise, the experiment also proved viable: thanks to sparse attention, the growth in consumption did not turn into a catastrophe, and both training and inference fit on an RTX 4090 with 48 GB of VRAM. During 8k training another practical problem arose, CUDA memory fragmentation, but it was worked around by configuring expandable_segments, after which peak consumption dropped from 46.8 to 38.5 GB.

Generation speed does fall as the prompt grows, but at the full 8k context the model still holds about 38 tokens per second, so this is not just a research trick but a fully workable local scenario.

For the Russian-language open-source segment, this is an important signal: even old models need not simply be archived, and can be brought up to a modern working state if the architectural details are restored carefully and validation is not skimped on. In the case of ruGPT3XL, this is no longer a cosmetic update but a real gain in usefulness: the model is closer to the original in quality, gained support in popular tools, and learned to work with long context without serious loss on short tasks.