Fix AdafactorBigVision dropping the factored row normalization for tall matrices - #2749

What

AdafactorBigVision silently drops the factored row normalization for any 2D parameter whose larger dimension comes first (e.g. an [out, in] weight with out > in, such as an MLP up-projection). The update then depends on the orientation of the weight matrix.

import torch
from timm.optim.adafactor_bv import AdafactorBigVision

grad = torch.randn(30, 20, dtype=torch.double)

def step(shape, g):
    p = torch.zeros(shape, dtype=torch.double, requires_grad=True)
    AdafactorBigVision([p], lr=1.0, momentum=None, eps=1e-30, weight_decay=0.0).step()  # after p.grad = g
    return p.detach()

# update(W) and update(W.T) should be transposes of each other; they are not:
# max|update(30,20) - correct| = 0.771, and update(30,20) == grad * col_factor only (row_factor == 1)

Why

In _single_tensor_adafactor, the row factor is exp_avg_sq_r / row_col_mean, where row_col_mean is meant to be the overall mean of the row accumulator:

exp_avg_sq_r.lerp_(grad_sqr.mean(dim=dr, keepdim=True), one_minus_beta2_t)   # keepdim -> full ndim retained
...
reduce_dc = dc - 1 if dc > dr else dc
row_col_mean = exp_avg_sq_r.mean(dim=reduce_dc, keepdim=True)

The dc - 1 if dc > dr shift is ported from the big_vision / optax source, where the row accumulator is stored with the reduced axis removed (so a later index has to be shifted down by one). timm keeps that axis via keepdim=True, so exp_avg_sq_r retains full rank and the shift is wrong: when dc > dr it points at an axis that is already size 1. Then row_col_mean == exp_avg_sq_r, row_factor = (exp_avg_sq_r / row_col_mean).rsqrt() is identically 1, and the row normalization vanishes.

Verified by running the optimizer: for grad shape (30, 20) the update equals grad * col_factor with row_factor == 1 exactly (diff 0.0), deviates from the correctly-factored update by max abs 0.771, and breaks transpose-equivariance by 0.771 — while the (20, 30) orientation (which hits the dc < dr no-shift branch) matches the reference exactly. No test caught it because test_adafactor only checks generic convergence, never the factored value.

Fix

exp_avg_sq_r keeps full ndim, so reduce over dc directly (the retained axis) — no optax axis-removal shift. This restores a nonzero row factor and makes the factored update transpose-equivariant again. The dc < dr orientation is unchanged (there reduce_dc already equaled dc).

Tests

Adds test_adafactor_bv_factored_row_normalization, asserting update(W.T) == update(W).T for a factored shape (both orientations). It fails on main and passes with the fix. Full tests/test_optim.py has no new failures (the two pre-existing failures — test_kron[kron] torch.compile/inductor on this platform, and the csgdp full-run convergence flake — reproduce identically on clean main).

The multi-tensor (foreach=True) Adafactor path is an unimplemented stub, so this single-tensor path is the only affected site.