refactor(muon): remove torch.compile from muon_update

Benchmarks show torch.compile provides negligible speedup (0-5%)
and is sometimes slower for muon_update, since computation is
dominated by matmul which already uses cuBLAS in eager mode.

Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>

Author: delock

deepspeed/runtime/zero/muon/original_muon.py

@@ -29,7 +29,6 @@
 
 import torch
 import deepspeed.comm as dist  # replace torch's distributed package with deepspeed.comm to resolve deepspeed check
-from deepspeed.runtime import compiler
 from deepspeed.accelerator import get_accelerator
 
 
@@ -135,7 +134,7 @@ def _zeropower_via_gram_newtonschulz(G, steps: int):
 NS_METHODS = {"standard", "gram"}
 
 
-def _muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram"):
+def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram"):
     orig_dtype = grad.dtype
     momentum.lerp_(grad, 1 - beta)
     update = grad.lerp_(momentum, beta) if nesterov else momentum
@@ -151,9 +150,6 @@ def _muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method
     return update
 
 
-muon_update = compiler.compile()(_muon_update)
-
-
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz