Dion2: Add selection_scope="global_capped": global selection at local comm cost

[alint77]

Follow-up to #97/#98. global selects better but ships the full shard (no comm saving vs Muon); local gets the fraction-scaled comm but selects per-shard, which converges worse (~0.09 nat in #98's A/B at small fraction). This adds a third scope that gets (nearly) both.

How it works

1. All-gather the row L1 norms only — one fp32 per row, ~KBs vs MBs for the rows.
2. Every rank derives the same exact-count global winner set — top-ceil(fraction·global) by stable argsort (deterministic tie-break), so the total arrivals per matrix are bounded by construction.
3. Each rank packs its winner rows into a fixed-size chunk per destination, pooled across that destination's per_rank matrices (one matrix's overflow uses another's spare room), back-to-back, with the per-matrix counts bit-cast as an int32 header into the chunk's first row. The receiver parses the sender's header — no recomputed plan, no host syncs, static shapes throughout. Winners that overflow the pooled budget are deferred: they skip error-feedback decay, so their momentum accumulates and they win a slot on a later step. Non-winners are never applied and never decayed.

capacity_factor controls the pooled budget: None (default) = auto per group, 1 + 2·√((1−1/world)/(per_rank·k)) (~0.02% expected deferral under a binomial winner-scatter model); a float ≥ 1.0 pins it (1.0 = pooling only, zero extra comm). Groups where packing cannot pay for itself (budget+1 reaches the full shard, header doesn't fit, col-sharded/non-2D) route to exact global. dion2.CAPPED_STATS exposes per-step deferred/winner row counts for measuring the real deferral rate.

Results (NorDion2, 5120-hidden/8-layer, 4×H100, fraction=0.25)

Measured on the previous fixed-slot transport; comm volume is unchanged modulo the ~0.1% header row (fresh trace pending).

per step (rank0)	global_capped	global	local
a2a payload	655 MB	2621 MB	655 MB
norms all-gather (the extra)	0.5 MB / 0.1 ms	—	—
opt.step GPU span	34.6 ms	35.5 ms	31.0 ms
step wall / MFU	462 ms / ~51%	467 ms / ~51%	459 ms / ~51%

Tests in tests/test_dion2_selection_scope.py (28/28 on 2×/4× H100): bit-exact pack→a2a→NS→unpack roundtrip (pooled deferral, blackout matrix, uneven counts); cross-matrix skew absorbed by pooling with scale-normalized eviction order; exact-count rank consistency; header bit-cast roundtrip; overflow deferral across steps; uneven/empty shards; fraction=1.0 == global; bf16 momentum.

Draft because

• The decisive convergence A/B is the fraction=0.25 sweep (≥2 seeds) — at fraction=0.5 all scopes tie (see review thread), so the scope only earns its keep where local actually loses.

🤖 Generated with Claude Code

[JohnLangford]

This is John's reviewbot.

Read the full diff plus dion2.py, nordion2.py, megabatch_base.py, opt_utils.py, and the existing scope tests. The core mechanism is sound and the implementation is careful. Notes below, most-substantive first.

Correctness — solid, with one semantic nuance to document

Verified the load-bearing invariants:

• Rank consistency holds by construction. all_gather_into_tensor delivers a bit-identical gathered on every rank, and thresh = vals[:, -1:] is a value (the k_total-th largest), so it's independent of topk tie-breaking. Every rank derives the same winner set. Good.
• No double/under momentum accumulation: dion2_pre_accumulate_norms does M += G, and dion2_pre_orthogonalize skips it exactly when priority is not None. Empty shards, uneven divisions, and fraction=1.0 (→ everything wins → equals global) all check out.
• No new deadlock risk: the extra all_gather is collective, selection_scope is group-uniform, and the entry condition is deterministic across ranks, so the collective order stays consistent.

But: "global selection" here is top-(k·world_size), not top-(fraction·global). They coincide only when world_size | global. For a non-divisible shard, e.g. global=17, world=4, fraction=0.35: global scope selects ceil(0.35·17)=6 rows, but global_capped uses budget k·world = ceil(0.35·5)·4 = 8. So global_capped selects a strictly larger set (top-8) than selection_scope="global" (top-6) — same count as local. This is fine (more orthogonalization, still a valid descent), but:

• The PR's claim "selection equals the exact global top-k when winners fit" should say top-(k·world_size); a test asserting equality against selection_scope="global" will disagree on non-divisible sizes.
• Your convergence A/B should keep this in mind — on the reported divisible 5120/4 case the counts match, but a fair vs-global comparison on ragged shapes will differ in count, not just deferral.

Design — the accumulation split is a footgun

Momentum accumulation is now owned by two compiled functions and kept correct only by two matching if priority is None guards living in different functions (and duplicated across dion2.py/nordion2.py). A future edit to one side silently double- or under-accumulates with no error. Consider either an explicit contract (e.g. pass an already_accumulated flag and assert), or push the winner/priority computation into dion2_pre_orthogonalize by handing it just the thresh tensor: it already stacks M and would compute slice_norms there anyway, so slice_norms += (slice_norms >= thresh) * offset keeps accumulation single-owner and drops the duplicate torch.stack(M). (Accumulation still has to precede the gather, so pre_accumulate_norms stays — but it becomes the only accumulator.)

Smaller items

• Dead parameter. device_rank is threaded into dion2_capped_priority_async but never used — each rank correctly reads its own norms directly rather than slicing gathered[device_rank]. Drop it. (If the intent was a gathered[device_rank] == norms self-check, it's unnecessary; if it wasn't, the unused arg is just confusing.)
• Missing robustness guard. dion2_capped_priority_async assumes local_size <= padded_local; if that FSDP2-contiguous assumption is ever violated, F.pad(norms, (0, padded_local - local_size)) gets a negative pad and silently trims norms instead of erroring. megabatch_orthogonalize_async has an explicit padded_local_size < original_local_size raise for the analogous case — mirror it here.
• Threshold ties. own >= thresh marks all values tied at the threshold as winners, so a hard tie can inflate the winner set beyond k·world. Measure-zero for continuous fp32 norms and bounded by the per-rank k cap anyway, so benign — just worth a one-line comment.
• Persistent deferral. A rank whose shard is structurally high-norm (not the "arbitrary cut" common case) defers the same overflow winners repeatedly; their momentum grows until they rotate in. Self-correcting in steady state, but there's no safeguard, and it's exactly the case a random-init A/B won't surface. Worth probing an imbalanced layout.

Tests

The committed change ships with no tests — test_dion2_selection_scope.py is untouched and the scratch harness isn't ported (acknowledged as draft). Please add, at minimum: (a) selection equals top-(k·world) global when winners fit; (b) an overflow winner is deferred and selected on the next step; (c) uneven/remainder shards; (d) fraction=1.0 == global; (e) empty shard; and (f) a rank-consistency assertion (all ranks agree on the winner set). Col-sharded (select_dim=-1) is untested here too, same as local.

Security / backward-compat

No concerns. Numeric optimizer path, no I/O or untrusted input; collective sizes derive from tensor shapes. The new enum value and trailing optional priority arg preserve existing behavior and callers use kwargs.

Nice work — the offset-priority trick to fold "winners-first, then fill by norm" into a single topk is clean, and the norms-only gather is a genuinely cheap way to buy global selection. Main asks before undraft: port the tests, tighten the accumulation contract, and document the top-(k·world) vs top-(fraction·global) distinction.

[alint77]

Pushed 6ce9ac5 addressing the review — all asks are in:

• Accumulation contract: dion2_pre_orthogonalize now takes the thresh tensor instead of a priority tensor; dion2_pre_accumulate_norms is the single accumulation owner and the offset machinery is gone.
• top-(k·world_size) documented in both optimizers (and demonstrated live by the empty-shard test, where capped(0.5) == global(1.0)).
• device_rank dropped, negative pad raises instead of silently trimming, >=-tie inflation commented.
• Tests ported into tests/test_dion2_selection_scope.py — your (a)–(f) plus the pad guard, 21/21 on 2×/4× H100.

One thing the thresh refactor surfaced that we both missed: as committed, the scope was trajectory-identical to local. The threshold is a value cut, so a rank's winners are always its top-norm rows — boosting a prefix of the sort order can never change a per-rank top-k (verified by 20k-trial brute force over random/uneven/empty shards). And since fillers received EF decay like real selections, the momentum dynamics matched too. Empirically: same-seed 100-step runs of old-capped and local produce loss curves identical to ~1e-6 (kernel-nondeterminism drift only), while global separates from step 5.

Fixed in the same commit: spare slots now carry zero rows (zeros pass through Newton-Schulz as exact zeros → no update) and fillers skip EF decay, so only global winners are applied and deferral is real. The new same-seed run diverges from both local and global from the first logged step, as it should. The deferral test is the one that separates the algorithms: with all 4 winners on rank 0 and k=2 slots, step 1 changes exactly rank 0's top-2 and nothing on rank 1 (local would apply rank 1's top-2), and step 2 applies exactly the deferred pair.

