GPT-OSS в MLA: как исправить проблему с RoPE-K и добиться почти lossless-конверсии 20B модели

Проблема: почему GPT-OSS ломается в MLA

Представьте ситуацию. Вы скачали GPT-OSS - открытую 20-миллиардную модель. Хотите запустить ее через MLA (Model Loader Architecture) для экспериментов. Стандартная конвертация через стандартные скрипты. И получаете на выходе... мусор. Perplexity (PPL) взлетает с 5.3 до 25+. Модель генерирует абракадабру.

Почему? Потому что в GPT-OSS используется нестандартная реализация RoPE (Rotary Positional Encoding) - RoPE-K. А MLA ожидает классическую RoPE. Архитектурный диссонанс. Модель теряет понимание позиций токенов в контексте.

RoPE-K vs классическая RoPE: в чем разница?

RoPE-K - это модификация, где ключевые (key) heads используют другую частоту вращения, чем value heads. В классической RoPE они идентичны. В GPT-OSS разработчики решили разделить их для лучшего улавливания дальних зависимостей.

MLA же заточена под стандарт. При конвертации она просто "склеивает" key и value heads, теряя эту дифференциацию. Результат - модель перестает понимать, какой токен где находится в последовательности.

Решение: патч для MLA вместо переписывания GPT-OSS

Можно переписать GPT-OSS под стандартную RoPE. Но это требует ретренинга - месяцы работы и терабайты данных. Или можно адаптировать MLA к пониманию RoPE-K. Второй вариант проще, быстрее, и сохраняет оригинальное качество модели.

Суть фикса: расширить механизм позиционного кодирования в MLA, чтобы он поддерживал раздельные частоты для key и value heads. Добавить флаг в конфиг, который говорит: "эта модель использует RoPE-K, будь внимателен".

1 Анализ оригинальной реализации GPT-OSS

Сначала смотрим, как устроен RoPE-K в исходниках GPT-OSS. Ищем файл с attention механизмом. Видим примерно такую структуру:

# В оригинальном GPT-OSS (упрощенно)
class RotaryEmbeddingK(nn.Module):
    def __init__(self, dim, base=10000):
        super().__init__()
        self.dim = dim
        self.base = base
        # Разные инварианты для key и value
        self.inv_freq_key = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.inv_freq_value = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim * 1.1))  # Вот она, разница!
        
    def forward(self, x, seq_dim=1):
        # key и value получают разное позиционное кодирование
        sin_key, cos_key = self._apply_rotary(x, self.inv_freq_key)
        sin_value, cos_value = self._apply_rotary(x, self.inv_freq_value)
        return sin_key, cos_key, sin_value, cos_value

Ключевая строка: self.inv_freq_value = ... * 1.1. Value heads используют частоту на 10% выше. Казалось бы, мелочь. Но для нейросети - катастрофа, если игнорировать.

2 Модификация MLA: добавляем поддержку RoPE-K

Теперь патчим MLA. Находим файл с реализацией Rotary Embedding. Обычно это modeling_rope.py или подобное. Добавляем новую конфигурационную опцию и модифицируем forward pass.

# В MLA: патч для поддержки RoPE-K
class PatchedRotaryEmbedding(nn.Module):
    def __init__(self, dim, base=10000, rope_type="standard", scaling_factor=1.0):
        super().__init__()
        self.dim = dim
        self.base = base
        self.rope_type = rope_type  # "standard" или "gpt_oss_k"
        self.scaling_factor = scaling_factor
        
        # Стандартная инициализация
        self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        
        # Дополнительно для RoPE-K
        if rope_type == "gpt_oss_k":
            # Точно такой же коэффициент, как в GPT-OSS
            self.inv_freq_value = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim * 1.1))
        
    def forward(self, x, seq_dim=1, head_type="key"):
        if self.rope_type == "standard" or head_type == "key":
            inv_freq = self.inv_freq
        elif self.rope_type == "gpt_oss_k" and head_type == "value":
            inv_freq = self.inv_freq_value
        else:
            inv_freq = self.inv_freq
            
        # Остальная логика применения ротационного кодирования
        t = torch.arange(x.shape[seq_dim], device=x.device).type_as(inv_freq)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos, sin

3 Интеграция патча в конвертер

Теперь нужно модифицировать сам процесс конвертации. MLA обычно использует скрипты на базе transformers. Добавляем флаг в аргументы:

# В скрипте конвертации (convert_gpt_oss_to_mla.py)
parser.add_argument("--rope-type", 
                    type=str, 
                    default="standard",
                    choices=["standard", "gpt_oss_k"],
                    help="Тип Rotary Positional Encoding")

