@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 19332c0479

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Model] Systematic support for fp32 head, pooling models part (#23810)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Sep 9, 2025 at 22:29 UTC 19332c0479bbc8e5389afd53e9d15a028bdaf5f8

14 files changed +166 -61

tests/models/language/pooling/mteb_utils.py

+31 -6

@@ -9,6 +9,7 @@ import mteb

9 import numpy as np

10 import pytest

11 import requests

12 +import torch

14 from tests.models.utils import (EmbedModelInfo, RerankModelInfo,

15 check_embeddings_close)

@@ -165,16 +166,19 @@ def mteb_test_embed_models(hf_runner,

166 vllm_extra_kwargs=None,

167 hf_model_callback=None,

168 atol=MTEB_EMBED_TOL):

169 + # A model family has many models with the same architecture,

170 + # and we don't need to test each one.

171 if not model_info.enable_test:

169 - # A model family has many models with the same architecture,

170 - # and we don't need to test each one.

172 pytest.skip("Skipping test.")

173

173 - example_prompts = ["The chef prepared a delicious meal."]

174 + # Test embed_dims, isnan and whether to use normalize

175 + example_prompts = ["The chef prepared a delicious meal." * 1000]

176

177 + # Allow vllm to test using the given dtype, such as float32

178 vllm_extra_kwargs = vllm_extra_kwargs or {}

179 vllm_extra_kwargs["dtype"] = model_info.dtype

180

181 + # Allow vllm to test using hf_overrides

182 if model_info.hf_overrides is not None:

183 vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

184

@@ -186,21 +190,32 @@ def mteb_test_embed_models(hf_runner,

190

191 model_config = vllm_model.llm.llm_engine.model_config

192

193 + # Confirm whether vllm is using the correct architecture

194 if model_info.architecture:

195 assert model_info.architecture in model_config.architectures

196 +

197 + # Confirm whether vllm uses the correct default_pooling_type, which

198 + # relates to whether chunked prefill and prefix caching are enabled

199 assert (model_config._model_info.default_pooling_type ==

200 model_info.default_pooling_type)

201

202 vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),

203 MTEB_EMBED_TASKS)

204 vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype

197 - vllm_outputs = vllm_model.embed(example_prompts)

205

206 + # Test embed_dims, isnan and whether to use normalize

207 + vllm_outputs = vllm_model.embed(example_prompts,

208 + truncate_prompt_tokens=-1)

209 + assert not torch.any(torch.isnan(torch.tensor(vllm_outputs)))

210 +

211 + # Accelerate mteb test by setting

212 + # SentenceTransformers mteb score to a constant

213 if model_info.mteb_score is None:

214 with hf_runner(model_info.name,

215 is_sentence_transformer=True,

216 dtype="float32") as hf_model:

217

218 + # e.g. setting default parameters for the encode method of hf_runner

219 if hf_model_callback is not None:

220 hf_model_callback(hf_model)

221

@@ -299,14 +314,16 @@ def mteb_test_rerank_models(hf_runner,

314 hf_model_callback=None,

315 vllm_mteb_encoder=VllmMtebEncoder,

316 atol=MTEB_RERANK_TOL):

317 + # A model family has many models with the same architecture,

318 + # and we don't need to test each one.

319 if not model_info.enable_test:

303 - # A model family has many models with the same architecture,

304 - # and we don't need to test each one.

320 pytest.skip("Skipping test.")

321

322 + # Allow vllm to test using the given dtype, such as float32

323 vllm_extra_kwargs = vllm_extra_kwargs or {}

324 vllm_extra_kwargs["dtype"] = model_info.dtype

325

326 + # Allow vllm to test using hf_overrides

327 if model_info.hf_overrides is not None:

328 vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

329

