Muon и стабильность обучения

Зачем это нужно

Почти все большие модели последних лет обучены AdamWAdamWEN AdamW (Adam with decoupled weight decay)Adam, в котором затухание весов применяется отдельно от адаптивного шага. Оптимизатор по умолчанию для обучения LLM.: вы разбирали его в уроке 1.4.2. В 2024–2025 годах у него появился серьёзный соперник. Оптимизатор MuonMuon (оптимизатор)EN Muon optimizerОптимизатор для матриц весов (Келлер Джордан, 2024) — momentum, затем ортогонализация обновления итерациями Ньютона–Шульца. Moonshot масштабировала его до Kimi K2. придумал Келлер Джордан для рекордов скорости обучения NanoGPT. Команда Moonshot AI (Kimi) показала, что он масштабируется: примерно вдвое эффективнее AdamW по вычислениям и на нём обучена Kimi K2 — модель на триллион параметров.

При масштабировании вылезла проблема — взрывающиеся attention-логиты. Решение Moonshot, QK-Clip, позволило обучить K2 на 15.5T токенов без единого всплеска lossвсплеск lossEN loss spikeРезкий скачок функции потерь во время обучения. В больших моделях может испортить многодневный прогон, поэтому техотчёты отдельно сообщают об их отсутствии.. Всплески — главный страх любого большого прогона: сутки работы 2048 GPU стоят около 100 тысяч долларов аренды (по 2 долл. за GPU-час, как в оценке DeepSeek-V3), и откат к чекпоинту сжигает их впустую.

Урок в двух частях: как устроен Muon и почему он работает, затем — откуда берутся всплески и чем их лечат.

Обновление матрицы — это тоже матрица

AdamW смотрит на параметры как на длинный список чисел. Каждое число получает свой адаптивный шаг, и неважно, что 4096 × 4096 этих чисел образуют матрицу $W$ линейного слоя.

Muon исходит из другого. Обновление матрицы весов — матрица $\Delta W$, и у неё есть сингулярные числа (сингулярное разложениесингулярное разложение (SVD)EN singular value decomposition (SVD)A = UΣVᵀ — любую матрицу можно разложить на поворот, растяжение вдоль осей и ещё один поворот. Эквивалентно сумме слоёв ранга 1 σᵢuᵢvᵢᵀ., урок 1.1.3): $\Delta W = U \Sigma V^\top = \sum_i \sigma_i u_i v_i^\top$. Каждое слагаемое — «направление» изменения, а $\sigma_i$ — его сила.

Эмпирический факт: у градиентов и momentum в трансформерах сингулярные числа очень неравны. Несколько направлений доминируют, остальные в десятки и сотни раз слабее — матрица «почти низкого ранга». Шаг по такому градиенту в основном усиливает то, что уже сильно, а редкие направления почти не обновляются. Хотя для обучения они могут быть важны.

Идея Muon: заменить обновление на его ортогонализацию (ортогонализация матрицыортогонализация матрицыEN matrix orthogonalizationЗамена матрицы G = UΣVᵀ на UVᵀ — все сингулярные числа становятся равными 1, а «направления» матрицы сохраняются.):

$$M = U \Sigma V^\top \quad\longrightarrow\quad O = U V^\top .$$

Направления те же, а все силы становятся равны 1. Сильные направления приглушаются, слабые усиливаются.

Взгляд Су Цзяньлиня: Adam и Muon — родственники

В блоге «Научное пространство» Су Цзяньлинь (соавтор работы Moonlight) объясняет Muon через знаковый спуск. Adam при $\beta_1 = \beta_2 = 0$ (и $\varepsilon = 0$) превращается в $\mathrm{sign}(g)$: каждой координате — шаг одинаковой величины, только знак свой. $UV^\top$ — это «матричный знак» $\mathrm{msign}(M)$: каждому сингулярному направлению — шаг одинаковой силы. Формально Muon — наискорейший спуск, если размер шага мерить спектральной нормойспектральная нормаEN spectral normНаибольшее сингулярное число σ₁ — во сколько раз матрица максимально может растянуть вектор. матрицы, а не поэлементно.

