Увеличение контекста ruGPT3XL до 8k: методика, PPL и Sparse Attention

Когда 2048 токенов — это просто смешно

Вы загружаете ruGPT3XL, даете ей техническое задание на пять страниц, а она забывает, о чем речь в первом абзаце. Знакомо? Оригинальная модель застряла в 2023 году с контекстом в 2048 токенов. В 2026 году это выглядит как попытка прочитать «Войну и мир» через замочную скважину.

Просто увеличить параметр max_position_embeddings в конфиге — самый верный способ получить модель, которая на длинных текстах будет генерировать бессвязный бред. Перплексия взлетит до небес. Почему? Потому что модель никогда не видела такие длинные последовательности во время предобучения. Ее позиционные эмбеддинги и паттерны внимания заточены под короткие окна.

Зачем нам Sparse Attention, если есть Flash Attention 4?

Flash Attention 4 — отличная вещь. Она ускоряет вычисление полного (dense) внимания. Но ее сложность все равно O(n²). Увеличиваем контекст с 2k до 8k — и требуемая память вырастет в 16 раз. На практике для ruGPT3XL (1.3B параметров) это означает, что даже на A100 80GB впадаешь в OOM еще на этапе forward pass.

Sparse Attention — это не просто оптимизация. Это принципиально иная архитектура, которая заставляет модель смотреть не на все токены, а только на избранных «соседей». Сложность становится линейной или субквадратичной. Именно это позволяет работать с окнами в 8k, 32k и даже больше без взрывного роста вычислений. Если хотите глубоко в теорию, у нас есть разбор субквадратичного внимания для 10M токенов.

1 Подготовка: берем модель, пачкаем руки

Сначала вытаскиваем модель. Не сгодится просто from_pretrained. Нам нужны веса и конфиг отдельно, чтобы их ковырять.

pip install torch==2.5.1 transformers==5.1.0 triton==3.0.0 --extra-index-url https://download.pytorch.org/whl/cu124

from transformers import GPT2LMHeadModel, GPT2Config
import torch

# Загружаем оригинальную ruGPT3XL
model_name = "sberbank-ai/rugpt3xl"
config = GPT2Config.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

print(f"Исходный контекст: {config.n_positions}")  # Должно быть 2048

2 Ломаем — не строим: меняем позиционные эмбеддинги

Позиционные эмбеддинги в GPT-2 архитектуре — это просто lookup-таблица. Чтобы модель понимала позиции до 8192, нужно расширить эту таблицу. Самый простой способ — интерполяция.

new_max_pos = 8192
orig_emb = model.transformer.wpe.weight.data

# Линейная интерполяция старых эмбеддингов до новых размеров
# Используем torch.nn.functional.interpolate
orig_emb = orig_emb.unsqueeze(0).transpose(1, 2)  # [1, dim, old_seq]
new_emb = torch.nn.functional.interpolate(orig_emb, size=new_max_pos, mode='linear')
new_emb = new_emb.transpose(1, 2).squeeze(0)  # [new_seq, dim]

# Заменяем в модели
model.transformer.wpe = torch.nn.Embedding(new_max_pos, config.n_embd)
model.transformer.wpe.weight.data = new_emb
config.n_positions = new_max_pos

Но линейная интерполяция — это костыль. Для качественного результата нужно дообучение. Об этом позже.

3 Ядро хака: пишем Sliding Window Attention на Triton

Здесь начинается магия. Мы заменим стандартный dense attention на sliding window. Окно размером 512 токенов: каждый токен видит только 512 предыдущих. Сложность падает с O(n²) до O(n*w), где w — размер окна.

Писать свое ядро внимания на Triton с нуля — боль. Поэтому возьмем за основу оптимизированное ядро из открытых репозиториев и адаптируем. (Кстати, если хотите понять, как такие хаки ломают фундамент LLM, посмотрите статью про DroPE).

import triton
import triton.language as tl

@triton.jit
def sliding_window_attention_kernel(
    Q, K, V, O,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    ...  # Параметры стадий и окон
):
    pid = tl.program_id(0)
    # Логика sliding window: ограничиваем диапазон ключей для каждого запроса
    window_start = tl.maximum(0, q_index - window_size)
    window_end = q_index + 1
    # Далее стандартные вычисления внимания, но только в окне
    ...

# Обертка для вызова ядра
def sliding_window_attention(q, k, v, window_size=512):
    """Вычисляет внимание со скользящим окном."""
    # Проверяем формы, выделяем выходной тензор
    o = torch.empty_like(q)
    # Запускаем ядро Triton
    grid = (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1])
    sliding_window_attention_kernel[grid](q, k, v, o, ...)
    return o

Полный код ядра занимает 150+ строк. Его суть — ограничить матрицу внимания диагональной полосой шириной window_size. Это та же идея, что и в Sliding Window из арсенала Karpathy, но реализованная на уровне CUDA ядра для максимальной скорости.

4 Интеграция: встраиваем Sparse Attention в ruGPT3XL

