[flyingdoubleG]

An expert's K-FAC covariances (CovarianceCollector) only sum the positions routed to it: collection_mask(module) selects the expert's _positions grid. The eigendecomposition then divides every module's sums by total_processed, which counts all gradient-carrying positions in the dataset. As a result, each expert's activation and gradient "means" are too small by a factor of p = (positions routed to it) / total_processed. Its uncorrected eigenvalue grid, λ_G ⊗ λ_A · total_processed / num_documents, ends up p times too small.

damped_inverse and the other inversions damp relative to each module's own mean eigenvalue, so this does not wash out. An expert's preconditioned gradient comes out roughly 1/p times too large compared with the dense modules. With top-2 of 64 experts that is about 32×. EK-FAC is unaffected: the eigenvectors don't depend on the scale, and the corrected eigenvalues are per-document sums divided by num_documents.

Fix

• CovarianceCollector counts the positions each module's covariances summed (num_positions). In a process group it all-reduces them at teardown.
• fit_factored_hessians passes those counts to eigendecompose_owned and save_uncorrected_eigenvalues. Each module is now divided by its own count instead of total_processed.
• An expert that received no positions has zero covariances, so its count is clamped to 1.
• Dense modules are unchanged, since their count equals total_processed.

Check

Each document has 2 tokens, so it has one gradient-carrying position. A module's per-document gradient there is g aᵀ, which makes the EK-FAC corrected eigenvalues sum to exactly E[|a|² |g|²]. The K-FAC grid sums to E[|a|²] E[|g|²], so the two should roughly agree. K-FAC / exact trace on a tiny Mixtral (4 experts, top-1, 64 docs), from the new test:

module	before	after
attention q/k/v/o	1.00–1.08	1.00–1.08 (unchanged)
experts, layer 0	0.13–0.44	0.89–1.05
experts, layer 1	0.06–0.40	1.00–1.09

The "before" ratios match each expert's routed fraction.

Tests

• tests/test_moe.py::test_kfac_normalizes_experts_by_their_routed_positions runs fit_factored_hessians on that model and asserts the ratio is in (0.7, 1.4) for every module. It fails on main and passes with this change.
• test_gloo_ranks_match_one_process_on_cpu now also checks that the all-reduced num_positions equal the single-process counts.

Not covered

compute_eigendecomposition (the file-based path used by tkfac/shampoo and approximate unrolling's segment aggregation) still divides by total_processed. Those collectors would need to save their counts alongside the covariance shards; I can follow up if you'd like.

Tested:

• pytest tests/test_moe.py tests/ekfac_tests: 56 passed, 15 skipped on CPU (torch 2.9.1).
• The four gloo process-group tests can't start with that Windows build of torch, so I ran them separately under torch 2.11: 14 passed.
• black, ruff and pyright: no new errors.

[luciaquirke]

Fabulous fix, thank you very much & LGTM!