Итерации Ньютона–Шульца

SVD большой матрицы на каждом шаге — дорого и плохо ложится на GPU. Нужен способ получить $UV^\top$ одними матричными умножениями. Это итерации Ньютона–Шульцаитерации Ньютона–ШульцаEN Newton–Schulz iterationМногочлен от матрицы X ← aX + b(XXᵀ)X + c(XXᵀ)²X, который подтягивает все сингулярные числа к 1. Заменяет дорогое SVD одними матричными умножениями.:

$$X_{k+1} = a X_k + b\,(X_k X_k^\top) X_k + c\,(X_k X_k^\top)^2 X_k .$$

Что здесь написано. Если $X = U \Sigma V^\top$, то $X X^\top X = U \Sigma V^\top V \Sigma U^\top U \Sigma V^\top = U \Sigma^3 V^\top$ (в SVD столбцы $U$ и $V$ ортонормированы, поэтому $V^\top V = I$ и $U^\top U = I$ — единичные матрицы, они «исчезают»): матрицы $U$ и $V$ не меняются, а каждое сингулярное число возводится в куб. Значит, один шаг применяет к каждому $\sigma$ один и тот же многочлен:

$$\sigma \mapsto a\sigma + b\sigma^3 + c\sigma^5 .$$

Задача сводится к выбору многочлена, который тянет любое $\sigma$ из $(0, 1]$ к единице. Перед итерациями матрицу делят на её норму Фробениусанорма ФробениусаEN Frobenius normКорень из суммы квадратов всех элементов матрицы — «L2-норма матрицы, вытянутой в вектор». Равна √(Σ σᵢ²). (корень из суммы квадратов всех элементов, урок 1.1.3), чтобы все $\sigma \le 1$.

Классическая кубика $a = 1.5$, $b = -0.5$, $c = 0$: $f(\sigma) = 1.5\sigma - 0.5\sigma^3$, $f(1) = 1$. Пример: $\sigma = 0.5$ даёт 0.6875, затем 0.869, 0.975, 0.999. Беда в маленьких $\sigma$: каждый шаг умножает их лишь на 1.5, и $\sigma = 0.01$ за 5 шагов дорастает до 0.076.

Квинтика Muon $a = 3.4445$, $b = -4.7750$, $c = 2.0315$. Коэффициенты подобраны численно так, чтобы малые числа росли как можно быстрее (почти в 3.4 раза за шаг), даже ценой того, что $\sigma$ не сходится к 1 точно, а колеблется около неё:

$\sigma_0$	шаг 1	шаг 2	шаг 3	шаг 4	шаг 5
0.01	0.034	0.118	0.400	1.093	0.699
0.1	0.340	0.992	0.707	1.107	0.712
0.5	1.189	0.896	0.824	0.938	0.765
1.0	0.701	1.114	0.721	1.090	0.696

После 5 шагов все $\sigma \ge 0.01$ попадают в коридор примерно от 0.7 до 1.2. Оказалось, этого хватает: точная ортогонализация ничего не добавляет к качеству, а 5 шагов дёшевы. По оценке автора, дополнительные FLOPs Muon при типичном обучении языковых моделей — меньше процента. Итерации устойчиво работают даже в BF16.

Алгоритм Muon целиком

Для каждой матрицы весов $W$ размера $A \times B$ на шаге $t$:

1. $M_t = \mu M_{t-1} + G_t$ — momentum ($\mu = 0.95$, $G_t$ — градиент). В варианте Нестерова (Nesterov) дальше берут $G_t + \mu M_t$: шаг как бы «заглядывает вперёд» по направлению накопленной скорости.
2. $O_t = \mathrm{NewtonSchulz}_5(M_t)$ — ортогонализация.
3. $W_t = W_{t-1} - \eta\left(0.2\sqrt{\max(A, B)} \cdot O_t + \lambda W_{t-1}\right)$ — шаг с weight decay (вариант Moonlight).

