refactor(muon): change argument adam_wd to weight_decay and handle params' type

build/torch26-cxx11-cu118-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9c77e5647b6056bfaee25050cca7948c40859db0a88fa4fcf40b67a85c947d8c
+oid sha256:f1f5df341112d93e43c0801e285abd66e79bfbe399d228f8be09ff26ece7421b
 size 1787272

build/torch26-cxx11-cu118-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx11-cu124-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:94ea66089cc8d9eda72b017733a9e05e4fee5a2f04c50658b690d2c19f0d3068
+oid sha256:2921aa2aa2587e261dc9ca4e5f60303b0d1c9a305d1584918a8c56b6dc79ebfb
 size 1824224

build/torch26-cxx11-cu124-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx11-cu126-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:46e01e1d957ada2d485b30cd60bc3ef7230b8857dffc59f2e7924339761ec577
+oid sha256:a93530e6981fdac23236dd7e3657c5b47513cda4accec78293234ce5f233400b
 size 1824224

build/torch26-cxx11-cu126-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx11-rocm62-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx11-rocm62-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a825a0cd31d8c1b91aa9db4b24248d7fc0a506615f625a385b40e6002025c7dd
+oid sha256:caa40905ac8f209fecccae42c6892c3766ad5c7069382e60d2339e73da6ee7d6
 size 1749744

build/torch26-cxx11-rocm62-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx98-cu118-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:579e9ddf66a4f17ead9232c2f32e6327fe6a3f16dd235e2e73e6cb282de1797e
+oid sha256:6919551ed599e7e0dc1a750d1972bdb31605f57583b3617054cb70dd40d54d26
 size 1787192

build/torch26-cxx98-cu118-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx98-cu124-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:beacb4ba2d56463b6d444875728b3462cb3ff6c1449e3c9693cd665bfbbbbb73
+oid sha256:f07cc2637669130fc9e209cb2c4358caba1c4c2d5837a108043b073d7897c3a7
 size 1824184

build/torch26-cxx98-cu124-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch26-cxx98-cu126-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9b04b011803d328d8dcd2edcf4c3840ddbb1bb2f093464c208f0ba2faf4f16bc
+oid sha256:8b9ef8fa2dd4d80cb3c1c3c2a72b99e0d76b3e676acd551f3a9ff4cdd21773eb
 size 1824184

build/torch26-cxx98-cu126-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch27-cxx11-cu118-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad6c725009f2e776b99d3134c75f15e11dd7fe75fe4ba1fa94779018c7871f8c
+oid sha256:8413f32011996384f13a985a99b4e2f863f8e4717acdb8439b63a10f77db6f15
 size 1787368

build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch27-cxx11-cu126-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch27-cxx11-cu126-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:50cb5819ff08a2179d78cd98164d07fd3cef1b66ee7703d599a310dfb140b9d1
+oid sha256:c9d303b11a0a82e9c51c7b32c7555bd351ec375b1879bf46eb64ea4aff32100f
 size 1824256

build/torch27-cxx11-cu126-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch27-cxx11-cu128-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch27-cxx11-cu128-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9c75e42265f382addc71327ad5628e8a2414da5872791c975e384708c4acd549
+oid sha256:b4334fe8d7157c2a9c85217cb981692daf9eb4c6d3f205d0fd41d4b717daefa1
 size 1883352

build/torch27-cxx11-cu128-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_036642a_dirty
-ops = torch.ops._optimizer_036642a_dirty
+from . import _optimizer_64757cb_dirty
+ops = torch.ops._optimizer_64757cb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_036642a_dirty::{op_name}"
+    return f"_optimizer_64757cb_dirty::{op_name}"

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_036642a_dirty.abi3.so → _optimizer_64757cb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a2363d4311d6a75fbcc03e6d4a71c73dae4d54e00a30135d25198d4078c6b0f
+oid sha256:272fcc69e3774fa43e222efefceeaca97a8c84ee3f1fe528a7478a8e80a70976
 size 1749648

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad

torch-ext/optimizer/muon.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor
+from torch.distributed._tensor import DTensor, Replicate
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -103,7 +103,7 @@ def _compute_u(state, steps, rank, compute_stream):
 
 
 @torch.no_grad()
-def _scatter(p, state, lr, wd, rank, comm_stream):
+def _scatter(p, state, lr, weight_decay, rank, comm_stream):
     u = state.computed_u
     mesh = p.device_mesh
 
