nano-embedding

nano-vllm/nanovllm/layers/embed_head.py · bb823b3 · L9–66 · 固定 commit 上游

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

9class VocabParallelEmbedding(nn.Module):
10
11    def __init__(
12        self,
13        num_embeddings: int,
14        embedding_dim: int,
15    ):
16        super().__init__()
17        self.tp_rank = dist.get_rank()
18        self.tp_size = dist.get_world_size()
19        assert num_embeddings % self.tp_size == 0
20        self.num_embeddings = num_embeddings
21        self.num_embeddings_per_partition = self.num_embeddings // self.tp_size
22        self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank
23        self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition
24        self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim))
25        self.weight.weight_loader = self.weight_loader
26
27    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
28        param_data = param.data
29        shard_size = param_data.size(0)
30        start_idx = self.tp_rank * shard_size
31        loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
32        param_data.copy_(loaded_weight)
33
34    def forward(self, x: torch.Tensor):
35        if self.tp_size > 1:
36            mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
37            x = mask * (x - self.vocab_start_idx)
38        y = F.embedding(x, self.weight)
39        if self.tp_size > 1:
40            y = mask.unsqueeze(1) * y
41            dist.all_reduce(y)
42        return y
43
44
45class ParallelLMHead(VocabParallelEmbedding):
46
47    def __init__(
48        self,
49        num_embeddings: int,
50        embedding_dim: int,
51        bias: bool = False,
52    ):
53        assert not bias
54        super().__init__(num_embeddings, embedding_dim)
55
56    def forward(self, x: torch.Tensor):
57        context = get_context()
58        if context.is_prefill:
59            last_indices = context.cu_seqlens_q[1:] - 1
60            x = x[last_indices].contiguous()
61        logits = F.linear(x, self.weight)
62        if self.tp_size > 1:
63            all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None
64            dist.gather(logits, all_logits, 0)
65            logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None
66        return logits