Только для скрытых матриц. Эмбеддинги, выходной слой, веса нормализаций и все одномерные параметры обучают AdamW. Строгого вывода здесь нет, есть интуиция и опыт. Muon рассчитан на слой, который умножает плотный вектор активаций на матрицу. Эмбеддинг устроен иначе: на вход приходит номер токена, и каждая строка матрицы — отдельный вектор своего токена. Для неё естественнее обновляться самостоятельно (AdamW так и делает, поэлементно), чем вместе со всей матрицей выравнивать сингулярные числа. Для выходного слоя Келлер Джордан тоже рекомендует AdamW: в экспериментах так получалось лучше.

Попутный бонус: у Muon один буфер состояния (momentum), а у AdamW два ($m$ и $v$). В пересчёте на FP32 это 4 байта на параметр экономии — вспомните таблицу памяти из урока 5.4.1.

Что Moonshot добавила для масштаба

Работа Moonlight (Moonshot AI, февраль 2025) начиналась с вопроса: Muon побеждал на маленьких моделях, а выдержит ли он миллиарды параметров и триллионы токенов? Понадобились два изменения.

1. Weight decay. Без него в длинном обучении веса и выходы слоёв растут, их значения выходят за удобный для BF16 диапазон, и качество начинает отставать от AdamW. С обычным decoupled weight decay (как в AdamW) Muon выигрывает и на длинной дистанции.
2. Согласование RMS обновления. RMS — корень из среднего квадрата элементов (как в RMSNorm, урок 4.5.3), типичный «размер» одного элемента. Ортогональная матрица $A \times B$ имеет RMS элементов $1/\sqrt{\max(A, B)}$ — у матриц разной формы шаг получался бы разным. Множитель $0.2\sqrt{\max(A, B)}$ делает RMS обновления равным 0.2, как типичное обновление AdamW. Тогда можно взять learning rate и weight decay от AdamW и не подбирать их заново.

Ещё нужна была распределённая версия. При ZeRO-1 (урок 5.4.3) состояние оптимизатора разрезано по GPU, а ортогонализации нужна вся матрица. Moonshot собирает градиент матрицы, считает Ньютона–Шульца и оставляет себе только свой кусок — и выложила эту реализацию в открытый доступ.

Итог: по экспериментам с законами масштабирования Muon достигает того же loss примерно за половину вычислений AdamW. На этом обучена Moonlight — MoE на 16B параметров, 3B активных, 5.7T токенов.

Практика

Задание в браузере

Реализуйте ортогонализацию через SVD (эталон), шаг и цикл Ньютона–Шульца, а затем обновление Muon с momentum и масштабом Moonlight. Кнопка «Запустить» покажет, как градиент с силами направлений 10, 3, 1 и 0.05 превращается в обновление, где все силы близки к единице.

Итерации Ньютона–Шульца и обновление Muonnumpy · нужен интернет6 тестов

import numpy as np

MUON_COEFFS = (3.4445, -4.7750, 2.0315)   # коэффициенты квинтики из Muon (Келлер Джордан)
CUBIC_COEFFS = (1.5, -0.5, 0.0)           # классическая кубическая итерация Ньютона–Шульца


def orthogonalize_svd(G):
    """Эталон: G = U Σ Vᵀ  →  U Vᵀ.

    Используйте np.linalg.svd(G, full_matrices=False) — он вернёт U, S, Vt.
    """
    # TODO: ваш код
    ...


def ns_step(X, a, b, c):
    """Один шаг итерации: X ← aX + b(XXᵀ)X + c(XXᵀ)²X.

    Подсказка: A = X @ X.T; B = b·A + c·(A @ A); ответ a·X + B @ X.
    """
    # TODO: ваш код
    ...


