Fix convolution flattening in megabatches and NorMuon normalization

[Zherui Yang (Adversarr)]

Problem and fix

With flatten=True, stacked convolution weights [N, out, ...] were reshaped
to [N, numel], so Newton–Schulz orthogonalized across layers. Preserve the
megabatch axis and orthogonalize [N, out, prod(rest)] instead. This also fixes
stacked Linear parameters with flatten=True.

The existing LR adjustment already uses the individual parameter's shape:
fan_out=out, fan_in=prod(rest). The corrected NS geometry now agrees with
that scaling; the LR formulas themselves are unchanged. The single-parameter
NS dispatch, flatten=False, and row-split geometry are preserved.

Following John (@JohnLangford)'s review and proposed patch, this PR also:

• Makes the megabatch-axis flag required and documents the intentional
  flatten=False dimensionality handling.
• Fixes NorMuon variance and normalization to use the same flattened matrix.
  The variance buffer holds one value per output channel. Local, replicated,
  and output-channel-sharded convolution parameters are supported; convolution
  column shards are explicitly rejected. Empty row shards are handled.
• Rejects flatten=True for 3D+ parameters in Dion2 and NorDion2/Dion3 before
  any parameter update: their submatrix selection/error feedback is not yet
  flatten-aware. AdamW/Lion fallback groups are exempt. flatten=False has
  different batch-of-spatial-matrices semantics, not equivalent conv flattening.

Compatibility

NorMuon retains its existing shard-local norm-preserving rescale. The sharded
oracle is an equivalent 2D parameter under the same sharding, not an
unsharded optimizer; this PR does not make normalization world-size invariant.

Old NorMuon flattened-convolution optimizer checkpoints are explicitly rejected
before replacing live state. A layout-version marker detects legacy states even
when singleton dimensions make the old and new variance shapes coincide. Load
model weights and construct a fresh optimizer to restart; no automatic migration
or silent variance reset is performed. Unaffected state layouts are preserved.

Correct convolution NS can cost more than the old layer-mixing operation; old
step times are not equivalent-work throughput baselines. CHANGELOG entries and
optimizer docstrings describe these behavior changes.

Validation

Validated with PyTorch 2.11.0+cu128:

• Full tracked CPU suite: 372 passed, 313 skipped, 18 baseline failures.
  Failure IDs match the original PR revision exactly; all 241 previous passes
  remain passing, with 131 new tests passing.
• Full tracked suite on 2 RTX 3090 GPUs: 674 passed, 13 skipped, 16 baseline
  failures. Failure IDs match a full original-revision GPU run exactly; all
  520 previous passes remain passing, with 136 new tests passing and 18 formerly
  skipped distributed tests executing successfully.
• Focused production/FSDP2 checks on 4 A100 GPUs: 7 passed, no skips. Both
  Muon and NorMuon execute actual FSDP2 with 2 and 4 workers.
• After the final BF16 test-oracle adjustment, the complete modified megabatch
  module passes 207/207 tests, no skips, on both RTX 3090 and A100.
• The broader 4-A100 sweep completed with 677 passed, 26 failed, no skips
  before that test-only adjustment. Eight failures were the BF16 comparison
  assumption fixed above. Combining the full run with the complete modified
  module rerun gives 685 passed, 18 baseline failures, no skips (an aggregate,
  not a second full-suite invocation). All 138 added tests pass, and no
  verified baseline passing test is lost.

The 16 GPU baseline failures are existing wrapper-subclass compiler failures
and optional Gram/CuTe schema failures, reproduced on the original PR revision
e9b9042. A100 also reproduces two unchanged Triton-vs-cuBLAS accuracy assertions
on that original revision. These kernels/tests are outside this patch and were
not disabled or loosened. CPU-only runs additionally encounter CUDA-only Triton
execution failures. Full-suite counts above precede the final test-only oracle
adjustment; that modified module was rerun completely afterward on both GPUs.

Coverage includes optimizer-level Muon grouping, Conv1d/2d/3d versus equivalent
matrices, multiple steps and LR-adjustment modes, state/checkpoint compatibility,
actual FSDP2 forward/backward with BF16 computation and FP32 master parameters,
uneven/empty row shards,
production Polar Express, and CUDA-graph replay. FP64 tests compare against
independent unbatched matrices with tight tolerances. BF16 GPU layer-isolation
tests preserve the BMM shape: direct PE GEMM/BMM calls differ by up to 0.0703125
on A100 even without the optimizer/flatten helper. Identical batched dispatch
remains exact; no production numerical tolerance or kernel was changed.

