Chuẩn Hóa Căn Bậc Hai Trung Bình (RMSNorm - Zhang & Sennrich) LLaMA

Mục tiêu kiến thức

• LayerNorm truyền thống thực hiện 2 thao tác:
• Trừ giá trị trung bình μ (re-centering).
• Chia cho độ lệch chuẩn σ (re-scaling).
• Zhang & Sennrich (NeurIPS 2019) chứng minh rằng lợi ích chính của LayerNorm nằm ở việc chia cho độ lớn (re-scaling) chứ không phải phép dịch tâm trừ μ.
• RMSNorm loại bỏ hoàn toàn phép trừ μ:

RMS(x) = √(1d ∑_i=1^d x_i² + ε) x_i = x_iRMS(x) · γ_i

• Lợi ích: Tiết kiệm 7% đến 50% thời gian chạy của tầng Normalization trên phần cứng, được ứng dụng trong hầu hết các mô hình hiện đại (LLaMA-1/2/3, Gemma, Mistral, Qwen).

Yêu cầu

Xây dựng lớp RMSNorm(torch.nn.Module):

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass

• Khởi tạo self.weight: nn.Parameter kích thước (dim,) khởi tạo toàn số 1.
• Áp dụng chuẩn hóa trên trục chiều cuối cùng (dim=-1).

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.