def newton_schulz(G, steps=5, coeffs=MUON_COEFFS, eps=1e-7):
    """Приближённая ортогонализация G одними матричными умножениями.

    1) если строк больше, чем столбцов, транспонируйте (и не забудьте транспонировать ответ обратно);
    2) поделите на норму Фробениуса np.linalg.norm(X) + eps — тогда все сингулярные числа ≤ 1;
    3) сделайте steps шагов ns_step с коэффициентами coeffs.
    Форма ответа совпадает с формой G.
    """
    # TODO: ваш код
    ...


def muon_update(grad, buf, beta=0.95, nesterov=True, steps=5):
    """Обновление Muon для одной матрицы весов (вариант Moonlight).

    buf ← beta·buf + grad                         (momentum)
    g   ← grad + beta·buf  если nesterov, иначе buf
    O   ← newton_schulz(g, steps)
    update = 0.2 · sqrt(max(A, B)) · O, где (A, B) = grad.shape
    Возвращает (update, новый buf).
    """
    # TODO: ваш код
    ...


def rms(M):
    """Среднеквадратичное значение элементов матрицы."""
    return float(np.sqrt(np.mean(np.asarray(M) ** 2)))


# Проверьте себя кнопкой «Запустить»:
rng = np.random.default_rng(0)
# «Градиент» почти низкого ранга: три сильных направления и слабый шум, как часто бывает у матриц весов
U, _ = np.linalg.qr(rng.normal(size=(64, 64)))
V, _ = np.linalg.qr(rng.normal(size=(32, 32)))
sigma = np.array([10.0, 3.0, 1.0] + [0.05] * 29)
G = U[:, :32] @ np.diag(sigma) @ V.T


def strengths(M):
    """Насколько M «растягивает» каждое из направлений (u_i, v_i) исходного G."""
    return np.round([U[:, i] @ M @ V[:, i] for i in (0, 1, 2, 3)], 3)


O = newton_schulz(G)
print("сила направлений в G:        ", strengths(G))
if O is not None:
    print("после Ньютона–Шульца, 5 шагов:", strengths(O))

ЗапуститьПроверитьПодсказкаСбросить

Muon против AdamW на PyTorch

Возьмём маленькую GPT из урока 5.4.2 (без MTP) и обучим её дважды: только AdamW и Muon для скрытых матриц + AdamW для остального. Learning rate одинаковый — благодаря согласованию RMS. Заодно будем следить за максимальным attention-логитом и нормой градиента: это два главных датчика стабильности.

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

TEXT = """У лукоморья дуб зелёный;
Златая цепь на дубе том:
И днём и ночью кот учёный
Всё ходит по цепи кругом;
Идёт направо — песнь заводит,
Налево — сказку говорит.
Там чудеса: там леший бродит,
Русалка на ветвях сидит;
Там на неведомых дорожках
Следы невиданных зверей;
Избушка там на курьих ножках
Стоит без окон, без дверей;
Там лес и дол видений полны;
Там о заре прихлынут волны
На брег песчаный и пустой,
И тридцать витязей прекрасных
Чредой из вод выходят ясных,
И с ними дядька их морской;
"""
chars = sorted(set(TEXT))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in TEXT])
V, D, T = len(chars), 64, 64


def newton_schulz(G, steps=5, eps=1e-7):
    """Ортогонализация обновления: квинтика Келлера Джордана, 5 шагов."""
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G / (G.norm() + eps)
    tall = X.shape[0] > X.shape[1]
    if tall:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X
    return X.T if tall else X


