Skip to content

Commit e65c9a5

Browse files
committed
incorporate the extra scale on the muon update proposed by kimi team with conjunction with Jianlin Su
1 parent 3510da1 commit e65c9a5

File tree

3 files changed

+30
-3
lines changed

3 files changed

+30
-3
lines changed

README.md

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,3 +99,15 @@ for _ in range(100):
9999
url = {https://kellerjordan.github.io/posts/muon/}
100100
}
101101
```
102+
103+
```bibtex
104+
@misc{kimiteam2025kimik2openagentic,
105+
title = {Kimi K2: Open Agentic Intelligence},
106+
author = {Kimi Team and Yifan Bai and Yiping Bao and Guanduo Chen and Jiahao Chen and Ningxin Chen and Ruijue Chen and Yanru Chen and Yuankun Chen and Yutian Chen and Zhuofu Chen and Jialei Cui and Hao Ding and Mengnan Dong and Angang Du and Chenzhuang Du and Dikang Du and Yulun Du and Yu Fan and Yichen Feng and Kelin Fu and Bofei Gao and Hongcheng Gao and Peizhong Gao and Tong Gao and Xinran Gu and Longyu Guan and Haiqing Guo and Jianhang Guo and Hao Hu and Xiaoru Hao and Tianhong He and Weiran He and Wenyang He and Chao Hong and Yangyang Hu and Zhenxing Hu and Weixiao Huang and Zhiqi Huang and Zihao Huang and Tao Jiang and Zhejun Jiang and Xinyi Jin and Yongsheng Kang and Guokun Lai and Cheng Li and Fang Li and Haoyang Li and Ming Li and Wentao Li and Yanhao Li and Yiwei Li and Zhaowei Li and Zheming Li and Hongzhan Lin and Xiaohan Lin and Zongyu Lin and Chengyin Liu and Chenyu Liu and Hongzhang Liu and Jingyuan Liu and Junqi Liu and Liang Liu and Shaowei Liu and T. Y. Liu and Tianwei Liu and Weizhou Liu and Yangyang Liu and Yibo Liu and Yiping Liu and Yue Liu and Zhengying Liu and Enzhe Lu and Lijun Lu and Shengling Ma and Xinyu Ma and Yingwei Ma and Shaoguang Mao and Jie Mei and Xin Men and Yibo Miao and Siyuan Pan and Yebo Peng and Ruoyu Qin and Bowen Qu and Zeyu Shang and Lidong Shi and Shengyuan Shi and Feifan Song and Jianlin Su and Zhengyuan Su and Xinjie Sun and Flood Sung and Heyi Tang and Jiawen Tao and Qifeng Teng and Chensi Wang and Dinglu Wang and Feng Wang and Haiming Wang and Jianzhou Wang and Jiaxing Wang and Jinhong Wang and Shengjie Wang and Shuyi Wang and Yao Wang and Yejie Wang and Yiqin Wang and Yuxin Wang and Yuzhi Wang and Zhaoji Wang and Zhengtao Wang and Zhexu Wang and Chu Wei and Qianqian Wei and Wenhao Wu and Xingzhe Wu and Yuxin Wu and Chenjun Xiao and Xiaotong Xie and Weimin Xiong and Boyu Xu and Jing Xu and Jinjing Xu and L. H. Xu and Lin Xu and Suting Xu and Weixin Xu and Xinran Xu and Yangchuan Xu and Ziyao Xu and Junjie Yan and Yuzi Yan and Xiaofei Yang and Ying Yang and Zhen Yang and Zhilin Yang and Zonghan Yang and Haotian Yao and Xingcheng Yao and Wenjie Ye and Zhuorui Ye and Bohong Yin and Longhui Yu and Enming Yuan and Hongbang Yuan and Mengjie Yuan and Haobing Zhan and Dehao Zhang and Hao Zhang and Wanlu Zhang and Xiaobin Zhang and Yangkun Zhang and Yizhi Zhang and Yongting Zhang and Yu Zhang and Yutao Zhang and Yutong Zhang and Zheng Zhang and Haotian Zhao and Yikai Zhao and Huabin Zheng and Shaojie Zheng and Jianren Zhou and Xinyu Zhou and Zaida Zhou and Zhen Zhu and Weiyu Zhuang and Xinxing Zu},
107+
year = {2025},
108+
eprint = {2507.20534},
109+
archivePrefix = {arXiv},
110+
primaryClass = {cs.LG},
111+
url = {https://arxiv.org/abs/2507.20534},
112+
}
113+
```

adam_atan2_pytorch/muon_adam_atan2.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from __future__ import annotations
2+
import math
23
from typing import Callable
34

45
import torch
@@ -66,6 +67,7 @@ def __init__(
6667
cautious_factor = 1., # set to 0. for zeroing out any updates not in same direction as gradient as in https://arxiv.org/abs/2411.16085
6768
a = 1.27,
6869
b = 1.,
70+
muon_rms_factor = 0.2,
6971
muon_steps = 5,
7072
muon_beta1 = 0.95,
7173
muon_newton_schulz5_coefs = (3.4445, -4.7750, 2.0315),
@@ -101,14 +103,15 @@ def __init__(
101103
muon_steps = muon_steps,
102104
muon_newton_schulz5_coefs = muon_newton_schulz5_coefs,
103105
muon_eps = muon_eps,
106+
muon_rms_factor = muon_rms_factor,
104107
)
105108

106109
if remove_muon_params_from_params:
107110
params = list(set(params) - set(muon_params))
108111

109112
param_groups = [
110113
dict(params = params, lr = lr),
111-
dict(params = muon_params, lr = muon_lr, beta1 = muon_beta1, use_muon = True)
114+
dict(params = muon_params, lr = muon_lr, beta1 = muon_beta1, rms_factor = muon_rms_factor, use_muon = True)
112115
]
113116

114117
super().__init__(param_groups, defaults)
@@ -134,6 +137,11 @@ def step(
134137

135138
param_init_lr = init_lr if not use_muon else init_muon_lr
136139

140+
# set lr scale to 1. if muon update
141+
142+
if use_muon:
143+
a = 1.
144+
137145
# maybe decoupled weight decay
138146

139147
if self.decoupled_wd:
@@ -193,7 +201,7 @@ def step(
193201
# maybe cautious update - algorithm 2 in https://arxiv.org/abs/2411.16085
194202
else:
195203

196-
muon_steps, muon_coefs, muon_eps = group['muon_steps'], group['muon_newton_schulz5_coefs'], group['muon_eps']
204+
muon_steps, muon_coefs, muon_eps, muon_rms_factor = group['muon_steps'], group['muon_newton_schulz5_coefs'], group['muon_eps'], group['muon_rms_factor']
197205

198206
# Muon from Keller Jordan
199207
# https://kellerjordan.github.io/posts/muon/
@@ -205,6 +213,13 @@ def step(
205213
eps = muon_eps
206214
)
207215

216+
# incorporate the match adam RMS from Kimi team
217+
# https://kexue.fm/archives/11267
218+
219+
muon_update_scale = math.sqrt(max(exp_avg.shape[-2:])) * muon_rms_factor
220+
221+
update = update * muon_update_scale
222+
208223
if cautious_factor < 1.:
209224
align_mask = (update * grad) > 0
210225
scale = torch.where(align_mask, torch.ones_like(grad), cautious_factor)

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "adam-atan2-pytorch"
3-
version = "0.2.2"
3+
version = "0.2.3"
44
description = "Adam-atan2 for Pytorch"
55
authors = [
66
{ name = "Phil Wang", email = "lucidrains@gmail.com" }

0 commit comments

Comments
 (0)