Muon optimizer: expert batching, parallel caching, A2A overlap [skip-build]

- Batched expert NS path for plain-tensor MoE params (skip expansion)
- Expert expansion cache to eliminate per-step detach overhead
- _setup_parallel() extraction for parallel metadata reuse
- Prelaunch first chunk A2A gather to overlap with expert NS compute
- Profiler annotations and clarify distributed_muon as test-only

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

torch-ext/optimizer/muon.py

@@ -10,14 +10,15 @@ from torch.profiler import record_function
 
 from .adamw import step_adamw
 from .async_utils import run_pipeline
-from .core import (_muon_state, adjust_lr_for_muon,
-                   get_default_muon_param_groups, is_expert_param, update_g,
-                   update_p)
+from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho,
+                   get_default_muon_param_groups, is_expert_param, update_p)
 from .distributed.utils import (_is_shard, construct_shard_mesh,
                                 get_slices_of_dtensor)
 from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO,
-                            _zeropower_via_newtonschulz5)
-from .pipeline import muon_chunk_pipeline
+                            _zeropower_via_newtonschulz5,
+                            zeropower_via_newtonschulz5,
+                            zeropower_via_newtonschulz5_batched)
+from .pipeline import muon_chunk_pipeline, prelaunch_first_gather
 from .qk_clip import compute_scales, get_qk_clip_info, qk_clip
 
 logger = logging.getLogger(__name__)
@@ -49,6 +50,18 @@ def _expand_expert_params(names, params, expert_keys):
         is_expert = is_expert_param(n, expert_keys)
         is_dtensor = isinstance(p.data, DTensor)
 
+        if is_expert:
+            if is_dtensor:
+                logger.debug(
+                    "[expand_expert] %s: expert DTensor, shape=%s, "
+                    "placements=%s, mesh=%s, local_shape=%s", n, p.shape,
+                    p.placements, p.device_mesh.mesh_dim_names,
+                    p.to_local().shape)
+            else:
+                logger.debug(
+                    "[expand_expert] %s: expert plain tensor, shape=%s", n,
+                    p.data.shape)
+
         if not is_expert:
             assert p.data.ndim <= 2, (
                 f"Param {n} has ndim={p.data.ndim} but does not match "
@@ -169,7 +182,6 @@ class Muon(torch.optim.Optimizer):
                      Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
         use_distributed_muon: Use distributed muon by Liu et al. (2024).
                               For testing purpose only.
-        small_param_numel_threshold: Threshold for classifying parameters as small and falling back to distributed Muon
         expert_keys: List of strings to identify expert-parallel parameters.
                      If any key appears in a parameter's name, its outermost
                      dimension is treated as the expert dimension and expanded
@@ -194,7 +206,6 @@ class Muon(torch.optim.Optimizer):
                  warmup_step=5,
                  chunk_size=-1,
                  use_distributed_muon=False,
-                 small_param_numel_threshold=65536,
                  expert_keys=None):
         defaults = dict(
             lr=lr,
@@ -229,8 +240,9 @@ class Muon(torch.optim.Optimizer):
         self.warmup_step = warmup_step
         self.chunk_size = chunk_size
         self.use_distributed_muon = use_distributed_muon
-        self.small_param_numel_threshold = small_param_numel_threshold
         self.expert_keys = expert_keys
+        self._parallel_cache: dict[tuple[str, ...], dict] = {}
+        self._expert_expand_cache: dict[tuple[int, ...], dict] = {}
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -334,8 +346,8 @@ class Muon(torch.optim.Optimizer):
             if g is None:
                 continue
 
-            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                             steps=group["ns_steps"])
+            u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                            steps=group["ns_steps"])
 
             adjusted_lr = adjust_lr_for_muon(lr, p.shape)
             update_p(p, u, lr, adjusted_lr, weight_decay)
@@ -356,52 +368,269 @@ class Muon(torch.optim.Optimizer):
         weight_decay: float,
         qk_logits: list[torch.Tensor | DTensor] | None,
     ):