class Muon(torch.optim.Optimizer):
    """Muon для 2D-матриц: momentum (Nesterov) → Ньютон–Шульц → масштаб 0.2·√max(A, B) → шаг с weight decay."""

    def __init__(self, params, lr=3e-3, momentum=0.95, weight_decay=0.1):
        super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                # self.state[p] — словарь состояния для параметра p; при первом шаге создаём буфер momentum из нулей
                buf = self.state[p].setdefault("buf", torch.zeros_like(p))
                buf.mul_(group["momentum"]).add_(p.grad)               # методы с «_» меняют тензор на месте: buf = μ·buf + g
                g = p.grad.add(buf, alpha=group["momentum"])          # Nesterov
                update = newton_schulz(g) * 0.2 * math.sqrt(max(p.shape))
                p.mul_(1 - group["lr"] * group["weight_decay"])       # decoupled weight decay, как в AdamW
                p.add_(update, alpha=-group["lr"])


class Block(nn.Module):
    def __init__(self, d, n_heads=4):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(d), nn.RMSNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d, bias=False), nn.Linear(d, d, bias=False)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d, bias=False), nn.GELU(), nn.Linear(4 * d, d, bias=False))
        self.h, self.max_logit = n_heads, 0.0

    def forward(self, x):
        B, L, d = x.shape
        q, k, v = self.qkv(self.n1(x)).view(B, L, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        logits = q @ k.transpose(-2, -1) / math.sqrt(d // self.h)            # (B, h, L, L)
        self.max_logit = max(self.max_logit, logits.max().item())           # мониторинг: максимальный логит
        mask = torch.ones(L, L, dtype=torch.bool).tril()
        a = logits.masked_fill(~mask, float("-inf")).softmax(-1) @ v
        x = x + self.proj(a.transpose(1, 2).reshape(B, L, d))
        return x + self.mlp(self.n2(x))


class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb, self.pos = nn.Embedding(V, D), nn.Embedding(T, D)
        nn.init.normal_(self.emb.weight, std=0.02)
        nn.init.normal_(self.pos.weight, std=0.02)
        self.blocks = nn.Sequential(Block(D), Block(D))
        self.norm = nn.RMSNorm(D)

    def forward(self, idx):
        h = self.blocks(self.emb(idx) + self.pos(torch.arange(idx.shape[1])))
        return self.norm(h) @ self.emb.weight.T


def run(use_muon, steps=200, lr=3e-3):
    torch.manual_seed(0)
    model = GPT()
    hidden = [p for n, p in model.named_parameters() if p.ndim == 2 and "blocks" in n]   # матрицы внутри блоков
    other = [p for n, p in model.named_parameters() if not (p.ndim == 2 and "blocks" in n)]
    if use_muon:   # Muon — только для скрытых матриц; эмбеддинги и нормы — AdamW (как в Moonlight и K2)
        opts = [Muon(hidden, lr=lr), torch.optim.AdamW(other, lr=lr, weight_decay=0.1)]
    else:
        opts = [torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.1)]
    g = torch.Generator().manual_seed(1)
    for step in range(steps + 1):
        ix = torch.randint(len(data) - T - 1, (32,), generator=g)
        x = torch.stack([data[i : i + T] for i in ix])
        y = torch.stack([data[i + 1 : i + T + 1] for i in ix])
        loss = F.cross_entropy(model(x).reshape(-1, V), y.reshape(-1))
        for o in opts:
            o.zero_grad()
        loss.backward()
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)       # клиппинг и заодно норма градиента
        for o in opts:
            o.step()
        if step % 50 == 0:
            ml = max(b.max_logit for b in model.blocks)
            print(f"{'Muon ' if use_muon else 'AdamW'} шаг {step:3d}: loss {loss.item():.3f}  |g| {gnorm:.2f}  max-логит {ml:5.1f}")
            for b in model.blocks:
                b.max_logit = 0.0


run(use_muon=False)
run(use_muon=True)

AdamW шаг   0: loss 3.840  |g| 1.63  max-логит   2.0
AdamW шаг  50: loss 2.221  |g| 0.62  max-логит   9.5
AdamW шаг 100: loss 0.795  |g| 1.41  max-логит  10.9
AdamW шаг 150: loss 0.175  |g| 1.06  max-логит  12.6
AdamW шаг 200: loss 0.094  |g| 0.59  max-логит  13.0
Muon  шаг   0: loss 3.840  |g| 1.63  max-логит   2.0
Muon  шаг  50: loss 1.639  |g| 1.43  max-логит   2.6
Muon  шаг 100: loss 0.323  |g| 1.62  max-логит   2.7
Muon  шаг 150: loss 0.097  |g| 0.64  max-логит   3.9
Muon  шаг 200: loss 0.075  |g| 0.56  max-логит   5.3

Обучение двух моделей занимает около минуты на CPU; точные числа на вашей машине могут немного отличаться. Max-логит — максимум за последние 50 шагов. На 100-м шаге у Muon loss 0.32, у AdamW — 0.80: при тех же гиперпараметрах Muon учится заметно быстрее. Не обобщайте одну игрушку на все случаи, но направление совпадает с выводами Moonlight. Логиты в нашей модели не взрываются ни у одного оптимизатора — эта беда проявляется на масштабе, к ней и переходим.

Всплески loss: откуда они берутся

Всплеск loss — резкий скачок функции потерь посреди спокойного обучения. Иногда loss возвращается сам за десятки шагов. Иногда уходит в nan, и прогон приходится откатывать. В статье PaLM (Google, 2022) описано около 20 всплесков при обучении модели на 540B, хотя клиппинг градиентов был включён. На маленьких моделях их не было.

Типичные механизмы:

• Рост attention-логитов. Если $q \cdot k / \sqrt{d}$ становится огромным, softmax вырождается почти в one-hot. Градиенты через такой softmax становятся резкими и неустойчивыми.
• Дрейф выходных логитов. Логиты словаря растут все вместе. Softmax от этого не меняется, но числа уходят туда, где низкая точность (урок 5.4.1) начинает врать.
• Адаптивный шаг на редких данных. У параметров, которые долго не получали градиента (редкие токены), второй момент Adam $v$ становится крошечным. Первый же заметный градиент даёт гигантский шаг. Поэтому в LLM берут $\beta_2 = 0.95$, а не 0.999: $v$ быстрее подстраивается.
• Сочетание батча и состояния. Авторы PaLM проверили «плохие данные»: те же батчи, поданные с более раннего чекпоинта, всплеска не давали. Виновато именно сочетание конкретных данных и конкретного состояния параметров.

Аптечка

Приём	Что делает	Где встречается
Warmup learning rate	первые шаги — маленький lr, пока статистики Adam не накопились	почти везде (урок 1.4.2)
Клиппинг нормы градиента	если $\lVert g \rVert$ больше порога (обычно 1.0), градиент масштабируется	почти везде
$\beta_2 = 0.95$, weight decay 0.1	быстрая адаптация $v$, ограничение роста весов	DeepSeek-V3, Llama, Qwen
z-loss	штраф $10^{-4} \cdot \log^2 Z$ за дрейф нормировочной константы softmax ($Z = \sum_j e^{z_j}$ — знаменатель softmax)	PaLM
QK-нормализацияQK-нормализация (QK-Norm)EN QK-NormНормализация (обычно RMSNorm) запросов и ключей перед скалярным произведением. Ограничивает рост attention-логитов; используется, например, в Qwen3.	RMSNorm на $q$ и $k$ перед скалярным произведением — логиты ограничены	Qwen3 и другие
Soft-capping логитов	$c \cdot \tanh(x / c)$ — плавно обрезает логиты	Gemma 2
Откат и пропуск батчей	вернуться на ~100 шагов назад и пропустить 200–500 батчей	PaLM
QK-Clip	уменьшить $W_q$, $W_k$ головы, если её max-логит превысил порог	Kimi K2

И главное — мониторинг. В журнал каждого прогона пишут не только loss, но и норму градиента, максимальный attention-логит по слоям, норму весов и отношение «размер обновления / размер весов». Всплеск почти всегда предваряется ростом одного из этих чисел.