@@ -319,9 +336,15 @@ def mteb_test_rerank_models(hf_runner,

336

337 model_config = vllm_model.llm.llm_engine.model_config

338

339 + # Confirm whether vllm is using the correct architecture

340 if model_info.architecture:

341 assert (model_info.architecture in model_config.architectures)

342 +

343 + # Score API is only enabled for num_labels == 1

344 assert model_config.hf_config.num_labels == 1

345 +

346 + # Confirm whether vllm uses the correct default_pooling_type, which

347 + # relates to whether chunked prefill and prefix caching are enabled

348 assert (model_config._model_info.default_pooling_type ==

349 model_info.default_pooling_type)

350

@@ -330,6 +353,8 @@ def mteb_test_rerank_models(hf_runner,

353 languages=MTEB_RERANK_LANGS)

354 vllm_dtype = model_config.dtype

355

356 + # Accelerate mteb test by setting

357 + # SentenceTransformers mteb score to a constant

358 if model_info.mteb_score is None:

359 st_main_score, st_dtype = mteb_test_rerank_models_hf(

360 hf_runner, model_info.name, hf_model_callback)

tests/models/language/pooling/test_bge_reranker_v2_gemma.py

@@ -14,6 +14,7 @@ from .mteb_utils import VllmMtebEncoder, mteb_test_rerank_models

14 RERANK_MODELS = [

15 LASTPoolingRerankModelInfo("BAAI/bge-reranker-v2-gemma",

16 architecture="GemmaForSequenceClassification",

17 + mteb_score=0.33757,

18 hf_overrides={

19 "architectures":

20 ["GemmaForSequenceClassification"],

vllm/config/__init__.py

+52 -1

@@ -745,7 +745,7 @@ class ModelConfig:

745

746 self.pooler_config = self._init_pooler_config()

747

748 - self.dtype = _get_and_verify_dtype(

748 + self.dtype: torch.dtype = _get_and_verify_dtype(

749 self.model,

750 self.hf_config,

751 self.dtype,

@@ -1751,6 +1751,32 @@ class ModelConfig:

1751 # `llm as reranker` models defaults to not using pad_token.

1752 return getattr(self.hf_config, "use_pad_token", True)

1753

1754 + @property

1755 + def head_dtype(self) -> torch.dtype:

1756 + """

1757 + "head" refers to the last Linear layer(s) of an LLM,

1758 + such as the lm_head in a generation model,

1759 + or the score or classifier in a classification model.

1760 +

1761 + The default head_dtype based on runner_type.\n

1762 + - The pooling model defaults to using fp32 head,

1763 + you can use --hf-overrides '{"head_dtype": "model"}' to disable it.\n

1764 + - The generate model defaults to not using fp32 head,

1765 + you can use --hf-overrides '{"head_dtype": "float32"}' to enable it.

1766 + """

1767 + head_dtype = _get_head_dtype(config=self.hf_config,

1768 + dtype=self.dtype,

1769 + runner_type=self.runner_type)

1770 +

1771 + if head_dtype not in current_platform.supported_dtypes:

1772 + logger.warning_once(

1773 + "The current platform does not support [%s] head dtype, "

1774 + "fallback to model dtype [%s].", head_dtype, self.dtype)

1775 + return self.dtype

1776 +

1777 + logger.debug_once("head dtype: %s", head_dtype)

1778 + return head_dtype

1779 +

1780 def get_and_verify_max_len(self, max_model_len: int):

1781 # Consider max_model_len in tokenizer_config only when

1782 # pooling models use absolute position_embedding.

@@ -2893,6 +2919,31 @@ def _get_and_verify_dtype(

2919 return torch_dtype

2920

2921

2922 +def _get_head_dtype(config: PretrainedConfig, dtype: torch.dtype,

2923 + runner_type: str) -> torch.dtype:

2924 + head_dtype: Optional[Union[str,

2925 + torch.dtype]] = getattr(config, "head_dtype",

2926 + None)

2927 +

2928 + if head_dtype == "model":

2929 + return dtype

2930 + elif isinstance(head_dtype, str):

2931 + head_dtype = head_dtype.lower()

2932 + if head_dtype not in _STR_DTYPE_TO_TORCH_DTYPE:

2933 + raise ValueError(f"Unknown dtype: {head_dtype!r}")

2934 + return _STR_DTYPE_TO_TORCH_DTYPE[head_dtype]

2935 + elif isinstance(head_dtype, torch.dtype):

2936 + return head_dtype

2937 + elif head_dtype is None:

2938 + if torch.float32 not in current_platform.supported_dtypes:

2939 + return dtype

2940 + if runner_type == "pooling":

2941 + return torch.float32

2942 + return dtype

2943 + else:

2944 + raise ValueError(f"Unknown dtype: {head_dtype}")

2945 +

2946 +

2947 def _get_and_verify_max_len(

2948 hf_config: PretrainedConfig,

2949 tokenizer_config: Optional[dict],

vllm/model_executor/layers/pooler.py

+23 -15

@@ -5,7 +5,7 @@ from collections.abc import Mapping, Set

5 from dataclasses import dataclass

6 from enum import IntEnum

7 from itertools import groupby

8 -from typing import Callable, Optional, TypeVar, Union, cast

8 +from typing import Callable, Optional, TypeVar, Union

10 import torch

11 import torch.nn as nn

@@ -362,14 +362,13 @@ class PoolerIdentity(PoolerActivation):

362 class PoolerNormalize(PoolerActivation):

363

364 def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:

365 - x = F.normalize(pooled_data.float(), p=2, dim=-1)

366 - return x.to(pooled_data.dtype)

365 + return F.normalize(pooled_data, p=2, dim=-1)

366

367

368 class PoolerMultiLabelClassify(PoolerActivation):

369

370 def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:

372 - return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)

371 + return F.sigmoid(pooled_data)

372

373

374 class PoolerClassify(PoolerActivation):

@@ -394,9 +393,9 @@ class PoolerClassify(PoolerActivation):

393 pooled_data.shape[-1])

394

395 if num_labels < 2:

397 - return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)

396 + return F.sigmoid(pooled_data)

397

399 - return F.softmax(pooled_data.float(), dim=-1).to(pooled_data.dtype)

398 + return F.softmax(pooled_data, dim=-1)

399

400

401 class LambdaPoolerActivation(PoolerActivation):

@@ -432,8 +431,9 @@ class EmbeddingPoolerHead(PoolerHead):

431 from vllm.model_executor.models.adapters import _load_st_projector

432

433 vllm_config = get_current_vllm_config()

435 - self.projector = _load_st_projector(

434 + self.projector: Optional[nn.Module] = _load_st_projector(

435 vllm_config.model_config) if vllm_config else None

436 + self.head_dtype = vllm_config.model_config.head_dtype

437

438 def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],

439 pooling_metadata: PoolingMetadata):

@@ -442,16 +442,11 @@ class EmbeddingPoolerHead(PoolerHead):

442 pooled_data = torch.stack(pooled_data)

443 # pooled_data shape: [batchsize, hidden_dimension]

444

445 + pooled_data = pooled_data.to(self.head_dtype)

446 +

447 # Apply ST projector

448 if self.projector is not None:

447 - projector = cast(nn.Module, self.projector)

448 -

449 - def _proj(x: torch.Tensor) -> torch.Tensor:

450 - orig_dtype = x.dtype

451 - y = projector(x.to(torch.float32))

452 - return y.to(orig_dtype)

453 -

454 - pooled_data = _proj(pooled_data)

449 + pooled_data = self.projector(pooled_data)

450 # pooled_data shape: [batchsize, embedding_dimension]

451

452 pooling_params = get_pooling_params(pooling_metadata)

@@ -494,8 +489,18 @@ class RewardPoolerHead(PoolerHead):

489 def __init__(self) -> None:

490 super().__init__(activation=PoolerClassify(static_num_labels=False))

491

492 + from vllm.config import get_current_vllm_config

493 + vllm_config = get_current_vllm_config()

494 + self.head_dtype = vllm_config.model_config.head_dtype

495 +

496 def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],

497 pooling_metadata: PoolingMetadata):

498 +

499 + if isinstance(pooled_data, list):

500 + pooled_data = [p.to(self.head_dtype) for p in pooled_data]

501 + else:

502 + pooled_data = pooled_data.to(self.head_dtype)

503 +

504 pooling_params = get_pooling_params(pooling_metadata)

505

506 # for softmax

@@ -641,6 +646,7 @@ class ClassifierPooler(Pooler):

646 self.act_fn = act_fn or PoolerClassify()

647 self.logit_bias: Optional[

648 float] = vllm_config.model_config.pooler_config.logit_bias

649 + self.head_dtype = vllm_config.model_config.head_dtype

650

651 def get_supported_tasks(self) -> Set[PoolingTask]:

652 return {"classify", "score"}

@@ -655,6 +661,8 @@ class ClassifierPooler(Pooler):

661 pooled_data = torch.stack(pooled_data)

662 # pooled_data shape: [batchsize, hidden_size]

663

664 + pooled_data = pooled_data.to(self.head_dtype)

665 +

666 if self.classifier is not None:

667 pooled_data = self.classifier(pooled_data)

668 # pooled_data shape: [batchsize, num_labels]

vllm/model_executor/models/adapters.py

+4 -6

@@ -62,7 +62,7 @@ def _load_st_projector(model_config: "ModelConfig") -> Optional[nn.Module]:

62 linear = nn.Linear(layer_config.get("in_features", 768),

63 layer_config.get("out_features", 768),

64 bias=layer_config.get("bias", True),

65 - dtype=torch.float32)

65 + dtype=model_config.head_dtype)

67 if not _load_dense_weights(linear, folder, model_config):

68 continue

@@ -70,7 +70,7 @@ def _load_st_projector(model_config: "ModelConfig") -> Optional[nn.Module]:

70 layers.append(linear)

71 if act_name := layer_config.get("activation_function"):

72 layers.append(get_act_fn(act_name))

73 - return nn.Sequential(*layers).to(dtype=torch.float32)

73 + return nn.Sequential(*layers).to(dtype=model_config.head_dtype)

74 except Exception:

75 logger.exception("ST projector loading failed")

@@ -105,15 +105,13 @@ def _load_dense_weights(linear: nn.Linear, folder: str,

105 if weight_key in state_dict:

106 weight_loader = getattr(linear.weight, "weight_loader",

107 default_weight_loader)

108 - weight_loader(linear.weight,

109 - state_dict[weight_key].to(torch.float32))

108 + weight_loader(linear.weight, state_dict[weight_key])

109

110 bias_key = weight_key.replace("weight", "bias")

111 if linear.bias is not None and bias_key in state_dict:

112 bias_loader = getattr(linear.bias, "weight_loader",

113 default_weight_loader)

115 - bias_loader(linear.bias,

116 - state_dict[bias_key].to(torch.float32))

114 + bias_loader(linear.bias, state_dict[bias_key])

115 return True

116 except Exception:

117 logger.exception("Failed to load %s", filename)

vllm/model_executor/models/bert.py

+3 -1

@@ -562,7 +562,9 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,

562 self.bert = BertPoolingModel(vllm_config=vllm_config,

563 prefix=maybe_prefix(prefix, "bert"),

564 embedding_class=BertEmbedding)

565 - self.classifier = nn.Linear(config.hidden_size, config.num_labels)

565 + self.classifier = nn.Linear(config.hidden_size,

566 + config.num_labels,

567 + dtype=vllm_config.model_config.head_dtype)

568

569 pooler_config = vllm_config.model_config.pooler_config

570 assert pooler_config is not None

vllm/model_executor/models/bert_with_rope.py

+8 -8

@@ -637,14 +637,14 @@ class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):

637 self.new = GteNewModel(vllm_config=vllm_config,

638 prefix=prefix,

639 add_pooling_layer=True)

640 - self.classifier = RowParallelLinear(config.hidden_size,

641 - config.num_labels,

642 - input_is_parallel=False,

643 - bias=True,

644 - quant_config=quant_config,

645 - prefix=maybe_prefix(

646 - prefix, "classifier"),

647 - return_bias=False)

640 + self.classifier = ReplicatedLinear(

641 + config.hidden_size,

642 + config.num_labels,

643 + bias=True,

644 + quant_config=quant_config,

645 + params_dtype=vllm_config.model_config.head_dtype,

646 + prefix=maybe_prefix(prefix, "classifier"),

647 + return_bias=False)

648

649 pooler_config = vllm_config.model_config.pooler_config

650 assert pooler_config is not None

vllm/model_executor/models/gpt2.py

+6 -4

@@ -339,7 +339,10 @@ class GPT2ForSequenceClassification(nn.Module):

339 config = vllm_config.model_config.hf_config

340 self.transformer = GPT2Model(vllm_config=vllm_config,

341 prefix=maybe_prefix(prefix, "gpt2"))

342 - self.score = nn.Linear(config.n_embd, config.num_labels, bias=False)

342 + self.score = nn.Linear(config.n_embd,

343 + config.num_labels,

344 + bias=False,

345 + dtype=vllm_config.model_config.head_dtype)

346

347 pooler_config = vllm_config.model_config.pooler_config

348 assert pooler_config is not None

@@ -348,7 +351,7 @@ class GPT2ForSequenceClassification(nn.Module):

351 "encode":

352 Pooler.for_encode(pooler_config),

353 "classify":

351 - Pooler.for_classify(pooler_config, classifier=None),

354 + Pooler.for_classify(pooler_config, classifier=self.score),

355 })