@@ -131,10 +131,14 @@ def _scatter(p, state, lr, wd, rank, comm_stream):
             placements=p.placements,
             device_mesh=mesh,
         )
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         p.data.add_(u, alpha=-lr)
 
 
+def default_is_muon(x, name):
+    return x.ndim >= 2 and "embed_tokens" not in name and "lm_head" not in name
+
+
 class Muon(torch.optim.Optimizer):
     """
     Muon - MomentUm Orthogonalized by Newton-schulz
@@ -159,18 +163,18 @@ class Muon(torch.optim.Optimizer):
         adamw_lr: The learning rate for the internal AdamW.
         adamw_betas: The betas for the internal AdamW.
         adamw_eps: The epsilon for the internal AdamW.
-        adamw_wd: The weight decay for the internal AdamW.
+        adamw_weight_decay: The weight decay for the internal AdamW.
     """
 
     def __init__(
         self,
         model,
-        is_muon_func,
+        is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
         nesterov=True,
         ns_steps=5,
-        adamw_wd=0.1,
+        weight_decay=0.1,
         adamw_betas=(0.9, 0.95),
         adamw_eps=1e-8,
         none_grad=True,
@@ -178,7 +182,7 @@ class Muon(torch.optim.Optimizer):
     ):
         defaults = dict(
             lr=lr,
-            wd=adamw_wd,
+            weight_decay=weight_decay,
             momentum=momentum,
             nesterov=nesterov,
             ns_steps=ns_steps,
@@ -272,7 +276,7 @@ class Muon(torch.optim.Optimizer):
 
         return param_to_state, ordered_params
 
-    def base(self, params, group, lr, wd, momentum):
+    def base(self, params, group, lr, weight_decay, momentum):
         # generate weight updates in distributed fashion
         for p in params:
             g = p.grad
@@ -299,7 +303,7 @@ class Muon(torch.optim.Optimizer):
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
 
             # apply weight decay
-            p.data.mul_(1 - lr * wd)
+            p.data.mul_(1 - lr * weight_decay)
 
             # apply update
             p.data.add_(u, alpha=-adjusted_lr)
@@ -317,15 +321,15 @@ class Muon(torch.optim.Optimizer):
             g = buf
         return g
 
-    def _update_p(self, p, u, lr, wd):
+    def _update_p(self, p, u, lr, weight_decay):
         # scale update
         adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
         # apply weight decay
-        p.data.mul_(1 - lr * wd)
+        p.data.mul_(1 - lr * weight_decay)
         # apply update
         p.data.add_(u, alpha=-adjusted_lr)
 
-    def parallel(self, params, group, lr, wd, momentum):
+    def parallel(self, params, group, lr, weight_decay, momentum):
         """
         Perform a parallel optimization step using Muon.
         """
@@ -364,7 +368,9 @@ class Muon(torch.optim.Optimizer):
             for p in ordered_params[start_idx : start_idx + chunk_size]:
                 state = param_to_state[id(p)]
                 adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
-                _scatter(p, state, adjusted_lr, wd, self.rank, self.comm_stream)
+                _scatter(
+                    p, state, adjusted_lr, weight_decay, self.rank, self.comm_stream
+                )
 
         chunk_size = params[0].device_mesh.mesh.numel()
 
@@ -398,23 +404,48 @@ class Muon(torch.optim.Optimizer):
 
             params = [p for p in group["params"] if self.state[p]["use_muon"]]
             lr = group["lr"]
-            wd = group["wd"]
+            weight_decay = group["weight_decay"]
             momentum = group["momentum"]
 
-            if isinstance(params[0].data, DTensor):
+            param_dtensors = []
+            param_tensors = []
+
+            for p in params:
+                if p is None or p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    if all(
+                        isinstance(placement, Replicate) for placement in p.placements
+                    ):
+                        param_tensors.append(p)
+                    else:
+                        param_dtensors.append(p)
+                elif isinstance(p.data, torch.Tensor):
+                    param_tensors.append(p)
+                else:
+                    raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+            if self.debug:
+                print(
+                    f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
+                    flush=True,
+                )
+
+            if len(param_dtensors) > 0:
                 self.parallel(
-                    params,
+                    param_dtensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
-            else:
+
+            if len(param_tensors) > 0:
                 self.base(
-                    params,
+                    param_tensors,
                     group,
                     lr=lr,
-                    wd=wd,
+                    weight_decay=weight_decay,
                     momentum=momentum,
                 )
 
@@ -426,7 +457,7 @@ class Muon(torch.optim.Optimizer):
             lr = group["lr"]
             beta1, beta2 = group["adamw_betas"]
             eps = group["adamw_eps"]
-            weight_decay = group["wd"]
+            weight_decay = group["weight_decay"]
 
             for p in params:
                 g = p.grad