MuonClip: как Kimi K2 обошлась без всплесков

При обучении на Muon в масштабе Kimi K2 (1T параметров, 32B активных) Moonshot столкнулась с взрывом max-логита: в одном из прогонов среднего размера максимальные attention-логиты быстро превысили 1000. Известные средства не подходили:

• QK-Norm невозможен в MLA (урок 5.2.2): на инференсе ключи не материализуются целиком — кэшируется только латентный вектор, а матрицы «поглощены». Нормировать нечего.
• Soft-capping ограничивает логит на выходе, но произведение $q \cdot k$ внутри всё равно растёт, и проблема просто прячется.

QK-ClipQK-ClipEN QK-ClipПриём из Kimi K2 — если максимальный attention-логит головы превысил порог τ, веса W_q и W_k этой головы умножают на √(τ / max). Вместе с Muon — оптимизатор MuonClip. действует на веса, а не на логиты. После каждого шага оптимизатора для каждой головы $h$ берут максимальный логит $S^h_{\max}$ на батче. Если он больше порога $\tau$:

$$\gamma_h = \frac{\tau}{S^h_{\max}}, \qquad W_q^h \leftarrow \sqrt{\gamma_h}\, W_q^h, \qquad W_k^h \leftarrow \sqrt{\gamma_h}\, W_k^h .$$

Логит — произведение запроса и ключа, поэтому на тех же входах он уменьшится ровно в $\gamma_h$ раз и станет равен $\tau$. Пример: $\tau = 100$, голова выдала $S_{\max} = 400$. Тогда $\gamma = 0.25$, и $W_q$, $W_k$ умножаются на $\sqrt{0.25} = 0.5$. Головы ниже порога не трогают. В MLA масштабируют только части, принадлежащие голове: её сжатые части запроса и ключа — на $\sqrt{\gamma_h}$, её RoPE-часть запроса — на $\gamma_h$ целиком. Общий для всех голов RoPE-ключ не трогают, чтобы не задеть соседние головы.

Muon вместе с QK-Clip называют MuonClip. В K2 $\tau = 100$. На графике из отчёта max-логит быстро упирается в 100, держится там, а примерно после 30% обучения сам опускается в спокойную зону — QK-Clip перестаёт срабатывать. Итог: 15.5T токенов без единого всплеска loss. Механику QK-Clip на NumPy вы реализуете в уроке china/08.

🖥 На вашей машине (4×GPU)

1. Перенесите класс Muon в свою GPT из урока 4.6.2 и обучите на русском корпусе на одной RTX 4080 дважды — с AdamW и с Muon — при равном числе токенов. Постройте обе кривые val loss.
2. Добавьте в лог max-логит по слоям и норму градиента. Поднимите learning rate в 3–10 раз: какой оптимизатор раньше даёт всплеск? Помогает ли QK-Norm (nn.RMSNorm на q и k по размерности головы)?
3. В свежих версиях PyTorch есть встроенный torch.optim.Muon. С adjust_lr_fn="match_rms_adamw" он делает то же, что наш класс, только итерации Ньютона–Шульца считает в BF16 (снова смешанная точность). По умолчанию там масштаб Келлера Джордана $\sqrt{\max(1, A/B)}$ — с ним learning rate от AdamW не подойдёт.

Удобная раскладка по картам: AdamW и Muon на базовом lr — на картах 0 и 1, оба оптимизатора на lr ×5 — на картах 2 и 3 (CUDA_VISIBLE_DEVICES=…, как в уроке 4.6.3). Все четыре кривые будут готовы одновременно. Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Китайский угол: Moonshot, Qwen и культура стабильности

