Skip to content

Commit a886c3d

Browse files
committed
refactor(muon): remove unused individually-compiled NS functions
torch.compile on muon_update already traces into the NS functions it calls, so the separately compiled wrappers are redundant. Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
1 parent 0d7f125 commit a886c3d

1 file changed

Lines changed: 0 additions & 3 deletions

File tree

deepspeed/runtime/zero/muon/original_muon.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -151,9 +151,6 @@ def _muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method
151151
return update
152152

153153

154-
# Compiled (torch.compile) versions for GPU paths
155-
zeropower_via_newtonschulz5 = compiler.compile()(_zeropower_via_newtonschulz5)
156-
zeropower_via_gram_newtonschulz = compiler.compile()(_zeropower_via_gram_newtonschulz)
157154
muon_update = compiler.compile()(_muon_update)
158155

159156

0 commit comments

Comments
 (0)