[BUG] RMSNorm implementation seems wrong

[laclouis5]

Describe the bug

the RMN norm implementation seems to use the variance (torch.var) instead of the root mean square:

pytorch-image-models/timm/layers/fast_norm.py

Lines 111 to 130 in 790decc

	def rms_norm(
	x: torch.Tensor,
	normalized_shape: List[int],
	weight: Optional[torch.Tensor] = None,
	eps: float = 1e-5,
	):
	norm_ndim = len(normalized_shape)
	if torch.jit.is_scripting():
	# ndim = len(x.shape)
	# dims = list(range(ndim - norm_ndim, ndim)) # this doesn't work on pytorch <= 1.13.x
	# NOTE -ve dims cause torchscript to crash in some cases, out of options to work around
	assert norm_ndim == 1
	v = torch.var(x, dim=-1).unsqueeze(-1) # ts crashes with -ve dim + keepdim=True
	else:
	dims = tuple(range(-1, -norm_ndim - 1, -1))
	v = torch.var(x, dim=dims, keepdim=True)
	x = x * torch.rsqrt(v + eps)
	if weight is not None:
	x = x * weight
	return x

This yields a different result from the PyTorch RMSNorm implementation in my tests.

Expected behavior
Same results between native Pytorch operator and timm operator.

Desktop (please complete the following information):

• OS: macOS
• This repository version: 1.0.12
• PyTorch version: 2.5

[rwightman]

@laclouis5 yes, it is wrong, I have a fix in the works w/ aimv2 model support as I need this to be correct ... the current impl is a valid norm, layer norm without the centering (mean subtraction) and no bias, it doesn't match the standard RMSNorm formulation.

[laclouis5]

@rwightman Thanks for the quick update.

Will the fix correct the current RMSNorm and implement the alternative behavior in another class / parameter or we have to expect timm's RMSNorm to diverge from PyTorch implementation?

[rwightman]

@laclouis5 I'll make the RmsNorm class match the standard / PyTorch form. Might try and dispatch to torch F.rms_norm when it's there (only >= 2.5 I think). And will fix the timm impl for older torch versions.

The current RmsNorm, I'll rename to SimpleNorm I think, for lack of a better name. It might have a name somewhere but I couldn't find something standard. I have only 2 models I trained with it, they worked well compared to LayerNorm.

[rwightman]

@laclouis5 it's merged now, for the aimv2 testing the corrected RmsProp was exact match in pytorch 2.4 and 2.5 to reference.

Slightly disappointing, I was hoping the impl they added in torch 2.5 would be fast but there doesn't appear to be a native cuda kernel. LayerNorm is faster.

[laclouis5]

Great!