Using adjust_lr in combination with split_sizes will give the wrong learning rate.

[skaae]

I have noticed that if i enable split_sizes in NorMuon in order to split a gated MLP and also enable adjust_lr="rmsnorm" my training runs will initially converge faster but then diverge.

I think there might be a bug here:

dion/dion/normuon.py

Line 329 in 0a059a4

adjusted_lr = adjust_lr_rms_norm(lr, X[0].shape, flatten=flatten)

where the the learning rate adjustment does not take into account the reduced sizes specified with split_sizes? This would result in too high learning rate for split matrices.
I don't completely follow the normuon_normalization_stacked but i don't think the gradients are normalized here?

BR Søren