Negative controls reproduce both defects: the pre-fix Muon step presents
[2, 288] instead of [2, 8, 36] to NS, and the NS-only fix still fails the
NorMuon convolution/matrix update oracle.

The normalization work is adapted from JohnLangford's proposed
61e3404,
with attribution in the commit. The distributed tests use all-rank collectives
and same-sharding references instead of the proposed unsharded comparison.

Existing training ablations — Muon only

These plots compare pre-fix Muon, AdamW, and fixed Muon. They are training-loss
examples, not NorMuon/Dion2/Dion3 convergence results or evidence of a universal
optimizer ranking. No full training campaign was rerun for the review follow-up.

Speech Commands v0.02: 35-keyword audio classification with Conv1d.

Food-101: 101-class food classification with the ConvNeXt-T UNet classifier.

ModelNet40: 40-class classification of voxelized point clouds with sparse Conv3d.

[Zherui Yang (Adversarr)]

@microsoft-github-policy-service agree

[John (JohnLangford)]

Review

Verdict: the bug is real, the diagnosis is right, and the 4-line fix is correct at all four call sites. I'd merge after a CHANGELOG entry, one flaky assertion tightened, and one optimizer-level test added.

Correctness — confirmed

Pre-fix, megabatch_base.py:902-908 with flatten=True turned a stacked [N, O, I, kh, kw] into [N, O*I*kh*kw], so Newton–Schulz orthogonalized across parameters. Same for stacked Linear with flatten=True → [N, O*I].

Independent confirmation that the new geometry is the intended one: adjust_lr_rms_norm / adjust_lr_spectral_norm (megabatch_base.py:1016-1034) are called with X[0].shape — a single parameter — and under flatten use fan_out=shape[0], fan_in=prod(shape[1:]). The learning rate has always been computed for O x (I*kh*kw) geometry; only NS disagreed. This fix makes them consistent, which is good evidence it is a correction rather than a behavior change. Worth putting in the PR description — it is the strongest argument for the patch.

All four call sites check out: :801 (sharded), :840 (replicated) and :873 (single-process stacked) pass True, correctly so even when per_rank == 1, since the stack axis exists unconditionally on those paths; :860 (N == 1) correctly keeps the default.

The elif X.ndim >= 4 branch still uses raw X.ndim rather than param_ndim. That is correct — flatten=False means "trailing two dims are matrices" and NS takes at most 3D — but please add a one-line comment saying the asymmetry is deliberate. The new param_ndim local sitting directly above it invites a future "consistency fix" that would be a regression.

The zero-pad interaction under sharding is sound: padding along a non--2 comm_dim inserts zero slices that become zero columns after flattening, and NS output column j is p(XX^T) x_j with zero columns contributing nothing to XX^T, so the real columns come back bit-identical to the unpadded result. _sharded_cpu_worker exercises exactly this (rows=5 over world_size 2 and 4, axes 0 and 1). Nice.

Design

has_megabatch_dim: bool = False defaults to the buggy behavior. Please make it a required keyword argument — there are only four call sites — so a future caller cannot silently reintroduce this. Also note in the docstring that the flag is ignored on the split_sizes path; that is correct (row blocks split on dim -2 regardless) and a test pins it, but a reader has to derive it.

A more elegant alternative, which I would reject: the flag only exists because the N == 1 branch hands NS an unstacked parameter. Make that branch torch.stack(U) too and has_megabatch_dim becomes unconditionally true and deletable. Don't do it — it changes single-parameter numerics from a 2D gemm to a [1, M, N] bmm (different cuBLAS kernel, no longer bit-identical to today) and adds a dynamo shape specialization. test_linear_parameters_remain_independent_matrices pins the pre-fix single/batched dispatch bitwise, so you clearly already considered this. Keeping the flag is the right call.

Tests — thorough, three issues

1. Likely flake. test_real_polar_express_matches_independent_convolutions compares eager(x.reshape(O, -1)) (a 2D gemm) against a row of a [3, O, F] batched run (a bmm) at rtol=0, atol=0. cuBLAS routinely selects different kernels and accumulation orders for gemm vs bmm, so this is architecture-dependent and will eventually fail on someone's GPU. That half needs a tolerance; the batched-vs-batched half can legitimately stay exact. test_zero_rank_deficient_noncontiguous_and_layer_independence reverses batch order under rtol=0, atol=0 too — safer on CPU float64, but the same class of assumption.

