arcee.py 15.7 KB
Newer Older
Raghav Ravishankar's avatar
Raghav Ravishankar committed
1
2
3
4
5
6
7
8
9
10
11
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Copyright 2023-2025 vLLM Team
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Inference-only Arcee (AFM) model – adds support for ReLU^2 feed-forward
# activation.

from collections.abc import Iterable
12
from itertools import islice
13
from typing import Any
Raghav Ravishankar's avatar
Raghav Ravishankar committed
14
15
16
17
18
19
20
21
22

import torch
from torch import nn
from transformers import LlamaConfig

from vllm.compilation.decorators import support_torch_compile
from vllm.distributed import get_pp_group
from vllm.model_executor.layers.activation import ReLUSquaredActivation
from vllm.model_executor.layers.layernorm import RMSNorm
23
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
Raghav Ravishankar's avatar
Raghav Ravishankar committed
24
25
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
26
27
28
    ParallelLMHead,
    VocabParallelEmbedding,
)
29
from vllm.model_executor.model_loader.weight_utils import (
30
31
32
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
33
34
from vllm.sequence import IntermediateTensors

35
36
37
38
39
40
41
from .interfaces import (
    EagleModelMixin,
    SupportsEagle,
    SupportsEagle3,
    SupportsLoRA,
    SupportsPP,
)
42
43
44
45
46
47
48
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
49
50
51
52
53
54


class ArceeMLP(nn.Module):
    """Feed-forward layer for Arcee using ReLU^2 activation
    (no gating as in LLaMA)."""

55
56
57
58
59
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
60
        quant_config: Any | None = None,
61
62
63
64
        bias: bool = False,
        prefix: str = "",
        reduce_results: bool = True,
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
        super().__init__()
        # Single linear projection up to intermediate size
        # (no separate gate projection)
        self.up_proj = ColumnParallelLinear(
            input_size=hidden_size,
            output_size=intermediate_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.up_proj",
        )
        # Down projection back to hidden size
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            reduce_results=reduce_results,
            prefix=f"{prefix}.down_proj",
        )
        if hidden_act != "relu2":
