diff --git a/dion/dion2.py b/dion/dion2.py index 2a853eb..bee8835 100644 --- a/dion/dion2.py +++ b/dion/dion2.py @@ -1,5 +1,6 @@ import math import torch +import torch.distributed as dist from collections import defaultdict from torch import Tensor from torch.distributed import ProcessGroup @@ -10,6 +11,7 @@ from .megabatch_base import ( DistributedOrthoBase, megabatch_orthogonalize_async, + muon_update_newton_schulz, adjust_lr_spectral_norm, adjust_lr_rms_norm, ) @@ -48,7 +50,17 @@ class Dion2(DistributedOrthoBase): assembled whole matrix -- layout-invariant/reproducible and better- converging. "local": per-shard top-k (union) -- cheaper comm but a sharding-dependent approximation that converges slightly worse; opt - in when comm-bound at large scale. No-op off the row-sharded path. + in when comm-bound at large scale. "global_capped": exact global + top-``ceil(fraction*global)`` selection at ~local comm cost -- + all-gather the row norms only (~KBs), then pack each rank's winner + rows into fixed-size chunks POOLED across the megabatch's matrices + (an int32 count header rides in each chunk). A rank whose winners + overflow its pooled chunk defers the overflow via error feedback; + see ``capacity_factor``. Groups where the packing cannot save comm + route to "global". No-op off the row-sharded path. + capacity_factor: "global_capped" pooled-chunk slack. None (default) = + auto per group, ``1 + 2*sqrt((1-1/world)/(per_rank*k))`` (covers + ~2 sigma of winner-count fluctuation); a float >= 1.0 pins it. Dion2 optimizer by Ahn et al.: TBD """ @@ -72,15 +84,21 @@ def __init__( verbose: bool = False, triton_post_ortho: bool = False, selection_scope: str = "global", + capacity_factor: Optional[float] = None, ): # Validate hyperparameters if lr < 0.0: raise ValueError(f"Invalid learning rate: {lr}") if not (0.0 < fraction <= 1.0): raise ValueError(f"fraction must be in (0, 1], got {fraction}") - if selection_scope not in ("local", "global"): + if selection_scope not in ("local", "global", "global_capped"): raise ValueError( - f"selection_scope must be 'local' or 'global', got {selection_scope!r}" + f"selection_scope must be 'local', 'global', or 'global_capped', " + f"got {selection_scope!r}" + ) + if capacity_factor is not None and capacity_factor < 1.0: + raise ValueError( + f"capacity_factor must be None (auto) or >= 1.0, got {capacity_factor}" ) if ef_decay < 0.0: raise ValueError(f"Invalid ef_decay: {ef_decay}") @@ -104,6 +122,7 @@ def __init__( algorithm="dion2", step=0, selection_scope=selection_scope, + capacity_factor=capacity_factor, ) super().__init__( params, distributed_mesh, "dion2", defaults, @@ -129,6 +148,9 @@ def _create_ortho_tasks( Mega-batched Dion2 task creation: groups ALL same-shape parameters into a single task to minimize communication rounds and kernel launches. """ + # New optimizer step: reset the rank-local capped-deferral counters so + # they accumulate exactly this step's packed megabatches. + CAPPED_STATS.clear() for group in param_groups: assert group["algorithm"] == self._algo_name assert all( @@ -154,6 +176,7 @@ def _create_ortho_tasks( verbose=self.verbose, triton_post_ortho=self._triton_post_ortho, selection_scope=group["selection_scope"], + capacity_factor=group["capacity_factor"], ) shape_groups: dict[tuple, list] = defaultdict(list) @@ -217,6 +240,7 @@ def dion2_update_megabatch_async( verbose: bool = False, triton_post_ortho: bool = False, selection_scope: str = "global", # "global" (exact whole-matrix top-k, default) or "local" (per-shard top-k; cheaper comm, sharding-variant) + capacity_factor: Optional[float] = None, # "global_capped" pooled-chunk slack; None = auto (z=2 model) ) -> Generator[None, None, None]: """ Mega-batched Dion2 update: processes ALL same-shape parameters in one @@ -237,6 +261,19 @@ def dion2_update_megabatch_async( sharding-dependent approximation of the true top-k (world-size variant). Cheaper comm (the win grows with model size), but converges slightly worse; opt in when comm-bound at large scale. + - ``"global_capped"``: exact global top-``ceil(fraction*global)`` selection + at ~"local" comm cost, via packed transport. The slice L1 norms are + all-gathered (one fp32 per slice, ~KBs) and every rank derives the same + exact-count winner set (stable-sort tie-break). Each rank packs its + winner rows into a fixed ``budget``-row chunk per destination, POOLED + across the destination's ``per_rank`` matrices, with an int32 count + header bit-cast into the chunk's first row -- the receiver parses the + sender's plan instead of recomputing it. Winners that overflow the + pooled budget are deferred: they skip error-feedback decay (only applied + rows are decayed, by the masked post), so their momentum accumulates and + they win later. Non-winners are never applied and never decayed -- a + per-rank top-k fill would be trajectory-identical to "local". Groups + whose budget would reach the full shard route to "global" instead. Off the row-sharded path (per-head, single-GPU, batch-sharded) each rank already holds whole matrices, so local and global selection coincide and @@ -269,12 +306,6 @@ def dion2_update_megabatch_async( # comm_dim for sharded communication: use select_dim (which equals normalized shard_dim) comm_dim = select_dim if is_sharded else None - # Decide whether selection happens before communication ("local", and only - # meaningful on the row-sharded path) or after the whole matrix is assembled - # ("global"). Off the sharded path each rank holds whole matrices, so the two - # are identical and we keep the cheaper pre-comm selection. - global_scope = selection_scope == "global" and comm_dim is not None - # On the sharded path X[0] must still be a DTensor, so .shape[comm_dim] # is the unsharded global size. The megabatch fn uses this to compute # the rank-consistent pad size for its alltoall. Catch the case where a @@ -290,6 +321,94 @@ def dion2_update_megabatch_async( else: global_dim_size = None + # --- "global_capped" packed-path decision (row-sharded 2D only) --- + # Pooled chunk budget: per (rank -> dest) pair, B = ceil(c * per_rank * k) + # rows shared across the dest's per_rank matrices. c defaults to the + # z=2 fluctuation model (see dion2_capped_pack); if the budget reaches the + # full shard, packing saves nothing and the group routes to exact "global". + capped_packed = ( + selection_scope == "global_capped" + and comm_dim is not None + and process_group is not None + and select_dim == -2 + and X[0].ndim == 2 + ) + if capped_packed: + padded_local = (global_dim_size + world_size - 1) // world_size + k = max(1, int(math.ceil(fraction * padded_local))) + per_rank = (N + world_size - 1) // world_size + if capacity_factor is None: + c = 1.0 + 2.0 * math.sqrt((1.0 - 1.0 / world_size) / (per_rank * k)) + else: + c = float(capacity_factor) + budget = int(math.ceil(c * per_rank * k)) + # Route to exact "global" when packing cannot pay for itself -- the + # packed chunk sends 1 + budget rows (count header included), so + # break-even against sending the full shard is budget + 1 -- or when + # the int32 count header cannot fit in the chunk's first row (2 bf16 + # slots per matrix). + if budget + 1 >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]: + capped_packed = False + + # Decide whether selection happens before communication ("local", and only + # meaningful on the row-sharded path) or after the whole matrix is assembled + # ("global"). Off the sharded path each rank holds whole matrices, so the two + # are identical and we keep the cheaper pre-comm selection. "global_capped" + # groups that fell out of the packed path (degenerate budget, col-sharded, + # non-2D) take the exact global path. + global_scope = comm_dim is not None and ( + selection_scope == "global" + or (selection_scope == "global_capped" and not capped_packed) + ) + + if capped_packed: + # --- "global_capped": winner-only packed transport --- + # Accumulate momentum (single owner on this path), all-gather the slice + # norms, pack the global winners into pooled fixed-size chunks, and + # come back with full-size shards that are zero except at the applied + # winner rows -- exactly the global path's masked-post format. + norms = dion2_pre_accumulate_norms( + G=to_local(G), M=to_local(M), select_dim=select_dim + ) + U_ortho = yield from dion2_capped_packed_async( + M_local=to_local(M), + norms=norms, + padded_local=padded_local, + fraction=fraction, + global_size=global_dim_size, + device_rank=device_rank, + world_size=world_size, + process_group=process_group, + per_rank=per_rank, + budget=budget, + kw=k * world_size, + newton_schulz_func=newton_schulz_func, + flatten=flatten, + epsilon=epsilon, + ) + if adjust_lr is None: + adjusted_lr = lr + elif adjust_lr == "spectral_norm": + adjusted_lr = adjust_lr_spectral_norm(lr, X[0].shape, flatten=flatten) + elif adjust_lr == "rms_norm": + adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten) + else: + raise ValueError(f"Unknown adjust_lr: {adjust_lr}") + _masked_post = dion2_post_orthogonalize_masked + if triton_post_ortho: + from .dion2_triton import dion2_post_orthogonalize_masked_triton as _masked_post + _masked_post( + X=to_local(X), + M=to_local(M), + U=U_ortho, + base_lr=lr, + adjusted_lr=adjusted_lr, + weight_decay=weight_decay, + ef_decay=ef_decay, + select_dim=select_dim, + ) + return + if global_scope: # --- Global selection: send the full shard, select after assembly --- # No pre-comm selection; momentum gets the gradient and the whole shard @@ -323,7 +442,10 @@ def dion2_update_megabatch_async( # U_ortho rows are exactly zero except at the globally-selected positions # this rank owns. Apply error-feedback decay to those rows of M and the # masked weight update, both keyed off the nonzero mask (no indices). - dion2_post_orthogonalize_masked( + _masked_post = dion2_post_orthogonalize_masked + if triton_post_ortho: + from .dion2_triton import dion2_post_orthogonalize_masked_triton as _masked_post + _masked_post( X=to_local(X), M=to_local(M), U=U_ortho, @@ -426,6 +548,371 @@ def dion2_update_megabatch_async( ) +@_inductor_workaround +@torch.compile(fullgraph=True) +def dion2_pre_accumulate_norms( + G: List[Tensor], M: List[Tensor], select_dim: int +) -> Tensor: + """ + Phase A of the "global_capped" scope: update momentum with the gradient and + return the stacked slice L1 norms ``[N, local_size]`` (fp32). The caller + all-gathers these tiny norms across ranks (an eager collective that cannot + live inside this compiled graph; see ``dion2_capped_select_async``). This + is the ONLY momentum accumulation on the capped path -- the packed + transport never calls ``dion2_pre_orthogonalize``, so the gathered norms + always reflect the post-accumulation momentum. + """ + dtype = M[0].dtype + G = [g.to(dtype=dtype) for g in G] + torch._foreach_add_(M, G) + norm_dim = -1 if select_dim == -2 else -2 + if M[0].size(select_dim) == 0: + return torch.zeros((len(M), 0), dtype=torch.float32, device=M[0].device) + return torch.stack(M, dim=0).norm(p=1, dim=norm_dim).float() + + +# Deferral instrumentation for the "global_capped" packed path. RANK-LOCAL +# semantics: cleared once per optimizer step (at ortho-task creation) and +# ACCUMULATED across that step's packed megabatches (shape groups), as GPU +# scalar tensors -- reading them costs no sync until the consumer calls +# .item(). Empty dict => no packed group ran this step (do not reuse stale +# values). "winner_rows" counts THIS rank's winners and can legitimately be +# zero; aggregate across ranks before forming a deferral ratio. +CAPPED_STATS: dict = {} + + +def dion2_capped_select_async( + norms: Tensor, # [N, local_size] fp32 from dion2_pre_accumulate_norms + padded_local: int, + fraction: float, + global_size: int, + device_rank: int, + world_size: int, + process_group: Optional[ProcessGroup], +) -> Generator[None, None, Tuple[Tensor, Tensor]]: + """ + "global_capped" winner selection: all-gather the slice L1 norms (one fp32 + per slice -- ~KBs vs MBs for the rows) and return, for this rank, + + - ``winner`` bool ``[N, local_size]``: membership in the exact global + top-``k_total`` set, ``k_total = ceil(fraction * global_size)``. The + set is EXACT-COUNT: a stable descending argsort over the gathered + norms breaks ties deterministically by flat (rank, row) position, so + ties cannot inflate the winner set past the budget (the receiver-side + buffer bound in the packed transport relies on this). + - ``thresh`` ``[N, 1]``: the k_total-th largest norm per matrix, used to + scale-normalize cross-matrix deferral priorities in the pack step. + + ``k_total <= global_size`` always, so the -1 shard padding can never win. + Rank consistency: every rank sorts the same bit-identical gathered tensor; + the receiver clip in ``dion2_capped_assemble`` guards the (pathological) + divergent case. Async generator in the megabatch style: yields at the + collective, returns the pair. + """ + N, local_size = norms.shape + if local_size > padded_local: + # F.pad with a negative amount would silently TRIM norms; fail loudly + # instead (mirrors the analogous guard in megabatch_orthogonalize_async). + raise ValueError( + f"Local norm count {local_size} exceeds padded size {padded_local}. " + "This should not happen with FSDP2 contiguous sharding." + ) + if local_size < padded_local: + norms_pad = torch.nn.functional.pad( + norms, (0, padded_local - local_size), value=-1.0 + ) + else: + norms_pad = norms + gathered = torch.empty( + world_size * N * padded_local, dtype=torch.float32, device=norms.device + ) + work = dist.all_gather_into_tensor( + gathered, norms_pad.reshape(-1).contiguous(), + group=process_group, async_op=True, + ) + yield + work.wait() + + # [world, N, padded_local] -> flat per-matrix [N, world*padded_local] with + # rank-major flat position (r * padded_local + row). + flat = gathered.view(world_size, N, padded_local).permute(1, 0, 2).reshape(N, -1) + k_total = min(max(1, int(math.ceil(fraction * global_size))), global_size) + order = torch.argsort(flat, dim=-1, descending=True, stable=True) + winner_flat = torch.zeros_like(flat, dtype=torch.bool) + winner_flat.scatter_(1, order[:, :k_total], True) + thresh = flat.gather(1, order[:, k_total - 1 : k_total]) + own = winner_flat.view(N, world_size, padded_local)[:, device_rank, :local_size] + return own.contiguous(), thresh + + +@_inductor_workaround +@torch.compile(fullgraph=True) +def dion2_capped_pack( + M: List[Tensor], # N tensors [local, cols] (momentum, already accumulated) + norms: Tensor, # [N, local] fp32 + winner: Tensor, # [N, local] bool, from dion2_capped_select_async + thresh: Tensor, # [N, 1] fp32 + world_size: int, + per_rank: int, + budget: int, # pooled row budget B per (rank -> dest) chunk +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Sender side of the packed "global_capped" transport. Matrices are assigned + to destination ranks in blocks of ``per_rank`` (megabatch order); for each + destination this rank packs its winner rows of that destination's matrices + back-to-back into a fixed ``budget``-row chunk -- slots are POOLED across + the per_rank matrices, so one matrix's overflow uses another's spare room. + + If a destination's total winners exceed ``budget``, the lowest-priority + rows are deferred (error feedback: they keep full momentum and win later). + Priority is the SCALE-NORMALIZED ``norm / thresh`` -- raw norms are not + comparable across matrices, and a large-scale matrix must not evict a + small-scale matrix's winners from the shared chunk. + + Returns ``(payload [world, budget, cols] bf16, counts [world, per_rank] + int32, src_index [world, budget] long, deferred_rows scalar)``. Empty + slots gather a dedicated zero row (never read by the receiver, which + parses only ``counts`` rows per segment). ``src_index`` maps each slot to + its flat local row (sentinel = the zero row) and is kept by the sender to + scatter the orthogonalized rows home; the packing is stable and + position-preserving -- the reverse path relies on it. + """ + N = len(M) + local = M[0].size(-2) + cols = M[0].size(-1) + n_pad = world_size * per_rank + device = M[0].device + + M_stacked = torch.stack(M, dim=0) # [N, local, cols] + if n_pad > N: + # Virtual matrices padding the group to world_size * per_rank: no + # winners, zero rows. + M_stacked = torch.cat( + [M_stacked, torch.zeros(n_pad - N, local, cols, dtype=M_stacked.dtype, device=device)] + ) + winner = torch.cat( + [winner, torch.zeros(n_pad - N, local, dtype=torch.bool, device=device)] + ) + norms = torch.cat( + [norms, torch.zeros(n_pad - N, local, dtype=norms.dtype, device=device)] + ) + thresh = torch.cat( + [thresh, torch.ones(n_pad - N, 1, dtype=thresh.dtype, device=device)] + ) + + prio = norms / thresh.clamp_min(1e-12) + score = torch.where(winner, prio, torch.full_like(prio, float("-inf"))) + score_d = score.view(world_size, per_rank * local) + b_eff = min(budget, per_rank * local) + # Stable argsort (not topk): on tied normalized priorities the flat + # (matrix, row) order is the deterministic secondary key, so WHICH row + # defers on overflow is reproducible across runs/devices -- matching the + # stable tie contract of the winner selection itself. + order = torch.argsort(score_d, dim=-1, descending=True, stable=True) + idxs = order[:, :b_eff] + vals = torch.gather(score_d, 1, idxs) + kept_flat = torch.zeros_like(score_d, dtype=torch.bool) + kept_flat.scatter_(1, idxs, vals.isfinite()) # winners only; spare slots stay empty + kept = kept_flat.view(world_size, per_rank, local) + + counts = kept.sum(dim=-1) # [world, per_rank] + seg_border = counts.cumsum(dim=1) - counts # exclusive, within chunk + # Slot of each kept row: matrix segment border + within-matrix ordinal + # (row order preserved -- stable packing). + ordinal = kept.cumsum(dim=-1) - 1 + slot = seg_border.unsqueeze(-1) + ordinal + slot_safe = torch.where(kept, slot, torch.full_like(slot, budget)) # trash col + + fid = torch.arange(n_pad * local, device=device).view(world_size, per_rank, local) + sentinel = n_pad * local # index of the appended zero row + src_index = torch.full( + (world_size, budget + 1), sentinel, dtype=torch.long, device=device + ) + src_index.scatter_( + 1, slot_safe.reshape(world_size, -1), fid.reshape(world_size, -1) + ) + src_index = src_index[:, :budget] + + M_flat = torch.cat( + [M_stacked.reshape(n_pad * local, cols), + torch.zeros(1, cols, dtype=M_stacked.dtype, device=device)] + ) + payload = M_flat[src_index].to(torch.bfloat16) # [world, budget, cols] + deferred = winner.sum() - kept.sum() + return payload, counts.to(torch.int32), src_index, deferred + + +@torch.compile(fullgraph=True) +def dion2_capped_assemble( + payload: Tensor, # [world, budget, cols] bf16 (received chunks, no header) + counts: Tensor, # [world, per_rank] int32 (parsed headers) + kw: int, # static NS row count per matrix (k * world_size) +) -> Tuple[Tensor, Tensor]: + """ + Receiver side: assemble each of this rank's ``per_rank`` matrices from the + variable-length segments of all senders' chunks into a static, zero-padded + ``[per_rank, kw, cols]`` Newton-Schulz input. Rows land in (sender rank, + within-segment) order; NS is permutation-equivariant and the reverse path + inverts the same mapping, so the order is only required to be stable. + + Arrivals beyond ``kw`` for a matrix (impossible with exact-count winners; + possible only under a divergent sender) are clipped to a trash slot rather + than corrupting the buffer. Returns ``(ns_input, ns_index)`` where + ``ns_index [per_rank, kw]`` maps NS rows back to flat payload positions + (sentinel = zero row) for ``dion2_capped_repack``. + """ + world_size, budget, cols = payload.shape + per_rank = counts.size(1) + device = payload.device + cnt = counts.to(torch.long) + + seg_border = cnt.cumsum(dim=1) - cnt # [world, per_rank] within sender chunk + dst_border = cnt.cumsum(dim=0) - cnt # [world, per_rank] within NS rows + max_seg = min(budget, kw) + j = torch.arange(max_seg, device=device) + valid = j < cnt.unsqueeze(-1) # [world, per_rank, max_seg] + src_pos = ( + torch.arange(world_size, device=device).view(-1, 1, 1) * budget + + seg_border.unsqueeze(-1) + + j + ) + src_pos = torch.where(valid, src_pos, torch.full_like(src_pos, world_size * budget)) + dst_pos = dst_border.unsqueeze(-1) + j + dst_pos = torch.where( + valid & (dst_pos < kw), dst_pos, torch.full_like(dst_pos, kw) # receiver clip + ) + + ns_index = torch.full( + (per_rank, kw + 1), world_size * budget, dtype=torch.long, device=device + ) + ns_index.scatter_( + 1, + dst_pos.permute(1, 0, 2).reshape(per_rank, -1), + src_pos.permute(1, 0, 2).reshape(per_rank, -1), + ) + ns_index = ns_index[:, :kw] + + pay_flat = torch.cat( + [payload.reshape(world_size * budget, cols), + torch.zeros(1, cols, dtype=payload.dtype, device=device)] + ) + ns_input = pay_flat[ns_index] # [per_rank, kw, cols], zero-padded tail + return ns_input, ns_index + + +@torch.compile(fullgraph=True) +def dion2_capped_repack( + ns_out: Tensor, # [per_rank, kw, cols] orthogonalized + ns_index: Tensor, # [per_rank, kw] from dion2_capped_assemble + world_size: int, + budget: int, +) -> Tensor: + """Inverse of assemble: scatter NS rows back to their payload positions. + Padding/sentinel rows collapse onto a trash row; payload slots that carried + no real row stay exactly zero (so the sender scatters zeros = no update).""" + per_rank, kw, cols = ns_out.shape + out = torch.zeros( + world_size * budget + 1, cols, dtype=ns_out.dtype, device=ns_out.device + ) + out.scatter_(0, ns_index.reshape(-1, 1).expand(-1, cols), ns_out.reshape(-1, cols)) + return out[: world_size * budget].view(world_size, budget, cols) + + +@torch.compile(fullgraph=True) +def dion2_capped_unpack( + recv: Tensor, # [world, budget, cols] orthogonalized rows, sender's layout + src_index: Tensor, # [world, budget] from dion2_capped_pack + n_params: int, + local: int, + total_rows: int, # n_pad * local (sentinel index) +) -> List[Tensor]: + """Sender-side finish: scatter the returned rows to their home (matrix, + row) positions. Output is a list of full-size ``[local, cols]`` shards that + are exactly zero except at this rank's applied winner rows -- the format + ``dion2_post_orthogonalize_masked`` consumes (same as the global scope).""" + world_size, budget, cols = recv.shape + U_flat = torch.zeros( + total_rows + 1, cols, dtype=recv.dtype, device=recv.device + ) + U_flat.scatter_(0, src_index.reshape(-1, 1).expand(-1, cols), recv.reshape(-1, cols)) + return list(U_flat[: n_params * local].view(n_params, local, cols).unbind(0)) + + +def dion2_capped_packed_async( + M_local: List[Tensor], + norms: Tensor, + padded_local: int, + fraction: float, + global_size: int, + device_rank: int, + world_size: int, + process_group: Optional[ProcessGroup], + per_rank: int, + budget: int, + kw: int, + newton_schulz_func: Callable, + flatten: bool, + epsilon: Tensor, +) -> Generator[None, None, List[Tensor]]: + """ + "global_capped" packed transport: winner selection (norms-only all-gather) + -> pooled pack -> fixed-size a2a with an int32 count header bit-cast into + each chunk's first row -> assemble + Newton-Schulz -> repack -> reverse a2a + -> scatter home. All collective sizes are construction-time constants; the + data-dependence lives entirely in GPU-side gathers/scatters (no host + syncs). The header makes the sender's packing plan the single source of + truth -- the receiver never recomputes winner counts. + """ + winner, thresh = yield from dion2_capped_select_async( + norms, padded_local, fraction, global_size, + device_rank, world_size, process_group, + ) + payload, counts, src_index, deferred = dion2_capped_pack( + M_local, norms, winner, thresh, world_size, per_rank, budget + ) + CAPPED_STATS["deferred_rows"] = CAPPED_STATS.get("deferred_rows", 0) + deferred + CAPPED_STATS["winner_rows"] = CAPPED_STATS.get("winner_rows", 0) + winner.sum() + + # Header row: int32 counts bit-cast into bf16 (2 bf16 per int32). Counts + # are NOT representable as bf16 *values* (exact only to 256); the bit-cast + # round-trips exactly. + cols = payload.size(-1) + header = torch.zeros( + world_size, 1, cols, dtype=torch.bfloat16, device=payload.device + ) + header[:, 0, : per_rank * 2] = counts.contiguous().view(torch.bfloat16) + chunks = torch.cat([header, payload], dim=1) # [world, 1 + budget, cols] + + send = [c.contiguous() for c in chunks.unbind(0)] + recv = [torch.empty_like(c) for c in send] + work = dist.all_to_all(recv, send, group=process_group, async_op=True) + yield + work.wait() + + recv_t = torch.stack(recv) + counts_recv = recv_t[:, 0, : per_rank * 2].contiguous().view(torch.int32) + ns_input, ns_index = dion2_capped_assemble( + recv_t[:, 1:].contiguous(), counts_recv, kw + ) + ns_out = muon_update_newton_schulz( + ns_input, newton_schulz_func=newton_schulz_func, + flatten=flatten, epsilon=epsilon, + ) + back = dion2_capped_repack(ns_out, ns_index, world_size, budget) + + send2 = [c.contiguous() for c in back.unbind(0)] + recv2 = [torch.empty_like(c) for c in send2] + work = dist.all_to_all(recv2, send2, group=process_group, async_op=True) + yield + work.wait() + + local = M_local[0].size(-2) + n_pad = world_size * per_rank + return dion2_capped_unpack( + torch.stack(recv2), src_index, len(M_local), local, n_pad * local + ) + + @_inductor_workaround @torch.compile(fullgraph=True) def dion2_pre_orthogonalize( @@ -454,6 +941,9 @@ def dion2_pre_orthogonalize( to the available count. The pad up to ``k_override`` is not done here -- the downstream megabatch pads U to ``local_comm_size=k_override`` for a uniform alltoall, and indices stay at the real selected count. + + (The "global_capped" scope does not pass through here at all -- it uses + the packed transport, ``dion2_capped_packed_async``, and the masked post.) """ dtype = M[0].dtype @@ -496,8 +986,6 @@ def dion2_pre_orthogonalize( return U_selected, indices_list M_stacked = torch.stack(M, dim=0) - - # Compute L1 norm along norm_dim (sum of absolute values) slice_norms = M_stacked.norm(p=1, dim=norm_dim) # Batched topk: indices shape (batch_size, k_topk). k_topk <= num_select is @@ -523,15 +1011,19 @@ def dion2_pre_orthogonalize( # Apply error feedback decay to selected slices in the original M tensors. # We reuse the already-gathered slices and write them back (scaled) using - # scatter_, which places values into positions specified by the index tensor. + # scatter_, which places values into positions specified by the index + # tensor. The .to(dtype) guards bf16 momentum: scatter_ requires src dtype + # to match M exactly, and it keeps a single fp32-multiply-then-round if + # ef_decay ever becomes a dimensioned tensor (a 0-dim fp32 scalar loses + # type promotion to bf16, but a dimensioned one would win it). indices_list = list(indices.unbind(dim=0)) - selected_list = list(selected_stacked.unbind(dim=0)) - for m, idx, selected in zip(M, indices_list, selected_list): + ef_src_list = list((selected_stacked * ef_decay).to(dtype).unbind(dim=0)) + for m, idx, ef_src in zip(M, indices_list, ef_src_list): if select_dim == -2: idx_exp = idx.unsqueeze(-1).expand(*idx.shape, m.size(-1)) else: idx_exp = idx.unsqueeze(-2).expand(*idx.shape[:-1], m.size(-2), idx.shape[-1]) - m.scatter_(dim=select_dim, index=idx_exp, src=selected * ef_decay) + m.scatter_(dim=select_dim, index=idx_exp, src=ef_src) # Convert to bf16 and unstack for communication U_selected = list(selected_stacked.to(dtype=torch.bfloat16).unbind(dim=0)) diff --git a/dion/dion2_triton.py b/dion/dion2_triton.py index cf6a59a..a62163d 100644 --- a/dion/dion2_triton.py +++ b/dion/dion2_triton.py @@ -231,3 +231,369 @@ def dion2_post_orthogonalize_triton( index_map.stride(0), SELECT_ROWS=SELECT_ROWS, ) + + +# --------------------------------------------------------------------------- +# Masked posts (selection_scope="global"/"global_capped") +# +# The masked post-orthogonalize has no index list: U arrives full-shard-shaped +# and exactly zero outside the selected rows, so selection membership is +# derived from the nonzero rows of U itself. These kernels fuse the entire +# masked post into one (Dion2) or two (NorDion2) row-programmed passes, +# replacing the compiled path's separate full-matrix sweeps (weight-decay +# foreach, mask reduction, EF decay, NorMuon normalization, weight update). +# Like the indexed kernel above, updates compute in fp32 with a single +# rounding on store, and in-place writes rule out @triton.autotune. +# +# Scalars (lr, weight decay, ...) arrive as 0-dim CPU tensors, so .item() is +# free (no GPU sync) and values are passed to the kernels by value. +# --------------------------------------------------------------------------- + + +@triton.heuristics( + { + "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))), + "num_warps": lambda args: 8 if args["N"] >= 2048 else 4, + } +) +@triton.jit +def _dion2_post_ortho_masked_kernel( + X_ptr, + M_ptr, + U_ptr, + a, + neg_b, + ef, + Mrows, + N, + x_stride_b, + x_stride_m, + x_stride_n, + m_stride_b, + m_stride_m, + m_stride_n, + u_stride_b, + u_stride_m, + u_stride_n, + BLOCK_N: tl.constexpr, +): + """Fused masked post for Dion2: one program per (batch, row). + + Pass 1 reduces |u| over the row to derive the selected flag (U is exactly + zero on non-selected rows). Pass 2 applies, in fp32 with one rounding: + selected: x = a*x + neg_b*u and m = ef*m + unselected: x = a*x and m untouched (ef would be *1) + Skipping M entirely on unselected rows is bitwise identical to the + compiled path's multiply-by-one and saves the M read+write there. + """ + pid = tl.program_id(0) + b = pid // Mrows + row = pid % Mrows + X_ptr += b * x_stride_b + row * x_stride_m + M_ptr += b * m_stride_b + row * m_stride_m + U_ptr += b * u_stride_b + row * u_stride_m + + absacc = tl.zeros((BLOCK_N,), dtype=tl.float32) + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32) + absacc += tl.abs(u) + sel = tl.sum(absacc) > 0.0 + + if sel: + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32) + u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32) + m = tl.load(M_ptr + idx * m_stride_n, mask=msk, other=0.0).to(tl.float32) + tl.store(X_ptr + idx * x_stride_n, a * x + neg_b * u, mask=msk) + tl.store(M_ptr + idx * m_stride_n, ef * m, mask=msk) + else: + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32) + tl.store(X_ptr + idx * x_stride_n, a * x, mask=msk) + + +def dion2_post_orthogonalize_masked_triton( + X: List[Tensor], + M: List[Tensor], + U: List[Tensor], + base_lr: Tensor, + adjusted_lr: Tensor, + weight_decay: Tensor, + ef_decay: Tensor, + select_dim: int, +): + """Triton-fused version of ``dion2_post_orthogonalize_masked``. + + Falls back to the compiled masked post when the fast path does not apply + (no triton/CUDA, column selection, or tensor-subclass params). Only row + selection (``select_dim == -2``) is fused: that is the only mode the + global/global_capped scopes produce today. + """ + use_fallback = ( + not TRITON_AVAILABLE + or select_dim != -2 + or not X + or not X[0].is_cuda + or any(is_traceable_wrapper_subclass(x) for x in X) + ) + if use_fallback: + from .dion2 import dion2_post_orthogonalize_masked + + dion2_post_orthogonalize_masked( + X=X, M=M, U=U, base_lr=base_lr, adjusted_lr=adjusted_lr, + weight_decay=weight_decay, ef_decay=ef_decay, select_dim=select_dim, + ) + return + + a = (1 - base_lr * weight_decay).item() + neg_b = (-adjusted_lr).item() + # Match the compiled path exactly: it multiplies M by ef_decay rounded to + # M's dtype (``ef_decay.to(M[0].dtype)``), so round before upcasting. + ef = ef_decay.to(M[0].dtype).item() + + for x, m, u in zip(X, M, U): + if x.numel() == 0: + continue + if not (x.is_contiguous() and m.is_contiguous() and u.is_contiguous()): + raise ValueError( + "dion2_post_orthogonalize_masked_triton requires contiguous tensors" + ) + orig_shape = x.shape + Mrows, N = orig_shape[-2], orig_shape[-1] + B = x.numel() // (Mrows * N) + x_flat = x.reshape(B, Mrows, N) + m_flat = m.reshape(B, Mrows, N) + u_flat = u.reshape(B, Mrows, N) + + _dion2_post_ortho_masked_kernel[(B * Mrows,)]( + x_flat, m_flat, u_flat, + a, neg_b, ef, + Mrows, N, + x_flat.stride(0), x_flat.stride(1), x_flat.stride(2), + m_flat.stride(0), m_flat.stride(1), m_flat.stride(2), + u_flat.stride(0), u_flat.stride(1), u_flat.stride(2), + ) + + +@triton.heuristics( + { + "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))), + "num_warps": lambda args: 8 if args["N"] >= 2048 else 4, + } +) +@triton.jit +def _nordion2_post_ortho_masked_stats_kernel( + U_ptr, + V_ptr, + M_ptr, + P_ptr, + ef, + lerp_w, + N, + u_stride_m, + u_stride_n, + v_stride_m, + m_stride_m, + m_stride_n, + p_stride_r, + BLOCK_N: tl.constexpr, +): + """NorDion2 masked post, phase A: one program per row. + + Computes the row's sum of squares of U (fp32), derives the selected flag, + updates the per-neuron variance V (lerp toward mean-square, selected rows + only -- unselected rows are stored back unchanged), applies error-feedback + decay to selected rows of momentum M, and emits per-row partials into P: + P[0, row] = sum_sq(u_row) (-> Frobenius norm of U) + P[1, row] = sum_sq(u_row) / denom^2 (-> Frobenius norm of U/denom) + P[2, row] = denom = sqrt(v_new) + 1e-8 (fp32, pre-rounding V store) + The matrix-level Frobenius rescale needs all rows, hence the two-phase + split; the epilogue reduces P and phase B applies the weight update. + """ + row = tl.program_id(0) + U_ptr += row * u_stride_m + M_ptr += row * m_stride_m + V_ptr += row * v_stride_m + + ssq = tl.zeros((BLOCK_N,), dtype=tl.float32) + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32) + ssq += u * u + sumsq = tl.sum(ssq) + sel = sumsq > 0.0 + + v_old = tl.load(V_ptr).to(tl.float32) + neuron = sumsq / N.to(tl.float32) + v_lerped = v_old + lerp_w * (neuron - v_old) + v_new = tl.where(sel, v_lerped, v_old) + denom = tl.sqrt(v_new) + 1e-8 + tl.store(V_ptr, v_new) + + tl.store(P_ptr + 0 * p_stride_r + row, sumsq) + tl.store(P_ptr + 1 * p_stride_r + row, sumsq / (denom * denom)) + tl.store(P_ptr + 2 * p_stride_r + row, denom) + + if sel: + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + m = tl.load(M_ptr + idx * m_stride_n, mask=msk, other=0.0).to(tl.float32) + tl.store(M_ptr + idx * m_stride_n, ef * m, mask=msk) + + +@triton.heuristics( + { + "BLOCK_N": lambda args: max(128, min(1024, triton.next_power_of_2(args["N"]))), + "num_warps": lambda args: 8 if args["N"] >= 2048 else 4, + } +) +@triton.jit +def _nordion2_post_ortho_masked_apply_kernel( + X_ptr, + U_ptr, + P_ptr, + scale_ptr, + mat_idx, + a, + neg_b, + N, + x_stride_m, + x_stride_n, + u_stride_m, + u_stride_n, + p_stride_r, + BLOCK_N: tl.constexpr, +): + """NorDion2 masked post, phase B: one program per row. + + Applies, in fp32 with a single rounding on store: + selected: x = a*x + neg_b * scale * u / denom + unselected: x = a*x (weight decay only; u is exactly zero anyway) + where scale = ||U||_F / max(||U/denom||_F, 1e-8) comes from the epilogue + reduction over phase A's partials. + """ + row = tl.program_id(0) + X_ptr += row * x_stride_m + U_ptr += row * u_stride_m + + sumsq = tl.load(P_ptr + 0 * p_stride_r + row) + sel = sumsq > 0.0 + + if sel: + denom = tl.load(P_ptr + 2 * p_stride_r + row) + scale = tl.load(scale_ptr + mat_idx).to(tl.float32) + coef = neg_b * scale / denom + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32) + u = tl.load(U_ptr + idx * u_stride_n, mask=msk, other=0.0).to(tl.float32) + tl.store(X_ptr + idx * x_stride_n, a * x + coef * u, mask=msk) + else: + for off in range(0, N, BLOCK_N): + idx = off + tl.arange(0, BLOCK_N) + msk = idx < N + x = tl.load(X_ptr + idx * x_stride_n, mask=msk, other=0.0).to(tl.float32) + tl.store(X_ptr + idx * x_stride_n, a * x, mask=msk) + + +def nordion2_post_orthogonalize_masked_triton( + X: List[Tensor], + M: List[Tensor], + V: List[Tensor], + U: List[Tensor], + base_lr: Tensor, + adjusted_lr: Tensor, + weight_decay: Tensor, + ef_decay: Tensor, + muon_beta2: Tensor, + select_dim: int, +): + """Triton-fused version of ``nordion2_post_orthogonalize_masked``. + + Two row-programmed kernels per matrix plus one batched epilogue reduction + per megabatch group (the NorMuon Frobenius rescale is a matrix-level + reduction, so a grid-wide dependency forces the two-phase split; the + epilogue is a deterministic torch reduction rather than fp32 atomics). + Falls back to the compiled masked post when the fast path does not apply. + All matrices in a megabatch shape group share one shape; mixed shapes fall + back too. + """ + shapes = {tuple(x.shape) for x in X} + use_fallback = ( + not TRITON_AVAILABLE + or select_dim != -2 + or not X + or not X[0].is_cuda + or X[0].ndim != 2 + or len(shapes) != 1 + or any(is_traceable_wrapper_subclass(x) for x in X) + ) + if use_fallback: + from .nordion2 import nordion2_post_orthogonalize_masked + + nordion2_post_orthogonalize_masked( + X=X, M=M, V=V, U=U, base_lr=base_lr, adjusted_lr=adjusted_lr, + weight_decay=weight_decay, ef_decay=ef_decay, + muon_beta2=muon_beta2, select_dim=select_dim, + ) + return + + Mrows, N = X[0].shape + if Mrows == 0 or N == 0: + return + + a = (1 - base_lr * weight_decay).item() + neg_b = (-adjusted_lr).item() + ef = ef_decay.to(M[0].dtype).item() + lerp_w = (1 - muon_beta2).item() + + for x, m, v, u in zip(X, M, V, U): + if not ( + x.is_contiguous() and m.is_contiguous() and u.is_contiguous() + and v.is_contiguous() + ): + raise ValueError( + "nordion2_post_orthogonalize_masked_triton requires contiguous tensors" + ) + + n_mat = len(X) + device = X[0].device + # Per-row partials: [n_mat, 3, rows] fp32 (sumsq, sumsq/denom^2, denom). + P = torch.empty((n_mat, 3, Mrows), dtype=torch.float32, device=device) + + for i, (m, v, u) in enumerate(zip(M, V, U)): + _nordion2_post_ortho_masked_stats_kernel[(Mrows,)]( + u, v, m, P[i], + ef, lerp_w, + N, + u.stride(0), u.stride(1), + v.stride(0), + m.stride(0), m.stride(1), + P.stride(1), + ) + + # Epilogue: matrix-level Frobenius rescale factors, one deterministic + # reduction for the whole group. Matches the compiled path's + # norm_U / clamp(norm_U_new, 1e-8). + sums = P[:, :2, :].sum(dim=-1) # [n_mat, 2] + scale = sums[:, 0].sqrt() / sums[:, 1].sqrt().clamp_min(1e-8) + scale = scale.contiguous() + + for i, (x, u) in enumerate(zip(X, U)): + _nordion2_post_ortho_masked_apply_kernel[(Mrows,)]( + x, u, P[i], scale, i, + a, neg_b, + N, + x.stride(0), x.stride(1), + u.stride(0), u.stride(1), + P.stride(1), + ) diff --git a/dion/nordion2.py b/dion/nordion2.py index 80eba38..1e9709f 100644 --- a/dion/nordion2.py +++ b/dion/nordion2.py @@ -15,9 +15,12 @@ ) from .opt_utils import AsyncTask, to_local from .dion2 import ( + CAPPED_STATS, dion2_pre_orthogonalize, dion2_post_orthogonalize, dion2_pre_accumulate, + dion2_pre_accumulate_norms, + dion2_capped_packed_async, _make_select_and_orthogonalize, ) from .normuon import normuon_normalization_stacked, _normuon_normalization_core @@ -55,7 +58,11 @@ class NorDion2(DistributedOrthoBase): assembled whole matrix -- layout-invariant/reproducible and better- converging. "local": per-shard top-k (union) -- cheaper comm but a sharding-dependent approximation that converges slightly worse; opt - in when comm-bound at large scale. No-op off the row-sharded path. + in when comm-bound at large scale. "global_capped": exact global + top-``ceil(fraction*global)`` selection at ~local comm cost via a + norms-only all-gather + pooled packed transport; overflow winners + deferred through error feedback (see Dion2, incl. + ``capacity_factor``). No-op off the row-sharded path. NorDion2 optimizer applying Dion2 update to NorMuon """ @@ -79,15 +86,21 @@ def __init__( newton_schulz_func: Optional[Callable] = None, triton_post_ortho: bool = False, selection_scope: str = "global", + capacity_factor: Optional[float] = None, ): # Validate hyperparameters if lr < 0.0: raise ValueError(f"Invalid learning rate: {lr}") if not (0.0 < fraction <= 1.0): raise ValueError(f"fraction must be in (0, 1], got {fraction}") - if selection_scope not in ("local", "global"): + if selection_scope not in ("local", "global", "global_capped"): raise ValueError( - f"selection_scope must be 'local' or 'global', got {selection_scope!r}" + f"selection_scope must be 'local', 'global', or 'global_capped', " + f"got {selection_scope!r}" + ) + if capacity_factor is not None and capacity_factor < 1.0: + raise ValueError( + f"capacity_factor must be None (auto) or >= 1.0, got {capacity_factor}" ) if mu < 0.0: raise ValueError(f"Invalid momentum factor (mu): {mu}") @@ -114,6 +127,7 @@ def __init__( algorithm="nordion2", step=0, selection_scope=selection_scope, + capacity_factor=capacity_factor, ) super().__init__( params, distributed_mesh, "nordion2", defaults, @@ -155,6 +169,9 @@ def _create_ortho_tasks( Mega-batched NorDion2 task creation: groups ALL same-shape parameters into a single task to minimize communication rounds and kernel launches. """ + # New optimizer step: reset the rank-local capped-deferral counters so + # they accumulate exactly this step's packed megabatches. + CAPPED_STATS.clear() for group in param_groups: assert group["algorithm"] == self._algo_name assert all( @@ -180,6 +197,7 @@ def _create_ortho_tasks( newton_schulz_func=self._newton_schulz_func, triton_post_ortho=self._triton_post_ortho, selection_scope=group["selection_scope"], + capacity_factor=group["capacity_factor"], ) shape_groups: dict[tuple, list] = defaultdict(list) @@ -246,6 +264,7 @@ def nordion2_update_megabatch_async( newton_schulz_func: Optional[Callable] = None, triton_post_ortho: bool = False, selection_scope: str = "global", + capacity_factor: Optional[float] = None, ) -> Generator[None, None, None]: """ Mega-batched NorDion2 update: processes ALL same-shape parameters in one @@ -266,7 +285,6 @@ def nordion2_update_megabatch_async( # comm_dim for sharded communication: use select_dim comm_dim = select_dim if is_sharded else None - global_scope = selection_scope == "global" and comm_dim is not None # On the sharded path X[0] must still be a DTensor, so .shape[comm_dim] # is the unsharded global size. The megabatch fn uses this to compute @@ -283,6 +301,85 @@ def nordion2_update_megabatch_async( else: global_dim_size = None + # --- "global_capped" packed-path decision; mirrors dion2 (see there) --- + capped_packed = ( + selection_scope == "global_capped" + and comm_dim is not None + and process_group is not None + and X[0].ndim == 2 + ) + if capped_packed: + padded_local = (global_dim_size + world_size - 1) // world_size + k = max(1, int(math.ceil(fraction * padded_local))) + per_rank = (N + world_size - 1) // world_size + if capacity_factor is None: + c = 1.0 + 2.0 * math.sqrt((1.0 - 1.0 / world_size) / (per_rank * k)) + else: + c = float(capacity_factor) + budget = int(math.ceil(c * per_rank * k)) + # Route to exact "global" when packing cannot pay for itself -- the + # packed chunk sends 1 + budget rows (count header included), so + # break-even against sending the full shard is budget + 1 -- or when + # the int32 count header cannot fit in the chunk's first row (2 bf16 + # slots per matrix). + if budget + 1 >= per_rank * padded_local or 2 * per_rank > X[0].shape[-1]: + capped_packed = False + + global_scope = comm_dim is not None and ( + selection_scope == "global" + or (selection_scope == "global_capped" and not capped_packed) + ) + + if capped_packed: + # --- "global_capped": winner-only packed transport (see dion2) --- + # Returns full-size shards, zero except at applied winner rows -- the + # masked post's format. Deferred winners skip both EF decay and the + # muon_beta2 variance update (the masked post only touches nonzero + # rows), so their momentum accumulates until they win a slot. + norms = dion2_pre_accumulate_norms( + G=to_local(G), M=to_local(M), select_dim=select_dim + ) + U_ortho = yield from dion2_capped_packed_async( + M_local=to_local(M), + norms=norms, + padded_local=padded_local, + fraction=fraction, + global_size=global_dim_size, + device_rank=device_rank, + world_size=world_size, + process_group=process_group, + per_rank=per_rank, + budget=budget, + kw=k * world_size, + newton_schulz_func=newton_schulz_func, + flatten=flatten, + epsilon=epsilon, + ) + if adjust_lr is None: + adjusted_lr = lr + elif adjust_lr == "spectral_norm": + adjusted_lr = adjust_lr_spectral_norm(lr, X[0].shape, flatten=flatten) + elif adjust_lr == "rms_norm": + adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten) + else: + raise ValueError(f"Unknown adjust_lr: {adjust_lr}") + _masked_post = nordion2_post_orthogonalize_masked + if triton_post_ortho: + from .dion2_triton import nordion2_post_orthogonalize_masked_triton as _masked_post + _masked_post( + X=to_local(X), + M=to_local(M), + V=to_local(V), + U=U_ortho, + base_lr=lr, + adjusted_lr=adjusted_lr, + weight_decay=weight_decay, + ef_decay=momentum, + muon_beta2=muon_beta2, + select_dim=select_dim, + ) + return + if global_scope: # --- Global selection: send the full shard, select after assembly --- U_full = dion2_pre_accumulate(G=to_local(G), M=to_local(M)) @@ -313,7 +410,10 @@ def nordion2_update_megabatch_async( # rows of M, runs NorMuon per-neuron normalization on the selected rows # (updating the local variance buffer V in place), and applies the # masked weight update -- all keyed off the nonzero mask, no indices. - nordion2_post_orthogonalize_masked( + _masked_post = nordion2_post_orthogonalize_masked + if triton_post_ortho: + from .dion2_triton import nordion2_post_orthogonalize_masked_triton as _masked_post + _masked_post( X=to_local(X), M=to_local(M), V=to_local(V), diff --git a/tests/test_dion2_post_ortho_masked_triton.py b/tests/test_dion2_post_ortho_masked_triton.py new file mode 100644 index 0000000..b9e3121 --- /dev/null +++ b/tests/test_dion2_post_ortho_masked_triton.py @@ -0,0 +1,301 @@ +"""Tests for the Triton masked post-orthogonalize kernels. + +``dion2_post_orthogonalize_masked_triton`` and +``nordion2_post_orthogonalize_masked_triton`` fuse the masked posts used by +the "global" and "global_capped" selection scopes. Parity is checked against +the compiled masked posts on identical inputs: + +- Unselected rows must match bitwise for X (both compute a*x with one + rounding) and be untouched for M/V. +- Selected rows differ at FP-rounding level: the compiled Dion2 path rounds + ``neg_lr * u`` to U's dtype before the add (two/three roundings), the + Triton kernel computes ``a*x + neg_b*u`` in fp32 with a single rounding. + +The end-to-end test runs NorDion2 with selection_scope="global_capped" on +2 GPUs, triton_post_ortho True vs False, and compares final params. +""" + +import multiprocessing as mp +import os +import pytest +import torch + +from dion.dion2 import dion2_post_orthogonalize_masked +from dion.nordion2 import nordion2_post_orthogonalize_masked +from dion.dion2_triton import ( + TRITON_AVAILABLE, + dion2_post_orthogonalize_masked_triton, + nordion2_post_orthogonalize_masked_triton, +) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +CUDA_AVAILABLE = torch.cuda.is_available() +TRITON_AND_CUDA = CUDA_AVAILABLE and TRITON_AVAILABLE + +torch._dynamo.config.cache_size_limit = 64 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_masked_data( + n_mat, rows, cols, sel_rows_per_mat, seed=42, x_dtype=torch.float32, + with_v=False, +): + """Lists of X, M, (V,) U with U exactly zero outside selected rows. + + ``sel_rows_per_mat`` is a list (len n_mat) of per-matrix selected-row + counts; 0 means no rows selected on that matrix. + """ + torch.manual_seed(seed) + X, M, V, U, sels = [], [], [], [], [] + for i in range(n_mat): + X.append(torch.randn(rows, cols, device=DEVICE, dtype=x_dtype)) + M.append(torch.randn(rows, cols, device=DEVICE, dtype=x_dtype)) + if with_v: + V.append( + torch.rand(rows, 1, device=DEVICE, dtype=x_dtype) + 0.01 + ) + u = torch.zeros(rows, cols, device=DEVICE, dtype=torch.bfloat16) + k = sel_rows_per_mat[i] + sel = torch.randperm(rows, device=DEVICE)[:k] + if k: + u[sel] = torch.randn(k, cols, device=DEVICE, dtype=torch.bfloat16) + U.append(u) + mask = torch.zeros(rows, dtype=torch.bool, device=DEVICE) + mask[sel] = True + sels.append(mask) + scalars = dict( + base_lr=torch.tensor(0.01), + adjusted_lr=torch.tensor(0.02), + weight_decay=torch.tensor(0.1), + ef_decay=torch.tensor(0.95), + ) + if with_v: + scalars["muon_beta2"] = torch.tensor(0.95) + return X, M, V, U, sels, scalars + return X, M, U, sels, scalars + + +def _tols(dtype): + if dtype == torch.bfloat16: + return dict(atol=1e-2, rtol=1e-2) + return dict(atol=1e-6, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# Function-level parity: Dion2 masked +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton") +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "rows,cols,sel_counts", + [ + (64, 128, [16, 0, 64]), # partial / none / all selected + (1, 32, [1, 0, 1]), # single row + (128, 4099, [32, 7, 0]), # cols > BLOCK_N, non-pow2 + ], +) +def test_dion2_masked_parity(x_dtype, rows, cols, sel_counts): + X, M, U, sels, sc = _make_masked_data( + len(sel_counts), rows, cols, sel_counts, x_dtype=x_dtype + ) + X_ref = [x.clone() for x in X] + M_ref = [m.clone() for m in M] + + dion2_post_orthogonalize_masked( + X=X_ref, M=M_ref, U=U, select_dim=-2, + base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"], + weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"], + ) + dion2_post_orthogonalize_masked_triton( + X=X, M=M, U=U, select_dim=-2, + base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"], + weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"], + ) + + for x, x_ref, m, m_ref, sel in zip(X, X_ref, M, M_ref, sels): + # Unselected rows: both paths compute a*x with one rounding -> bitwise. + torch.testing.assert_close(x[~sel], x_ref[~sel], rtol=0, atol=0) + # M untouched on unselected rows (compiled multiplies by exactly 1). + torch.testing.assert_close(m[~sel], m_ref[~sel], rtol=0, atol=0) + torch.testing.assert_close(x, x_ref, **_tols(x_dtype)) + torch.testing.assert_close(m, m_ref, **_tols(x_dtype)) + + +@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton") +def test_dion2_masked_batched_3d(): + """Leading batch dim: (B, rows, cols) with per-(b,row) selection.""" + torch.manual_seed(0) + B, rows, cols = 3, 32, 64 + x = torch.randn(B, rows, cols, device=DEVICE) + m = torch.randn(B, rows, cols, device=DEVICE) + u = torch.zeros(B, rows, cols, device=DEVICE, dtype=torch.bfloat16) + for b in range(B): + sel = torch.randperm(rows)[: 4 * (b + 1)] + u[b, sel] = torch.randn(len(sel), cols, device=DEVICE, dtype=torch.bfloat16) + args = dict( + base_lr=torch.tensor(0.01), adjusted_lr=torch.tensor(0.02), + weight_decay=torch.tensor(0.1), ef_decay=torch.tensor(0.95), + select_dim=-2, + ) + x_ref, m_ref = x.clone(), m.clone() + dion2_post_orthogonalize_masked(X=[x_ref], M=[m_ref], U=[u], **args) + dion2_post_orthogonalize_masked_triton(X=[x], M=[m], U=[u], **args) + torch.testing.assert_close(x, x_ref, **_tols(torch.float32)) + torch.testing.assert_close(m, m_ref, **_tols(torch.float32)) + + +@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton") +def test_dion2_masked_empty_shard_skipped(): + args = dict( + base_lr=torch.tensor(0.01), adjusted_lr=torch.tensor(0.02), + weight_decay=torch.tensor(0.1), ef_decay=torch.tensor(0.95), + select_dim=-2, + ) + x = torch.empty(0, 64, device=DEVICE) + m = torch.empty(0, 64, device=DEVICE) + u = torch.empty(0, 64, device=DEVICE, dtype=torch.bfloat16) + # Must not raise. + dion2_post_orthogonalize_masked_triton(X=[x], M=[m], U=[u], **args) + + +# --------------------------------------------------------------------------- +# Function-level parity: NorDion2 masked +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton") +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "rows,cols,sel_counts", + [ + (64, 128, [16, 0, 64]), + (128, 4099, [32, 7, 0]), + (16, 20480, [4, 16, 0]), # production-wide rows + ], +) +def test_nordion2_masked_parity(x_dtype, rows, cols, sel_counts): + X, M, V, U, sels, sc = _make_masked_data( + len(sel_counts), rows, cols, sel_counts, x_dtype=x_dtype, with_v=True + ) + X_ref = [x.clone() for x in X] + M_ref = [m.clone() for m in M] + V_ref = [v.clone() for v in V] + + kw = dict( + base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"], + weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"], + muon_beta2=sc["muon_beta2"], select_dim=-2, + ) + nordion2_post_orthogonalize_masked(X=X_ref, M=M_ref, V=V_ref, U=U, **kw) + nordion2_post_orthogonalize_masked_triton(X=X, M=M, V=V, U=U, **kw) + + for x, x_ref, m, m_ref, v, v_ref, sel in zip( + X, X_ref, M, M_ref, V, V_ref, sels + ): + torch.testing.assert_close(x[~sel], x_ref[~sel], rtol=0, atol=0) + torch.testing.assert_close(m[~sel], m_ref[~sel], rtol=0, atol=0) + torch.testing.assert_close(v[~sel], v_ref[~sel], rtol=0, atol=0) + torch.testing.assert_close(x, x_ref, **_tols(x_dtype)) + torch.testing.assert_close(m, m_ref, **_tols(x_dtype)) + torch.testing.assert_close(v, v_ref, **_tols(x_dtype)) + + +@pytest.mark.skipif(not TRITON_AND_CUDA, reason="requires CUDA + triton") +def test_nordion2_masked_no_rows_selected_anywhere(): + """All-zero U: pure weight decay, V and M untouched, no NaN from the + Frobenius rescale (norm 0 / clamp).""" + X, M, V, U, sels, sc = _make_masked_data( + 2, 32, 64, [0, 0], with_v=True + ) + kw = dict( + base_lr=sc["base_lr"], adjusted_lr=sc["adjusted_lr"], + weight_decay=sc["weight_decay"], ef_decay=sc["ef_decay"], + muon_beta2=sc["muon_beta2"], select_dim=-2, + ) + X_ref = [x.clone() for x in X] + M_ref = [m.clone() for m in M] + V_ref = [v.clone() for v in V] + nordion2_post_orthogonalize_masked(X=X_ref, M=M_ref, V=V_ref, U=U, **kw) + nordion2_post_orthogonalize_masked_triton(X=X, M=M, V=V, U=U, **kw) + for x, x_ref in zip(X, X_ref): + assert torch.isfinite(x).all() + torch.testing.assert_close(x, x_ref, rtol=0, atol=0) + for m, m_ref in zip(M, M_ref): + torch.testing.assert_close(m, m_ref, rtol=0, atol=0) + for v, v_ref in zip(V, V_ref): + torch.testing.assert_close(v, v_ref, rtol=0, atol=0) + + +# --------------------------------------------------------------------------- +# End-to-end: NorDion2 global_capped, triton vs compiled post (2 GPUs) +# --------------------------------------------------------------------------- + +def _e2e_worker(rank, world_size, port, triton_post, out_path): + import torch.distributed as dist + from dion import NorDion2 + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + + torch.manual_seed(1234) # same init on all ranks (replicated DDP-style) + full = [torch.randn(64, 96) for _ in range(3)] + params = [ + torch.nn.Parameter( + torch.tensor_split(w, world_size, dim=0)[rank].clone().to(device) + ) + for w in full + ] + opt = NorDion2( + [dict(params=params)], + distributed_mesh=dist.group.WORLD, + lr=0.01, + fraction=0.25, + selection_scope="global_capped", + capacity_factor=1.5, + use_triton=False, + triton_post_ortho=triton_post, + ) + torch.manual_seed(99) + for _ in range(3): + grads = [torch.randn(64, 96) for _ in range(3)] + for p, g in zip(params, grads): + p.grad = torch.tensor_split(g, world_size, dim=0)[rank].clone().to(device) + opt.step() + opt.zero_grad() + + if rank == 0: + torch.save([p.detach().cpu() for p in params], out_path) + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.skipif( + not TRITON_AND_CUDA or torch.cuda.device_count() < 2, + reason="requires >=2 GPUs + triton", +) +def test_nordion2_capped_e2e_triton_vs_compiled(tmp_path): + ctx = mp.get_context("spawn") + outs = {} + for triton_post, port in ((False, 29711), (True, 29713)): + out = tmp_path / f"params_{triton_post}.pt" + procs = [ + ctx.Process( + target=_e2e_worker, args=(r, 2, port, triton_post, str(out)) + ) + for r in range(2) + ] + for p in procs: + p.start() + for p in procs: + p.join(300) + assert p.exitcode == 0 + outs[triton_post] = torch.load(out) + + for p_ref, p_tri in zip(outs[False], outs[True]): + torch.testing.assert_close(p_tri, p_ref, atol=1e-5, rtol=1e-4) diff --git a/tests/test_dion2_selection_scope.py b/tests/test_dion2_selection_scope.py index ea7e475..fcfc4e8 100644 --- a/tests/test_dion2_selection_scope.py +++ b/tests/test_dion2_selection_scope.py @@ -32,7 +32,7 @@ import torch.multiprocessing as mp from torch.distributed.tensor import DeviceMesh, distribute_tensor, Shard -from dion.dion2 import Dion2 +from dion.dion2 import Dion2, dion2_capped_select_async from dion.nordion2 import NorDion2 from dion.polar_express import polar_express @@ -241,3 +241,336 @@ def test_fraction_one_local_equals_global(OptCls, kw, tmp_path): dl = _run_sharded(OptCls, "local", 1.0, kw, 2, 29630, tmp_path) dg = _run_sharded(OptCls, "global", 1.0, kw, 2, 29640, tmp_path) torch.testing.assert_close(dl["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL) + + +# ---- selection_scope="global_capped" (pooled packed transport) ---- +# +# All-gather slice norms only; every rank derives the same EXACT-COUNT global +# top-ceil(fraction*global) winner set (stable-sort tie-break). Each rank +# packs its winner rows into fixed-size chunks pooled across the megabatch's +# matrices (int32 count header bit-cast into each chunk's first row); winners +# overflowing the pooled budget defer via error feedback (no EF decay -> +# momentum accumulates -> selected later). Only applied rows are decayed, by +# the same masked post the global scope uses. Groups where packing cannot +# save comm route to the exact "global" scope. + + +def _capped_worker(rank, world_size, port, OptCls, scope, kw, out_path, shape, + fraction, boosts, steps, dtype): + """Multi-step, multi-param worker with crafted gradients: small noise + everywhere, plus rows in ``boosts[i]`` (one dict per param) set to a + constant so their momentum L1 norm is exactly the boost magnitude. Steps + after the first use a zero gradient (isolates the error-feedback + dynamics). Same-shape params land in ONE megabatch group, so multiple + params exercise the pooled packed transport. Saves the full W of every + param after every step. ``dtype`` sets the param (and therefore momentum) + dtype -- bf16 exercises the mixed-precision paths fp32 can never reach.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + dev = torch.device(f"cuda:{rank}") + mesh = DeviceMesh("cuda", list(range(world_size))) + + rows, cols = shape + g = torch.Generator(device=dev).manual_seed(1234) + W0s, Gs, params = [], [], [] + for boost in boosts: + W0 = torch.randn(rows, cols, generator=g, device=dev).to(dtype) + G = (torch.randn(rows, cols, generator=g, device=dev) * 0.01).to(dtype) + for r, mag in boost.items(): + G[r] = mag / cols # exact row L1 norm == mag + W0s.append(W0) + Gs.append(G) + params.append(torch.nn.Parameter(distribute_tensor(W0, mesh, [Shard(0)]))) + + opt = OptCls( + [dict(params=params)], + distributed_mesh=mesh, + lr=0.1, + fraction=fraction, + newton_schulz_func=_ns, + selection_scope=scope, + **kw, + ) + Ws = [] + for s in range(steps): + for p, G in zip(params, Gs): + grad = G if s == 0 else torch.zeros_like(G) + p.grad = distribute_tensor(grad, mesh, [Shard(0)]) + opt.step() + Ws.append([p.detach().full_tensor().cpu() for p in params]) + if rank == 0: + torch.save({"W0": [w.cpu() for w in W0s], "Ws": Ws}, out_path) + dist.destroy_process_group() + + +def _run_capped(OptCls, scope, kw, port, tmp, *, shape=(16, 32), fraction=0.25, + boost=None, steps=1, dtype=torch.float32): + out = str(tmp / f"out_{port}.pt") + boosts = boost if isinstance(boost, list) else [boost or {}] + mp.spawn( + _capped_worker, + args=(2, port, OptCls, scope, kw, out, shape, fraction, boosts, + steps, dtype), + nprocs=2, + join=True, + ) + return torch.load(out) + + +def _changed_rows(W_after, W_before): + """Row indices where anything changed at all (bitwise).""" + return set((W_after != W_before).any(dim=-1).nonzero().flatten().tolist()) + + +@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs") +@pytest.mark.parametrize("OptCls,kw", [ + (Dion2, {}), + (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}), +]) +def test_capped_scope_matches_global_when_winners_fit(OptCls, kw, tmp_path): + """When every rank owns at most k winners and world_size divides the rows + (so the top-(k*world) budget equals global's top-ceil(fraction*global)), + capped must equal global-scope exactly. 16 rows / 2 ranks / fraction 0.25: + k=2, budget 4; the 4 boosted rows are split 2 per rank.""" + boost = {0: 3.0, 1: 2.9, 8: 2.8, 9: 2.7} + dc = _run_capped(OptCls, "global_capped", kw, 29680, tmp_path, boost=boost) + dg = _run_capped(OptCls, "global", kw, 29681, tmp_path, boost=boost) + torch.testing.assert_close(dc["Ws"][0][0], dg["Ws"][0][0], + rtol=BF16_RTOL, atol=BF16_ATOL) + + +@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs") +@pytest.mark.parametrize("OptCls,kw", [ + (Dion2, {"ef_decay": 0.5, "weight_decay": 0.0, "capacity_factor": 1.0}), + (NorDion2, {"mu": 0.5, "muon_beta2": 0.95, "weight_decay": 0.0, + "capacity_factor": 1.0}), +]) +def test_capped_scope_defers_overflow_winners(OptCls, kw, tmp_path): + """All 4 global winners live on rank 0 but its chunk budget is pinned to + k=2 rows (capacity_factor=1.0, single matrix): step 1 sends exactly its + top-2 {0,1} by normalized priority (rows 2,3 deferred WITHOUT + error-feedback decay; rank 1 sends nothing, so NO row of rank 1 changes -- + this is what distinguishes capped from local, which would fill and apply + rank 1's top-2). Step 2 (zero grad): rows 0,1 decayed to 1.5/1.45 while + the deferred 2,3 kept 2.8/2.7, so exactly {2,3} are applied. + weight_decay=0 makes non-applied rows bitwise unchanged.""" + boost = {0: 3.0, 1: 2.9, 2: 2.8, 3: 2.7} + d = _run_capped(OptCls, "global_capped", kw, 29682, tmp_path, + boost=boost, steps=2) + assert _changed_rows(d["Ws"][0][0], d["W0"][0]) == {0, 1} + assert _changed_rows(d["Ws"][1][0], d["Ws"][0][0]) == {2, 3} + + +@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs") +@pytest.mark.parametrize("OptCls,kw", [ + (Dion2, {}), + (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}), +]) +def test_capped_scope_uneven_shards_fraction_one_equals_global(OptCls, kw, tmp_path): + """Uneven division (17 rows over 4 ranks -> [5,5,5,2] + one short shard): + at fraction=1.0 the threshold is the -1 pad value, everything wins, and + capped must equal global (both apply all 17 rows; the megabatch zero-pads + the short shard's spare slots).""" + dc = _run_sharded(OptCls, "global_capped", 1.0, kw, 4, 29683, tmp_path, + shape=(17, 32)) + dg = _run_sharded(OptCls, "global", 1.0, kw, 4, 29684, tmp_path, + shape=(17, 32)) + torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL) + + +@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs") +def test_capped_scope_uneven_shards_runs(tmp_path): + """Partial-fraction sanity on uneven shards (one short shard) through the + real packed path: capacity_factor is pinned to 1.0 because the auto slack + on this tiny k would reach the full shard and route the group to global.""" + d = _run_sharded(Dion2, "global_capped", 0.35, {"capacity_factor": 1.0}, + 4, 29685, tmp_path, shape=(17, 32)) + assert torch.isfinite(d["W"]).all() + assert not torch.allclose(d["W"], d["W0"], atol=1e-6) + + +@pytest.mark.skipif(CUDA < 4, reason="needs 4 GPUs") +def test_capped_scope_empty_shard(tmp_path): + """3 rows over 4 ranks chunk to [1,1,1,0]: rank 3's shard is empty. For a + matrix this tiny the pooled budget reaches the full shard, so the group + must route to the exact global scope -- capped(f) == global(f), with the + exact top-ceil(fraction*global) budget (2 rows at f=0.5, not k*world).""" + dc = _run_sharded(Dion2, "global_capped", 0.5, {}, 4, 29686, tmp_path, + shape=(3, 32)) + assert torch.isfinite(dc["W"]).all() + dg = _run_sharded(Dion2, "global", 0.5, {}, 4, 29687, tmp_path, + shape=(3, 32)) + torch.testing.assert_close(dc["W"], dg["W"], rtol=BF16_RTOL, atol=BF16_ATOL) + + +@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs") +@pytest.mark.parametrize("scope", ["local", "global_capped"]) +@pytest.mark.parametrize("OptCls,kw", [ + (Dion2, {}), + (NorDion2, {"mu": 0.95, "muon_beta2": 0.95}), +]) +def test_scope_bf16_momentum(OptCls, kw, scope, tmp_path): + """bf16 params (=> bf16 momentum) must run through the selection scopes. + Guards the EF write-back dtype: the capped-scope ef_factor is a + *dimensioned* fp32 tensor, and ``bf16 * fp32[N,k,1]`` type-promotes to + fp32, which the in-place ``scatter_`` into bf16 momentum rejects with + "Expected self.dtype to be equal to src.dtype". A 0-dim fp32 scalar (the + pre-capped code, and the local branch) does NOT promote, so fp32-param + tests can never catch this -- it only fires in real bf16 training.""" + port = 29690 + (scope == "global_capped") * 2 + (OptCls is NorDion2) + d = _run_capped(OptCls, scope, kw, port, tmp_path, dtype=torch.bfloat16) + W = d["Ws"][0][0] + assert torch.isfinite(W.float()).all() + assert not torch.equal(W, d["W0"][0]) # the step actually applied an update + + +def _select_consistency_worker(rank, world_size, port): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + dev = torch.device(f"cuda:{rank}") + + # Distinct norms everywhere (no ties), different on each rank; one short + # shard (rank 1 has 3 of padded_local=4 slices) to exercise the -1 pad. + N, padded_local = 2, 4 + global_size = 7 # 4 + 3 real rows + fraction = 0.5 # k_total = ceil(0.5 * 7) = 4 per matrix + local = padded_local - rank # rank0: 4, rank1: 3 + norms = (torch.arange(N * local, dtype=torch.float32, device=dev) + .reshape(N, local) * 1.7 + 3.0 + rank * 0.37) + + gen = dion2_capped_select_async( + norms, padded_local, fraction, global_size, rank, world_size, None + ) + next(gen) + try: + next(gen) + raise AssertionError("generator should be exhausted after the collective") + except StopIteration as e: + winner, thresh = e.value + + assert winner.shape == (N, local) and winner.dtype == torch.bool + assert thresh.shape == (N, 1) + # (f) rank consistency: every rank must hold the bit-identical threshold. + both = [torch.empty_like(thresh) for _ in range(world_size)] + dist.all_gather(both, thresh.contiguous()) + assert torch.equal(both[0], both[1]) + # Exact-count winner set: exactly ceil(fraction*global) winners per matrix + # across all ranks -- the receiver buffer bound relies on this. + winners = winner.sum() + dist.all_reduce(winners) + assert winners.item() == N * 4 + dist.destroy_process_group() + + +@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs") +def test_capped_select_rank_consistent(): + mp.spawn(_select_consistency_worker, args=(2, 29688), nprocs=2, join=True) + + +def test_capped_select_negative_pad_raises(): + """local_size > padded_local would make F.pad silently TRIM the norms; the + guard must raise instead. (Raise happens before the collective, so no + process group is needed.) Pure-CPU unit test.""" + gen = dion2_capped_select_async(torch.zeros(1, 5), 3, 0.5, 8, 0, 2, None) + with pytest.raises(ValueError, match="exceeds padded size"): + next(gen) + + +def test_capped_header_bitcast_roundtrip(): + """The packed transport ships int32 per-matrix counts bit-cast into the + bf16 chunk's first row (2 bf16 slots per int32). The bit-cast must + round-trip exactly for counts far beyond 256 (where bf16 *values* lose + integer exactness -- the failure mode this encoding avoids). Pure CPU.""" + counts = torch.tensor([[320, 4097, 0], [70000, 1, 2 ** 20]], dtype=torch.int32) + world, per_rank = counts.shape + cols = 32 + header = torch.zeros(world, cols, dtype=torch.bfloat16) + header[:, : per_rank * 2] = counts.contiguous().view(torch.bfloat16) + back = header[:, : per_rank * 2].contiguous().view(torch.int32) + assert torch.equal(back, counts) + + +def test_capped_pack_roundtrip_identity(): + """pack -> (simulated a2a) -> assemble -> identity-NS -> repack -> + (simulated reverse a2a) -> unpack must return each rank's kept winner rows + bit-exactly (bf16) at their home positions and exact zeros elsewhere, with + sent rows always a subset of the winner set. Covers pooled deferral (4 + winners vs budget 3), a matrix with zero winners anywhere ("blackout"), + and uneven per-(rank, matrix) counts. Pure CPU, no process group.""" + from dion.dion2 import ( + dion2_capped_pack, dion2_capped_assemble, + dion2_capped_repack, dion2_capped_unpack, + ) + torch.manual_seed(0) + world, per_rank, N, local, cols, budget, kw = 2, 2, 4, 4, 8, 3, 4 + + Ms, normss = [], [] + for _ in range(world): + M = [torch.randn(local, cols) + 5 for _ in range(N)] + Ms.append(M) + normss.append(torch.stack([m.abs().sum(-1) for m in M])) + w0 = torch.zeros(N, local, dtype=torch.bool) + w1 = torch.zeros(N, local, dtype=torch.bool) + w0[0, :3] = True; w1[0, 0] = True # m0: rank0 pooled 3+1 > budget -> defer + w0[1, 1] = True; w1[1, 2] = True # m1: 1 + 1 + w0[2, :2] = True # m2: 2 + 0; m3: blackout + winners = [w0, w1] + threshs = [torch.ones(N, 1)] * world # uniform scale: prio == raw norm + + packs = [ + dion2_capped_pack(Ms[r], normss[r], winners[r], threshs[r], + world, per_rank, budget) + for r in range(world) + ] + backs = {} + for d in range(world): + recv = torch.stack([packs[r][0][d] for r in range(world)]) + cnt = torch.stack([packs[r][1][d] for r in range(world)]) + ns_in, ns_ix = dion2_capped_assemble(recv, cnt, kw) + backs[d] = dion2_capped_repack(ns_in, ns_ix, world, budget) + + assert int(packs[0][3]) == 1 and int(packs[1][3]) == 0 # deferred counts + for r in range(world): + recv2 = torch.stack([backs[d][r] for d in range(world)]) + U = dion2_capped_unpack(recv2, packs[r][2], N, local, + world * per_rank * local) + for i in range(N): + sent = U[i].abs().sum(-1) > 0 + assert bool((sent <= winners[r][i]).all()) # winners only + assert torch.equal(U[i][sent], Ms[r][i].to(torch.bfloat16)[sent]) + assert U[i][~sent].abs().sum() == 0 + + +@pytest.mark.skipif(CUDA < 2, reason="needs 2 GPUs") +def test_capped_pooled_slots_absorb_cross_matrix_skew(tmp_path): + """The T2 value proposition. 4 same-shape params form ONE megabatch group; + world=2 gives per_rank=2, so matrices {0,1} assemble on rank 0 and {2,3} + on rank 1, and rank 0's chunk to dest 0 POOLS matrices 0 and 1 with a + budget of per_rank*k = 4 rows (capacity_factor pinned to 1.0). + + Matrix 0 has ALL 4 of its global winners on rank 0 (needs 4 > k=2 rows of + "its own" slots); matrix 1 contributes 1 rank-0 winner at a 100x smaller + norm scale but the HIGHEST scale-normalized priority (3.0/2.7 = 1.111). + Rank 0's pooled winners = 5 > budget 4, so exactly one row defers -- and + it must be matrix 0's LOWEST normalized priority (row 3, 275/275 = 1.0), + NOT matrix 1's small-raw-norm winner, which a raw-norm packing order + would evict. Per-matrix walls (the pre-pooling design) could never have + sent matrix 0's third winner at all. weight_decay=0 makes unapplied rows + bitwise unchanged.""" + boosts = [ + {0: 300.0, 1: 290.0, 2: 280.0, 3: 275.0}, # m0: prios 1.09/1.05/1.02/1.0 + {0: 3.0, 8: 2.9, 9: 2.8, 10: 2.7}, # m1: prio(row0) = 1.111 + {}, + {}, + ] + kw = {"capacity_factor": 1.0, "weight_decay": 0.0} + d = _run_capped(Dion2, "global_capped", kw, 29694, tmp_path, boost=boosts) + # Matrix 0: rank 0 keeps its top-3 by normalized priority; row 3 deferred. + assert _changed_rows(d["Ws"][0][0], d["W0"][0]) == {0, 1, 2} + # Matrix 1: the small-scale winner survives the shared chunk; rank 1's + # winners {8,9,10} ride its own uncontended chunk. + assert _changed_rows(d["Ws"][0][1], d["W0"][1]) == {0, 8, 9, 10}