Diff Coverage

                group = self.param_groups[group_idx]

                # Compute momentum for this assignment's owned params only.
                if self._can_overlap_shard(hsdp_assign):
                    self._process_shard_params_pipelined(group, hsdp_assign, group_idx, local_work.get(group_idx))
                    continue
                self._process_serial_assignment(group, hsdp_assign, group_idx, broadcaster)

        for group_idx, work in local_work.items():
            while work["batches"]:
                self._run_local_ns_batch(self.param_groups[group_idx], work)
        broadcaster.wait_all()
        return loss

    def _process_serial_assignment(self, group, hsdp_assign, group_idx, broadcaster):

        if not hsdp_assign.is_shard:
            if ns_inputs:
                self._process_unshard_params(group, ns_inputs)
        else:
            self._process_shard_params(
                group, ns_inputs, [hsdp_assign], group_idx,
                buffer_cache={},
            )
        if self.post_update_fn is not None:
            self._run_post_update_fn(group, ns_inputs.keys())

        # Flush broadcasts after each assignment.
        broadcaster.flush_group(hsdp_assign)

        for group_idx, group in enumerate(self.param_groups):
            info = self._hsdp_assignment_batches.get(group_idx)
            if not info:
                # No HSDP info — all params are no_comm.
                unshard_params = self.unshard_params_by_group.get(group_idx, [])
                no_comm_ns[group_idx] = self._update_muon_momentum(
                    group, unshard_params
                )
            else:
                no_comm_params = info.get("no_comm", [])

    @staticmethod
    def _local_workspace_bytes(batches):
        """Reserve storage for the largest local NS batch."""
        return max((32 * sum(to_local_if_dtensor(param.data).numel() for param in params)
                    for params, _ in batches), default=0)

    def _prepare_local_work(self, group_linear_batches, no_comm_ns):
        """Move local NS batches into communication windows when overlap is eligible."""

        # Expert-axis shards live in local assignments, not necessarily in no_comm.
        for group_idx, assignments in group_linear_batches.items():
            group = self.param_groups[group_idx]
            if any(self._can_overlap_shard(assignment) for assignment in assignments):
                batches = [(list(batch), batch) for batch in self._unshard_ns_batches(no_comm_ns[group_idx])]
                remaining = []
                for assignment in assignments:
                    if not assignment.is_shard and not assignment.is_replicated:
                        params = [param for param in assignment.owned_params if param.grad is not None]
                        for shape_group in self._group_by_shape(params).values():
                            batches.extend((batch, None) for batch in
                                           self._split_into_memory_safe_batches(shape_group, shard_size=1))
                    else:
                        remaining.append(assignment)
                group_linear_batches[group_idx] = remaining
                local_work[group_idx] = {
                    "batches": deque(batches),
                    "workspace_bytes": self._local_workspace_bytes(batches),
                }
            elif no_comm_ns[group_idx]:

                }
            elif no_comm_ns[group_idx]:
                self._process_unshard_params(group, no_comm_ns[group_idx])
                if self.post_update_fn is not None:
                    self._run_post_update_fn(group, no_comm_ns[group_idx].keys())

        return local_work

    def _run_post_update_fn(

                    records=hsdp_assign.owned_records,
                    replicate_sizes=shard_sizes,
                )
                self._shard_compute_coord[group_idx].update(compute_by_index)
                self._shard_param_compute_coord.update(
                    (record.param, compute_by_index[record.index]) for record in hsdp_assign.owned_records
                )

            self.unshard_params_by_group[group_idx] = unshard_params

                gathered_inputs[p] = full_inp
        local_inputs.clear()
        gathered_list.clear()

        my_updates = self._compute_shard_ns_updates(my_params, gathered_inputs, group, total_shard_size)
        return my_updates, param_compute_coord

    def _compute_shard_ns_updates(self, my_params, gathered_inputs, group, total_shard_size):
        """Keep identical shape grouping and NS numerics in serial and pipelined execution."""
        my_updates: Dict[torch.nn.Parameter, torch.Tensor] = {}

                        shard_sizes,
                    )

            if valid_params:
                if len(shard_pgs) == 1 and not hsdp_assign.is_replicated:
                    safe_batches = self._shard_pipeline_batches(valid_params, total_shard_size)
                else:
                    safe_batches = self._split_into_memory_safe_batches(
                        valid_params, shard_size=total_shard_size, min_batch_size=total_shard_size,
                    )
                for sub_batch in safe_batches:
                    my_updates, param_compute_coord = self._gather_and_compute_shard_updates(

        checked = set()
        for assignments in self.shard_assignments_by_group.values():
            for assignment in assignments:
                if assignment.is_replicated or len(assignment.shard_pgs) != 1:
                    continue
                pg = assignment.shard_pgs[0]
                if pg is None or pg in checked:
                    continue
                checked.add(pg)
                device = to_local_if_dtensor(assignment.owned_params[0].data).device
                config = torch.tensor(
                    [int(self.fsdp_overlap), self.fsdp_overlap_memory_budget_bytes or -1,

                    raise ValueError("Muon FSDP overlap settings and callback eligibility must agree across ranks")

    def _pipeline_parameter_bytes(self, param, shard_size):
        """Estimate native NS and padded communication storage for one parameter."""
        meta = self._param_shard_metadata[param]
        padded_numel = math.prod(meta.stage_metas[0].pad_shape)
        padded_numel = ((padded_numel + 255) // 256) * 256
        return 32 * param.numel() + 2 * padded_numel * (shard_size + 1) + 1024

    def _can_overlap_shard(self, assignment):
        """Allow local NS overrides while retaining single-axis FSDP scheduling."""
        if not self.fsdp_overlap or not assignment.is_shard or assignment.is_replicated:

        """Allow local NS overrides while retaining single-axis FSDP scheduling."""
        if not self.fsdp_overlap or not assignment.is_shard or assignment.is_replicated:
            return False
        if len(assignment.shard_pgs) != 1 or assignment.shard_pgs[0] is None:
            return False
        if any(callback is not None for callback in (
                self.reshape_fn, self.momentum_update_fn,
                self.ns_transform_fn, self.post_update_fn)):
            return False

    def _pipeline_bucket_bytes(self, params, shard_size):
        """Bound shared communication storage plus the largest compute-owner workspace."""
        owners = getattr(self, "_shard_param_compute_coord", {})
        if any(param not in owners for param in params):
            return sum(self._pipeline_parameter_bytes(param, shard_size) for param in params)
        workspace = defaultdict(int)
        communication = 0
        for param in params:
            workspace[owners[param]] += 32 * param.numel()

        return result

    def _prepare_pipeline_buckets(self, batches, assignment, compute_coord, ns_inputs):
        """Cache static plans without enqueueing device work for later buckets."""
        key = (id(assignment), tuple(tuple(params) for params in batches))
        plans = self._pipeline_plans.get(key)
        if plans is None:
            _, local_coords, _, _ = self._get_shard_info(assignment)
            indices = {record.param: record.index for record in assignment.owned_records}
            plans = []
            for params in batches:
                coords = {param: compute_coord[indices[param]] for param in params}
                plans.append({
                    "params": params, "coords": coords,
                    "keep": {index for index, param in enumerate(params) if coords[param] == local_coords},
                    "stage_metas": [self._param_shard_metadata[param].stage_metas[0] for param in params],
                    "gather_plan": {},

                    "keep": {index for index, param in enumerate(params) if coords[param] == local_coords},
                    "stage_metas": [self._param_shard_metadata[param].stage_metas[0] for param in params],
                    "gather_plan": {},
                })
            self._pipeline_plans[key] = plans
        return [dict(plan, ns_inputs=ns_inputs, shard_metadata=self._param_shard_metadata) for plan in plans]

    @staticmethod
    def _launch_pipeline_gather(bucket, assignment, cache, previous, comm_stream):
        """Pack just this bucket on the communication stream, then launch its gather."""

        """Pack just this bucket on the communication stream, then launch its gather."""
        if previous is not None:
            comm_stream.wait_event(previous["ns_done"])
        if "prepared" not in bucket:
            local_inputs = build_pad_ns_inputs(bucket["params"], bucket["ns_inputs"], bucket["shard_metadata"])
            # Buckets retain every input until the caller stream has joined the last apply.
            # This keeps compute-stream allocations alive across communication reads.
            bucket["prepared"] = prepare_muon_gather(
                local_inputs, assignment.shard_pgs[0], assignment.layout_spec.shard_axes[0][1],
                bucket["stage_metas"], bucket["keep"], cache, bucket["gather_plan"],
            )
        bucket["pending"] = bucket["prepared"].launch(assignment.shard_pgs[0], cache)

            bucket["ns_done"].record(compute_stream)

    def _enqueue_pipeline_broadcast(self, bucket, assignment, backend, comm_stream):
        """Submit broadcasts after output production without waiting on the compute stream."""
        with backend.stream(comm_stream), torch.profiler.record_function("muon.pipeline.broadcast"):
            comm_stream.wait_event(bucket["ns_done"])
            if bucket["reuse_after"] is not None:
                comm_stream.wait_event(bucket["reuse_after"])
            for update in bucket["updates"].values():
                update.record_stream(comm_stream)
            active_coords = set(bucket["coords"].values())
            for key in tuple(bucket["cache"]):
                if key[1] not in active_coords:
                    del bucket["cache"][key]
            packed = self._pack_shard_updates(
                bucket["params"], bucket["updates"], bucket["coords"], assignment, bucket["cache"],
            )
            bucket["packed"] = packed
            pg = assignment.shard_pgs[0]
            bucket["works"] = [
                dist.broadcast(tensor, src=dist.get_global_rank(pg, coord[0]), group=pg, async_op=True)
                for coord, tensor in packed[0].items()
            ]

            ]

    def _enqueue_pipeline_apply(self, bucket, assignment, group, backend, update_stream):
        """Consume broadcast outputs on the update stream, leaving NS independent."""
        with backend.stream(update_stream), torch.profiler.record_function("muon.pipeline.apply"):
            for work in bucket["works"]:
                work.wait()
            for tensor in bucket["packed"][0].values():
                tensor.record_stream(update_stream)
            self._apply_shard_updates(bucket["packed"], group, assignment)
            bucket["apply_done"] = backend.Event()
            bucket["apply_done"].record(update_stream)

    def _process_shard_params_pipelined(self, group, assignment, group_idx, local_work=None):
        """Run G0, G1, B0, G2, B1 with two communication slots and budgeted host lookahead."""
        shard_size = dist.get_world_size(assignment.shard_pgs[0])

        budget = self.fsdp_overlap_memory_budget_bytes or 512 * 1024 * 1024
        prepared_bytes = self._pipeline_prepared_bytes(assignment.owned_params)
        peak_bucket_bytes = max(self._pipeline_bucket_bytes(batch, shard_size) for batch in batches)
        if len(batches) < 2 or prepared_bytes + 2 * peak_bucket_bytes > budget:
            if local_work is not None:
                while local_work["batches"]:
                    self._run_local_ns_batch(group, local_work)
            inputs = self._update_muon_momentum(group, assignment.owned_params)
            self._process_shard_params(group, inputs, [assignment], group_idx, buffer_cache={})
            return
        if local_work is not None:
            local_reserve = local_work["workspace_bytes"]
            if prepared_bytes + 2 * peak_bucket_bytes + local_reserve <= budget:
                prepared_bytes += local_reserve
            else:
                while local_work["batches"]:
                    self._run_local_ns_batch(group, local_work)
                local_work = None
        queue_depth = min(len(batches), (budget - prepared_bytes) // peak_bucket_bytes)
        device = to_local_if_dtensor(batches[0][0].data).device
        backend = getattr(torch, device.type)
        compute_stream = backend.current_stream(device)

            self._run_pipeline_buckets(
                buckets, assignment, group, shard_size, backend, compute_stream,
                comm_stream, update_stream, gather_caches, broadcast_caches, queue_depth, local_work,
            )
        except Exception:
            # Drain while cross-stream storage remains alive if a step fails midway.
            backend.synchronize(device)
            raise

    def _run_pipeline_buckets(
            self, buckets, assignment, group, shard_size, backend, compute_stream,
            comm_stream, update_stream, gather_caches, broadcast_caches, queue_depth, local_work,

            group: Dict[str, Any],
            hsdp_assign: HSDPGroupAssignment,
    ) -> None:
        """Fused broadcast and apply for shard-group updates."""
        packed = self._pack_shard_updates(valid_params, my_updates, param_compute_coord, hsdp_assign, {})
        shard_sizes, local_coords, shard_pgs, total_shard_size = self._get_shard_info(hsdp_assign)
        if total_shard_size > 1:
            self._batched_relay_broadcast(packed[0], shard_pgs, shard_sizes, local_coords)
        self._apply_shard_updates(packed, group, hsdp_assign)

    def _pack_shard_updates(self, valid_params, my_updates, param_compute_coord, hsdp_assign, buffer_cache):
        """Pack complete owner updates using slot-local communication storage."""
        _, local_coords, _, _ = self._get_shard_info(hsdp_assign)
        device = to_local_if_dtensor(valid_params[0].data).device

        coord_groups: Dict[Tuple[int, ...], List[torch.nn.Parameter]] = defaultdict(list)
        for p in valid_params:

                param_offsets.append((total_padded_numel, actual_numel, padded_numel))
                total_padded_numel += padded_numel

            coord_param_offsets[coord] = param_offsets
            cache_key = ("broadcast", coord)
            pack_buffer = buffer_cache.get(cache_key)
            if pack_buffer is None or pack_buffer.numel() < total_padded_numel:
                pack_buffer = torch.empty(total_padded_numel, dtype=torch.bfloat16, device=device)
                buffer_cache[cache_key] = pack_buffer
            pack_buffer = pack_buffer[:total_padded_numel]

            if is_compute_rank:
                for p, (offset, actual_numel, padded_numel) in zip(coord_params, param_offsets):
                    update = my_updates[p]

                        pack_buffer[offset + actual_numel:offset + padded_numel].zero_()

            pack_buffers[coord] = pack_buffer

        return pack_buffers, coord_groups, coord_param_offsets

    def _apply_shard_updates(self, packed, group, hsdp_assign):
        """Apply received updates on the caller's stream after communication has completed."""
        pack_buffers, coord_groups, coord_param_offsets = packed
        lr = group["lr"]
        weight_decay = group["weight_decay"]
        all_local_params = []
        all_update_shards = []
        layout_spec = hsdp_assign.layout_spec
        for coord, coord_params in coord_groups.items():
            pack_buffer = pack_buffers[coord]
            param_offsets = coord_param_offsets[coord]

def _build_shard_stage(gathered_shapes, record_idx, tensor_dim, shard_size, cur_rank):
    """Construct one padding stage and its reconstructed shape."""
    rank_shapes = tuple(tuple(gathered_shapes[rank_idx][record_idx]) for rank_idx in range(shard_size))
    tensor_dim_norm = tensor_dim % len(rank_shapes[0])
    split_sizes = tuple(shape[tensor_dim_norm] for shape in rank_shapes)
    pad_shape = tuple(max(shape[dim] for shape in rank_shapes) for dim in range(len(rank_shapes[0])))
    stage = ParamShardStageMeta(
        tensor_dim=tensor_dim_norm, shard_size=shard_size, cur_rank_in_shard_group=cur_rank,
        rank_shapes=rank_shapes, split_sizes=split_sizes, pad_shape=pad_shape,
    )
    output_shape = list(rank_shapes[cur_rank])
    if shard_size > 1:
        output_shape[tensor_dim_norm] = sum(split_sizes)
    return stage, tuple(output_shape)


def build_param_shard_metadata_for_group(
        hsdp_group: HSDPCommGroup,

        cur_rank = dist.get_rank(shard_pg) if shard_pg is not None and shard_size > 1 else 0
        next_shapes: List[Tuple[int, ...]] = []

        for record_idx, _ in enumerate(hsdp_group.records):
            stage, output_shape = _build_shard_stage(gathered_shapes, record_idx, tensor_dim, shard_size, cur_rank)
            stage_metas_per_param[record_idx].append(stage)
            next_shapes.append(output_shape)

        current_shapes = next_shapes

    param_to_meta: Dict[torch.Tensor, ParamShardMeta] = {}

        total_padded_numel * shard_size,
        dtype, device,
    )[:total_padded_numel * shard_size]

    return pack_buffer, gathered_buffer


def _pack_and_allgather(
        gather_inputs: List[torch.Tensor],

        shard_size: int,
        buffer_cache: Optional[Dict],
) -> torch.Tensor:
    """Pack local shards into one buffer, all-gather, return gathered view."""
    pack_buffer, gathered_buffer = _pack_gather_buffers(
        gather_inputs, param_meta, total_padded_numel, axis_idx,
        dtype, device, shard_size, buffer_cache,
    )
    dist.all_gather_into_tensor(gathered_buffer, pack_buffer, group=shard_pg)