Dion2/NorDion2 FSDP2 padding regression

[alint77]

Problem

Dion's whole point is a cheaper orthogonalization step: select only a
fraction (default 0.25) of each matrix's rows and run Newton-Schulz on that
submatrix. Under FSDP2 row-sharding this saving is silently erased.

In megabatch_orthogonalize_async (dion/megabatch_base.py), every rank's
selected rows are padded back up to the full per-rank shard before the
all-to-all:

padded_local_size = ceil(global_comm_dim_size / world_size)   # the FULL shard
# e.g. 5120-row weight, 4 ranks: selected 320 rows -> padded back to 1280

Because select_dim == shard_dim (both -2/rows), the selection dimension is
exactly the dimension the pad rule re-inflates. Result: the assembled matrix is
full size (~¾ zero rows), so comm volume and NS GEMM cost are identical to
NorMuon — Dion2's fraction buys nothing here.

Not a correctness bug (zero rows contribute nothing to UᵀU and are narrowed off
after the second all-to-all). It is a performance regression: Dion2 pays the
selection overhead (topk, gather, scatter, ~244-kernel tail) and collects none of
its compute/comm reward.

How it was identified

• GPU kernel breakdown of nordion2_gramns vs normuon_gramns (5120/8L, bs1
  traces): both spend ~63% in the same gramNS GEMMs; nordion2's NS is only
  modestly cheaper (49 vs 65 ms), not the ~4x expected from fraction=0.25.
• Trace shapes confirmed it: nordion2's all_to_all chunks and
  gemm_symmetric inputs are full [*, 5120, 5120], byte-identical comm
  (2500 MB / 6 calls) to normuon — the selection never shrank the matrix.
• git blame + history traced the padding to the megabatch path and showed the
  original Dion2 did not do this.

When it was introduced

The original Dion2 (Kwangjun Ahn) had a dedicated distributed path that
communicated and orthogonalized only the selected submatrix — the old
dion2_update_batch_async sized its all-to-all buffers to U_selected and
carried the literal comment # Only submatrix is orthogonalized!. No pad to
global size.

The regression was introduced by the megabatch refactor that moved Dion2
onto NorMuon's shared comm path, then hardened by later padding fixes:

date	commit	what
2025-11-13	5ca6fa5 Kwangjun Ahn (#15)	original Dion2; selected-submatrix-only comm + NS
2026-03-27	4266bd0 John Langford	Refactor Dion2 to use megabatching — Dion2 joins the shared path (regression starts here)
2026-03-27	e4cc9b2 John Langford	Dion2 + NorMuon share the megabatch base class
2026-05-04	e8e1c6a John (#74)	pad to ceil(global/world) to fix an FSDP2 empty-shard NCCL hang (makes the full-size inflation explicit)
2026-05-06	ff0fb99 (#75) / 4dc5728 (#77)	precompute + harden padded_local_size

Root cause

megabatch_orthogonalize_async is shared between full-shard callers (Muon/
NorMuon, where padding to ceil(global/world) is correct and necessary to avoid
the all-to-all hang) and selected callers (Dion2/NorDion2). The pad rule has no
notion that the selected input was already subsampled, so it re-inflates the
selection to the full shard. The original per-optimizer path avoided this by
sizing comm to k = ceil(fraction * rows) directly.

[alint77]

Discussed this extensively with @NoahAmsel. I am working on a fix atm @JohnLangford

[NoahAmsel]

Thanks for finding this @alint77. If I could summarize:

Imagine we have 4 gpus and a matrix with leading dimension 14. This matrix will be divided into 4 shards with respectively 4, 4, 4, and 2 rows. The megabatch_orthogonalize_async function pads the final shard with 2 rows of zeros so that we have 4, 4, 4, and 4. The target number of rows (4 in this example) is called padded_local_size and it's computed here.

The problem comes when we are using dion2 or nordion2. In that case, we subselect 50% of the rows. So now we have shards of size 2, 2, 2, and 1, respectively. However, padded_local_size is still 4, since the megabatch_orthogonalize_async does not know about dion2. Therefore, the blocks get padded with zeros to be size 4, 4, 4, and 4. This ruins the cost savings of dion2.

@alint77 is that correct? @JohnLangford?

[alint77]

Fix

A user-facing selection_scope arg (nanoplm: dion_selection_scope, default
"local") on Dion2/NorDion2 selects between two behaviors on the row-sharded
path. Off that path (per-head, single-GPU, batch-sharded) each rank already
holds whole matrices, so both modes coincide and the arg is a no-op.

"local" — per-shard top-k, communicate only the selected rows. Restores
the original Dion2 cost profile: comm and Newton-Schulz both scale with
fraction. Implementation: when the global comm dim divides evenly across ranks
(the common case), select a uniform k = ceil(fraction * global/world) on every
rank (so all-to-all sizes match) and pass global_comm_dim_size = k * world to
the megabatch, which then computes padded_local_size = k, sees U_selected
already at that size, and skips padding. Numerically identical to the previous
full-pad implementation (the dropped padded zero rows never affected U^T U),
just without paying for them. The selected set is the union of per-rank top-k, so
it depends on the sharding layout (world-size variant). Uneven/empty shards fall
back to the original full-pad path.

"global" — send the full shard (full comm, like NorMuon), but take the
top-k on the assembled whole matrix and Newton-Schulz only that submatrix. The
selection is exact and invariant to the sharding layout (reproducible across
world sizes). Implementation: skip pre-comm selection; wrap newton_schulz_func
in _make_select_and_orthogonalize, which selects + orthogonalizes + scatters
back into a full-size, otherwise-zero tensor (so it composes with split_sizes /
per-head, which call NS per block / per head). The post step
(dion2_post_orthogonalize_masked / nordion2_post_orthogonalize_masked)
derives the selected mask from the nonzero rows of the result -- no indices --
and applies error-feedback decay, NorMuon per-neuron normalization (NorDion2),
and the masked weight update.

NS compute scales with fraction in both modes; "local" also shrinks comm but
is world-size variant, "global" keeps full comm but is exact/reproducible.

Touched files: dion/dion/dion2.py, dion/dion/nordion2.py (modes + helpers);
src/nanoplm/pretraining/{config,optim}.py (plumbing).

Measured effect (nordion2, 5120/8L, fraction=0.25, world=4, bs1, step-60)

run	t7 ops	t7 ms (GPU)	nccl MB	Gram GEMM shape (attn)
OLD (broken, full-pad)	324	89.3	2500	[*, 5120, 5120]
local (this fix)	230	39.2	625	[*, 1280, 1280]
global (this fix)	248	42.8	2500	[*, 1280, 1280]

The fix does what the regression analysis predicted:

• Newton-Schulz GPU time more than halves (89 → 39/43 ms) because the Gram
  product now runs on the selected submatrix: the square attention weight's Gram
  drops from [5120, 5120] to [1280, 1280] (= fraction * 5120, ~16x fewer
  flops), and the fused-QKV blocks from [5120,5120] to [1280,1280].
• local also cuts comm 4x (2500 → 625 MB): the all-to-all chunks shrink
  from the full [*, 1280, 5120] shard to the selected [*, 320, 5120]
  (320 = ceil(0.25 * 1280)). global keeps full comm (2500 MB, full
  [*, 1280, 5120] chunks) as designed -- it must gather whole rows to select
  globally.
• Kernel count drops (324 → 230/248) as the padded-zero-row scatter/copy tail
  shrinks.

End-to-end both modes train the real 5120/8L/bs1/gns config for 100 steps at
~53% MFU (local final loss 2.66, global 2.64).

[alint77]

Scaling to 2 nodes / world_size=8

Rerun on 2×H100 nodes (world_size=8, same 5120/8L/fraction=0.25/bs1 config).
Per-rank FSDP2 shard halves (5120/8=640 rows), so per-rank comm and per-rank NS
compute both scale with 1/world_size (data-parallel over parameters).

run	t7 ops	t7 ms (GPU)	nccl ms	nccl MB	t7 occ%
OLD broken (ws4)	324	89.3	11.1	2500	99.1
local ws4	230	39.2	3.4	625	98.6
global ws4	248	42.8	10.8	2500	98.4
local ws8	234	21.7	9.0	325	97.2
global ws8	252	38.7	29.8	1300	59.5

Both modes still run NS on the selected submatrix at ws8 (Gram [1280,1280] for
the square attention weight, vs [5120,5120] broken). a2a chunks confirm:

• local ws8 sends the selected rows per rank: [*,160,5120] (Wo, 0.25×640),
  [*,480,5120] (QKV, 0.25×1920), [*,640,5120] (MLP-up, 0.25×2560).
• global ws8 sends the full per-rank shard: [*,640,5120], [*,1920,5120],
  [*,2560,5120].

The ws8 result exposes the local-vs-global tradeoff at multi-node scale:

• local ws8 keeps scaling: GPU compute halves again (39→22 ms, 1/world_size),
  comm halves again (625→325 MB), occupancy stays high (97%). The small selected-
  row a2a hides behind NS compute.
• global ws8 stalls on comm: GPU compute halves (43→39 ms) but the full-shard
  a2a now crosses the inter-node network, and nccl time jumps to 29.8 ms --
  exceeding the 23 ms of NS work -- so the GPU idles 40% of the step
  (occupancy collapses 98%→59.5%). Full comm is the bottleneck at 2 nodes.

[alint77]

FIX: NorDion2 world-size sweep — opt.step scaling (5120/8L, bs1, unfused QKV)

NorDion2, gram-NS, selection_scope=local, fraction=0.25, fused_qkv=false,
global_batch_tokens=1, across world_size 1 / 2 / 4 on one 4×H100 node.
Profiler steps 58–64. Traces: output/profiling_runs_scopefix_bs1/nordion2_local_nofusedqkv{,_1gpu,_2gpu}/.

metric	1 GPU	2 GPU	4 GPU
step time (steady, ms)	590	510	475
H100 MFU	43%	50%	53.5%
throughput (tok/s)	~28k	~64k	~137k
opt.step GPU wall (ms)	117.9	61.8	31.4
opt.step % of total step	20.0%	12.1%	6.6%
a2a comm (MB/step)	0	642	164
a2a NCCL time (ms)	0	15.9	3.9

a2a shapes [N, k_rows, cols]: 2 GPU [17,640,5120]+[8,2560,5120]; 4 GPU
[9,320,5120]+[4,1280,5120] (half the per-rank k; k = ceil(fraction·rows/world)).

Findings

• Newton-Schulz distributes ~linearly: ≈115 ms / world_size (115 → 61 → 31).
  Each rank orthogonalizes only its shard's selected fraction; 1 GPU has no sharding
  so it eats the full undistributed 115 ms.
• opt.step share collapses 20% → 12% → 6.6% — NS wall shrinks ~1/N while the
  forward/backward per step stays comparable.
• Comm stays cheap and shrinks past 2 GPUs: 0 → 642 MB → 164 MB (smaller shards
  and smaller per-rank k). Adding ranks cuts both NS compute and comm.
• Throughput scales ~1 : 2.3 : 4.9 (bs1 adds data-parallel width, so step time
  barely moves); the MFU gain is exactly the opt.step share shrinking.

[alint77]

NorDion2 opt.step Amdahl — fraction-scaled vs world-scaled

What fraction of the optimizer step is cut by dion_fraction, what is cut by
world_size, and what is the floor neither knob touches. Measured on the
padding-fixed local-scope path (5120/8L, bs1, gram-NS, fused_qkv=false),
profiler steps 58–64, stream-7 GPU compute inside the Optimizer.step#NorDion2.step
window. Four runs forming a 2×2 (world_size × fraction):

• ws2 / f=0.25 — nordion2_local_nofusedqkv_2gpu
• ws2 / f=0.50 — nordion2_local_nofusedqkv_2gpu_frac05
• ws4 / f=0.25 — nordion2_local_nofusedqkv
• ws4 / f=0.50 — nordion2_local_nofusedqkv_4gpu_frac05

Traces: output/profiling_runs_scopefix_bs1/<run>/.../profiler_traces/chrome_trace.json.

The two scaling classes

The opt.step splits into ops that touch only the k = ceil(fraction · padded_local)
selected rows, and ops that touch the entire local shard (padded_local = global_rows / world) regardless of fraction:

• fraction-scaled — Newton-Schulz (Gram matmuls + polynomial), the all-to-all, and
  the selected-row post-ops (error-feedback decay, NorMuon normalization, variance EMA,
  weight-update scatter). Cut by both fraction and 1/world.
• world-scaled (world-only) — momentum M += G, the L1-norm scan that ranks rows
  for selection, the top-k, and weight-decay X *= (1 − lr·wd). Cut by 1/world
  only — these process the whole shard no matter what fraction is.

Bucketing is by name-stable base kernel names (trailing torch.compile hashes
stripped), not raw hashed names, so the same logical op is comparable across runs that
recompile to different kernel names. NS = quackgemm matmuls + eye_gemm /
gemm_out setup; world-only = triton_for_fused (fixed 5-iter poly loop, launch-bound)

• triton_poi_fused_copy_ (momentum add over the full shard) + sbtopk::gatherTopK.

opt.step bucket table — full 2×2 (world × fraction)

bucket (ms)	ws2 / f=.25	ws2 / f=.50	ws4 / f=.25	ws4 / f=.50	scales with
NS orthogonalization	35.9 (59%)	58.8 (62%)	18.1 (58%)	29.1 (60%)	fraction × 1/world
select + post + a2a-pack (selected rows)	12.2 (20%)	22.8 (24%)	6.7 (21%)	12.6 (26%)	fraction × 1/world
world-only (momentum, L1-norm, top-k, wd)	13.0 (21%)	13.0 (14%)	6.5 (21%)	6.6 (14%)	1/world ONLY
opt.step GPU total	61.1	94.6	31.2	48.2
a2a comm volume (MB/step)	642	1285	164	328	fraction × 1/world
a2a shape [N, k, cols] (attn)	[17, 640, 5120]	[17, 1280, 5120]	[9, 320, 5120]	[9, 640, 5120]	k = fraction·rows/world

Scaling along world (fraction fixed): ws2 → ws4

bucket	f=.25: ws2→ws4	factor	f=.50: ws2→ws4	factor	prediction
NS	35.9 → 18.1	0.50×	58.8 → 29.1	0.50×	1/world
select + post + a2a-pack	12.2 → 6.7	0.54×	22.8 → 12.6	0.55×	1/world
world-only	13.0 → 6.5	0.50×	13.0 → 6.6	0.51×	1/world
opt.step total	61.1 → 31.2	0.51×	94.6 → 48.2	0.51×
a2a MB	642 → 164	0.26×	1285 → 328	0.26×	1/world × (k shrinks)

All three buckets halve (0.50–0.55×) at both fractions when world doubles — both
scaling classes track 1/world. a2a volume drops faster (0.26×) because k shrinks
and the shard shrinks. No floor appears along world at this scale: more ranks
shrink the whole opt.step, at every fraction.

Scaling along fraction (world fixed): f=0.25 → f=0.50

bucket	ws2: f.25→f.50	factor	ws4: f.25→f.50	factor	prediction
NS	35.9 → 58.8	1.64×	18.1 → 29.1	1.61×	~fraction
select + post + a2a-pack	12.2 → 22.8	1.86×	6.7 → 12.6	1.89×	~fraction
world-only	13.0 → 13.0	1.00×	6.5 → 6.6	1.01×	flat (1.0×)
opt.step total	61.1 → 94.6	1.55×	31.2 → 48.2	1.54×
a2a MB	642 → 1285	2.00×	164 → 328	2.00×	exactly fraction

The decisive test, now at both world sizes. Doubling fraction doubles the
fraction-scaled buckets (NS 1.6×, post ~1.9×, a2a exactly 2.00× — the a2a shape
k doubles: 640→1280 at ws2, 320→640 at ws4) and leaves the world-only bucket dead
flat (1.00× / 1.01×) at both ws2 and ws4. Those ops process the whole shard regardless
of fraction, so their absolute time is unchanged; only their share of opt.step drops
(21% → 14%) as the fraction-scaled part grows around them.

NS scales ~1.6× (sub-2×) because the Gram matrix is k × k — quadratic in k, so
doubling k is ~1.5–2× depending on tile regime (the quackgemm_default kernel
crosses a CUTLASS tile threshold at f=0.5 and jumps, a shape-regime artifact, not a
clean 2×).

Narrative

Two knobs, two floors. The opt.step has a fraction-removable part and a
fraction-irremovable part. The fraction-scaled part (NS + a2a + selected-row post,
~79% at f=0.25/ws4) is everything that touches only the k selected rows — shrinking
fraction shrinks it, and it is also the only part that benefits from the padding fix
(the broken path padded k back to the full shard, so this part ran ~4× too big).
The world-only part (~21% at f=0.25/ws4) — momentum accumulation, the row-ranking L1
norm, top-k, and weight decay over the full shard — is the fraction-Amdahl floor:
no matter how small you make fraction, these stay, because they must process every
row to decide which rows to select and to keep momentum/weights current.

World scaling has no such floor (at this scale). Both classes halve when world
doubles (0.50–0.54×), so adding ranks shrinks the entire opt.step — the world-only
floor is 1/world, not a constant. Its share is flat across world (21% at ws2 and
ws4, f=0.25) but drops when fraction grows (21% → 14%, ws4 f=.25 → f=.50).

Where the leverage is. At f=0.25/ws4, NS is 58% of opt.step and the whole
fraction-scaled part is 79% — so fraction is the high-leverage knob today (halving it
removes ~40% of opt.step, floored at the 21% world-only part). World is a clean
~1/N win on the whole opt.step with no floor. The padding fix matters precisely
because it puts the fraction-scaled part back onto the fraction × 1/world curve it
was supposed to be on — on the broken path that part was 1/world only (full-shard),
so fraction did nothing for it.

Caveat on the world-only floor. The world-only bucket is partly launch/latency-bound
(the 5-iteration NS poly loop triton_for_fused is ~20 ms and flat in fraction but is
counted under NS setup here; the momentum M+=G and weight-decay are memory-bandwidth
bound). Its 21% share is real but its 1/world scaling will flatten toward a launch
floor at very small shards / very high world — the same Amdahl limit, just along the
world axis, that bites later than the fraction one.

[alint77]

Before the fix: dion main (padding regression)

Same sweep on unpatched dion main, which pads each rank's selected rows back to the
full shard before the all-to-all — so NS and comm run on k = padded_local rows
(the whole shard) instead of k = ceil(fraction·padded_local), i.e. 1/fraction = 4×
too much. Traces: output/profiling_runs_scopefix_bs1/nordion2_broken_nofusedqkv{,_1gpu,_2gpu}/.

opt.step component	1 GPU	2 GPU	4 GPU
opt.step ms (broken → fixed)	114 → 118	174 → 62	86 → 31
speedup from fix	~1.0×	2.82×	2.75×
NS ms (broken → fixed)	81 → 85	131 → 44	67 → 22
a2a MB (broken → fixed)	0 → 0	2569 → 642	655 → 164
a2a NCCL ms (broken → fixed)	0 → 0	44.0 → 15.9	13.4 → 3.9
step MFU (broken → fixed)	44% → 43%	41% → 50%	48% → 53.5%

• The bug makes the optimizer scale negatively: broken opt.step goes 114 → 174 → 86 ms
  — 2 GPU is slower than 1 GPU, because padding forces full-shard NS plus a 2.6 GB
  all-to-all. The fix restores ~1/N scaling (118 → 62 → 31).
• 1 GPU is unaffected (broken ≈ fixed): no sharding ⇒ no padding, so the regression
  is isolated to the sharded (2/4 GPU) path.
• The fix cuts NS and comm by ~fraction (4×) on the sharded path — 2.8× faster opt.step,
  and +5–9 MFU points end-to-end.