Muon

class torch.optim.Muon(params, lr=0.001, weight_decay=0.1, momentum=0.95, nesterov=True, ns_coefficients=(3.4445, -4.775, 2.0315), eps=1e-07, ns_steps=5, adjust_lr_fn=None)

Реализует алгоритм Muon.

$$\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{ (lr)},\ \lambda \text{ (weight decay)},\ \mu \text{ (momentum)},\ \textit{nesterov}\in\{True,False\},\\ &\hspace{13mm}(a,b,c)\ \text{ (NS coefficients)},\ \varepsilon \text{ (epsilon)},\ k \text{ (NS steps)},\ \theta_0 \text{ (params)},\ f(\theta) \text{ (objective)} \\ &\textbf{initialize} : B_0 \leftarrow 0 \text{ (momentum buffer)} \\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{for}\ t=1\ \textbf{to}\ \ldots\ \textbf{do} \\[0.25ex] &\hspace{5mm} g_t \leftarrow \nabla_{\theta} f_t(\theta_{t-1}) \\[0.25ex] &\hspace{5mm} B_t \leftarrow \mu B_{t-1} + g_t \\[0.25ex] &\hspace{5mm} \widetilde{B}_t \leftarrow \begin{cases} g_t + \mu B_t, & \text{if nesterov}=True \\ B_t, & \text{if nesterov}=False \end{cases} \\[1.0ex] &\hspace{5mm} O_t \leftarrow \mathrm{NS}^{(a,b,c)}_{k}\!\big(\widetilde{B}_t;\ \varepsilon\big) \\[0.5ex] &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma\,\lambda\,\theta_{t-1} \quad\text{(decoupled weight decay)} \\[0.25ex] &\hspace{5mm} \gamma \leftarrow \mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big) \\[0.25ex] &\hspace{5mm} \theta_t \leftarrow \theta_t - \gamma\, O_t \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\mathbf{return}\ \theta_t \\[-1.ex] &\rule{110mm}{0.4pt}s \end{aligned}$$

Здесь $\mathrm{NS}^{(a,b,c)}_{k}(\cdot;\varepsilon)$ обозначает $k$ итераций оператора ортогонализации Ньютона—Шульца, параметризованного коэффициентами $(a,b,c)$ и численной стабилизацией $\varepsilon$.

Цель функции $\mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big)$ — обеспечить согласованный масштаб ортогонализированного обновления для прямоугольных матриц.

В исходной реализации Keller масштаб обновления умножается на $\sqrt{\max\!\left(1, \frac{A}{B}\right)}$, где $A$ и $B$ — размеры оптимизируемой матрицы, соответствующие числу выходных и входных связей для матрицы весов Linear.

Реализация Moonshot ориентирована на соответствие $RMS$ оптимизатора AdamW. Коррекция вычисляется следующим образом: $\gamma \leftarrow {0.2}\gamma\,\sqrt{\max\!\left({A}, {B}\right)}$ Метод взят из статьи Muon масштабируется для обучения LLM. Результаты исследований показывают, что с этой коррекцией Muon может напрямую использовать скорость обучения и затухание весов, настроенные для AdamW.

Jeremy Bernstein в статье Вывод Muon предлагает условие масштабирования по спектральной норме, при котором обновление умножается на $\sqrt{\frac{A}{B}}$. Это похоже на «исходную» реализацию Keller, но без ограничения снизу значением 1.

Мы предоставляем следующие варианты коррекции скорости обучения: «original», соответствующий реализации Keller; «match_rms_adamw», соответствующий реализации Moonshot; и «spectral_unclamped», соответствующий реализации Bernstein. Если adjust_lr_fn не задан, по умолчанию используется «original».

Дополнительные сведения об алгоритме см. в статьях Muon: оптимизатор для скрытых слоёв нейронных сетей, Muon масштабируется для обучения LLM и Вывод Muon.

Параметры:

• params (iterable) – итерируемый объект с параметрами или именованными параметрами для оптимизации либо итерируемый объект словарей, задающих группы параметров. При использовании named_parameters все параметры во всех группах должны иметь имена. Обратите внимание, что Muon — это оптимизатор для двумерных параметров скрытых слоёв нейронной сети. Другие параметры, например bias и embedding, следует оптимизировать стандартным методом, таким как AdamW.
• lr (float, Tensor, необязательный) – скорость обучения (по умолчанию: 1e-3).
• weight_decay (float, необязательный) – затухание весов (штраф L2). (по умолчанию: 0.1)
• momentum (float, необязательный) – коэффициент импульса (по умолчанию: 0.95)
• nesterov (bool, необязательный) – включает импульс Нестерова. Применимо только при ненулевом значении momentum
• ns_coefficients (tuple из трёх чисел с плавающей запятой, необязательный) – коэффициенты (a,b,c) для полинома ортогонализации Ньютона—Шульца (по умолчанию: (3.4445, -4.775, 2.0315))
• eps (float, необязательный) – слагаемое, добавляемое к знаменателю для численной устойчивости. (по умолчанию: 1e-07)
• ns_steps (int, необязательный) – число итераций Ньютона—Шульца. (по умолчанию: 5)
• adjust_lr_fn (str, необязательный) – функция коррекции скорости обучения. Одно из значений: “original”, “match_rms_adamw” или “spectral_unclamped”. Если значение не задано, по умолчанию используется “original”. (по умолчанию: None)