2. No optimizer-level test. Everything targets megabatch_orthogonalize_async directly. The bug was only reachable through Muon(..., flatten=True) with two or more same-shape conv params, and the highest-value regression test is exactly that: two identical nn.Conv2d modules, one Muon(flatten=True) step, compared against stepping each in its own optimizer. That is the test that catches a future regression in the grouping logic rather than in NS.

3. Cost. _sharded_cpu_worker runs 4 ndims x 3 row counts x 3 counts x 2 axes x 2 return_stacked = 144 full all-to-all rounds, at world_size 2 and 4, plus 16 more in the replicated worker. That is minutes of gloo traffic to test a reshape. The interesting axes are {ndim} x {divisible, non-divisible, empty-shard} x {axis}; count and return_stacked are orthogonal and already covered single-process. Also, the existing distributed tests here (test_normuon_split_lr.py, test_dion2_selection_scope.py) write results to a file and assert in the parent — asserting inside the worker still surfaces via mp.spawn, but with a much worse failure message.

Minor: monkeypatching module-level dist.all_to_all onto all_to_all_single is clever and properly contained (subprocess + finally), but it validates the packing against a re-implemented collective — keep the NCCL variant as the authoritative check.

Throughput

The fix raises orthogonalization FLOPs by design: old was one Gram of an N x numel matrix, about N^2 * numel; new is N batched O x (I*kh*kw) orthogonalizations, about N * O^2 * I*kh*kw — roughly O^2/N times more work for conv megabatches. That is the correct amount of work (it equals the non-megabatched cost) and the megabatch still buys its comm and kernel-launch amortization, but it is worth saying in the description that the old fast path was fast because it was wrong, so conv users should expect step time to rise. The flattens stay views (torch.stack / torch.cat outputs are contiguous), so no extra copies.

Backward compatibility

This silently changes numerics for anyone running flatten=True with megabatching. It needs a CHANGELOG.md [Unreleased] → Fixed entry, ideally carrying the LR-adjustment consistency argument. The flatten docstring repeated in muon.py:42, normuon.py:44, nordion2.py:48, dion2.py:40 ("3D+ tensors are flattened to 2D") should say each parameter is flattened to 2D.

Adjacent bug this PR does not fix

NorMuon and NorDion2 with conv + flatten=True are still wrong after this patch. variance_neuron is allocated as zeros_like(param[..., 0:1]) (normuon.py:125), so for [O, I, kh, kw] it is [O, I, kh, 1], and _normuon_normalization_core (normuon.py:420-444) takes the Frobenius norm over dims (-2, -1) = (kh, kw) and the neuron variance over dim=-1 = kw. So once NS correctly produces O x (I*kh*kw) orthogonality, the normalization renormalizes per kh x kw slice and partially destroys it — while adjust_lr_*(flatten=True) is scaling for the flattened geometry. normuon_normalization_stacked's own signature comment says U: [N, rows, cols], which a 5D conv tensor is not.

That is out of scope here, but it interacts directly with this PR's stated goal. Either handle flatten in the normalization or reject flatten=True for those two optimizers until it is handled — and it would be useful to know which optimizer produced the W&B curves above.

Security: nothing of concern — no I/O, no deserialization, no new dependencies.

— posted as John's reviewbot

[John (JohnLangford)]

Follow-on: the NorMuon half of the same bug

I put together the NorMuon fix for the adjacent bug I flagged above, branched directly on your head commit e9b9042, in case you'd rather land the convolution story as one PR than leave a known-broken path behind this one. Entirely your call — take it, take part of it, or ignore it.

git fetch https://github.com/JohnLangford/dion.git feat/normuon-conv-flatten
git cherry-pick 61e34048253c0dce8c26792b64ae94ae6e7de953

Branch: https://github.com/JohnLangford/dion/tree/feat/normuon-conv-flatten

Why it belongs here rather than in a separate PR

The two halves only work together, in both directions.

Your patch alone leaves NorMuon(flatten=True) on convolutions wrong: Newton-Schulz now correctly produces out x (in*kh*kw) orthogonality, and then the normalization renormalizes per kh x kw slice and partially undoes it.

My patch alone is equally incomplete: without your fix, NS still flattens a stacked megabatch to [N, numel], so flatten=True convolutions stay broken regardless of what the normalization does. Concretely, on main a 2-rank replicated NorMuon step on a (8, 4, 3, 3) conv diverges from the single-process result by ~2.8e-4; with your commit plus this one it matches to 1e-14.

One of the new tests (test_replicated_conv_flatten_matches_single_process) fails on main for exactly the reason your PR fixes, so it doubles as optimizer-level regression coverage for this PR — the level I suggested was missing above. It runs on gloo/CPU.

What it does