-        """ Implementation of Distributed Muon by Liu et al. """
+        """Batched Distributed Muon — for testing/correctness verification only.
 
-        # Momentum is already applied by _step_muon before this method.
-        for n, p in zip(names, params):
-            g = p.grad
-            if g is None:
-                continue
-
-            # Gather G
-            if isinstance(p.data, DTensor):
-                g_full = g.full_tensor()
-                p_full = p.data.full_tensor()
-            else:
-                g_full = g
-                p_full = p
-
-            u_full = _zeropower_via_newtonschulz5(g_full.to(COMM_DTYPE),
-                                                  steps=group["ns_steps"])
-
-            adjusted_lr = adjust_lr_for_muon(lr, p_full.shape)
-            update_p(p_full, u_full, lr, adjusted_lr, weight_decay)
+        Uses all-gather to reconstruct full tensors, computes Newton-Schulz on
+        the full grad, then slices back to local shards.  This is simpler but
+        slower than the parallel pipeline (all2all) path, so it serves as a
+        reference implementation for verifying correctness.
+        """
+        with record_function("distributed_muon"):
+            # Momentum is already applied by _step_muon before this method.
+            ns_steps = group["ns_steps"]
 
-            qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits)
+            # Separate plain tensors (no communication) from DTensors.
+            plain_names, plain_params = [], []
+            dtensor_names, dtensor_params = [], []
+            for n, p in zip(names, params):
+                if p.grad is None:
+                    continue
+                if isinstance(p.data, DTensor):
+                    dtensor_names.append(n)
+                    dtensor_params.append(p)
+                else:
+                    plain_names.append(n)
+                    plain_params.append(p)
+
+            # Process plain tensors per-param (no communication).
+            for n, p in zip(plain_names, plain_params):
+                u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE),
+                                                 steps=ns_steps)
+                adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                update_p(p, u, lr, adjusted_lr, weight_decay)
+
+                qk_clip_state = get_qk_clip_info(self.clip_config, n,
+                                                 qk_logits)
+                scales_full = compute_scales(
+                    p, qk_clip_state) if qk_clip_state is not None else None
+                if scales_full is not None:
+                    qk_clip(p, scales_full, qk_clip_state.head_dim)
+
+            if not dtensor_params:
+                return
+
+            # Group DTensors by (placements, mesh) for batched all-gather.
+            placement_groups: dict[tuple,
+                                   tuple[list,
+                                         list]] = defaultdict(lambda: ([], []))
+            for n, p in zip(dtensor_names, dtensor_params):
+                key = (p.placements, p.device_mesh)
+                placement_groups[key][0].append(n)
+                placement_groups[key][1].append(p)
+
+            logger.info(
+                "distributed_muon: %d placement groups, %d total dtensors",
+                len(placement_groups), len(dtensor_params))
+
+            for (placements, mesh), (grp_names,
+                                     grp_params) in placement_groups.items():
+                shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+                    placements, mesh)
+                rank = dist.get_rank(shard_pg)
+                world_size = dist.get_world_size(shard_pg)
+
+                logger.info("  group: %d params, placements=%s, world_size=%d",
+                            len(grp_params), placements, world_size)
+
+                # Separate params that can be batched (all shard dims evenly
+                # divisible) from those needing per-param full_tensor
+                # (e.g. MoE gate weights with fewer rows than shard ranks).
+                # all_gather_into_tensor requires equal buffer sizes across
+                # ranks, so uneven splits must use DTensor full_tensor().
+                batch_names, batch_params = [], []
+                single_names, single_params = [], []
+                for n, p in zip(grp_names, grp_params):
+                    even = all(p.shape[pl.dim] %
+                               shard_mesh.mesh.shape[dim_idx] == 0
+                               for dim_idx, pl in enumerate(shard_placements))
+                    if even:
+                        batch_names.append(n)
+                        batch_params.append(p)
+                    else:
+                        single_names.append(n)
+                        single_params.append(p)
+
+                # Process uneven-split params per-param via full_tensor().
+                for n, p in zip(single_names, single_params):
+                    with record_function("distributed_muon::newton_schulz"):
+                        g_full = p.grad.full_tensor().to(COMM_DTYPE)
+                        u_full = _zeropower_via_newtonschulz5(g_full,
+                                                              steps=ns_steps)
+                        del g_full
+                    with record_function("distributed_muon::update"):
+                        adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                        p._local_tensor.mul_(1 - lr * weight_decay)
+                        local_indices = get_slices_of_dtensor(
+                            p, rank, shard_mesh, shard_placements)
+                        u_local = u_full[local_indices]
+                        p._local_tensor.add_(u_local, alpha=-adjusted_lr)
+                        del u_full
+
+                        qk_clip_state = get_qk_clip_info(
+                            self.clip_config, n, qk_logits)
+                        scales_full = compute_scales(
+                            p, qk_clip_state
+                        ) if qk_clip_state is not None else None
+                        if scales_full is not None:
+                            ratio = p.shape[0] // scales_full.shape[0]
+                            idx0 = local_indices[0]
+                            if isinstance(idx0, slice):
+                                start = idx0.start or 0
+                                idx0 = torch.arange(start,
+                                                    idx0.stop,
+                                                    device=scales_full.device)
+                            row_scales = scales_full[idx0 // ratio]
+                            p._local_tensor.mul_(row_scales.view(-1, 1))
+
+                if not batch_params:
+                    continue
 
-            scales_full = compute_scales(
-                p_full, qk_clip_state) if qk_clip_state is not None else None
+                logger.info("  batched=%d, single=%d", len(batch_params),
+                            len(single_params))
+
+                # Concat all local grad shards into a single flat buffer.
+                with record_function("distributed_muon::gather"):
+                    grad_locals = [
+                        p.grad.to_local().to(COMM_DTYPE).flatten()
+                        for p in batch_params
+                    ]
+                    numels = [g.numel() for g in grad_locals]
+                    grad_concat = torch.cat(grad_locals)
+                    del grad_locals
+
+                    # Single all-gather (replaces N separate full_tensor).
+                    grad_gathered = torch.empty(
+                        grad_concat.numel() * world_size,
+                        dtype=COMM_DTYPE,
+                        device="cuda",
+                    )
+                    dist.all_gather_into_tensor(grad_gathered,
+                                                grad_concat,
+                                                group=shard_pg)
+
+                total_numel = grad_concat.numel()
+                del grad_concat
+
+                # Precompute per-param offsets within the concat buffer.
+                offsets = []
+                off = 0
+                for ne in numels:
+                    offsets.append(off)
+                    off += ne
+
+                # Per-param: reconstruct full grad → NS → local update.
+                for i, (n, p) in enumerate(zip(batch_names, batch_params)):
+                    with record_function("distributed_muon::newton_schulz"):
+                        g_full = torch.empty(p.shape,
+                                             dtype=COMM_DTYPE,
+                                             device="cuda")
+                        for r in range(world_size):
+                            r_start = r * total_numel + offsets[i]
+                            shard = grad_gathered[r_start:r_start + numels[i]]
+                            indices = get_slices_of_dtensor(
+                                p, r, shard_mesh, shard_placements)
+                            g_full[indices] = shard.reshape(
+                                g_full[indices].shape)
+
+                        u_full = _zeropower_via_newtonschulz5(g_full,
+                                                              steps=ns_steps)
+                        del g_full
+
+                    with record_function("distributed_muon::update"):
+                        adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                        p._local_tensor.mul_(1 - lr * weight_decay)
+                        local_indices = get_slices_of_dtensor(
+                            p, rank, shard_mesh, shard_placements)
+                        u_local = u_full[local_indices]
+                        p._local_tensor.add_(u_local, alpha=-adjusted_lr)
+                        del u_full
+
+                        qk_clip_state = get_qk_clip_info(
+                            self.clip_config, n, qk_logits)
+                        scales_full = compute_scales(
+                            p, qk_clip_state
+                        ) if qk_clip_state is not None else None
+                        if scales_full is not None:
+                            ratio = p.shape[0] // scales_full.shape[0]
+                            idx0 = local_indices[0]
+                            if isinstance(idx0, slice):
+                                start = idx0.start or 0
+                                idx0 = torch.arange(start,
+                                                    idx0.stop,
+                                                    device=scales_full.device)
+                            row_scales = scales_full[idx0 // ratio]
+                            p._local_tensor.mul_(row_scales.view(-1, 1))
+
+    def _setup_parallel(self, names, params, group, qk_logits):
+        """Compute (or retrieve cached) parallel pipeline metadata.
+
+        Returns:
+            (ordered_params, param_to_state, rank, chunk_size)
+        """
+        cache_key = tuple(names)
 
-            if scales_full is not None:
-                qk_clip(p_full, scales_full, qk_clip_state.head_dim)
+        if cache_key not in self._parallel_cache:
+            # First call: compute metadata and populate cache.
+            param_to_state, ordered_params = self.init_state_and_assign_params(
+                names, params, group, qk_logits)
 
-            if isinstance(p.data, DTensor):
-                ndims = len(p.device_mesh.mesh.shape)
-                p_replicate = DTensor.from_local(
-                    p_full,
-                    device_mesh=p.device_mesh,
-                    placements=[Replicate() for _ in range(ndims)],
-                )
+            shard_pg = param_to_state[id(ordered_params[0])].process_group
+            rank = dist.get_rank(group=shard_pg)
 
-                p_sharded = p_replicate.redistribute(
-                    device_mesh=p.device_mesh,
-                    placements=p.placements,
+            if self.chunk_size == -1:
+                shard_ranks = dist.get_world_size(shard_pg)
+                chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+            elif self.chunk_size > 0:
+                chunk_size = self.chunk_size
+            else:
+                raise ValueError(
+                    "chunk_size must be -1 or a positive integer.")
+
+            ordered_names = [
+                param_to_state[id(p)].name for p in ordered_params
+            ]
+            name_to_state = {
+                param_to_state[id(p)].name: param_to_state[id(p)]
+                for p in ordered_params
+            }
+            self._parallel_cache[cache_key] = {
+                'ordered_names': ordered_names,
+                'name_to_state': name_to_state,
+                'rank': rank,
+                'chunk_size': chunk_size,
+            }
+        else:
+            # Cached path: rebuild param_to_state with current id(p) keys.
+            cache = self._parallel_cache[cache_key]
+            rank = cache['rank']
+            chunk_size = cache['chunk_size']
+
+            name_to_param = dict(zip(names, params))
+            ordered_params = [name_to_param[n] for n in cache['ordered_names']]
+
+            param_to_state = {}
+            for p, n in zip(ordered_params, cache['ordered_names']):
+                cached_state = cache['name_to_state'][n]
+                param_to_state[id(p)] = _muon_state(
+                    worker_rank=cached_state.worker_rank,
+                    process_group=cached_state.process_group,
+                    rank_indices=cached_state.rank_indices,
+                    rank_numels=cached_state.rank_numels,
+                    name=n,
+                    qk_clip_state=get_qk_clip_info(self.clip_config, n,
+                                                   qk_logits),
                 )
 
-                p.copy_(p_sharded)
+        return ordered_params, param_to_state, rank, chunk_size
 
-    def parallel(self, names, params, group, lr, weight_decay, qk_logits):
+    def parallel(self,
+                 names,
+                 params,
+                 group,
+                 lr,
+                 weight_decay,
+                 qk_logits,
+                 prelaunch_gather=None):
         """
         Perform a parallel optimization step using Muon.
 
@@ -410,31 +639,23 @@ class Muon(torch.optim.Optimizer):
         interleaves multiple chunks so that communication and computation
         overlap across chunks (the same overlap previously achieved by the
         warmup + main-loop index scheduling).
+
+        If ``prelaunch_gather`` is provided, it is passed to the first
+        chunk's generator to skip re-launching the already in-flight
+        A2A gather.
         """
 
         # Momentum is already applied by _step_muon before this method.
 
-        param_to_state, ordered_params = self.init_state_and_assign_params(
-            names, params, group, qk_logits)
-
-        # Compute local rank for this group's shard process group.
-        shard_pg = param_to_state[id(ordered_params[0])].process_group
-        rank = dist.get_rank(group=shard_pg)
-
-        if self.chunk_size == -1:
-            shard_ranks = dist.get_world_size(param_to_state[id(
-                ordered_params[0])].process_group)
-            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
-        elif self.chunk_size > 0:
-            chunk_size = self.chunk_size
-        else:
-            raise ValueError("chunk_size must be -1 or a positive integer.")
+        ordered_params, param_to_state, rank, chunk_size = (
+            self._setup_parallel(names, params, group, qk_logits))
 
         def pipelines():
+            first = True
             for start in range(0, len(ordered_params), chunk_size):
                 chunk = ordered_params[start:start + chunk_size]
                 if chunk:
-                    yield muon_chunk_pipeline(
+                    kwargs = dict(
                         params=chunk,
                         param_to_state=param_to_state,
                         rank=rank,
@@ -443,9 +664,11 @@ class Muon(torch.optim.Optimizer):
                         weight_decay=weight_decay,
                         none_grad=group["none_grad"],
                     )
+                    if first and prelaunch_gather is not None:
+                        kwargs['prelaunch_gather'] = prelaunch_gather
+                    first = False
+                    yield muon_chunk_pipeline(**kwargs)
 
-        with record_function("muon::barrier"):
-            dist.barrier()
         with record_function("muon::pipeline"):
             run_pipeline(pipelines(), max_concurrent=self.warmup_step + 1)
 
@@ -457,16 +680,152 @@ class Muon(torch.optim.Optimizer):
         names = group["names"]
 
         # Apply momentum to all params before routing/expansion.
+        # Batched using _foreach_* ops (compiled, fullgraph=True).
         with record_function("muon::momentum"):
-            for n, p in zip(names, params):
-                g = p.grad
-                if g is None:
+            active_params = [p for p in params if p.grad is not None]
+            if active_params:
+                # Ensure momentum buffers exist (avoid zeros_like when already present).
+                for p in active_params:
+                    if "momentum_buffer" not in self.state[p]:
+                        self.state[p]["momentum_buffer"] = torch.zeros_like(
+                            p.grad)
+
+                # Extract local tensors for compiled batch function.
+                local_grads = [
+                    p.grad._local_tensor
+                    if isinstance(p.grad, DTensor) else p.grad
+                    for p in active_params
+                ]
+                local_bufs = [
+                    self.state[p]["momentum_buffer"]._local_tensor
+                    if isinstance(self.state[p]["momentum_buffer"], DTensor)
+                    else self.state[p]["momentum_buffer"]
+                    for p in active_params
+                ]
+
+                # Wrap momentum as tensor for torch.compile.
+                batch_pre_ortho(local_grads, local_bufs,
+                                torch.tensor(momentum), group["nesterov"])
+
+                # For non-nesterov, the result is the momentum buffer.
+                if not group["nesterov"]:
+                    for p in active_params:
+                        p.grad = self.state[p]["momentum_buffer"]
+
+        # Identify batched experts for deferred NS.
+        # Detection is cheap (condition checks only); actual NS compute is
+        # deferred so it can overlap with the first chunk's A2A gather.
+        deferred_expert_work = []
+        if self.expert_keys:
+            batched_expert_indices = []
+            for i, (n, p) in enumerate(zip(names, params)):
+                if not (is_expert_param(n, self.expert_keys)
+                        and p.grad is not None):
                     continue
-                g = update_g(self.state, p, g, group, momentum)
-                p.grad = g
+                # Eligible: plain tensor, or DTensor with no non-dim-0 shards.
+                if isinstance(p.data, DTensor):
+                    has_tp = any(
+                        _is_shard(pl) and pl.dim != 0 for pl in p.placements)
+                    if has_tp:
+                        continue
+                batched_expert_indices.append(i)
+
+            if batched_expert_indices:
+                # Save refs for deferred NS; free grads from param list.
+                for i in batched_expert_indices:
+                    p = params[i]
+                    g = p.grad
+                    local_g = (g._local_tensor
+                               if isinstance(g, DTensor) else g)
+                    local_data = (p.data._local_tensor if isinstance(
+                        p.data, DTensor) else p.data)
+                    deferred_expert_work.append((local_data, local_g))
+                    p.grad = None
+
+                # Remove batched experts from lists before expansion.
+                keep = sorted(
+                    set(range(len(params))) - set(batched_expert_indices))
+                names = [names[i] for i in keep]
+                params = [params[i] for i in keep]
+
+        def _run_deferred_expert_ns():
+            """Execute deferred batched expert NS."""
+            if not deferred_expert_work:
+                return
+            with record_function("muon::batched_expert_ns"):
+                ns_steps = group["ns_steps"]
+                for local_data, local_g in deferred_expert_work:
+                    u = zeropower_via_newtonschulz5_batched(
+                        local_g.to(COMM_DTYPE), steps=ns_steps)
+                    adjusted_lr = adjust_lr_for_muon(lr, local_g.shape[1:])
+                    local_data.mul_(1 - lr * weight_decay)
+                    local_data.add_(u, alpha=-adjusted_lr)
 
         # Expand expert params by splitting on dim 0.
-        names, params = _expand_expert_params(names, params, self.expert_keys)
+        logger.debug("[_step_muon] before expand: %d params, expert_keys=%s",
+                     len(params), self.expert_keys)
+        if self.expert_keys:
+            cache_key = tuple(id(p) for p in params)
+            cache = self._expert_expand_cache.get(cache_key)
+
+            if cache is None:
+                # Cold path: full expansion + build cache metadata.
+                exp_names, exp_params = _expand_expert_params(
+                    names, params, self.expert_keys)
+
+                # Build per-expert-group info for hot-path grad updates.
+                grad_info = []
+                exp_idx = 0
+                for orig_idx, (n, p) in enumerate(zip(names, params)):
+                    if not is_expert_param(n, self.expert_keys):
+                        exp_idx += 1
+                        continue
+
+                    is_dt = isinstance(p.data, DTensor)
+                    num_experts = (p.to_local() if is_dt else p.data).shape[0]
+
+                    # Detect TP mesh from the first expanded expert param.
+                    tp_mesh = None
+                    tp_pls = None
+                    sample = exp_params[exp_idx]
+                    if isinstance(sample.data, DTensor):
+                        tp_mesh = sample.data.device_mesh
+                        tp_pls = list(sample.data.placements)
+
+                    grad_info.append((orig_idx, num_experts, exp_idx, is_dt,
+                                      tp_mesh, tp_pls))
+                    exp_idx += num_experts
+
+                self._expert_expand_cache[cache_key] = {
+                    'names': exp_names,
+                    'params': exp_params,
+                    'grad_info': grad_info,
+                }
+                names, params = exp_names, exp_params
+            else:
+                # Hot path: reuse cached params, only update expert grads.
+                for (orig_idx, num_experts, exp_start, is_dt, tp_mesh,
+                     tp_pls) in cache['grad_info']:
+                    p = params[orig_idx]
+                    g = p.grad
+                    local_grad = (g.to_local()
+                                  if is_dt and isinstance(g, DTensor) else g)
+                    for i in range(num_experts):
+                        expert_p = cache['params'][exp_start + i]
+                        sg = local_grad[i]
+                        if tp_mesh is not None:
+                            expert_p.grad = DTensor.from_local(
+                                sg, device_mesh=tp_mesh, placements=tp_pls)
+                        else:
+                            expert_p.grad = sg
+                    p.grad = None
+
+                names = cache['names']
+                params = cache['params']
+        else:
+            names, params = _expand_expert_params(names, params,
+                                                  self.expert_keys)
+        logger.debug("[_step_muon] after expand: %d params", len(params))
 
         param_dtensors = []
         name_dtensors = []
@@ -474,10 +833,10 @@ class Muon(torch.optim.Optimizer):
         param_tensors = []
         name_tensors = []
 
-        param_dtensors_small = []
-        name_dtensors_small = []
-
+        # distributed_muon is a reference implementation for testing only.
+        # The parallel pipeline (all2all) path below is the production path.
         if self.use_distributed_muon:
+            _run_deferred_expert_ns()
             self.distributed_muon(names=names,
                                   params=params,
                                   group=group,
@@ -486,8 +845,6 @@ class Muon(torch.optim.Optimizer):
                                   qk_logits=qk_logits)
             return
 
-        # For simplicity, we use distributed Muon for small parameters
-        # whose number of elements is below a threshold.
         for n, p in zip(names, params):
             if p is None or p.grad is None:
                 continue
@@ -495,23 +852,28 @@ class Muon(torch.optim.Optimizer):
                 if all(
                         isinstance(placement, Replicate)
                         for placement in p.placements):
+                    logger.debug(
+                        "[route] %s → base (DTensor all-Replicate), "
+                        "shape=%s, placements=%s", n, p.shape, p.placements)
                     param_tensors.append(p)
                     name_tensors.append(n)
-                elif p.data.numel() <= self.small_param_numel_threshold:
-                    param_dtensors_small.append(p)
-                    name_dtensors_small.append(n)
                 else:
+                    logger.debug(
+                        "[route] %s → parallel (DTensor), shape=%s, "
+                        "placements=%s, mesh=%s", n, p.shape, p.placements,
+                        p.device_mesh.mesh_dim_names)
                     param_dtensors.append(p)
                     name_dtensors.append(n)
             elif isinstance(p.data, torch.Tensor):
+                logger.debug("[route] %s → base (plain tensor), shape=%s", n,
+                             p.data.shape)
                 param_tensors.append(p)
                 name_tensors.append(n)
             else:
                 raise TypeError(f"Unsupported parameter type: {type(p.data)}")
 
-        logger.debug(
-            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors, "
-            f"{len(param_dtensors_small)} Small DTensors")
+        logger.debug(f"[Muon] {len(param_dtensors)} DTensors → parallel, "
+                     f"{len(param_tensors)} Tensors → base")
 
         def group_dtensors(dtensors, names):
             # To support different placements, we group parameters by placements
@@ -527,21 +889,6 @@ class Muon(torch.optim.Optimizer):
                                            p.device_mesh])][1].append(p)
             return placement_to_params
 
-        if len(param_dtensors_small) > 0:
-            if not dist.is_initialized():
-                raise RuntimeError(
-                    "Parallel Muon requires torch.distributed to be initialized."
-                )
-
-            self.distributed_muon(
-                params=param_dtensors_small,
-                names=name_dtensors_small,
-                group=group,
-                lr=lr,
-                weight_decay=weight_decay,
-                qk_logits=qk_logits,
-            )
-
         if len(param_dtensors) > 0:
             if not dist.is_initialized():
                 raise RuntimeError(
@@ -549,7 +896,26 @@ class Muon(torch.optim.Optimizer):
                 )
 
             dtensor_group = group_dtensors(param_dtensors, name_dtensors)
+
+            # Pre-launch the first chunk's A2A gather so that the NCCL
+            # communication overlaps with the (deferred) batched expert NS
+            # compute on the default CUDA stream.
+            prelaunch = None
+            if deferred_expert_work:
+                first_names, first_params = next(iter(dtensor_group.values()))
+                ordered, pts, rnk, csz = self._setup_parallel(
+                    first_names, first_params, group, qk_logits)
+                first_chunk = ordered[:csz]
+                if first_chunk:
+                    prelaunch = prelaunch_first_gather(first_chunk, pts, rnk,
+                                                       group["none_grad"])
+
+            _run_deferred_expert_ns()
+
+            first_group = True
             for _, (names, params) in dtensor_group.items():
+                pg = prelaunch if first_group else None
+                first_group = False
                 self.parallel(
                     names,
                     params,
@@ -557,7 +923,10 @@ class Muon(torch.optim.Optimizer):
                     lr=lr,
                     weight_decay=weight_decay,
                     qk_logits=qk_logits,
+                    prelaunch_gather=pg,
                 )
+        else:
+            _run_deferred_expert_ns()
 
         if len(param_tensors) > 0:
             self.base(
@@ -586,10 +955,18 @@ class Muon(torch.optim.Optimizer):
             with torch.enable_grad():
                 loss = closure()
 
-        for group in self.param_groups:
+        logger.debug("[Muon.step] expert_keys=%s, %d param groups",
+                     self.expert_keys, len(self.param_groups))
+
+        for i, group in enumerate(self.param_groups):
             if group["use_muon"]:
+                logger.debug("[Muon.step] group %d: use_muon=True, %d params",
+                             i, len(group["params"]))
                 self._step_muon(group, qk_logits=qk_logits)
             else:
+                logger.debug(
+                    "[Muon.step] group %d: use_muon=False (AdamW), %d params",
+                    i, len(group["params"]))
                 step_adamw(self.state, group)
 
         return loss