Xây dựng Tầng Feed-Forward SwiGLU của Transformer Hiện Đại (LLaMA)

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

• Hiểu bước tiến hóa kiến trúc FFN từ Transformer cổ điển sang Transformer hiện đại:
• FFN cổ điển (Vaswani 2017): FFN(x) = ReLU(x W₁ + b₁) W₂ + b₂.
• FFN SwiGLU (Noam Shazeer 2020, dùng trong LLaMA, Mistral, Gemma):

SwiGLU(x) = ( SiLU(x W_gate) ⊙ (x W_up) ) W_down Trong đó:

• SiLU(z) = z · σ(z) = z1 + e^-z (còn gọi là hàm Swish).
• W_gate ∈ R^{D × D_ffn}: Chiếu cổng (gate projection).
• W_up ∈ R^{D × D_ffn}: Chiếu mở rộng (up projection).
• W_down ∈ R^{D_ffn × D}: Chiếu thu gọn (down projection).
• Toàn bộ các lớp tuyến tính đều không sử dụng bias (bias=False).

Yêu cầu

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

class SwiGLUFeedForward(torch.nn.Module):
    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        pass

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

• Khởi tạo 3 lớp tuyến tính không bias: w_gate, w_up, w_down.
• Nhận đầu vào x ∈ R^{B × S × D}, trả về tensor cùng kích thước (B, S, D).

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ả.