nano-rmsnorm

nano-vllm/nanovllm/layers/layernorm.py · bb823b3 · L5–50 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

5class RMSNorm(nn.Module):
6
7    def __init__(
8        self,
9        hidden_size: int,
10        eps: float = 1e-6,
11    ) -> None:
12        super().__init__()
13        self.eps = eps
14        self.weight = nn.Parameter(torch.ones(hidden_size))
15
16    @torch.compile
17    def rms_forward(
18        self,
19        x: torch.Tensor,
20    ) -> torch.Tensor:
21        orig_dtype = x.dtype
22        x = x.float()
23        var = x.pow(2).mean(dim=-1, keepdim=True)
24        x.mul_(torch.rsqrt(var + self.eps))
25        x = x.to(orig_dtype).mul_(self.weight)
26        return x
27
28    @torch.compile
29    def add_rms_forward(
30        self,
31        x: torch.Tensor,
32        residual: torch.Tensor,
33    ) -> tuple[torch.Tensor, torch.Tensor]:
34        orig_dtype = x.dtype
35        x = x.float().add_(residual.float())
36        residual = x.to(orig_dtype)
37        var = x.pow(2).mean(dim=-1, keepdim=True)
38        x.mul_(torch.rsqrt(var + self.eps))
39        x = x.to(orig_dtype).mul_(self.weight)
40        return x, residual
41
42    def forward(
43        self,
44        x: torch.Tensor,
45        residual: torch.Tensor | None = None,
46    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
47        if residual is None:
48            return self.rms_forward(x)
49        else:
50            return self.add_rms_forward(x, residual)