85
86
87
88
            raise ValueError(
                f"Unsupported activation: {hidden_act}. "
                "Only 'relu2' is supported for AFM."
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        # Define ReLU^2 activation: (ReLU(x))^2 elementwise
        self.act_fn = ReLUSquaredActivation()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.up_proj(x)  # Project to intermediate size
        x = self.act_fn(x)  # Apply ReLU^2 activation elementwise
        x, _ = self.down_proj(x)  # Project back down to hidden size
        return x


class ArceeDecoderLayer(nn.Module):
    """Transformer decoder block for Arcee, with self-attention and
    ReLU^2 MLP."""

103
104
105
    def __init__(
        self,
        config: LlamaConfig,
106
107
        cache_config: Any | None = None,
        quant_config: Any | None = None,
108
109
        prefix: str = "",
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
110
111
        super().__init__()
        self.hidden_size = config.hidden_size
112
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
113
114
        # Determine if attention bias is needed (some variants use bias terms)
        attention_bias = getattr(config, "attention_bias", False) or getattr(
115
116
            config, "bias", False
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
117
118
119
120
121
122
        bias_o_proj = attention_bias
        if hasattr(config, "qkv_bias"):
            attention_bias = config.qkv_bias

        # Self-Attention (using LLaMA's attention structure)
        from vllm.model_executor.models.llama import (
123
124
125
            LlamaAttention,  # import here to avoid circular import
        )

Raghav Ravishankar's avatar
Raghav Ravishankar committed
126
127
128
129
        self.self_attn = LlamaAttention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
130
131
132
            num_kv_heads=getattr(
                config, "num_key_value_heads", config.num_attention_heads
            ),
Raghav Ravishankar's avatar
Raghav Ravishankar committed
133
134
135
136
137
138
139
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            bias=attention_bias,
            bias_o_proj=bias_o_proj,
            cache_config=cache_config,
            prefix=f"{prefix}.self_attn",
            attn_type=getattr(
140
141
                config, "attn_type", "decoder"
            ),  # assume decoder (causal) unless specified
Raghav Ravishankar's avatar
Raghav Ravishankar committed
142
143
144
145
146
147
148
149
150
151
152
        )
        # MLP with ReLU^2 activation
        self.mlp = ArceeMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            bias=getattr(config, "mlp_bias", False),
            prefix=f"{prefix}.mlp",
        )
        # Layer normalization layers (RMSNorm as in LLaMA)
153
154
155
156
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
157
158

    def forward(
159
160
161
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
162
        residual: torch.Tensor | None,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
163
164
165
166
167
168
169
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self-Attention block
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            # Fused residual add + layernorm if supported
170
171
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
172
        # Feed-forward block
173
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
174
175
176
177
178
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


@support_torch_compile
179
class ArceeModel(nn.Module, EagleModelMixin):
Raghav Ravishankar's avatar
Raghav Ravishankar committed
180
181
182
    """The transformer model backbone for Arcee (embedding layer + stacked
    decoder blocks + final norm)."""

183
184
185
186
187
188
189
    def __init__(
        self,
        *,
        vllm_config,
        prefix: str = "",
        layer_type: type[nn.Module] = ArceeDecoderLayer,
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
190
191
192
193
194
195
196
197
198
        super().__init__()
        config: LlamaConfig = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        self.quant_config = quant_config
        self.config = config
        self.vocab_size = config.vocab_size

        # Word embeddings (parallelized if using pipeline parallel)
199
200
201
        if get_pp_group().is_first_rank or (
            config.tie_word_embeddings and get_pp_group().is_last_rank
        ):
Raghav Ravishankar's avatar
Raghav Ravishankar committed
202
203
204
205
206
207
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
            )
        else:
208
            self.embed_tokens = PPMissingLayer()  # placeholder on non-embedding ranks
Raghav Ravishankar's avatar
Raghav Ravishankar committed
209
210
211
212

        # Build decoder layers across pipeline ranks
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
213
214
215
216
217
218
            lambda prefix: layer_type(
                config=config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=prefix,
            ),
Raghav Ravishankar's avatar
Raghav Ravishankar committed
219
220
221
222
223
224
225
226
227
228
            prefix=f"{prefix}.layers",
        )
        # Final RMSNorm on the last pipeline stage
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()

        # Prepare factory for empty intermediate tensors
        # (for pipeline scheduling)
229
230
231
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
232

233
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
234
235
236
237
        return self.embed_tokens(input_ids)

    def forward(
        self,
238
        input_ids: torch.Tensor | None,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
239
        positions: torch.Tensor,
240
241
242
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
243
244
        # Embedding lookup (on first pipeline rank)
        if get_pp_group().is_first_rank:
245
246
247
            hidden_states = (
                inputs_embeds
                if inputs_embeds is not None
248
                else self.embed_input_ids(input_ids)
249
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
250
251
252
            residual = None
        else:
            assert intermediate_tensors is not None, (
253
254
                "IntermediateTensors must be provided for non-first pipeline ranks"
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
255
256
257
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

258
        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
259
        for idx, layer in enumerate(
260
261
            islice(self.layers, self.start_layer, self.end_layer)
        ):
Raghav Ravishankar's avatar
Raghav Ravishankar committed
262
            hidden_states, residual = layer(positions, hidden_states, residual)
263
264
265
            self._maybe_add_hidden_state(
                aux_hidden_states, idx + 1, hidden_states, residual
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
266
267
268

        if not get_pp_group().is_last_rank:
            # Send intermediate results to the next pipeline stage
269
270
271
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
272
273
274
275
276
277
        # On last rank: apply final layer norm
        hidden_states, _ = self.norm(hidden_states, residual)
        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
        return hidden_states

278
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
279
280
281
282
283
284
285
286
287
288
289
290
291
        """Load weights, mapping q/k/v projections to fused qkv_proj."""
        stacked_params_mapping = [
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
        ]

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
292
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
293
294
                continue

295
296
297
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
298
                param = params_dict[scale_name]
299
300
301
302
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
303
304
305
306
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue

307
            if "scale" in name or "zero_point" in name:
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
                remapped_name = maybe_remap_kv_scale_name(name, params_dict)
                if remapped_name is None:
                    continue
                name = remapped_name

            mapped = False
            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)

                if name.endswith(".bias") and name not in params_dict:
                    mapped = True
                    break

                if is_pp_missing_parameter(name, self):
                    mapped = True
                    break

                param = params_dict[name]
                weight_loader = param.weight_loader  # type: ignore[attr-defined]
                weight_loader(param, loaded_weight, shard_id)
                loaded_params.add(name)
                mapped = True
                break

            if mapped:
                continue

            if name.endswith(".bias") and name not in params_dict:
                continue

            if is_pp_missing_parameter(name, self):
                continue

            param = params_dict[name]
345
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
346
347
348
349
350
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params

Raghav Ravishankar's avatar
Raghav Ravishankar committed
351

352
353
354
class ArceeForCausalLM(
    nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
):
Raghav Ravishankar's avatar
Raghav Ravishankar committed
355
356
    """Arcee Model for causal language modeling, integrated with vLLM
    runtime."""
357

358
    # Map fused module names to their submodule components
Raghav Ravishankar's avatar
Raghav Ravishankar committed
359
360
361
362
363
364
365
366
367
368
369
    # (for quantization and LoRA)
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
    }

    def __init__(self, *, vllm_config, prefix: str = "") -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.config = config

        # Initialize the inner Transformer model (ArceeModel)
370
        self.model = ArceeModel(vllm_config=vllm_config, prefix=f"{prefix}.model")
Raghav Ravishankar's avatar
Raghav Ravishankar committed
371
372
373
374
375
376
        # On the last pipeline stage, set up the LM head and logits processor
        if get_pp_group().is_last_rank:
            # Determine vocabulary size (including any LoRA extra tokens
            # for padded LM head)

            self.lm_head = ParallelLMHead(
377
                config.vocab_size,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
378
379
380
381
382
383
384
                config.hidden_size,
                quant_config=vllm_config.quant_config,
                bias=getattr(config, "lm_head_bias", False),
                prefix=f"{prefix}.lm_head",
            )
            if config.tie_word_embeddings:
                # Tie output weights with input embedding matrix
385
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
386
            logit_scale = getattr(config, "logit_scale", 1.0)
387
            self.logits_processor = LogitsProcessor(
388
                config.vocab_size, scale=logit_scale
389
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
390
391
392
        else:
            # Placeholder for lm_head on non-last ranks
            self.lm_head = PPMissingLayer()
393

Raghav Ravishankar's avatar
Raghav Ravishankar committed
394
        self.make_empty_intermediate_tensors = (
395
396
            self.model.make_empty_intermediate_tensors
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
397
398
399

    def forward(
        self,
zhuwenwen's avatar
zhuwenwen committed
400
        input_ids: torch.Tensor,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
401
        positions: torch.Tensor,
402
403
404
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
405
406
407
408
409
410
        model_output = self.model(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
411
412
        return model_output

413
    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
414
        # Compute final logits from hidden states (last pipeline rank only)
415
        logits = self.logits_processor(self.lm_head, hidden_states)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
416
417
        return logits

418
419
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
420

421
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
422
423
424
425
        """Load weights into the model (delegates to inner model and handles
        tied embeddings)."""
        loader = AutoWeightsLoader(
            self,
426
427
428
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
            skip_substrs=["gate_proj"],
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
429
430
431
        # AutoWeightLoader handles weight name remapping, including fusing
        # separate q_proj, k_proj, v_proj into qkv_proj
        return loader.load_weights(weights)