Пример

>>> # Muon only supports 2D params; use a standard optimizer
>>> # such as AdamW for biases, embeddings, and other non-2D
>>> # parameters.
>>> muon_params = [
...     p for p in model.parameters() if p.ndim == 2
... ]
>>> other_params = [
...     p for p in model.parameters() if p.ndim != 2
... ]
>>> optim_muon = torch.optim.Muon(
...     muon_params, lr=0.02, momentum=0.95
... )
>>> optim_adamw = torch.optim.AdamW(
...     other_params, lr=3e-4, weight_decay=0.01
... )
>>> optim_muon.zero_grad()
>>> optim_adamw.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optim_muon.step()
>>> optim_adamw.step()

add_param_group(param_group) [исходный код]

Добавляет группу параметров в Optimizers param_groups.

Это может быть полезно при дообучении предварительно обученной сети: замороженные слои можно сделать обучаемыми и добавлять в Optimizer по мере обучения.

Параметры:

param_group (dict) – определяет, какие тензоры следует оптимизировать, а также параметры оптимизации, специфичные для группы.

load_state_dict(state_dict) [исходный код]

Загружает состояние оптимизатора.

Параметры:

state_dict (dict) – состояние оптимизатора. Должно быть объектом, возвращённым вызовом state_dict().

Предупреждение

Обязательно вызывайте этот метод после инициализации torch.optim.lr_scheduler.LRScheduler, поскольку предварительный вызов перезапишет загруженные скорости обучения.

Примечание

Имена параметров (если они присутствуют в ключе “param_names” каждой группы параметров в state_dict()) не влияют на процесс загрузки. Чтобы использовать имена параметров в пользовательских случаях (например, если параметры в загруженном словаре состояния отличаются от параметров, инициализированных в оптимизаторе), необходимо реализовать пользовательский register_load_state_dict_pre_hook для соответствующей адаптации загружаемого словаря. Если param_names присутствуют в загруженном словаре состояния param_groups, они будут сохранены и переопределят текущие имена в состоянии оптимизатора, если таковые имеются. Если в загруженном словаре состояния их нет, param_names оптимизатора останется без изменений.

Пример

>>> optimizer = ...  # initialized optimizer matching the saved state
>>> scheduler1 = torch.optim.lr_scheduler.LinearLR(
...     optimizer,
...     start_factor=0.1,
...     end_factor=1,
...     total_iters=20,
... )
>>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
...     optimizer,
...     T_max=80,
...     eta_min=3e-5,
... )
>>> lr = torch.optim.lr_scheduler.SequentialLR(
...     optimizer,
...     schedulers=[scheduler1, scheduler2],
...     milestones=[20],
... )
>>> lr.load_state_dict(torch.load("./save_seq.pt"))
>>> # now load the optimizer checkpoint after loading the LRScheduler
>>> optimizer.load_state_dict(torch.load("./save_optim.pt"))

register_load_state_dict_post_hook(hook, prepend=False) [исходный код]

Регистрирует постхук load_state_dict, который вызывается после вызова load_state_dict(). Он должен иметь следующую сигнатуру:

hook(optimizer) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора.

Хук вызывается с аргументом self после вызова load_state_dict для self. Зарегистрированный хук можно использовать для постобработки после того, как load_state_dict загрузил state_dict.

Параметры:

• hook (Callable) – пользовательский хук для регистрации.
• prepend (bool) – если True, предоставленный постхук hook будет вызван перед всеми уже зарегистрированными постхуками для load_state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных постхуков. (по умолчанию: False)

Возвращает:

дескриптор, с помощью которого можно удалить добавленный хук, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook, prepend=False) [исходный код]

Регистрирует предхук load_state_dict, который вызывается перед вызовом load_state_dict(). Он должен иметь следующую сигнатуру:

hook(optimizer, state_dict) -> state_dict or None

Аргумент optimizer — это используемый экземпляр оптимизатора, а аргумент state_dict — поверхностная копия state_dict, переданного пользователем в load_state_dict. Хук может изменить state_dict на месте или, при желании, вернуть новый. Если возвращён state_dict, он будет загружен в оптимизатор.

Хук вызывается с аргументами self и state_dict перед вызовом load_state_dict для self. Зарегистрированный хук можно использовать для предварительной обработки перед вызовом load_state_dict.

Параметры:

• hook (Callable) – пользовательский хук для регистрации.
• prepend (bool) – если True, предоставленный предхук hook будет вызван перед всеми уже зарегистрированными предхуками для load_state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных предхуков. (по умолчанию: False)