• variance_neuron is sized from the group's flatten: [out, 1, 1, 1] instead of [out, in, kh, 1], so a "neuron" is an output channel. Built with chained narrow, which preserves a DTensor's sharding.
• The normalization runs on the [N, out, in*kh*kw] view. normuon_normalization_stacked is untouched — it already took [N, rows, cols], so this is a view on both sides and the compiled kernel and CUDA-graph capture are unaffected.
• _get_or_initialize_state takes the owning param group (it needs flatten at allocation time). Required, not optional: a buffer that silently defaults to the wrong shape is not recoverable later. All seven call sites already had group in scope.
• A 3D+ parameter sharded anywhere but dim 0 is rejected. Under flatten every other dim holds columns, so such a shard splits a neuron's row across ranks while the variance is reduced locally. FSDP2 shards dim 0, so this rejects a configuration that was already producing wrong numbers.
• Dion2/NorDion2 raise NotImplementedError for flatten=True with a 3D+ parameter. Their low-rank selection indexes rows at dim -2, which for a convolution is a kernel dimension, so the rows selected, error-fed back and scattered were never neurons. Making that flatten-aware means reshaping the momentum and error-feedback buffers and reworking the row-sharded local scope — a redesign, not a patch — so it is rejected for now rather than run silently wrong. That path is broken on main today, independently of this PR.
• split_sizes and num_heads needed no new guard: _resolve_split_sizes and _resolve_num_heads already reject flatten=True.

The oracle throughout is exact equivalence: a [out, *rest] parameter under flatten=True must step bitwise identically to a [out, prod(rest)] weight fed the reshaped gradient, since they are the same matrix. That holds for Conv1d/2d/3d, megabatches of 1/2/3/5, all three adjust_lr modes, and the replicated distributed path.

Two things to know

It is a checkpoint break for 3D+ parameters under flatten=True: variance_neuron changes shape, so an old checkpoint will not load. Those checkpoints came from the broken path, and 2D parameters and flatten=False are untouched. There is a CHANGELOG entry saying so.

One test is unexecuted. test_sharded_conv_flatten_matches_single_process is NCCL/CUDA-gated because the FSDP2-sharded path needs all_to_all, which gloo does not implement — the same wall your _sharded_cpu_worker adapter works around. I have no GPU here, so it skipped. Everything else (49 tests) passes on CPU, and the full existing suite shows no regressions (identical pass/fail sets before and after: 226 passed, 18 pre-existing Triton/CUDA failures on a CPU-only box). If you take this, that one test wants a run on a 2-GPU box.

— posted as John's reviewbot

[Zherui Yang (Adversarr)]

Thanks John (@JohnLangford) — I incorporated the review feedback and adapted your
NorMuon patch into this PR, with co-author credit and a link to the source commit.

• has_megabatch_dim is now required and keyword-only; the singleton call is
  explicit, and the intentional flatten=False branch is documented.
• Added actual Muon optimizer/grouping regressions and reduced the redundant
  Gloo matrix. FP64 tests compare independent matrices with tight tolerances;
  BF16 layer-isolation checks preserve BMM shape, since direct PE GEMM/BMM
  calls differ by up to 0.0703125 on A100 without any optimizer/flatten helper.
  The identical batched dispatch still has an exact assertion.
• NorMuon now allocates one variance per output channel and normalizes the
  flattened matrix. Non-row convolution shards are rejected, and empty row
  shards work. Old incompatible optimizer states fail explicitly before load
  mutates live state, including singleton kernels whose buffer shapes coincide.
• Dion2 and NorDion2/Dion3 reject flattened 3D+ parameters before any weight
  update, including grad-less ones. AdamW/Lion fallback groups are exempt.

For the sharded test, I used an equivalent equally sharded 2D matrix as the
oracle, because NorMuon retains its existing shard-local norm rescale. All ranks
participate in gradient gathers and save their local comparison results. The
tests exercise actual FSDP2 forward/backward, not just DTensor construction.

Validation: actual FSDP2 with both 2 and 4 A100 GPUs, sharded state reload,
production PE, and CUDA-graph replay all passed (7 focused GPU tests, no skips).
The final modified megabatch module passes all 207 tests on both A100 and
RTX 3090. Full CPU and local GPU suites preserve every baseline passing test;
their remaining 18/16 failures respectively reproduce on the original PR
revision. The broader A100 checks also reproduce two unchanged Triton accuracy
assertions on the original source. Its full run plus the final modified-module
rerun passes all 138 added tests, with only those 18 verified baseline failures
remaining. No unrelated kernel/tolerance changes were
made. I also isolated compiler caches between the new test modules so the shape
sweeps do not exhaust the budget of later tests.

