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
35
from vllm.sequence import IntermediateTensors

from .interfaces import SupportsLoRA, SupportsPP
36
37
38
39
40
41
42
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
43
44
45
46
47
48


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

49
50
51
52
53
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
54
        quant_config: Any | None = None,
55
56
57
58
        bias: bool = False,
        prefix: str = "",
        reduce_results: bool = True,
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
        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":
79
80
81
82
            raise ValueError(
                f"Unsupported activation: {hidden_act}. "
                "Only 'relu2' is supported for AFM."
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
        # 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."""

97
98
99
    def __init__(
        self,
        config: LlamaConfig,
100
101
        cache_config: Any | None = None,
        quant_config: Any | None = None,
102
103
        prefix: str = "",
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
104
105
        super().__init__()
        self.hidden_size = config.hidden_size
106
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
107
108
        # Determine if attention bias is needed (some variants use bias terms)
        attention_bias = getattr(config, "attention_bias", False) or getattr(
109
110
            config, "bias", False
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
111
112
113
114
115
116
        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 (
117
118
119
            LlamaAttention,  # import here to avoid circular import
        )

Raghav Ravishankar's avatar
Raghav Ravishankar committed
120
121
122
123
        self.self_attn = LlamaAttention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
124
125
126
            num_kv_heads=getattr(
                config, "num_key_value_heads", config.num_attention_heads
            ),
Raghav Ravishankar's avatar
Raghav Ravishankar committed
127
128
129
130
131
132
133
            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(
134
135
                config, "attn_type", "decoder"
            ),  # assume decoder (causal) unless specified
Raghav Ravishankar's avatar
Raghav Ravishankar committed
136
137
138
139
140
141
142
143
144
145
146
        )
        # 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)
147
148
149
150
        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
151
152

    def forward(
153
154
155
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
156
        residual: torch.Tensor | None,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
157
158
159
160
161
162
163
    ) -> 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
164
165
            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
166
        # Feed-forward block
167
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
168
169
170
171
172
173
174
175
176
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


@support_torch_compile
class ArceeModel(nn.Module):
    """The transformer model backbone for Arcee (embedding layer + stacked
    decoder blocks + final norm)."""

177
178
179
180
181
182
183
    def __init__(
        self,
        *,
        vllm_config,
        prefix: str = "",
        layer_type: type[nn.Module] = ArceeDecoderLayer,
    ) -> None:
Raghav Ravishankar's avatar
Raghav Ravishankar committed
184
185
186
187
188
189
190
191
192
        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)
193
194
195
        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
196
197
198
199
200
201
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
            )
        else:
202
            self.embed_tokens = PPMissingLayer()  # placeholder on non-embedding ranks
Raghav Ravishankar's avatar
Raghav Ravishankar committed
203
204
205
206

        # Build decoder layers across pipeline ranks
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
207
208
209
210
211
212
            lambda prefix: layer_type(
                config=config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=prefix,
            ),
Raghav Ravishankar's avatar
Raghav Ravishankar committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
            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()

        # For optional capturing of intermediate hidden states
        # (not used by default)
        self.aux_hidden_state_layers: tuple[int, ...] = tuple()

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

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

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

        aux_hidden_states: list[torch.Tensor] = []
        for idx, layer in enumerate(
258
259
            islice(self.layers, self.start_layer, self.end_layer)
        ):
Raghav Ravishankar's avatar
Raghav Ravishankar committed
260
261
            if idx in self.aux_hidden_state_layers:
                aux_hidden_states.append(
262
263
                    hidden_states + residual
                )  # capture pre-layer hidden state if needed
Raghav Ravishankar's avatar
Raghav Ravishankar committed
264
265
266
267
            hidden_states, residual = layer(positions, hidden_states, residual)

        if not get_pp_group().is_last_rank:
            # Send intermediate results to the next pipeline stage
268
269
270
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
271
272
273
274
275
276
        # 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

277
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
278
279
280
281
282
283
284
285
286
287
288
289
290
        """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
291
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
292
293
                continue

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

306
            if "scale" in name or "zero_point" in name:
307
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
                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]
344
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
345
346
347
348
349
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params

Raghav Ravishankar's avatar
Raghav Ravishankar committed
350
351
352
353

class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
    """Arcee Model for causal language modeling, integrated with vLLM
    runtime."""
354

355
    # Map fused module names to their submodule components
Raghav Ravishankar's avatar
Raghav Ravishankar committed
356
357
358
359
360
361
362
363
364
365
366
    # (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)
367
        self.model = ArceeModel(vllm_config=vllm_config, prefix=f"{prefix}.model")
Raghav Ravishankar's avatar
Raghav Ravishankar committed
368
369
370
371
372
373
        # 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(
374
                config.vocab_size,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
375
376
377
378
379
380
381
                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
382
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
Raghav Ravishankar's avatar
Raghav Ravishankar committed
383
            logit_scale = getattr(config, "logit_scale", 1.0)
384
            self.logits_processor = LogitsProcessor(
385
                config.vocab_size, scale=logit_scale
386
            )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
387
388
389
        else:
            # Placeholder for lm_head on non-last ranks
            self.lm_head = PPMissingLayer()
390

Raghav Ravishankar's avatar
Raghav Ravishankar committed
391
        self.make_empty_intermediate_tensors = (
392
393
            self.model.make_empty_intermediate_tensors
        )
Raghav Ravishankar's avatar
Raghav Ravishankar committed
394
395
396

    def forward(
        self,
397
        input_ids: torch.Tensor | None,
Raghav Ravishankar's avatar
Raghav Ravishankar committed
398
        positions: torch.Tensor,
399
400
401
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
402
403
404
405
406
407
        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
408
409
        return model_output

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

415
416
    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
417

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