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{輸入} : \gamma \text{ (學習率)},\ \lambda \text{ (權重衰減)},\ \mu \text{ (動量)},\ \textit{nesterov}\in\{True,False\},\\ &\hspace{13mm}(a,b,c)\ \text{ (NS 係數)},\ \varepsilon \text{ (epsilon)},\ k \text{ (NS 步數)},\ \theta_0 \text{ (引數)},\ f(\theta) \text{ (目標函式)} \\ &\textbf{初始化} : B_0 \leftarrow 0 \text{ (動量緩衝區)} \\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{對於}\ t=1\ \textbf{到}\ \ldots\ \textbf{執行} \\[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{如果 nesterov}=True \\ B_t, & \text{如果 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{(解耦權重衰減)} \\[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{返回}\ \theta_t \\[-1.ex] &\rule{110mm}{0.4pt}s \end{aligned}

這裡， $\mathrm{NS}^{(a,b,c)}_{k}(\cdot;\varepsilon)$ 表示具有數值穩定 $\varepsilon$ 的、由係數 $(a,b,c)$ 引數化的 Newton–Schulz 正交化運算元 $k$ 次迭代。

$\mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big)$ 的目的在於使正交化更新後的 $RMS$ 在處理矩形矩陣時保持一致。

Keller 的原始實現透過 $\sqrt{\max\!\left(1, \frac{A}{B}\right)}$ 縮放更新，其中 $A$ 和 $B$ 是被最佳化矩陣的維度。

Moonshot 的實現也致力於匹配 AdamW 的 $RMS$ 。調整計算如下： $\gamma \leftarrow {0.2}\gamma\,\sqrt{\max\!\left({A}, {B}\right)}$ 該方法改編自 Muon is Scalable for LLM Training。研究結果表明，透過此調整，Muon 可以直接複用為 AdamW 調整過的學習率和權重衰減。

我們提供了兩種學習率調整的選項：“original”（原始），遵循 Keller 的實現；以及“match_rms_adamw”，參考 Moonshot 的實現。這為使用者提供了選擇兩者的靈活性。如果未指定 adjust_lr_fn，則預設使用“original”。

有關該演算法的更多詳細資訊，請參閱 Muon: An optimizer for hidden layers in neural networks 和 Muon is Scalable for LLM Training。

引數

params (iterable) – 要最佳化的引數或命名引數的可迭代物件，或定義引數組的可迭代字典。當使用命名引數時，所有組中的引數都應命名。請注意，Muon 是用於神經網路隱藏層二維引數的最佳化器。偏差和嵌入等其他引數應由標準方法（如 AdamW）進行最佳化。
lr (float, Tensor, optional) – 學習率（預設為 1e-3）。
weight_decay (float, optional) – 權重衰減（L2 正則化）。（預設為 0.1）
momentum (float, optional) – 動量因子（預設為 0.95）。
nesterov (bool, optional) – 啟用 Nesterov 動量。僅當動量非零時適用。
ns_coefficients (tuple of three floats, optional) – Newton–Schulz 正交化多項式的係數（a,b,c）（預設為 (3.4445, -4.775, 2.0315)）。
eps (float, optional) – 新增到分母以提高數值穩定性的項。（預設為 1e-07）。
ns_steps (int, optional) – Newton–Schulz 迭代步數。（預設為 5）。
adjust_lr_fn (str, optional) – 用於調整學習率的函式。可以是“original”或“match_rms_adamw”。如果未指定，則預設為使用“original”。（預設為 None）。

add_param_group(param_group)[source]#

向 Optimizer 的 param_groups 新增一個引數組。

這在微調預訓練網路時可能很有用，因為隨著訓練的進行，可以使凍結的層變得可訓練並新增到 Optimizer 中。

引數: param_group (dict) – 指定哪些 Tensor 應該被最佳化，以及組特定的最佳化選項。

load_state_dict(state_dict)[source]#

載入最佳化器狀態。

引數: state_dict (dict) – 最佳化器狀態。應為呼叫 state_dict() 返回的物件。

警告

請確保在初始化 torch.optim.lr_scheduler.LRScheduler 後呼叫此方法，因為在此之前呼叫會覆蓋載入的學習率。

注意

引數的名稱（如果它們存在於每個引數組的“param_names”鍵下，在 state_dict() 中）不會影響載入過程。為了使用引數名稱進行自定義（例如，當載入的狀態字典中的引數與最佳化器中初始化的引數不同時），應實現自定義的 register_load_state_dict_pre_hook 來相應地調整載入的字典。如果載入的狀態字典 param_groups 中存在 param_names，它們將被儲存並覆蓋最佳化器狀態中當前存在的名稱。如果它們不存在於載入的狀態字典中，最佳化器 param_names 將保持不變。

示例

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

register_load_state_dict_post_hook(hook, prepend=False)[source]#

註冊一個 load_state_dict 後置鉤子，它將在呼叫 load_state_dict() 後被呼叫。它應該具有以下簽名：

hook(optimizer) -> None

引數 optimizer 是正在使用的最佳化器例項。

呼叫 load_state_dict 到 self 上後，鉤子將使用引數 self 呼叫。註冊的鉤子可用於在 load_state_dict 載入了 state_dict 後執行後處理。

引數

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 True，則提供的後置 hook 將在 load_state_dict 上所有已註冊的後置鉤子之前執行。否則，提供的 hook 將在所有已註冊的後置鉤子之後執行。(預設: False)

返回

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_load_state_dict_pre_hook(hook, prepend=False)[source]#

註冊一個 load_state_dict 前置鉤子，它將在呼叫 load_state_dict() 之前被呼叫。它應該具有以下簽名：

hook(optimizer, state_dict) -> state_dict or None

引數 optimizer 是正在使用的最佳化器例項，引數 state_dict 是使用者傳遞給 load_state_dict 的 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.RemoveableHandle

register_state_dict_post_hook(hook, prepend=False)[source]#

註冊一個 state_dict 後置鉤子，它將在呼叫 state_dict() 後被呼叫。

它應具有以下簽名

hook(optimizer, state_dict) -> state_dict or None

鉤子將使用引數 self 和 state_dict 呼叫，在 self 上生成 state_dict 後。鉤子可以就地修改 state_dict，或者選擇性地返回一個新的。註冊的鉤子可用於在返回 state_dict 之前對其進行後處理。

引數

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 True，則提供的後置 hook 將在 state_dict 上所有已註冊的後置鉤子之前執行。否則，提供的 hook 將在所有已註冊的後置鉤子之後執行。(預設: False)

返回

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_state_dict_pre_hook(hook, prepend=False)[source]#

註冊一個 state_dict 前置鉤子，它將在呼叫 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.RemoveableHandle

register_step_post_hook(hook)[source]#

註冊一個最佳化器步驟後鉤子，它將在最佳化器步驟之後被呼叫。

它應具有以下簽名

hook(optimizer, args, kwargs) -> None

引數 optimizer 是正在使用的最佳化器例項。

引數: hook (Callable) – 使用者定義的待註冊鉤子。
返回: 一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子
返回型別: torch.utils.hooks.RemovableHandle

register_step_pre_hook(hook)[source]#

註冊一個最佳化器步驟預鉤子，它將在最佳化器步驟之前被呼叫。

它應具有以下簽名

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()[source]#

將最佳化器的狀態作為 dict 返回。

它包含兩個條目

state：一個包含當前最佳化狀態的 Dict。其內容
在不同的最佳化器類中會有所不同，但有一些共同的特點。例如，狀態是按引數儲存的，而引數本身不儲存。 state 是一個對映引數 ID 到一個包含每個引數對應狀態的 Dict 的字典。
param_groups：一個包含所有引數組的 List，其中每個
引數組是一個 Dict。每個引數組包含最佳化器特有的元資料，例如學習率和權重衰減，以及組中引數的 ID 列表。如果引數組使用 named_parameters() 初始化，則名稱內容也會儲存在狀態字典中。

注意：引數 ID 可能看起來像索引，但它們只是將狀態與 param_group 關聯的 ID。從 state_dict 載入時，最佳化器會按順序匹配 param_group 的 params（int ID）和最佳化器的 param_groups（實際的 nn.Parameter），以匹配狀態，而無需額外驗證。

返回的狀態字典可能看起來像

{
    '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)[source]#: 執行一次最佳化步驟。

zero_grad(set_to_none=True)[source]#

重置所有最佳化過的 torch.Tensor 的梯度。

引數

set_to_none (bool, optional) –

將梯度設定為 None，而不是設定為零。預設值：True

這通常會降低記憶體佔用，並能適度提高效能。但是，它會改變某些行為。例如：

當用戶嘗試訪問梯度並對其進行手動運算時，None 屬性或全零的 Tensor 會產生不同的行為。
如果使用者請求 zero_grad(set_to_none=True) 然後執行 backward，對於未收到梯度的引數，其 .grad 保證為 None。
torch.optim 最佳化器在梯度為 0 或 None 時行為不同（一種情況是以 0 梯度執行步長，另一種情況是跳過該步長）。

Muon#

文件

教程

資源