One consequence for the A/B: per-step applied rows are k·world minus deferrals, so the effective fraction dips slightly below nominal whenever winners overflow a rank — your persistent-deferral probe is the adversarial version of this and still worth running.

Perf unchanged (fresh 5120/8L trace, table in the description): a2a 655 MB vs global's 2621 per step, opt.step GPU 34.6 ms vs 35.5 (global) / 31.0 (local), MFU parity. Description updated to match the new semantics.

Remaining ask before undraft is the convergence A/B vs global — now a meaningful experiment, since the mechanism genuinely differs from local. Over to you for the runs.

[JohnLangford]

This is John's reviewbot.

Re-reviewed at 6ce9ac5, which responds to the earlier round. Nearly everything I raised is addressed, and I verified the fixes — including running the ported suite (21/21 pass on 4×H100) and prototyping the alternative end-to-end. Details below.

Round-1 items — all resolved, verified

• Single-owner accumulation: priority→thresh, and dion2_pre_accumulate_norms is now the sole accumulator with the offset trick gone (norm computed once for selection, threshold compared). Cleaner. (The if thresh is None accumulation guard is still split across two functions — fine and documented, just the one spot a future edit could desync.)
• Dead device_rank removed; negative-pad guard added with a unit test; tie semantics and the top-k·world vs top-ceil(fraction·global) budget both documented in docstrings and tests. 👍
• Tests ported and genuinely good — test_capped_scope_defers_overflow_winners asserting exact changed-row sets {0,1} then {2,3} across a zero-grad follow-up step is a nice EF-dynamics check. Coverage now spans winners-fit==global, overflow deferral, uneven shards, fraction=1.0, empty shard, rank-consistency, and the pad guard.

New winners-only masking — correct

Zeroing non-winner fillers and skipping their EF decay is sound: a zero input row survives Newton-Schulz as an exact zero (the global masked-post already relies on this), and since the assembled Gram is block-diagonal between winners and zero-fillers, the winners' orthogonalization is identical to sending winners alone. So with no overflow this now yields exactly the global top-k·world update — a real improvement over the previous local-equivalent fill.

Empirical findings from testing this (may inform the pending A/B)

I ran the mechanism on H100s. Two things worth having before the convergence run:

1. Deferral is common, not a tail case. Across 150 shape/world/fraction configs, the cap defers ≥1 winner in 74% of them even with random (balanced) sharding; single-step coverage of the true global top-k·world runs 71–91%. End-to-end, wherever it defers, the capped weight sits 1.8e-2–9e-2 off global that step (vs bf16 round-off when it doesn't). So the A/B is meaningful — and to isolate deferral, use a shape where world_size divides the sharded dim (else the budget-count difference confounds it).

2. A precision boundary that isn't deferral. The norms-only path selects on fp32 norms while global selects on bf16 shard norms; at the k-th-norm boundary these pick a different row ~20–90% of the time (rising with k, as the boundary gap shrinks). So global_capped won't match global even at zero deferral, and part of any A/B gap is precision, not the algorithm. There's now also a smaller internal version: thresh is fp32 (pre_accumulate_norms does .float()), but the winner mask compares it against slice_norms recomputed in pre_orthogonalize in the momentum dtype in a separate compiled graph — a reduction-order mismatch there could misclassify a boundary slice (self-correcting via EF, so benign). Cheap hardening: reuse the already-gathered norms for the winner mask instead of recomputing, so selection and thresholding share one reduction.

The more elegant alternative — now validated end-to-end

I prototyped the asymmetric (variable-length) reshard we discussed and it works on real hardware (2 & 4 ranks): drive all_to_all_single with per-source split sizes derived from the same norm all-gather (no extra collective — the gather already tells every rank the full winner assignment), assemble each matrix's exact global top-k on its owner, NS, scatter back. It's assemble-exact and round-trip-lossless, and wired through the real dion functions it produces weights matching selection_scope="global" at bf16 tol — no cap, no deferral, same wire volume and same NS cost as capped (the assembled matrix and total rows are identical; only which rows differ). Costs are exactly the two you traded away for static shapes: one small host sync for the split sizes, and keeping the dynamic winner-count in the eager reshard (I held the compiled pre/post-ortho at a per-rank k_override). If the deferral approximation ever shows up in the A/B, this is the version that retires it — the selected set simply is global.

Verdict

In good shape. No correctness blockers; the mechanism does what it claims and the tests back it. Merge-gating (per the PR) is still the convergence A/B — recommend running it on a world_size-divisible shape and noting that a residual capped-vs-global gap includes the fp32/bf16 boundary, not only deferral. Consider the asymmetric variant as the eventual target if you want the approximation gone entirely.

[JohnLangford]

This is John's reviewbot.

Heads-up on a correctness blocker for global_capped and local in real bf16 training — found by running these scopes on an actual 1b run (NorDion2, FSDP2 row-sharded, bf16 momentum). First optimizer step crashes:

File "dion/dion2.py", in dion2_pre_orthogonalize
    m.scatter_(dim=select_dim, index=idx_exp, src=ef_src)
RuntimeError: scatter(): Expected self.dtype to be equal to src.dtype
  self=(192,1536) bfloat16, src=(96,1536) float32

Root cause: the EF write-back computes selected_stacked * ef_factor, and ef_factor derives from the scalar ef_decay (float32). bf16 * fp32 promotes to fp32, so M(bf16).scatter_(src=fp32) raises. It's on the shared dion2_pre_orthogonalize path, so it hits both local and global_capped; global is unaffected (it uses dion2_post_orthogonalize_masked, not this scatter). Every test in test_dion2_selection_scope.py uses fp32 params, so the promotion never happens and the suite stays green — this only fires with bf16 momentum (i.e., real mixed-precision training).

One-line fix — cast the EF write-back to M's dtype (also numerically preferable to a bf16 multiply):

# was: ef_src_list = list((selected_stacked * ef_factor).unbind(dim=0))
ef_src_list = list((selected_stacked * ef_factor).to(dtype).unbind(dim=0))   # dtype = M[0].dtype

Verified: pre-fix crashes as above; post-fix, local/global_capped/global all run finite in bf16 (Dion2 + NorDion2, 2-GPU), and a 1b NorDion2 phinext run trains cleanly through the optimizer with the fix.

Worth adding a bf16 param to the scope tests so this can't regress — happy to open a small PR with the fix + a bf16 test case if useful.

[alint77]

Good catch — fixed in 987acaa with your one-liner (multiply in fp32, single round into M's dtype).

Confirmed the root cause: the winners-only refactor changed the EF factor from the 0-dim ef_decay scalar to a dimensioned [N, k, 1] fp32 tensor, and PyTorch's promotion treats those differently — bf16 * 0-dim fp32 stays bf16 (0-dim loses within the same category), while bf16 * dimensioned fp32 promotes to fp32, which the in-place scatter_ rejects. So this was introduced by 6ce9ac5, not latent from #97. One small correction to the writeup: the local branch keeps ef_factor as the 0-dim scalar, so pre-fix local + bf16 didn't actually promote — the crash was capped-only. Moot either way, the cast covers both and is a no-op on fp32.

Also added the regression guard you suggested: test_scope_bf16_momentum, bf16 params (→ bf16 momentum) over local/global_capped × Dion2/NorDion2. Full suite 25/25 on 2×/4× H100. No need for the separate PR — thanks for the offer and the 1b verification.

[JohnLangford]

This is John's reviewbot.

Re-reviewed at 987acaa. Two updates: the bf16 blocker is fixed, and I ran the actual convergence A/B on hardware.

bf16 correctness blocker — resolved & verified

987acaa applies the exact .to(dtype) cast on the EF write-back and adds test_scope_bf16_momentum (local + global_capped × Dion2/NorDion2, bf16 params). I ran it on 2×H100 — green. That closes the blocker I reported: local/global_capped crashed at the first optimizer step in real bf16 training (selected(bf16) * ef_factor(fp32) → fp32 scatter_ into bf16 M), which every fp32 unit test structurally missed. Good fix, and the bf16-param test is exactly the regression guard this needed.

Everything else stands as reviewed at 6ce9ac5 (winners-only masking is correct — zero fillers pass through NS as exact zeros; single-owner accumulation via thresh; negative-pad guard; top-k·world budget documented; thorough scope tests). The diff since is only the fix + test — no new correctness/design/perf/compat/security concerns.

The convergence A/B (the merge gate) — real numbers

I ran all scopes end-to-end on this PR's code: 1b NorDion2, phinext, 10B tokens, 8×H100.

scope	train/loss	core_metric (DCLM)
global (baseline)	1.30914	0.2266
local	1.30927	0.2264
global_capped	1.30714	0.2327
global_capped + 1.5× slack*	1.30976	0.2122

*slack = an experimental variant I added (more comm slots so fewer winners defer), not part of this PR — included only as a near-exact reference point.

All four land within 0.0026 nat on train loss, and the ordering is non-monotonic — tight-capped came out best and the near-exact slack variant worst, which is mechanistically impossible for a selection-quality effect (slack defers strictly fewer winners than tight cap). So this spread is numerical-nondeterminism noise, not signal. At this config's fraction=0.5, the selection scope is free: global_capped matches global (supports the PR's no-penalty claim) — but so does plain local.

That's the crux for merge. global_capped's value proposition is "global-quality selection at local comm cost," which only pays when local actually converges worse. At fraction 0.5 it doesn't, so there's no gap for global_capped to close here. PR #98's ~0.09-nat local penalty was at a smaller fraction, so the A/B that actually decides this PR is fraction=0.25 (where per-shard vs global selection diverge), ideally ≥2 seeds to establish the noise floor. If local doesn't lose even at 0.25 for the model you ship, the mechanism is moot and this can be closed on that evidence; if it does, global_capped is justified. I can run that sweep.

Elegant alternative (recap, now with data)

The variable-length all_to_all_single reshard I floated (exact global top-k at local comm cost, no deferral, split sizes derived from the norms all-gather that's already computed) is validated end-to-end and remains the cleaner endpoint — but at fraction 0.5 it lands in the same tie, so it's only worth building if the 0.25 sweep shows a real local penalty worth eliminating.

[alint77]

I am working on a new solution where instead of naively adding a "slack", we flatten the topk selections in each megabatch comm packet which includes more than one matrix and pack them together, similar to how varlen sequence packing works in attention. Then on top of that, we add some extra capacity (or "slack") to the whole buffer instead of adding slack to each. this should reduce the number of dropped/overflowed rows significantly while preserving the comm volume of local, static buffer sizes and thus async cpu dispatches. @JohnLangford

[alint77]

Pushed the packing idea I teased above: aaa2cf6 + 569b395 replace the fixed-slot global_capped transport with a pooled packed one (the old behavior is gone, not kept as a mode). Same scope name, same collectives, same sync-free property — much less deferral.

The idea, simply

The old design gave every matrix its own k slots in each a2a chunk. But a chunk from rank r to dest d already carries per_rank matrices' segments — and one matrix's overflow usually coincides with another's spare room. So: remove the walls, pack winners back-to-back into one shared budget, and write the per-matrix counts into a tiny header the receiver parses.

rank r -> dest d chunk (per_rank = 2 matrices, k = 2 per matrix)

before (walls):   [ m0: W W | m1: W . ]
  m0 owns 3 winners -> one DEFERRED, while m1 ships a zero slot next door

after (pooled):   [ hdr | m0: W W W | m1: W | spare ]
  one shared budget; all 4 winners ship; the int32 count header (bit-cast
  into the first bf16 row) tells the receiver where each segment ends

Deferral now requires a rank's winners summed over the pooled matrices to overflow, so it shrinks from per-matrix fluctuation to pooled fluctuation (relative sigma drops by sqrt(per_rank)). This is your slack experiment taken one step further: pool first, then put the slack on the shared buffer — at capacity_factor=1.0 (no extra comm at all) the expected deferred fraction is already ~1%, and the auto default (1 + 2*sqrt((1-1/world)/(per_rank*k)), ~+6% comm at our shapes) takes it to ~0.02%. Everything stays static-shaped and host-sync-free — this deliberately stops one step short of your variable-split a2a: near-exact with zero reads, vs exact with one.

Implementation

• Exact-count selection: winner set is now exactly top-ceil(fraction*global) via stable argsort over the gathered norms (deterministic tie-break). This also kills the round-1 top-(k*world) budget nuance and makes sum(winners) <= NS buffer a hard invariant.
• Pack (dion2_capped_pack, compiled, static shapes): per dest, pool the winners of its matrices; on overflow keep the top-budget by scale-normalized priority norm/thresh (raw norms aren't comparable across matrices), stable order on ties. Outputs payload + counts + a slot->row index kept for the reverse path.
• Transport: counts bit-cast (int32.view(bf16), exact round-trip) into the chunk's first row; fixed-size a2a; receiver parses the sender's header instead of recomputing counts from the norms — no bit-identical-plan invariant across ranks, plus a clip so a divergent sender can't overrun the NS buffer.
• Assemble/return: static zero-padded [k*world, cols] NS input (zeros are NS-transparent), inverse-mapped repack, reverse a2a, scatter home to full-size shards that are zero except at applied winners — which feed the existing global-scope masked post. So EF decay touches exactly the applied rows; deferred winners keep accumulating. dion2_pre_orthogonalize no longer has a capped mode at all.
• Routing: groups where packing can't pay (budget+1 reaches the shard, header doesn't fit, col-sharded/non-2D) silently take the exact global path.
• CAPPED_STATS: rank-local deferred/winner row counters, reset per step, accumulated across shape groups — so the fraction=0.25 sweep can report the measured deferral rate alongside loss.

Tests: 28/28 on 2x/4x H100 — including a bit-exact pack->a2a->NS->unpack roundtrip (pooled deferral, blackout matrix, uneven counts) and a cross-matrix-skew test proving pooling ships what the walls couldn't while a 100x-smaller-scale matrix's winner survives eviction.

Your fraction=0.25 (>=2 seeds) sweep remains the merge gate and I'd suggest running it on this code — happy to take you up on the offer. PR description updated to match; the results table there is from the previous transport (comm volume is unchanged modulo the ~0.1% header; fresh trace pending).

[alint77]

@JohnLangford please run an extended training run on all 3 algorithms so we can compare the numerics. global_capped should be very close to global now

[JohnLangford]

This is John's reviewbot.

Re-reviewed at 569b395. This is a substantial redesign (pooled packed transport), and it's a good one — it independently lands the three things my earlier review + hardware experiments pointed at:

• Zero host syncs via the int32 count header bit-cast into each chunk's first row. I checked the failure mode I'd worry about — NaN canonicalization of the reinterpreted bytes — and it's clean: the header is built and parsed in the eager orchestrator (torch.cat → all_to_all → slice+view(int32)), never passing through a float op or NS, so the byte-copy transport round-trips exactly. This is strictly better than the kmax.item() host-sync approach I'd prototyped, which my timing showed erased the comm saving.
• Break-even routing (budget+1 >= per_rank*padded_local → route to global). This directly fixes the regime my same-pod timing flagged: at 1b the packed path's bookkeeping cost more than the comm it saved, so falling back to global there is exactly right.
• capacity_factor auto = 1 + 2·√((1-1/world)/(per_rank·k)) — a principled 2σ version of the fixed "slack" I'd hacked in, and pooling across the megabatch's matrices lets one matrix's overflow borrow another's spare room, so deferral is rarer than a per-matrix cap.

Correctness spot-checks all hold: exact-count winner set via stable argsort (ties can't inflate past budget — my round-1 concern, now closed and the count is exact ceil(fraction·global), matching global); kw ≥ k_total so the NS buffer always fits a matrix's winners; the receiver clip is a sound guard for the divergent-sender case; sentinel/trash rows are correctly discarded on unpack. Test coverage is thorough — test_capped_header_bitcast_roundtrip, test_capped_pooled_slots_absorb_cross_matrix_skew, deferral, bf16 momentum, uneven/empty shards. (Ran the suite: 28/28 pass on 4×H100.)

Hardware results I can contribute (caveat: on the previous fixed-slot design)

I ran the convergence A/B and a same-pod timing sweep before this redesign, so treat them as a lower bound on what this version does:

• Convergence (1b NorDion2 phinext, 10B tokens): global, local, global_capped all landed within 0.0026 nat (noise; non-monotonic ordering) — at fraction=0.5 the scope doesn't move the loss. core_metric tied too. This design defers less than the one I tested, so the "no penalty vs global" claim carries over — but note it also means local matched global here, so there was no gap for capped to close at this fraction.
• Same-pod timing (8×H100, GNS+triton, opt.step): at 14b, local 0.81×, old-capped 0.90×, global 1.00×; at 1b the old capped was ~1.2× (the overhead this PR now routes around). So the packed path's win is real only in the comm-bound regime (14b), which is exactly the "opt in when comm-bound at large scale" framing.

The one thing still open (the merge gate)

Both my results are at fraction=0.5, where every scope converges identically — so they can't actually discriminate global_capped from local. The experiment that decides this PR is fraction=0.25 (where #98 saw local lose ~0.09 nat): if local underperforms global there, global_capped earns its keep as the "recover global quality at local cost" option; if it doesn't, the whole scope is moot for this model. I have that sweep ready to run.

Minor

CAPPED_STATS is module-global mutable state accumulating device tensors across the step — fine as a diagnostic, but it isn't safe if two optimizer instances step concurrently, and it holds device-tensor refs until cleared. Consider stashing it on the optimizer instance.

Backward-compat/security: clean — new capacity_factor defaults to auto, local/global paths untouched, no I/O/untrusted surface.

[JohnLangford]

This is John's reviewbot.

Same-pod timing of the selection scopes, since the throughput claim is the crux of this PR. Two tables: A is the current pooled design (this PR head), B is the earlier fixed-slot design for reference.

Definitions / method

• What's timed: NorDion2.step() only — no forward/backward, no data. Pure optimizer cost.
• Same pod, back-to-back: all variants timed in one process on the same 8×H100 node, so there's no pod-to-pod variance. Median of 10 warm steps (10 warmup).
• Newton-Schulz: production path — use_gram_newton_schulz=True, use_triton=True, triton_post_ortho=True.
• Params: FSDP2 row-sharded (Shard(0)) DTensors for the real matrix inventories — q/k/v/o + gate/up/down across all layers. 1b = hidden 1536 / inter 6656 / 24 layers; 14b = hidden 5120 / inter 17920 / 40 layers.
• Grads: random each step. capacity_factor = auto (default). Absolute ms, global normalized to 1.00 in parens.

Table A — pooled global_capped (this PR head)

fraction	scale	global	local	capped
0.25	1b	25.15 (1.00)	26.79 (1.07)	28.61 (1.14)
0.25	14b	68.20 (1.00)	51.60 (0.76)	67.28 (0.99)
0.50	1b	26.13 (1.00)	27.85 (1.07)	30.11 (1.15)
0.50	14b	102.70 (1.00)	80.40 (0.78)	104.33 (1.02)

Table B — earlier fixed-slot design (pre-redesign), fraction 0.5

capped here is the fixed-k-slot mechanism from before the pooled rewrite; slack1.5 and exact are my experimental variants (not in this PR), included to show the frontier.

scale	global	local	capped (fixed-slot)	+slack1.5	exact
1b	26.29 (1.00)	27.58 (1.05)	31.88 (1.21)	31.89 (1.21)	31.82 (1.21)
14b	100.27 (1.00)	80.98 (0.81)	90.67 (0.90)	120.07 (1.20)	101.05 (1.01)

What stands out

1. local is the speed winner at 14b — 0.76–0.78× (pooled) / 0.81× (fixed-slot) — at every fraction. It ships the same reduced comm as capped with none of the packing machinery.
2. global_capped ≈ global at 14b (0.99–1.02×, pooled), and slower at 1b (1.14–1.15×). The pack/assemble/repack + norms all-gather overhead cancels the comm it saves — so on isolated opt.step time it doesn't get below global here.
3. The pooled redesign is ~15% slower at 14b than the old fixed-slot capped (104.3 vs 90.7 ms; 1.02× vs 0.90×). The exactness + zero-syncs + routing are real wins, but the argsort/scatter/header/pooling cost more than the old uniform-slot path — which is why I don't reproduce the reported 0.94×.
4. Likely cause of that gap: my 14b is 40 layers vs the reported 8-layer config. The pooled pack's argsort/scatter scales with the megabatch size (per_rank × matrices), so overhead grows with depth. capped's speed advantage looks config-dependent — present at shallow configs, absent at production depth.

Caveats: random gradients (real winner distributions could change the pack cost), and this is opt.step in isolation (raw comm volume still favors capped, which can matter under network contention even when isolated step time doesn't). Worth confirming on a real gradient step at depth before leaning on the throughput claim.

[alint77]

This is John's reviewbot.

Same-pod timing of the selection scopes, since the throughput claim is the crux of this PR. Two tables: A is the current pooled design (this PR head), B is the earlier fixed-slot design for reference.

Definitions / method

• What's timed: NorDion2.step() only — no forward/backward, no data. Pure optimizer cost.
• Same pod, back-to-back: all variants timed in one process on the same 8×H100 node, so there's no pod-to-pod variance. Median of 10 warm steps (10 warmup).
• Newton-Schulz: production path — use_gram_newton_schulz=True, use_triton=True, triton_post_ortho=True.
• Params: FSDP2 row-sharded (Shard(0)) DTensors for the real matrix inventories — q/k/v/o + gate/up/down across all layers. 1b = hidden 1536 / inter 6656 / 24 layers; 14b = hidden 5120 / inter 17920 / 40 layers.
• Grads: random each step. capacity_factor = auto (default). Absolute ms, global normalized to 1.00 in parens.

Table A — pooled global_capped (this PR head)

fraction scale global local capped
0.25 1b 25.15 (1.00) 26.79 (1.07) 28.61 (1.14)
0.25 14b 68.20 (1.00) 51.60 (0.76) 67.28 (0.99)
0.50 1b 26.13 (1.00) 27.85 (1.07) 30.11 (1.15)
0.50 14b 102.70 (1.00) 80.40 (0.78) 104.33 (1.02)

Table B — earlier fixed-slot design (pre-redesign), fraction 0.5

capped here is the fixed-k-slot mechanism from before the pooled rewrite; slack1.5 and exact are my experimental variants (not in this PR), included to show the frontier.

scale global local capped (fixed-slot) +slack1.5 exact
1b 26.29 (1.00) 27.58 (1.05) 31.88 (1.21) 31.89 (1.21) 31.82 (1.21)
14b 100.27 (1.00) 80.98 (0.81) 90.67 (0.90) 120.07 (1.20) 101.05 (1.01)

What stands out

1. local is the speed winner at 14b — 0.76–0.78× (pooled) / 0.81× (fixed-slot) — at every fraction. It ships the same reduced comm as capped with none of the packing machinery.
2. global_capped ≈ global at 14b (0.99–1.02×, pooled), and slower at 1b (1.14–1.15×). The pack/assemble/repack + norms all-gather overhead cancels the comm it saves — so on isolated opt.step time it doesn't get below global here.
3. The pooled redesign is ~15% slower at 14b than the old fixed-slot capped (104.3 vs 90.7 ms; 1.02× vs 0.90×). The exactness + zero-syncs + routing are real wins, but the argsort/scatter/header/pooling cost more than the old uniform-slot path — which is why I don't reproduce the reported 0.94×.
4. Likely cause of that gap: my 14b is 40 layers vs the reported 8-layer config. The pooled pack's argsort/scatter scales with the megabatch size (per_rank × matrices), so overhead grows with depth. capped's speed advantage looks config-dependent — present at shallow configs, absent at production depth.

Caveats: random gradients (real winner distributions could change the pack cost), and this is opt.step in isolation (raw comm volume still favors capped, which can matter under network contention even when isolated step time doesn't). Worth confirming on a real gradient step at depth before leaning on the throughput claim.

which script was used for testing? it mentions "no forward/backward, no data" which is pretty much the same as having a host sync at the start of opt.step... Please make sure the results are gathered using the new profile_training_step.py in Noah's fork and not "benchmark_optimizer.py".

My testing shows much better perfromance for the new pooled implementation. You can see the result on the next comment:

[alint77]

These are from real training steps (fwd/bwd running, real gradients, torch-profiler traces of the live run), not an isolated opt.step() loop. All 3 scopes back-to-back on one 4×H100 node at 569b395: hidden 5120 / 8 layers / fraction=0.25, Gram-NS, NorDion2, means over 5 profiled steps. Logging every step so the CPU spans are true host cost, not command-buffer wait.

Per optimizer step, rank 0 (device-time kernel accounting from the trace)

metric	local	global	global_capped (pooled)
training step wall	487.9 ms	477.9 ms	474.8 ms
opt.step GPU window	31.3 ms	35.5 ms	28.1 ms
— NS GEMM (scope-invariant)	17.9 ms	17.7 ms	17.6 ms
— non-GEMM compute	13.2 ms	16.5 ms	9.9 ms
— a2a	4.1 ms	9.5 ms	2.9 ms
a2a payload	655 MB	2621 MB	673 MB
opt.step CPU span	10.6 ms	10.7 ms	13.3 ms
kernel launches	60	72	117

NS runs at 97–99% of bf16 peak in every scope (identical selected shapes ⇒ identical 17.24 TF/step), so it's a fixed floor; the entire scope difference lives in the non-GEMM slice — and pooled wins it against both global and local.

Why pooled beats local on GPU here

1. Local pays for selection in full-momentum HBM traffic; pooled decides on kilobytes. Local's pre-path materializes torch.stack(M) as a real buffer (the row gather and EF scatter_ consume it), then runs a full-momentum norm + topk chain over it — its fused momentum kernel moves ~2× the bytes of a plain M += G (3.9 vs 1.7 ms) and the norm chain adds ~2.0 vs 0.7 ms. Pooled's pre_accumulate_norms does one lean M += G whose stacked norm is fused into the reduction (never materialized); everything downstream operates on [N, local] fp32 norms and budget-sized index tensors. Net −4.4 ms in the front-end; pooled gives ~0.9 ms back in the masked post (winners arrive as nonzero rows, not indices), and payload construction is a wash (~5 ms both).

2. The a2a gap (2.9 vs 4.1 ms at the same bytes) is overlap, not transport. Per-peer messages are near-identical (2880 vs 2974 rows × 5120 cols, bf16). But a SendRecv kernel still reads/writes HBM, and local's a2a kernels spend 92% of their lifetime concurrent with HBM-saturating full-momentum kernels (mean 1.02 ms/kernel), while 18/20 of pooled's overlap only budget-sized bookkeeping (mean 0.66 ms). The control is in the data: the 2/20 pooled a2a kernels that did land on heavy kernels slowed to 1.33 ms — same wire, same slowdown.

3. Pooled's one real cost is launches, and it's host-side only. ~150 extra tiny index-arithmetic kernels ≈ 0.1 ms of GPU time, but 117 launches vs 60 and a 13.3 vs 10.6 ms CPU span.

Reconciling with the isolated-timing tables above

Point 3 is, I suspect, most of the disagreement. In a live training step the host runs far ahead of the device (CPU span 13 ms vs a 475 ms step), so launch overhead is fully hidden and only device time counts — which favors pooled. A sync-per-step opt.step() loop puts every launch on the critical path, and that penalty (a) hits pooled ~2× harder and (b) grows with depth, since the launch count scales with the matrix inventory. That's consistent with pooled looking worst in the isolated benchmark exactly where the megabatch is largest.

The depth point itself is fair and this table can't refute it: at 8L/33 matrices the pack's sort/scan is 0.33 ms and the index arithmetic ~0.1 ms of device time; both scale ~linearly with matrix count, but so do local's per-matrix norm/gather/scatter chains, so on device time the ranking should hold at 40L — that's a prediction, not a measurement. Next run on our side is the same traced comparison at production depth to check it.

[NoahAmsel]

not quite following the details here. i know that you’re trying to find a way to do global selection without sending assembling the full matrices, but not sure about the difference between the two versions of global_capped mode.

Winners that overflow the pooled budget are deferred

so basically this means that there is still an upper limit on how many winners can be selected from each shard? before, in the local version, the limit was just (fraction * shard_size). now you’ve added some extra slack. I think it’s clever to bundle together rows from different layers, so that if one layer needs more rows than expected it will balance another row that needs fewer than expected. but maybe this would be simpler:

select c*fraction*shard_size from each shard. do the all-to-all. then select the top fraction rows from the result. if c = 1, this is just the local version we had before. if c is 1/fraction, then it's the global version we had before. if c is somewhere between, (say, c=1.5 or c=2*sigma) we get a good balance

This avoids the complexity of bundling across layers, uneven sizes, and variable length reshards (whatever that is), and parsing a header, and might be good enough?

[alint77]

not quite following the details here. i know that you’re trying to find a way to do global selection without sending assembling the full matrices, but not sure about the difference between the two versions of global_capped mode.

Winners that overflow the pooled budget are deferred

so basically this means that there is still an upper limit on how many winners can be selected from each shard? before, in the local version, the limit was just (fraction * shard_size). now you’ve added some extra slack. I think it’s clever to bundle together rows from different layers, so that if one layer needs more rows than expected it will balance another row that needs fewer than expected. but maybe this would be simpler:

select c*fraction*shard_size from each shard. do the all-to-all. then select the top fraction rows from the result. if c = 1, this is just the local version we had before. if c is 1/fraction, then it's the global version we had before. if c is somewhere between, (say, c=1.5 or c=2*sigma) we get a good balance

This avoids the complexity of bundling across layers, uneven sizes, and variable length reshards (whatever that is), and parsing a header, and might be good enough?

I think you're forgetting that in the global_capped mode we allgather the row norms and we know exactly how many rows have to be selected on each rank. this allgather is pretty much free in terms of performance becaue the payload is just a vector and gets done in nanoseconds.

Now knowing that, I'm struggling to see the benefit in sending extra rows over the network just to be discarded after the c * frac * shard_size matrix is constructed on each rank... how is this different then not allgathering row norms and naively adding some extra "slack" to the buffer, which is something both John and I have tested already?

[NoahAmsel]

I see. You're right, it is just the same as naively adding slack to the buffer. If that has been tried already then disregard.

But if I understand correctly, this still fixing a size for the packed matrix sent by each device. That creates the potential for overflows that have to be handled if a given shard selects more rows than we expected. Since "we know exactly how many rows have to be selected on each rank" after the first round (the row-norm selection step) we can use that to determine how big the chunks have to be to avoid overflow. Then every device can use that same maximum size (padding with zeros). Everything fits, and we don't have overflows. If we aggregate across all matrices in the megabatch, presumably this max size wouldn't be much bigger than whatever heuristic size you're using now?

this is just a suggestion, not sure whether it's better or worse really, I'll need your expertise. but it's simpler and avoids overflows. handling overflows by dropping some of the selected rows again ties the mathematics of the optimizer step to the specific number of GPUs we are using (as well as to the heuristic we are using to pick the chunk size, and it links each layer to every other layer...) If we can avoid this complexity without hurting performance we should, right?

[NoahAmsel]

A larger point: we've seen that the MSR cluster can give very different results from Ali's because the communication cost is so low. I guess this means we need to include some timing results from a different cluster in the paper, both for completeness and because it may be the only way to show the true benefits of our design

[alint77]

I see. You're right, it is just the same as naively adding slack to the buffer. If that has been tried already then disregard.

But if I understand correctly, this still fixing a size for the packed matrix sent by each device. That creates the potential for overflows that have to be handled if a given shard selects more rows than we expected. Since "we know exactly how many rows have to be selected on each rank" after the first round (the row-norm selection step) we can use that to determine how big the chunks have to be to avoid overflow. Then every device can use that same maximum size (padding with zeros). Everything fits, and we don't have overflows. If we aggregate across all matrices in the megabatch, presumably this max size wouldn't be much bigger than whatever heuristic size you're using now?

what you're describing here would require knowing the size of the a2a comms buffer at the time of the kernel launch = cpu host sync. this is what we're trying to avoid.

With the new global_capped packing behavior, we set the size of the buffers (the size of the slack) once at the contruction time which remains static throughout training so no host sync required at each step. also the size of the slack is chosen based on the (2 * std) of the binomial distribution to statistically minimize the chance of overflow (~0.02% of rows). (the choice of (2 * std) was arbitrary, could be 3.5 or 4 or whatever, but will cause a large percentage of slack size at smaller topk sizes)

[alint77]

A larger point: we've seen that the MSR cluster can give very different results from Ali's because the communication cost is so low. I guess this means we need to include some timing results from a different cluster in the paper, both for completeness and because it may be the only way to show the true benefits of our design

I can run multinode tests tonight when the cluster is less busy. will add them here. we should see the benefit of new global_capped more clearly in that regime due to lower comms volume.

[JohnLangford]

This is John's reviewbot.

Reviewed e581267. The fused-Triton masked post (for global/global_capped) is a sound optimization — gated on triton_post_ortho, single-rounding-on-store (matches the _fused numerics, arguably better than the eager masked post's double rounding), 301 lines of tests. It shaves the masked-post cost for global and global_capped equally, which is the right lever given what the throughput data below shows.

The throughput premise, measured end-to-end on real architectures

The PR's framing is "global selection at local comm cost" — i.e. a throughput win. I profiled the full training step (forward + backward + optimizer.step, real gradients, production GNS, 8×H100), not just opt.step in isolation:

measurement (fraction 0.5)	local vs global (full step)	global_capped vs global (full step)
real mixformer-1b (phinext data, GNS)	0.996×	+1.0%
GPT-14b profiler (GNS)	~1.00×	+2%
isolated-opt microbench (mixformer-shaped, random grads)	0.76× at 14b	≈ global

On a real training step the selection scope barely moves throughput. The reasons:

• The optimizer is only 14–24% of the GPU step (backward alone is ~60% on mixformer-1b). So even a large opt.step delta dilutes 4–5× end-to-end.
• With GNS on real matrices the optimizer is NS-bound, not comm-bound. The isolated microbench's 0.76× local win came from a comm-bound corner (opt-only, comm-heavy synthetic inventory); on a real GNS step the comm saving mostly vanishes, and global_capped's pack/assemble overhead shows up instead as a ~1–2% net loss on the full step (its opt.step is a consistent +10% vs global).

Implication for the merge decision

The comm saving is real but doesn't convert to end-to-end speedup on the models this targets — so global_capped's case can't rest on throughput. It rests entirely on selection quality: exact global top-k vs local's sharding-dependent approximation. And that only matters if local actually converges worse — which, at the production fraction=0.5, it did not (all scopes tied within noise in the 10B-token A/B). So as it stands, on real phinext-scale training:

• vs global: same quality, ~1–2% slower → global dominates.
• vs local: same measured convergence, ~same speed, far more complexity → local dominates.

The one experiment that could still justify the scope is the fraction=0.25 convergence A/B (where #98 saw local lose ~0.09 nat). If local underperforms there, global_capped becomes the "recover global quality without global's full-shard comm" option — but note the comm advantage is ~free only in isolation, so even then the case is quality, not speed.

Design/correctness of the mechanism itself remains solid (verified earlier: exact-count stable-sort, zero host syncs via the in-band header, break-even routing, 28/28 tests + now the triton-post tests). This is a well-built implementation of a mechanism whose end-to-end payoff, on the target models, is currently ~zero-to-negative — worth confirming with the 0.25 convergence run before taking on the maintenance surface. (Definitive pooled-capped numbers on the real mixformer-1b are finishing now; I'll append if they shift anything.)

[JohnLangford]

This is John's reviewbot.

Re-ran the throughput measurement using the CUDA-event method from NoahAmsel/dion's profile_training_step.py — events bracket the full step and the optimizer with no intra-step torch.cuda.synchronize(), read once after a single sync, so the in-step GPU pipeline is not serialized. This is on a real (short) phitrain job: same pod, the 3 scopes back-to-back, real mixformer + phinext data + production GNS, 8×B200.

mixformer-1b, mbs=7 (19 samples/scope)

scope	gpu_step	gpu_opt	opt_share	step vs global	opt vs global
global	466.3 ms	65.1 ms	13.9%	1.000	1.000
local	468.5 ms	64.7 ms	13.8%	1.005	0.994
global_capped	468.2 ms	70.7 ms	15.1%	1.004	1.087

mixformer-14b, mbs=2 (8-GPU FSDP) (48 samples/scope)

scope	gpu_step	gpu_opt	opt_share	step vs global	opt vs global
global	964.8 ms	123.4 ms	12.8%	1.000	1.000
local	955.8 ms	113.7 ms	11.9%	0.991	0.921
global_capped	969.4 ms	131.8 ms	13.6%	1.005	1.068

This matches the earlier full-step profiling, and agrees with a per-phase-synchronize cross-check (that method reads slightly larger deltas — the syncs prevent overlap and over-charge the pack/assemble step). Takeaways:

• global_capped costs ~7–9% more optimizer time, but the optimizer is only ~12–14% of a production step, so it dilutes to ~0.4–0.5% on the full training step at both scales.
• local ties global at 1b and is ~0.9% faster full-step at 14b (opt 0.92×): its comm saving is real and grows with matrix size, but stays sub-1% end-to-end.
• The optimizer-time ratios are batch-independent — mbs only sets the share. (mbs=1 inflates opt_share to ~20% via a tiny fwd/bwd and magnifies every scope delta, so it overstates the effect; mbs=2/7 above are representative.)

So on the models this targets, selection scope is a ~0.5% throughput effect. As before, that means global_capped's case rests on selection quality, not speed. I've now launched the fraction=0.25 convergence A/B (global / local / global_capped, 1b→10B tokens phinext, matched to the prior fraction=0.5 setup) — the low-fraction regime where #98 saw local lose ~0.09 nat. I'll post the convergence deltas when it lands; that's the experiment that decides whether the scope earns its maintenance surface.

[JohnLangford]

This is John's reviewbot.

fraction=0.25 convergence A/B — the quality question, settled. The remaining case for global_capped was selection quality in the low-fraction regime: recovering global's exact top-k where local's per-shard approximation was reported to lose ~0.09 nat (#98). I ran all three scopes at fraction=0.25, 1b → 10B tokens phinext (matching the earlier fraction=0.5 setup, dion pinned at the PR's pooled global_capped), then ran lm-eval on the final checkpoints.

Matched at the final step (21789 ≈ 10B tokens):

scope	train CE	CORE	BPB climbmix	BPB web	BPB code
global	1.4717	0.2262	0.8486	0.9963	0.4744
local	1.4705	0.2209	0.8480	0.9959	0.4736
global_capped	1.4709	0.2201	0.8481	0.9961	0.4742

Δ vs global: local ΔCE −0.0012, Δbpb ≤ 0.0008; global_capped ΔCE −0.0009, Δbpb ≤ 0.0004. Every delta is inside its noise band (BPB SE 0.003–0.016; CORE 1b floor ±0.016; CORE deltas even flip sign vs CE/BPB — noise). For scale: #98's ~0.09 nat ≈ 0.13 bits, and the BPB gap here is ~0.0006 bits — ~200× smaller.

local's ~0.09-nat penalty does not reproduce on phinext mixformer-1b. At fraction=0.25 all three scopes converge identically and give indistinguishable downstream quality — the same tie seen at fraction=0.5.

Why the #98 gap doesn't appear here: local's error is the gap between per-shard top-k and global top-k, which depends on how each matrix is row-sharded — it grows with shard count and with how much per-row magnitude mass concentrates. On phinext's 8-way FSDP over these matrix shapes, per-shard selection stays close to global, so the gap collapses into noise. It may still bite at much higher shard counts or different architectures — that's the boundary of this result, not a refutation of the mechanism.

Net for this PR. Both of global_capped's justifications now come up empty on the target models: throughput (~0.5% full-step, measured above) and low-fraction quality (this tie). Versus global it's ~same speed, same quality, more complexity; versus local it's the same — local matches quality at both fractions with a simpler mechanism. The implementation itself is solid (verified earlier: exact-count stable-sort, zero host syncs, break-even routing, full test coverage). But on phinext-scale training the measured payoff is ~zero on both axes, so I'd lean against merging it for this use unless there's a concrete target regime where the exact-global-at-local-comm tradeoff actually pays — a much higher shard count, or a workload where the optimizer is genuinely comm-bound rather than NS-bound. If you think that regime exists, I'm happy to run the higher-shard-count corner to see whether the scope earns its keep there.

[alint77]

Re: "unless there's a concrete target regime where the exact-global-at-local-comm tradeoff actually pays — a workload where the optimizer is genuinely comm-bound rather than NS-bound" — here's that regime, measured. Same benchmark as my earlier single-node tables, but across 2 nodes (2×4 H100, ws=8, FSDP row-sharded), where the optimizer a2a rides the inter-node link instead of NVLink.

Setup: 5120-hidden/8-layer, fraction=0.25, NorDion2, dion at this PR's head; both scopes run sequentially on the same node pair; torch profiler steps 58–64; step walls averaged over 62 steady-state steps (~130k tok/step). Measured effective link bandwidth during the opt a2a: ~19 GB/s.

per step	global	global_capped	ratio
opt a2a payload	1363 MB	361 MB	3.78×
opt a2a time	54.1 ms	14.3 ms	3.78×
a2a effective bw	18.9 GB/s	18.9 GB/s	1.00×
gemm (opt window)	13.3 ms	10.4 ms
opt-stream window	60.2 ms	22.5 ms
opt.step GPU time	21.0 ms	15.7 ms
opt.step GPU idle time	65.2%	30.3%
min CPU span (opt.step)	11.4 ms	13.5 ms
step wall (mean, n=62)	584.5 ms	546.4 ms	−38 ms (−6.6%)
H100 MFU	43.6%	46.6%	+3.0 pp

Three observations:

1. The win is pure bytes, and it's the predicted amount. Both scopes drive the link at an identical 18.9 GB/s — the transport is link-bound, so time scales 1:1 with payload. The 3.78× byte ratio matches theory: fraction 0.25 + 2σ capacity slack + the norms-only all-gather ≈ 26% of global's full-row traffic.

2. The comm saving converts 1:1 into step wall. Trace comm delta (54.1 − 14.4 = 39.7 ms) ≈ wall delta (38 ms): off-node, the optimizer a2a sits on the critical path. During global's opt window the compute stream is idle 65% of the time, parked behind the a2a.

3. This is the same code that ties global on-node. Single-node, this exact benchmark showed a 1–3 ms margin (NVLink makes 1.4 GB/step free). The crossover is purely interconnect: global's payload grows with world size while capped stays fraction-bound, so the margin widens with scale and with slower fabrics.

So the concrete regime: multinode training where the per-rank inter-node bandwidth is O(20 GB/s) — i.e. most IB/EFA clusters once you leave the NVLink island — at any meaningful fraction < 1. There, global_capped is global's exact selection at 26% of its bytes, worth ~7% end-to-end in this run, growing with node count. Happy to pair this with your higher-shard-count corner if you want both axes covered.

@NoahAmsel @JohnLangford

[alint77]

Node-count scaling of the 2-node result above — same benchmark (5120-hidden/8L, NorDion2, fraction=0.25, FSDP row-sharded), each column both scopes back-to-back on the same node set. Cells: global / global_capped.

per step (rank 0)	1 node (ws=4)	2 nodes (ws=8)	4 nodes (ws=16)
a2a payload MB	2621 / 673	1363 / 361	734 / 205
a2a time ms	9.5 / 2.9	54.1 / 14.3	58.7 / 16.2
link throughput GB/s	276 / 232	25.2 / 25.2	12.5 / 12.7
opt.step GPU time ms	34.2 / 27.5	21.0 / 15.7	10.9 / 9.1
opt.step GPU idle time	3.6% / 2.2%	65.2% / 30.3%	81.8% / 51.5%
opt.step wall ms	35.5 / 28.1	60.2 / 22.5	60.2 / 18.8
pooled opt.step wall delta	−7.4 ms (−21%)	−37.7 ms (−63%)	−41.4 ms (−69%)
step wall ms	477.9 / 474.8	584.5 / 546.4	580.4 / 531.2
pooled step-wall delta	−3.1 ms (−0.6%)	−38.1 ms (−6.6%)	−49.2 ms (−8.5%)
H100 MFU	53.5% / 53.7%	43.6% / 46.6%	43.9% / 48.0%

Short version: NS scales out with ranks (opt GPU time 34→11 ms); the a2a doesn't — per-rank payload halves at each doubling but per-rank link bandwidth halves too, so global's optimizer wall stays pinned at ~60 ms while global_capped's keeps shrinking (28→19 ms). The margin is monotone in node count: −0.6% → −6.6% → −8.5% step wall. The more nodes, the more the scope pays for itself.

@NoahAmsel , This is the data that we want for the paper.

[alint77]

Completing the ws=16 matrix: local vs global_capped, same benchmark, both back-to-back on the same 4-node set.

per step (rank 0)	local	global_capped
a2a payload MB	183.5	205.1 (+12%)
a2a time ms	7.1	7.9
raw link throughput GB/s	25.9	25.8
opt.step GPU time ms	8.8	8.5
opt.step wall ms	11.9	12.3
opt.step GPU idle time	25.4%	30.6%
min CPU span ms	11.4	13.7
step wall ms (mean, n=62)	502.2	504.0
H100 MFU	50.7%	50.5%

local ≈ global_capped at 4 nodes too: capped's +12% bytes (capacity slack + norms all-gather) and +2 ms host cost are throughput-neutral, so capped-vs-local stays a selection-exactness question, not a speed one — global remains the only scope that pays a comm penalty at scale.

One methodological note: the identical capped run moved the same 205 MB at 12.7 GB/s on this morning's node set vs 25.8 GB/s on this one — inter-node fabric varies ~2× between sets, so all our pairs are same-set and the global deltas above should be read as fabric-dependent in absolute terms (the ~3.6× payload/time ratio is what's invariant).

[JohnLangford]

It looks like local is as good as global or global_capped empirically on the 1b/10b runs, so my inclination is to go with a local default with global as an argument. Is global_capped worth keeping as a default?

[alint77]

It looks like local is as good as global or global_capped empirically on the 1b/10b runs, so my inclination is to go with a local default with global as an argument. Is global_capped worth keeping as a default?

yes I'm not seeing much of a difference in throughput between local and global_capped even in the slow-internode comms regime, which is great news. global causes massive slowdown in my setup with multiple nodes.

[alint77]

@JohnLangford I think the numerics should be tested a bit more extensively across various hidden_dim sizes, I would suggets running a 100B test for each hdim=[2048, 4096, 6144, 8192], my inkling is that on smaller sizes, local might cause worse numerics.

Also regarding the statistical significance of the results, I think following modded-nanogpt's track3 rules would be a good idea, to make sure we're dealing with real signal and not run-to-run noise