Skip to content

vllm.model_executor.models.nemotron_parse3_5_mtp

One-step auxiliary MTP proposer for Nemotron Parse 3.5.

Classes:

  • ParseAuxMTPModel –

    Draft one token using Parse 3.5's trained auxiliary projections.

ParseAuxMTPModel

Bases: Module

Draft one token using Parse 3.5's trained auxiliary projections.

Source code in vllm/model_executor/models/nemotron_parse3_5_mtp.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class ParseAuxMTPModel(nn.Module):
    """Draft one token using Parse 3.5's trained auxiliary projections."""

    has_own_embed_tokens = False
    has_own_lm_head = False

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "decoder.embed_tokens.": "model.embed_tokens.",
            "decoder.extra_heads.0.": "model.extra_head.",
            "decoder.extra_proj.0.": "model.extra_proj.",
        }
    )

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        del prefix
        self.config = vllm_config.model_config.hf_config
        decoder_config = self.config.decoder
        vocab_size = int(decoder_config.vocab_size)
        hidden_size = int(decoder_config.d_model)

        self.model = _ParseAuxMTPInner(vocab_size, hidden_size)
        self.lm_head = ParallelLMHead(vocab_size, hidden_size)
        self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
        self.logits_processor = LogitsProcessor(vocab_size, vocab_size)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        del positions, inputs_embeds, kwargs
        return self.model(input_ids, hidden_states)

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.logits_processor(self.lm_head, hidden_states)

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_tokens(input_ids)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        def mtp_weights() -> Iterable[tuple[str, torch.Tensor]]:
            for name, weight in weights:
                if name == "decoder.embed_tokens.weight" or name.startswith(
                    ("decoder.extra_heads.0.", "decoder.extra_proj.0.")
                ):
                    yield name, weight

        loader = AutoWeightsLoader(self)
        return loader.load_weights(mtp_weights(), mapper=self.hf_to_vllm_mapper)