Возвращает:

дескриптор, с помощью которого можно удалить добавленный хук, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_state_dict_post_hook(hook, prepend=False) [исходный код]

Регистрирует постхук словаря состояния, который вызывается после вызова state_dict().

Он должен иметь следующую сигнатуру:

hook(optimizer, state_dict) -> state_dict or None

Хук вызывается с аргументами self и state_dict после создания state_dict для self. Хук может изменить state_dict на месте или, при желании, вернуть новый. Зарегистрированный хук можно использовать для постобработки state_dict перед его возвратом.

Параметры:

• hook (Callable) – пользовательский хук для регистрации.
• prepend (bool) – если True, предоставленный постхук hook будет вызван перед всеми уже зарегистрированными постхуками для state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных постхуков. (по умолчанию: False)

Возвращает:

дескриптор, с помощью которого можно удалить добавленный хук, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_state_dict_pre_hook(hook, prepend=False) [исходный код]

Регистрирует предхук словаря состояния, который вызывается перед вызовом state_dict().

Он должен иметь следующую сигнатуру:

hook(optimizer) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора. Хук вызывается с аргументом self перед вызовом state_dict для self. Зарегистрированный хук можно использовать для предварительной обработки перед вызовом state_dict.

Параметры:

• hook (Callable) – пользовательский хук для регистрации.
• prepend (bool) – если True, предоставленный предхук hook будет вызван перед всеми уже зарегистрированными предхуками для state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных предхуков. (по умолчанию: False)

Возвращает:

дескриптор, с помощью которого можно удалить добавленный хук, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_step_post_hook(hook) [исходный код]

Регистрирует постхук шага оптимизатора, который вызывается после выполнения шага оптимизатора.

Он должен иметь следующую сигнатуру:

hook(optimizer, args, kwargs) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора.

Параметры:

hook (Callable) – пользовательский хук для регистрации.

Возвращает:

дескриптор, с помощью которого можно удалить добавленный хук, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_step_pre_hook(hook) [исходный код]

Зарегистрировать предварительный хук шага оптимизатора, который будет вызываться перед шагом оптимизатора.

Он должен иметь следующую сигнатуру:

hook(optimizer, args, kwargs) -> None or modified args and kwargs

Аргумент optimizer — это используемый экземпляр оптимизатора. Если предварительный хук изменяет args и kwargs, преобразованные значения возвращаются в виде кортежа, содержащего new_args и new_kwargs.

Параметры:

hook (Callable) – Пользовательский хук, который необходимо зарегистрировать.

Возвращает:

дескриптор, который можно использовать для удаления добавленного хука, вызвав handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

state_dict() [исходный код]

Возвращает состояние оптимизатора в виде dict.

Оно содержит две записи:

• state: a Dict holding current optimization state. Its content

  различается в зависимости от класса оптимизатора, но некоторые общие характеристики сохраняются. Например, состояние сохраняется для каждого параметра, а сам параметр НЕ сохраняется. state — это словарь, сопоставляющий идентификаторы параметров со словарём, содержащим состояние каждого параметра.

• param_groups: a List containing all parameter groups where each

  группа параметров — это словарь. Каждая группа параметров содержит метаданные, относящиеся к оптимизатору, например скорость обучения и затухание весов, а также список идентификаторов параметров, входящих в группу. Если группа параметров была инициализирована с помощью named_parameters(), в словаре состояния также будут сохранены имена.

ПРИМЕЧАНИЕ: Идентификаторы параметров могут выглядеть как индексы, но это всего лишь идентификаторы, связывающие состояние с param_group. При загрузке из state_dict оптимизатор сопоставляет элементы params param_group (целочисленные идентификаторы) и param_groups оптимизатора (фактические nn.Parameter s), чтобы сопоставить состояния БЕЗ дополнительной проверки.

Возвращаемый словарь состояния может выглядеть примерно так:

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}

Тип возвращаемого значения:

dict[str, Any]

step(closure=None) [исходный код]

Выполняет один шаг оптимизации.

zero_grad(set_to_none=True) [исходный код]

Сбрасывает градиенты всех оптимизируемых torch.Tensor s.

Параметры:

set_to_none (bool, необязательно) –

Вместо установки нулевого значения устанавливает для градиентов значение None. По умолчанию: True

Как правило, это уменьшает потребление памяти и может несколько повысить производительность. Однако это меняет некоторые особенности поведения. Например:

1. Если пользователь попытается получить доступ к градиенту и выполнить над ним операции вручную, атрибут со значением None и тензор, заполненный нулями, будут вести себя по-разному.
2. Если пользователь запросит zero_grad(set_to_none=True), а затем выполнит обратное распространение, для .grads параметров, не получивших градиент, гарантированно будет установлено значение None.
3. Оптимизаторы torch.optim ведут себя по-разному, если градиент равен 0 или None (в первом случае выполняется шаг с градиентом 0, а во втором шаг полностью пропускается).