RMSNorm

import torch
import torch.nn as nn

# Enable KuDNN.
torch._C._set_kdnn_enabled(True)

# Input data: (batch_size, seq_len, hidden_dim)
input = torch.randn(2, 5, 10)

# Construct the RMSNorm layer.
# Specify eps as 1e-5 and the data type as fp16 to normalize the last dimension.
rms_norm = nn.RMSNorm(10, eps = 1e-5, dtype = torch.float16)  
output = rms_norm(input)

# Verification: The output root mean square is close to the scaling value.
print("RMSNorm output shape: ", output.shape) # Same as the input.
print("RMSNorm last dimension RMS:", torch.sqrt(torch.mean(output.pow(2), dim=-1))) # Close to the scaling value.