Skip to content

vllm.model_executor.models.llama_eagle

logger module-attribute 

logger = init_logger(__name__)

EagleLlamaForCausalLM

Bases: LlamaForCausalLM

Source code in vllm/model_executor/models/llama_eagle.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
class EagleLlamaForCausalLM(LlamaForCausalLM):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        nn.Module.__init__(self)
        self.config = vllm_config. \
            speculative_config.draft_model_config.hf_config
        target_layer_num = vllm_config.model_config.get_num_layers(
            vllm_config.parallel_config)
        self.model = LlamaModel(vllm_config=vllm_config,
                                prefix="model",
                                start_layer_id=target_layer_num)

        logit_scale = getattr(self.config, "logit_scale", 1.0)
        self.logits_processor = LogitsProcessor(self.config.vocab_size,
                                                scale=logit_scale)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return self.model(input_ids, positions, hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=None,
        )

        model_weights = {}
        for name, loaded_weight in weights:
            if "lm_head" not in name:
                name = "model." + name
            model_weights[name] = loaded_weight
        loader.load_weights(model_weights.items())

config instance-attribute 

config = hf_config

logits_processor instance-attribute 

logits_processor = LogitsProcessor(
    vocab_size, scale=logit_scale
)

model instance-attribute 

model = LlamaModel(
    vllm_config=vllm_config,
    prefix="model",
    start_layer_id=target_layer_num,
)

__init__ 

__init__(*, vllm_config: VllmConfig, prefix: str = '')
Source code in vllm/model_executor/models/llama_eagle.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
    nn.Module.__init__(self)
    self.config = vllm_config. \
        speculative_config.draft_model_config.hf_config
    target_layer_num = vllm_config.model_config.get_num_layers(
        vllm_config.parallel_config)
    self.model = LlamaModel(vllm_config=vllm_config,
                            prefix="model",
                            start_layer_id=target_layer_num)

    logit_scale = getattr(self.config, "logit_scale", 1.0)
    self.logits_processor = LogitsProcessor(self.config.vocab_size,
                                            scale=logit_scale)

forward 

forward(
    input_ids: Tensor,
    positions: Tensor,
    hidden_states: Tensor,
) -> tuple[Tensor, Tensor]
Source code in vllm/model_executor/models/llama_eagle.py 
146
147
148
149
150
151
152
def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    return self.model(input_ids, positions, hidden_states)

load_weights 

load_weights(weights: Iterable[tuple[str, Tensor]])
Source code in vllm/model_executor/models/llama_eagle.py 
154
155
156
157
158
159
160
161
162
163
164
165
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
    loader = AutoWeightsLoader(
        self,
        skip_prefixes=None,
    )

    model_weights = {}
    for name, loaded_weight in weights:
        if "lm_head" not in name:
            name = "model." + name
        model_weights[name] = loaded_weight
    loader.load_weights(model_weights.items())

LlamaModel

Bases: Module

Source code in vllm/model_executor/models/llama_eagle.py 
 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
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@support_torch_compile
class LlamaModel(nn.Module):

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        start_layer_id: int = 0,
    ) -> None:
        super().__init__()
        self.config = vllm_config. \
            speculative_config.draft_model_config.hf_config
        self.vocab_size = self.config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.config.vocab_size,
            self.config.hidden_size,
            prefix=maybe_prefix(prefix, "embed_tokens"),
        )

        self.layers = nn.ModuleList([
            LlamaDecoderLayer(
                self.config,
                i == 0,
                prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
            ) for i in range(self.config.num_hidden_layers)
        ])
        self.fc = torch.nn.Linear(self.config.hidden_size * 2,
                                  self.config.hidden_size,
                                  bias=False)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        input_embeds = self.embed_tokens(input_ids)
        hidden_states = self.fc(
            torch.cat((input_embeds, hidden_states), dim=-1))
        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(
                positions,
                hidden_states,
                residual,
            )
        hidden_states = hidden_states + residual
        return hidden_states, hidden_states

    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:

                # if PP disabled then draft will share embed with target
                if get_pp_group().world_size == 1 and \
                    "embed_tokens." in name:
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

config instance-attribute 

config = hf_config

embed_tokens instance-attribute 

embed_tokens = VocabParallelEmbedding(
    vocab_size,
    hidden_size,
    prefix=maybe_prefix(prefix, "embed_tokens"),
)

fc instance-attribute 

fc = Linear(hidden_size * 2, hidden_size, bias=False)

layers instance-attribute 

layers = ModuleList(
    [
        LlamaDecoderLayer(
            config,
            i == 0,
            prefix=maybe_prefix(
                prefix, f"layers.{i + start_layer_id}"
            ),
        )
        for i in range(num_hidden_layers)
    ]
)

vocab_size instance-attribute 

vocab_size = vocab_size

__init__ 

__init__(
    *,
    vllm_config: VllmConfig,
    prefix: str = "",
    start_layer_id: int = 0,
) -> None
Source code in vllm/model_executor/models/llama_eagle.py 
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
def __init__(
    self,
    *,
    vllm_config: VllmConfig,
    prefix: str = "",
    start_layer_id: int = 0,
) -> None:
    super().__init__()
    self.config = vllm_config. \
        speculative_config.draft_model_config.hf_config
    self.vocab_size = self.config.vocab_size

    self.embed_tokens = VocabParallelEmbedding(
        self.config.vocab_size,
        self.config.hidden_size,
        prefix=maybe_prefix(prefix, "embed_tokens"),
    )

    self.layers = nn.ModuleList([
        LlamaDecoderLayer(
            self.config,
            i == 0,
            prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
        ) for i in range(self.config.num_hidden_layers)
    ])
    self.fc = torch.nn.Linear(self.config.hidden_size * 2,
                              self.config.hidden_size,
                              bias=False)

forward 

forward(
    input_ids: Tensor,
    positions: Tensor,
    hidden_states: Tensor,
) -> tuple[Tensor, Tensor]
Source code in vllm/model_executor/models/llama_eagle.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    input_embeds = self.embed_tokens(input_ids)
    hidden_states = self.fc(
        torch.cat((input_embeds, hidden_states), dim=-1))
    residual = None
    for layer in self.layers:
        hidden_states, residual = layer(
            positions,
            hidden_states,
            residual,
        )
    hidden_states = hidden_states + residual
    return hidden_states, hidden_states

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]
Source code in vllm/model_executor/models/llama_eagle.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def load_weights(self, weights: Iterable[tuple[str,
                                               torch.Tensor]]) -> set[str]:
    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        (".qkv_proj", ".q_proj", "q"),
        (".qkv_proj", ".k_proj", "k"),
        (".qkv_proj", ".v_proj", "v"),
        (".gate_up_proj", ".gate_proj", 0),
        (".gate_up_proj", ".up_proj", 1),
    ]
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    for name, loaded_weight in weights:
        for param_name, weight_name, shard_id in stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)
            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:

            # if PP disabled then draft will share embed with target
            if get_pp_group().world_size == 1 and \
                "embed_tokens." in name:
                continue

            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
    return loaded_params