Optimize pipeline: batched update, zero-copy scatter, prelaunch gather [skip-build]

- Batched update_params with _foreach_mul_/_foreach_add_
- Zero-copy scatter directly into param views
- Deferred QK clip buffer allocation
- Buffer packing optimization
- prelaunch_first_gather() for overlapping first chunk A2A with compute

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

torch-ext/optimizer/pipeline.py

@@ -6,8 +6,8 @@ import torch.distributed as dist
 from torch.distributed.tensor import DTensor
 from torch.profiler import record_function
 
-from .core import _muon_state, adjust_lr_for_muon, update_p
-from .newton_schulz import COMM_DTYPE, _zeropower_via_newtonschulz5
+from .core import _muon_state, adjust_lr_for_muon
+from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5
 from .qk_clip import compute_scales
 
 logger = logging.getLogger(__name__)
@@ -45,26 +45,33 @@ def _launch_gather(
         else:
             gathered_grads[id(p)] = None
 
-    # Build send buffer
-    per_dst: list[list[torch.Tensor]] = [[] for _ in range(num_ranks)]
+    # Build send buffer – batch grad copies via torch.cat
+    # (1-2 fused kernels vs N individual narrow().copy_() calls).
     send_counts = [0] * num_ranks
-
     for p in params:
         state = param_to_state[id(p)]
-        dst = state.worker_rank
-        assert dst < num_ranks
-        shard_elems = state.rank_numels[rank]
-        g = p.grad
-        g = g.to_local().to(COMM_DTYPE).contiguous()
-        assert g.numel() == shard_elems
-        per_dst[dst].append(g.view(-1))
-        send_counts[dst] += shard_elems
-
-    assert any(
-        len(v) > 0 for v in
-        per_dst), "At least one destination rank must receive a sharded tensor"
-    per_dst_flat = [t for dst in per_dst for t in dst]
-    send_buf = torch.cat(per_dst_flat, dim=0)
+        send_counts[state.worker_rank] += state.rank_numels[rank]
+
+    total_send = sum(send_counts)
+    if total_send > 0:
+        # Group grad slices by destination rank in a single pass.
+        dst_to_grads = [[] for _ in range(num_ranks)]
+        for p in params:
+            state = param_to_state[id(p)]
+            n = state.rank_numels[rank]
+            if n > 0:
+                g = p.grad.to_local()
+                dst_to_grads[state.worker_rank].append(g.reshape(-1))
+
+        # Flatten in dst order and cat once.
+        all_slices = []
+        for dst in range(num_ranks):
+            all_slices.extend(dst_to_grads[dst])
+        send_buf = torch.cat(all_slices)
+        if send_buf.dtype != COMM_DTYPE:
+            send_buf = send_buf.to(COMM_DTYPE)
+    else:
+        send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
     # Build recv buffer
     recv_counts = [0] * num_ranks
@@ -120,7 +127,8 @@ def _complete_gather(
 
             shard_view = gathered_grads[id(p)][indices]
             n = shard_view.numel()
-            assert n > 0
+            if n == 0:
+                continue
 
             sg = recv_buf.narrow(0, off + inner_off, n)
             sg = sg.reshape(shard_view.shape)
@@ -143,7 +151,7 @@ def _compute_ns(
     """
     computed_us: dict[int, torch.Tensor | None] = {}
     for p in owned_params:
-        u = _zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps)
+        u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps)
         gathered_grads[id(p)] = None  # free gathered grad
         computed_us[id(p)] = u
     return computed_us
@@ -163,46 +171,47 @@ def _launch_scatter(
     Returns:
         work: Async operation handle.
         recv_buf: Flat receive buffer (needed by ``_complete_scatter``).
-        scattered_us: ``{id(p): empty_local_tensor}`` for all params.
+        scattered_us: Empty dict, populated by ``_complete_scatter`` with
+            zero-copy views into ``recv_buf``.
         recv_counts: Per-source-rank element counts.
     """
-    # Allocate scattered-u buffers
+    # scattered_us is populated by _complete_scatter with zero-copy views
+    # into recv_buf, avoiding N empty_like allocations + N copy_ calls.
+    # Pre-seed entries for params whose local shard is empty (rank_numels == 0)
+    # so _update_params can iterate all params without KeyError.
     scattered_us: dict[int, torch.Tensor] = {}
     for p in params:
-        scattered_us[id(p)] = torch.empty_like(p.to_local(), dtype=COMM_DTYPE)
+        if param_to_state[id(p)].rank_numels[rank] == 0:
+            scattered_us[id(p)] = torch.empty_like(p.to_local(),
+                                                   dtype=COMM_DTYPE)
 
-    # Build send buffer (from computed_us on owner ranks)
-    per_dst: list[list[torch.Tensor]] = [[] for _ in range(num_ranks)]
+    # Build send buffer – batch via torch.cat
+    # (1 fused kernel vs N*num_ranks individual narrow().copy_() calls).
     send_counts = [0] * num_ranks
-
     if owned_params:
         for p in owned_params:
             state = param_to_state[id(p)]
-
-            assert computed_us[id(p)] is not None
-            u_full = computed_us[id(p)].to(COMM_DTYPE).contiguous()
-
-            total_sent = 0
             for dst_rank in range(num_ranks):
-                indices = state.rank_indices[dst_rank]
-                su = u_full[indices].flatten()
-
-                n = su.numel()
-                assert n > 0
+                send_counts[dst_rank] += state.rank_numels[dst_rank]
 
-                per_dst[dst_rank].append(su)
-                send_counts[dst_rank] += n
-                total_sent += n
-
-            assert total_sent == u_full.numel()
-
-    lengths = [len(v) for v in per_dst]
-    if all(l > 0 for l in lengths):
-        assert all(
-            l == lengths[0] for l in lengths
-        ), "All destination ranks must have the same number of sharded tensor"
-        per_dst_flat = [t for dst in per_dst for t in dst]
-        send_buf = torch.cat(per_dst_flat, dim=0)
+    total_send = sum(send_counts)
+    if total_send > 0:
+        # Cache u_full conversions to avoid redundant .to() per dst_rank.
+        u_fulls = {}
+        for p in owned_params:
+            u_fulls[id(p)] = computed_us[id(p)].to(COMM_DTYPE).contiguous()
+
+        # Collect slices in dst order (matches all-to-all send layout).
+        all_slices = []
+        for dst_rank in range(num_ranks):
+            for p in owned_params:
+                state = param_to_state[id(p)]
+                su = u_fulls[id(p)][state.rank_indices[dst_rank]].flatten()
+                if su.numel() > 0:
+                    all_slices.append(su)
+
+        send_buf = torch.cat(all_slices) if all_slices else torch.empty(
+            0, dtype=COMM_DTYPE, device="cuda")
     else:
         send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
@@ -218,7 +227,6 @@ def _launch_scatter(
         recv_counts[src] = total
 
     recv_total = sum(recv_counts)
-    assert recv_total > 0
     recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
     # Launch async all-to-all
@@ -242,7 +250,13 @@ def _complete_scatter(
     rank: int,
     scattered_us: dict[int, torch.Tensor],
 ) -> None:
-    """Copy recv buffer into scattered_us (in-place)."""
+    """Populate scattered_us with zero-copy views into recv_buf.
+
+    Instead of pre-allocating tensors and copying, we assign views directly
+    from ``recv_buf``.  This eliminates N ``empty_like`` + N ``copy_`` calls.
+    The underlying storage of ``recv_buf`` is kept alive through the views
+    until ``scattered_us`` is cleared after ``_update_params``.
+    """
     off = 0
     for src in range(len(recv_counts)):
         block = recv_counts[src]
@@ -255,11 +269,11 @@ def _complete_scatter(
             if state.worker_rank != src:
                 continue
             n = state.rank_numels[rank]
-            assert n > 0
+            if n == 0:
+                continue
 
-            flat_local = recv_buf.narrow(0, off + inner_off,
-                                         n).view_as(p.to_local())
-            scattered_us[id(p)].copy_(flat_local)
+            scattered_us[id(p)] = recv_buf.narrow(0, off + inner_off,
+                                                  n).view_as(p.to_local())
 
             inner_off += n
 
@@ -275,23 +289,40 @@ def _update_params(
     lr: float,
     weight_decay: float,
 ) -> None:
-    """Apply weight decay, Muon update, and optional QK clipping."""
-    for p in params:
-        state = param_to_state[id(p)]
-        u_dtensor = DTensor.from_local(
-            scattered_us[id(p)],
-            placements=p.placements,
-            device_mesh=p.device_mesh,
-        )
+    """Apply weight decay, Muon update, and optional QK clipping.
 
+    Uses batched ``_foreach_mul_`` for weight decay and batched
+    ``_foreach_add_`` for the Muon update, grouping parameters by
+    adjusted_lr to minimize kernel launches while preserving float32
+    precision for the alpha scaling.
+    """
+    if not params:
+        return
+
+    # Batched weight decay: p *= (1 - lr * wd) — single fused kernel.
+    p_locals = [p._local_tensor for p in params]
+    torch._foreach_mul_(p_locals, 1.0 - lr * weight_decay)
+
+    # Group params by adjusted_lr so _foreach_add_ can use a single
+    # alpha per group (preserves float32 precision for alpha scaling).
+    lr_groups: dict[float, tuple[list, list]] = {}
+    for p in params:
         adjusted_lr = adjust_lr_for_muon(lr, p.shape)
-        update_p(p, u_dtensor, lr, adjusted_lr, weight_decay)
+        if adjusted_lr not in lr_groups:
+            lr_groups[adjusted_lr] = ([], [])
+        lr_groups[adjusted_lr][0].append(p._local_tensor)
+        lr_groups[adjusted_lr][1].append(scattered_us[id(p)])
 
-        # QK clipping – applied directly on the local tensor to
-        # avoid DTensor sharding-propagation issues with _StridedShard.
-        scales_full = compute_scales(
-            p,
-            state.qk_clip_state) if state.qk_clip_state is not None else None
+    for adjusted_lr, (p_group, u_group) in lr_groups.items():
+        torch._foreach_add_(p_group, u_group, alpha=-adjusted_lr)
+
+    # QK clipping – applied directly on the local tensor to
+    # avoid DTensor sharding-propagation issues with _StridedShard.
+    for p in params:
+        state = param_to_state[id(p)]
+        if state.qk_clip_state is None:
+            continue
+        scales_full = compute_scales(p, state.qk_clip_state)
         if scales_full is not None:
             ratio = p.shape[0] // scales_full.shape[0]
             idx0 = state.rank_indices[rank][0]
@@ -304,6 +335,45 @@ def _update_params(
             p._local_tensor.mul_(row_scales.view(-1, 1))
 
 
+# ======================================================================
+# Pre-launch helper for overlapping first chunk's gather with other work.
+# ======================================================================
+
+
+@torch.no_grad()
+def prelaunch_first_gather(
+    params: list[DTensor],
+    param_to_state: dict[int, _muon_state],
+    rank: int,
+    none_grad: bool,
+) -> tuple[dist.Work, torch.Tensor, dict[int, torch.Tensor | None], list[int]]:
+    """Launch the first chunk's A2A gather early for overlap with other compute.
+
+    Call this *before* expensive GPU work (e.g. batched expert NS) so that
+    the NCCL all-to-all runs concurrently on the NCCL stream while the
+    default stream executes compute.
+
+    Returns the same 4-tuple that ``_launch_gather`` produces, which should
+    be passed as ``prelaunch_gather`` to :func:`muon_chunk_pipeline`.
+    """
+    process_group = param_to_state[id(params[0])].process_group
+    num_ranks = dist.get_world_size(group=process_group)
+    owned_params = [
+        p for p in params if param_to_state[id(p)].worker_rank == rank
+    ]
+
+    with record_function("muon::prelaunch_gather"):
+        work, recv_buf, gathered_grads, recv_counts = _launch_gather(
+            params, owned_params, param_to_state, rank, num_ranks,
+            process_group)
+
+    if none_grad:
+        for p in params:
+            p.grad = None
+
+    return work, recv_buf, gathered_grads, recv_counts
+
+
 # ======================================================================
 # Main generator – thin orchestrator that wires stages together.
 # ======================================================================
@@ -318,6 +388,7 @@ def muon_chunk_pipeline(
     lr: float,
     weight_decay: float,
     none_grad: bool,
+    prelaunch_gather: tuple | None = None,
 ) -> Generator[None, None, None]:
     """Process one chunk of parameters through the full Muon pipeline.
 
@@ -334,9 +405,12 @@ def muon_chunk_pipeline(
     runs concurrently on the NCCL stream — no separate ``comm_stream``
     is required.
 
+    If ``prelaunch_gather`` is provided, the gather was already launched
+    by :func:`prelaunch_first_gather` and we skip launching it again.
+
     Yields exactly **2** times:
 
-    1. After launching async all-to-all gather.
+    1. After launching async all-to-all gather (or immediately if pre-launched).
     2. After launching async all-to-all scatter.
     """
     process_group = param_to_state[id(params[0])].process_group
@@ -345,15 +419,19 @@ def muon_chunk_pipeline(
         p for p in params if param_to_state[id(p)].worker_rank == rank
     ]
 
-    # Stages 1-2: launch async gather.
-    with record_function("muon::launch_gather"):
-        work, recv_buf, gathered_grads, recv_counts = _launch_gather(
-            params, owned_params, param_to_state, rank, num_ranks,
-            process_group)
-
-        if none_grad:
-            for p in params:
-                p.grad = None
+    if prelaunch_gather is not None:
+        # Gather was pre-launched; none_grad already handled by caller.
+        work, recv_buf, gathered_grads, recv_counts = prelaunch_gather
+    else:
+        # Normal path: launch async gather.
+        with record_function("muon::launch_gather"):
+            work, recv_buf, gathered_grads, recv_counts = _launch_gather(
+                params, owned_params, param_to_state, rank, num_ranks,
+                process_group)
+
+            if none_grad:
+                for p in params:
+                    p.grad = None
 
     yield  # --- YIELD 1: other chunks can launch their gather ---