nano-lm-head

nano-vllm/nanovllm/models/qwen3.py · bb823b3 · L186–216 · 固定 commit 上游

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

186class Qwen3ForCausalLM(nn.Module):
187    packed_modules_mapping = {
188        "q_proj": ("qkv_proj", "q"),
189        "k_proj": ("qkv_proj", "k"),
190        "v_proj": ("qkv_proj", "v"),
191        "gate_proj": ("gate_up_proj", 0),
192        "up_proj": ("gate_up_proj", 1),
193    }
194
195    def __init__(
196        self,
197        config: Qwen3Config
198    ) -> None:
199        super().__init__()
200        self.model = Qwen3Model(config)
201        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
202        if config.tie_word_embeddings:
203            self.lm_head.weight.data = self.model.embed_tokens.weight.data
204
205    def forward(
206        self,
207        input_ids: torch.Tensor,
208        positions: torch.Tensor,
209    ) -> torch.Tensor:
210        return self.model(input_ids, positions)
211
212    def compute_logits(
213        self,
214        hidden_states: torch.Tensor,
215    ) -> torch.Tensor:
216        return self.lm_head(hidden_states)