356

357 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

@@ -367,8 +370,7 @@ class GPT2ForSequenceClassification(nn.Module):

370 position_ids=positions,

371 inputs_embeds=inputs_embeds,

372 intermediate_tensors=intermediate_tensors)

370 - logits = self.score(hidden_states)

371 - return logits

373 + return hidden_states

374

375

376 def _add_transformer_prefix(

vllm/model_executor/models/internlm2.py

+11 -8

@@ -423,13 +423,15 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):

423 delattr(self, attr)

424

425 config = vllm_config.model_config.hf_config

426 - self.v_head = RowParallelLinear(

427 - config.hidden_size,

428 - 1,

429 - bias=False,

430 - input_is_parallel=False,

431 - prefix=maybe_prefix(prefix, "v_head"),

432 - )

426 + self.head_dtype = vllm_config.model_config.head_dtype

427 +

428 + self.v_head = RowParallelLinear(config.hidden_size,

429 + 1,

430 + bias=False,

431 + input_is_parallel=False,

432 + params_dtype=self.head_dtype,

433 + prefix=maybe_prefix(prefix, "v_head"),

434 + return_bias=False)

435

436 pooler_config = vllm_config.model_config.pooler_config

437 assert pooler_config is not None

@@ -446,5 +448,6 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):

448 ) -> Union[torch.Tensor, IntermediateTensors]:

