tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-04-06 07:35:32 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 8be3907514 into 1144d186ed

Browse Source
This commit is contained in:
zichongli5 2026-02-05 10:43:32 +08:00 committed by GitHub
commit fb9a32a9f3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
nanochat/optim.py
View File

@ -247,6 +247,7 @@ class MuonAdamW(torch.optim.Optimizer):
momentum_buffer = state["momentum_buffer"]
# Second momentum buffer is factored, either per-row or per-column
# from NorMuon: https://arxiv.org/abs/2510.05491
if "second_momentum_buffer" not in state:
state_shape = (num_params, shape[-2], 1) if shape[-2] >= shape[-1] else (num_params, 1, shape[-1])
state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)