parser.add_argument("--rope-scaling-factor",
                    type=float,
                    default=1.0,
                    help="Коэффициент масштабирования для value heads (для GPT-OSS: 1.1)")

И применяем эти параметры при создании конфига модели:

config = MLAConfig(
    vocab_size=model.config.vocab_size,
    hidden_size=model.config.hidden_size,
    num_attention_heads=model.config.num_attention_heads,
    num_hidden_layers=model.config.num_hidden_layers,
    rope_type=args.rope_type,
    rope_scaling_factor=args.rope_scaling_factor,
    # ... остальные параметры
)

4 Проверка качества конверсии

После конвертации обязательно проверяем PPL на валидационном датасете. Используем тот же датасет, что и оригинальная GPT-OSS для сравнения.

# Замер PPL до и после
python evaluate_ppl.py \
  --model-path ./converted_gpt_oss_20b \
  --dataset wikitext-2 \
  --batch-size 4 \
  --max-length 2048

# Ожидаемые результаты:
# Без патча: PPL ≈ 25-30 (плохо)
# С патчем: PPL ≈ 5.5-6.0 (почти как оригинал 5.3)

Разница в 0.2-0.7 пунктов PPL - это и есть "почти lossless". Для большинства задач неразличимо. Если хотите добиться полной идентичности, нужно точнее подобрать коэффициент scaling_factor через небольшой grid search.

Типичные ошибки и как их избежать

А что с другими моделями?

GPT-OSS - не единственная модель с нестандартным RoPE. Аналогичные проблемы встречаются в:

• GLM-4.7-REAP (использует динамическое масштабирование RoPE)
• Некоторые версии Granite (MoE-архитектуры)
• Кастомные модели после обучения с нуля на специфичных данных

Принцип решения тот же: анализируем исходную реализацию, находим отличия от стандарта, патчим загрузчик, а не переучиваем модель. Это в разы дешевле и быстрее.

Практический workflow: от скачивания до работающей модели

1. Скачиваем GPT-OSS 20B с официального репозитория
2. Клонируем MLA с GitHub и применяем наш патч
3. Запускаем конвертацию с флагами --rope-type gpt_oss_k --rope-scaling-factor 1.1
4. Проверяем PPL на wikitext-2
5. Если PPL > 6.0, экспериментируем с scaling_factor
6. Тестируем на целевых задачах (генерация, классификация)

Весь процесс занимает 2-3 часа вместо недель ретренинга. И сохраняет 99% качества оригинальной модели.

Что будет, если проигнорировать проблему?

Можно попробовать использовать стандартную конвертацию и надеяться, что "авось пронесет". Не пронесет. Модель будет:

• Путать порядок событий в тексте
• Терять контекстную связность после 512 токенов
• Генерировать грамматически правильный, но семантически бессмысленный текст
• Показывать PPL в 4-5 раз хуже оригинала

По сути, вы получите очень дорогой (20B параметров!) случайный генератор текста. Впустую потраченное время на загрузку, конвертацию и ожидание, что вот сейчас "заработает".

RoPE-K - это баг или фича?

Спорный вопрос. С одной стороны, нестандартная реализация ломает совместимость. С другой - разработчики GPT-OSS добились этим улучшения на длинных контекстах (по их заявлениям).

Мой вердикт: это фича, но плохо документированная. Если бы в README GPT-OSS было четко указано "используем RoPE-K с scaling_factor=1.1 для value heads", половина проблем с конвертацией исчезла бы. Но документация - вечная боль open-source проектов.

Поэтому приходится лезть в исходники. И это нормально. Настоящий инженер не ждет готовых решений, а разбирается, как оно работает внутри. Как в том гайде по выжиманию VRAM - там тоже приходится копаться в низкоуровневых оптимизациях.

Будущее: будут ли такие проблемы с GPT-OSS:120B?

Если прогнозы о GPT-OSS:120B сбудутся, и модель выйдет, почти наверняка в ней будут свои архитектурные особенности. Возможно, еще более сложные модификации RoPE. Или вообще другой механизм позиционного кодирования.

Но принцип решения останется тем же: анализировать, понимать, адаптировать. Не пытаться впихнуть квадратный колышек в круглую дыру, а менять форму дыры под колышек.

И последний совет: всегда сохраняйте патчи отдельно от основной кодовой базы MLA. Когда выйдет новая версия MLA, вы сможете быстро применить свои модификации к обновленному коду. Не становитесь тем парнем, у которого "работает только на старой версии, потому что я ее сильно патчил".