Moonshot превратила чужую идею в индустриальный инструмент. Muon придумал Келлер Джордан, а Moonshot AI (китайское название — «Тёмная сторона Луны») довела его до триллиона параметров: weight decay, согласование RMS, распределённая реализация (код — в MoonshotAI/Moonlight), затем QK-Clip. Второй автор работы Moonlight — Су Цзяньлинь, создатель RoPE (урок 4.5.2). В его блоге «Научное пространство» (kexue.fm) есть целая серия о Muon: от вывода через матричный знак до объяснения QK-Clip — с разбором, почему QK-Norm не подходит для MLA. Muon используется и в DeepSeek-V4: в аннотации техотчёта он назван одним из ключевых улучшений — «для более быстрой сходимости и большей стабильности обучения».

DeepSeek-V3 в отчёте подчёркивает: за весь претрейн не было ни одного необратимого всплеска loss и ни одного отката к чекпоинту.

Qwen3-Next (сентябрь 2025) описывает свой набор мер: RMSNorm с нулевым центром и weight decay на весах нормализации (в Qwen3 некоторые из этих весов разрастались аномально), выходной гейт внимания, который убирает attention sink и массивные активации, нормированная инициализация роутера MoE.

Термины. Оптимизатор по-китайски — буквально «оптимизирующее устройство», с тем же суффиксом, что в словах «прибор» и «сосуд». Ортогонализация и стабильность обучения переводятся дословно, а всплеск loss называют «острым пиком потерь».

Проверь себя

Проверь себя6 вопросов·

1 Что Muon делает с momentum-матрицей $M = U\Sigma V^\top$ перед шагом?

Делит каждый элемент на корень из второго момента, как AdamЗаменяет на $UV^\top$ (приближённо): направления те же, все сингулярные числа ≈ 1Оставляет только первое сингулярное направлениеОбнуляет маленькие элементы

2 Кубическая итерация Ньютона–Шульца: $\sigma \mapsto 1.5\sigma - 0.5\sigma^3$. Во что превратится $\sigma = 0.5$ за один шаг?

3 К каким параметрам Muon обычно НЕ применяют (их обучают AdamW)? Выберите все.Можно выбрать несколько ответов

Эмбеддинги токеновВыходной слой (LM-голова)Веса нормализаций (RMSNorm)Матрицы $W_q$, $W_k$, $W_v$ и матрицы MLP внутри блоков

4 QK-Clip с порогом $\tau = 100$. Голова выдала максимальный логит 400. На какое число умножат $W_q$ этой головы?

5 Почему Moonshot не могла просто взять QK-Norm для Kimi K2?

QK-Norm запатентованВ MLA ключи на инференсе не материализуются целиком — кэшируется латентный вектор, матрицы поглощеныQK-Norm не работает с Muon в принципеQK-Norm увеличивает KV-кэш вдвое

6 Какие приёмы Moonlight понадобились, чтобы Muon работал на больших моделях? Выберите все.Можно выбрать несколько ответов

Weight decayМасштаб обновления $0.2\sqrt{\max(A, B)}$, чтобы RMS совпадал с AdamWЗамена Ньютона–Шульца на точное SVDПрименение Muon и к эмбеддингам

Проверить ответы

Итоги

• Muon = momentum → ортогонализация (5 итераций Ньютона–Шульца с квинтикой 3.4445, −4.7750, 2.0315) → масштаб → шаг. Только для скрытых матриц; остальное — AdamW. Состояния вдвое меньше, чем у AdamW.
• Итерация действует на сингулярные числа многочленом $a\sigma + b\sigma^3 + c\sigma^5$; квинтика за 5 шагов приводит их в коридор ≈ 0.7–1.2.
• Moonlight масштабировала Muon: weight decay + RMS обновления 0.2 → примерно вдвое эффективнее AdamW по вычислениям.
• Всплески loss лечат warmup, клиппингом, $\beta_2 = 0.95$, z-loss, QK-Norm, soft-capping и откатами. QK-Clip в Kimi K2 ограничивает рост логитов через веса головы — 15.5T токенов без единого всплеска.

Дальше: модуль 5.5 о длинном контексте. Начнём с расширения RoPE — PI, NTK и YaRN.