449 hidden_states = self.model(input_ids, positions, intermediate_tensors,

450 inputs_embeds)

449 - logits, _ = self.v_head(hidden_states)

451 + hidden_states = hidden_states.to(self.head_dtype)

452 + logits = self.v_head(hidden_states)

453 return logits

vllm/model_executor/models/jamba.py

+1 -1

@@ -613,7 +613,7 @@ class JambaForSequenceClassification(JambaForCausalLM):

613 config.hidden_size,

614 num_labels,

615 bias=score_bias,

616 - dtype=torch.float32,

616 + dtype=vllm_config.model_config.head_dtype,

617 )

618

619 pooler_config = vllm_config.model_config.pooler_config

vllm/model_executor/models/jina_vl.py

+8 -5

@@ -5,9 +5,9 @@ from typing import Optional

6 import torch

7 import torch.nn as nn

8 -from transformers import BatchFeature, PretrainedConfig

8 +from transformers import BatchFeature

10 -from vllm.config import VllmConfig

10 +from vllm.config import ModelConfig, VllmConfig

11 from vllm.inputs import TokensPrompt

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -28,13 +28,17 @@ logger = init_logger(__name__)

29 class JinaVLScorer(nn.Module):

31 - def __init__(self, config: PretrainedConfig):

31 + def __init__(self, model_config: "ModelConfig"):