Теперь нужно заменить forward pass в слоях внимания модели. Модифицируем класс GPT2Attention.

from transformers.models.gpt2.modeling_gpt2 import GPT2Attention
import torch.nn as nn

class SparseGPT2Attention(GPT2Attention):
    def __init__(self, config, window_size=512):
        super().__init__(config)
        self.window_size = window_size

    def _attn(self, q, k, v, attention_mask=None, head_mask=None):
        # Вместо стандартного matmul используем наше ядро
        attn_weights = sliding_window_attention(q, k, v, self.window_size)
        
        # Применяем маску если есть (важно для padding)
        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask
        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
        
        # Применяем dropout во время обучения
        attn_weights = self.attn_dropout(attn_weights)
        
        # Умножаем на значения
        attn_output = torch.matmul(attn_weights, v)
        return attn_output, attn_weights

# Заменяем внимание в модели
for i, layer in enumerate(model.transformer.h):
    old_attn = layer.attn
    new_attn = SparseGPT2Attention(config, window_size=512)
    # Копируем веса (Q, K, V проекции)
    new_attn.c_attn.weight.data = old_attn.c_attn.weight.data
    new_attn.c_attn.bias.data = old_attn.c_attn.bias.data
    # Копируем выходную проекцию
    new_attn.c_proj.weight.data = old_attn.c_proj.weight.data
    new_attn.c_proj.bias.data = old_attn.c_proj.bias.data
    layer.attn = new_attn

5 Дообучение: учим модель жить в большом мире

Теперь модель технически может обрабатывать 8k токенов, но качественно она это делать не умеет. Нужно дообучение на длинных последовательностях. Берем русскоязычный корпус (например, смесь Wikipedia, новостей и книг) и нарезаем на примеры по 8192 токена.

Важный момент — learning rate. Нельзя использовать тот же LR, что и для полного дообучения. Мы лишь адаптируем позиционные эмбеддинги и немного подстраиваем веса внимания под новую схему.

from transformers import Trainer, TrainingArguments

# Псевдоданные для примера
train_dataset = ...  # Датасет с последовательностями длиной 8192

# Используем последовательный параллелизм для обработки длинных контекстов
# Об этом подробно в статье про Ulysses Sequence Parallelism.

training_args = TrainingArguments(
    output_dir='./rugpt3xl-8k',
    per_device_train_batch_size=1,  # Один длинный пример на GPU
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=5e-5,  # Очень маленький LR!
    warmup_steps=500,
    logging_dir='./logs',
    fp16=True,  # Используем половинную точность
    gradient_checkpointing=True,  # Экономия памяти
    dataloader_num_workers=4,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

Дообучение на одной A100 80GB займет примерно 2-3 дня. Если нет своего железа, можно взять в аренду инстанс с GPU в Lambda Labs (спонсированная ссылка) — их цены на H100 в 2026 году самые адекватные на рынке.

Измеряем ущерб: PPL как главный судья

После дообучения нельзя просто сказать «работает». Нужно измерить перплексию (PPL) на отложенном наборе данных. PPL показывает, насколько модель «удивлена» реальным данным. Чем ниже, тем лучше.

Сравниваем три модели:

Видите разницу? Линейная интерполяция без Sparse Attention приводит к краху PPL на длинных контекстах. Наша методика дает небольшой рост PPL (что ожидаемо), но модель остается стабильной даже на 8192 токенах. И память растет линейно, а не квадратично.

Где все падает: частые грабли

• Ошибка в маске внимания. При sliding window нужно правильно строить маску, чтобы модель не заглядывала в будущее и учитывала padding. Одна ошибка — и модель видит токены из следующего абзаца, что убивает качество.
• Слишком большое окно. Выставляете window_size=2048, чтобы «почти как dense». Но тогда теряется смысл Sparse Attention, и память снова растет. Начинайте с 512.
• Недостаток данных для дообучения. 10 тысяч примеров по 8k токенов — это минимум. Лучше 50-100k. Экономия на данных — гарантия высокой PPL.
• Игнорирование особенностей русской морфологии. ruGPT3XL обучена на русских текстах. Sliding window может «разрезать» сложные грамматические конструкции. Проверяйте качество на художественных текстах, а не только на новостях.

Что дальше? Вместо выводов

Увеличить контекст ruGPT3XL до 8k — реально. Но это не финиш, а старт. Дальше можно экспериментировать с блочно-разреженным вниманием для иерархических документов, добавить Tuneable Attention для адаптивного выбора паттерна или даже применить методы из статьи про Zero-Shot Transferable Adapter, чтобы научить модель работать с разными длинами контекста на лету.

Главный совет на 2026 год: не зацикливайтесь на одной архитектуре внимания. Через полгода выйдет новый хак, который сделает все это legacy. Но понимание принципов — как PPL отражает качество, как sparse attention экономит память, как правильно дообучать — останется с вами и поможет быстро адаптировать любую модель под новые требования. Даже если это будет не ruGPT3XL, а какой-нибудь ruGPT-7 с изначальным контекстом в 100k.

Подписаться на канал