The description and CHANGELOG now explain the LR/NS geometry agreement,
checkpoint break, supported optimizer scope, and throughput caveat. The three
existing training plots are explicitly labeled Muon-only; this follow-up makes
no new convergence or universal optimizer-ranking claim.

[John (JohnLangford)]

The core fix looks right to me, and the validation is unusually thorough — I verified the NS geometry rewrite at all four call sites, the singleton and flatten=False paths really are unchanged, and the sharded path still gathers rows before NS so orthogonalization stays globally exact. Four things I'd like resolved or discussed.

A. The Dion2/NorDion2 rejection is broader than its stated cause, and fires too late. (blocking)

The justification is that submatrix selection isn't flatten-aware. That's true of the local selection scope, which does pre-comm top-k on the raw tensor. But in the global scope, _select_ns is passed as newton_schulz_func into megabatch_orthogonalize_async — so after this PR it receives the already-flattened [N, out, prod(rest)], and select_dim=-2 indexes real output channels. That path looks like it is repaired by your NS change, not broken by it. Could you either scope the guard to selection_scope == "local", or show the global path is still wrong?

Two smaller points on the same guard:

• It raises from _create_ortho_tasks, i.e. at the first step(), after a full forward/backward. flatten and p.ndim are both known at __init__/add_param_group — rejecting there (as NorMuon's shard guard already effectively does via prepopulate) fails in a second instead of a step, and removes the per-step cost.
• This hard-breaks anyone running Dion2 with flatten=True (the rewritten test_3d_params_flatten was exactly that). That's fine if the math is wrong, but the PR ships no test demonstrating the wrongness, only that it now raises. A short numeric negative control would earn the break, the way the other two defects got one.

B. Per-step revalidation in NorMuon._create_ortho_tasks.

The new loop calls _get_shard_info + _validate_variance_state for every ndim>2 param on every step. State can only change via load_state_dict (hooked), add_param_group (hooked), or direct mutation of group["flatten"]. This file optimizes host dispatch hard — the V-writeback comment cites ~20% of step CPU — so an O(#params) Python plus DTensor-placement walk each step is out of character. Validating once and caching would do it; the only case needing per-step detection (test_changing_live_flatten_is_rejected) is served by comparing a stored per-param flatten flag against group["flatten"], which is much cheaper than _get_shard_info.

C. load_state_dict robustness.

• saved["algorithm"] and saved["flatten"] will KeyError on any state dict lacking those keys — older checkpoints, hand-built dicts, third-party tooling. .get with a clear fallback would be safer.
• type(version) is not int rejects a version that round-trips as numpy.int64 or a 0-d tensor through some checkpoint backends. isinstance(version, int) and not isinstance(version, bool) avoids that.
• The step_dev comment directly above explains why this codebase deliberately keeps state values as tensors for DCP key-set symmetry. Worth a sentence on why a plain int is fine here, or just store the version as a 0-d tensor for consistency.

D. has_megabatch_dim can be removed entirely by flattening at the boundary.

Rather than teaching NS whether a stack axis is present, reshape each local U[i] from [out, *rest] to [out, prod(rest)] once on entry to megabatch_orthogonalize_async when flatten=True, and reshape back on exit. Then muon_update_newton_schulz loses both flatten and has_megabatch_dim and only ever sees [*, m, n]; normuon's flatten_rows reshape/restore dance disappears; comm_dim is always -2 on the flatten path instead of shard_dim - ndim; and neuron_variance_buffer's narrow loop reduces to a plain [out, 1]. Every reshape is a view on a contiguous local tensor, so it costs nothing.

Two things push back, and they're why I think the flag is defensible as written: the N == 1 path deliberately hands NS a 2D tensor to hit a GEMM rather than a BMM (your own notes put PE GEMM vs BMM at up to 0.07 apart in bf16, so that dispatch is load-bearing), and keeping variance_neuron at parameter rank/placement matters for DCP. Both survive the refactor — keep the singleton path 2D, keep the state buffer at param rank and reshape at use — so I'd still prefer it.

If the flag stays, please pass param_ndim: int rather than a bool. Both branches then read directly (flatten and param_ndim >= 3, start_dim = X.ndim - param_ndim + 1) and it can't be silently wrong for a future caller that stacks two axes.

Minor: _reject_flattened_convolutions lives in dion2.py and is imported under its private name by nordion2.py. It's shared policy and belongs in megabatch_base.py alongside the other _resolve_*/_validate_* guards.