32 super().__init__()

33 + config = model_config.hf_config

34 + head_dtype = model_config.head_dtype

35 self.dense = ColumnParallelLinear(config.hidden_size,

36 config.hidden_size,

37 + params_dtype=head_dtype,

38 bias=True)

39 self.out_proj = RowParallelLinear(config.hidden_size,

40 config.num_labels,

41 + params_dtype=head_dtype,

42 bias=True)

44 def forward(self, x, **kwargs):

@@ -88,11 +92,10 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,

92 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

93 super().__init__(vllm_config=vllm_config,

94 prefix=maybe_prefix(prefix, "qwen2_vl"))

91 - config = vllm_config.model_config.hf_config

95 pooler_config = vllm_config.model_config.pooler_config

96 assert pooler_config is not None

95 - self.score = JinaVLScorer(config)

98 + self.score = JinaVLScorer(vllm_config.model_config)

99 self.pooler = DispatchPooler({

100 "encode":

101 Pooler.for_encode(pooler_config),

vllm/model_executor/models/modernbert.py

+3 -1

@@ -306,7 +306,9 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):

306 self.config = config

307 self.model = ModernBertModel(vllm_config=vllm_config,

308 prefix=maybe_prefix(prefix, "modernbert"))

309 - self.classifier = nn.Linear(config.hidden_size, config.num_labels)

309 + self.classifier = nn.Linear(config.hidden_size,

310 + config.num_labels,

311 + dtype=vllm_config.model_config.head_dtype)

312 self.pooling = ModernBertPooler(config)

313

314 pooler_config = vllm_config.model_config.pooler_config

vllm/model_executor/models/qwen2_rm.py

@@ -53,15 +53,18 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):

