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

Add NorMuon paper link

Browse Source
This commit is contained in:
zichongli5 2026-02-03 21:53:43 -05:00 committed by GitHub
parent 542beb0c8c
commit 8be3907514
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

@ -243,6 +243,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)