53 self.quant_config = quant_config

54 self.model = Qwen2Model(vllm_config=vllm_config,

55 prefix=maybe_prefix(prefix, "model"))

56 + self.head_dtype = vllm_config.model_config.head_dtype

58 self.score = nn.Sequential(

59 ColumnParallelLinear(config.hidden_size,

60 config.hidden_size,

61 quant_config=quant_config,

62 + params_dtype=self.head_dtype,

63 return_bias=False),

64 nn.ReLU(),

65 RowParallelLinear(config.hidden_size,

66 config.num_labels,

67 + params_dtype=self.head_dtype,

68 quant_config=quant_config,

69 return_bias=False),

70 )

@@ -80,6 +83,7 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):

83 ) -> Union[torch.Tensor, IntermediateTensors]:

84 hidden_states = self.model(input_ids, positions, intermediate_tensors,

85 inputs_embeds)

86 + hidden_states = hidden_states.to(self.head_dtype)

87 logits = self.score(hidden_states)

88 return logits

vllm/model_executor/models/roberta.py

+11 -5

@@ -8,7 +8,7 @@ import torch

8 from torch import nn

9 from transformers import RobertaConfig

11 -from vllm.config import VllmConfig

11 +from vllm.config import ModelConfig, VllmConfig

12 from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,

13 DispatchPooler, Pooler)

14 from vllm.model_executor.layers.vocab_parallel_embedding import (

@@ -73,10 +73,16 @@ class RobertaEmbedding(nn.Module):

73 class RobertaClassificationHead(nn.Module):

74 """Head for sentence-level classification tasks."""

76 - def __init__(self, config: RobertaConfig):

76 + def __init__(self, model_config: "ModelConfig"):

77 super().__init__()

78 - self.dense = nn.Linear(config.hidden_size, config.hidden_size)

79 - self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

78 + config = model_config.hf_config

79 + head_dtype = model_config.head_dtype

80 + self.dense = nn.Linear(config.hidden_size,

81 + config.hidden_size,

82 + dtype=head_dtype)

83 + self.out_proj = nn.Linear(config.hidden_size,

84 + config.num_labels,

85 + dtype=head_dtype)

87 def forward(self, x: torch.Tensor) -> torch.Tensor:

88 # CLSPool has already been applied in `pooling`

@@ -184,7 +190,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

190 self.roberta = BertModel(vllm_config=vllm_config,

191 prefix=maybe_prefix(prefix, "bert"),

192 embedding_class=RobertaEmbedding)

187 - self.classifier = RobertaClassificationHead(config)

193 + self.classifier = RobertaClassificationHead(vllm_config.model_config)

194

195 pooler_config = vllm_config.model_config.pooler_config

196 assert pooler_config is not None