Standardise weight tying on `ParallelLMHead.tie_weights` (#52147)

vllm/model_executor/models/arctic.py

+1 -1

@@ -557,7 +557,7 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

557 prefix=maybe_prefix(prefix, "lm_head"),

558 )

559 if self.config.tie_word_embeddings:

560 - self.lm_head.weight = self.model.embed_tokens.weight

560 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

561 self.num_experts = config.num_local_experts

562 self.num_experts_per_tok = config.num_experts_per_tok

563

vllm/model_executor/models/bailing_moe.py

+7 -8

@@ -509,15 +509,14 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

509 self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)

510

511 if get_pp_group().is_last_rank:

512 + self.lm_head = ParallelLMHead(

513 + config.vocab_size,

514 + config.hidden_size,

515 + quant_config=quant_config,

516 + prefix=maybe_prefix(prefix, "lm_head"),

517 + )

518 if self.tie_word_embeddings:

513 - self.lm_head = self.model.word_embeddings

514 - else:

515 - self.lm_head = ParallelLMHead(

516 - config.vocab_size,

517 - config.hidden_size,

518 - quant_config=quant_config,

519 - prefix=maybe_prefix(prefix, "lm_head"),

520 - )

519 + self.lm_head = self.lm_head.tie_weights(self.model.word_embeddings)

520 self.logits_processor = LogitsProcessor(config.vocab_size)

521 else:

522 self.lm_head = PPMissingLayer()

vllm/model_executor/models/bloom.py

+6 -7

@@ -324,14 +324,13 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):

324 self.transformer = BloomModel(

325 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer")

326 )

327 + self.lm_head = ParallelLMHead(

328 + self.config.vocab_size,

329 + self.config.hidden_size,

330 + prefix=maybe_prefix(prefix, "lm_head"),

331 + )

332 if self.config.tie_word_embeddings:

328 - self.lm_head = self.transformer.word_embeddings

329 - else:

330 - self.lm_head = ParallelLMHead(

331 - self.config.vocab_size,

332 - self.config.hidden_size,

333 - prefix=maybe_prefix(prefix, "lm_head"),

334 - )

333 + self.lm_head = self.lm_head.tie_weights(self.transformer.word_embeddings)

334

335 self.logits_processor = LogitsProcessor(config.vocab_size)

336 self.make_empty_intermediate_tensors = (

vllm/model_executor/models/chameleon.py

+1 -1

@@ -977,7 +977,7 @@ class ChameleonForConditionalGeneration(

977 prefix=maybe_prefix(prefix, "lm_head"),

978 )

979 if config.tie_word_embeddings:

980 - self.lm_head.weight = self.model.embed_tokens.weight

980 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

981

982 logit_scale = getattr(config, "logit_scale", 1.0)

983 self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)

vllm/model_executor/models/ernie45_moe.py

+1 -1

@@ -543,7 +543,7 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe

543 self.lm_head = PPMissingLayer()

544

545 if self.config.tie_word_embeddings:

546 - self.lm_head.weight = self.model.embed_tokens.weight

546 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

547 self.logits_processor = LogitsProcessor(config.vocab_size)

548 self.make_empty_intermediate_tensors = (

549 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/ernie45_vl_moe.py

+1 -1

@@ -611,7 +611,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

611 self.lm_head = PPMissingLayer()

612

613 if self.config.tie_word_embeddings:

614 - self.lm_head.weight = self.model.embed_tokens.weight

614 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

615 self.logits_processor = LogitsProcessor(config.vocab_size)

616 self.make_empty_intermediate_tensors = (

617 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/ernie_mtp.py

+1 -1

@@ -177,7 +177,7 @@ class ErnieMTP(nn.Module):

177 )

178

179 if self.config.tie_word_embeddings:

180 - self.lm_head.weight = self.model.embed_tokens.weight

180 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

181

182 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

183 return self.model.embed_input_ids(input_ids)

vllm/model_executor/models/exaone.py

+1 -1

@@ -410,7 +410,7 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

410 prefix=maybe_prefix(prefix, "lm_head"),

411 )

412 if config.tie_word_embeddings:

413 - self.lm_head.weight = self.transformer.wte.weight

413 + self.lm_head = self.lm_head.tie_weights(self.transformer.wte)

414

415 logit_scale = getattr(config, "logit_scale", 1.0)

416 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/exaone4.py

+1 -1

@@ -410,7 +410,7 @@ class Exaone4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

410 prefix=maybe_prefix(prefix, "lm_head"),

411 )

412 if config.tie_word_embeddings:

413 - self.lm_head.weight = self.model.embed_tokens.weight

413 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

414

415 logit_scale = getattr(config, "logit_scale", 1.0)

416 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/exaone4_5_mtp.py

+1 -1

@@ -159,7 +159,7 @@ class Exaone4_5_MTP(ExaoneMoeMTP, SupportsMultiModal):

159 prefix=maybe_prefix(prefix, "lm_head"),

160 )

161 if config.tie_word_embeddings:

162 - self.lm_head.weight = self.model.embed_tokens.weight

162 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

163 self.logits_processor = LogitsProcessor(

164 self.unpadded_vocab_size, config.vocab_size

165 )

vllm/model_executor/models/exaone_moe.py

+1 -1

@@ -522,7 +522,7 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

522 quant_config=quant_config,

523 )

524 if config.tie_word_embeddings:

525 - self.lm_head.weight = self.model.embed_tokens.weight

525 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

526

527 logit_scale = getattr(config, "logit_scale", 1.0)

528 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/exaone_moe_mtp.py

+1 -1

@@ -159,7 +159,7 @@ class ExaoneMoeMTP(nn.Module):

159 prefix=maybe_prefix(prefix, "lm_head"),

160 )

161 if config.tie_word_embeddings:

162 - self.lm_head.weight = self.model.embed_tokens.weight

162 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

163 self.logits_processor = LogitsProcessor(

164 self.unpadded_vocab_size, config.vocab_size

165 )

vllm/model_executor/models/falcon.py

+7 -8

@@ -499,15 +499,14 @@ class FalconForCausalLM(nn.Module, SupportsPP):

499 if config.tie_word_embeddings is not None

500 else True

501 )

502 + self.lm_head = ParallelLMHead(

503 + config.vocab_size,

504 + config.hidden_size,

505 + quant_config=quant_config,

506 + prefix=maybe_prefix(prefix, "lm_head"),

507 + )

508 if self.tie_word_embeddings:

503 - self.lm_head = self.transformer.word_embeddings

504 - else:

505 - self.lm_head = ParallelLMHead(

506 - config.vocab_size,

507 - config.hidden_size,

508 - quant_config=quant_config,

509 - prefix=maybe_prefix(prefix, "lm_head"),

510 - )

509 + self.lm_head = self.lm_head.tie_weights(self.transformer.word_embeddings)

510 self.logits_processor = LogitsProcessor(config.vocab_size)

511 self.make_empty_intermediate_tensors = (

512 self.transformer.make_empty_intermediate_tensors

vllm/model_executor/models/funasr.py

+7 -8

@@ -919,15 +919,14 @@ class FunASRForConditionalGeneration(

919 )

920 logit_scale = getattr(config, "logit_scale", 1.0)

921

922 + self.lm_head = ParallelLMHead(

923 + config.vocab_size,

924 + config.hidden_size,

925 + quant_config=quant_config,

926 + prefix=maybe_prefix(prefix, "lm_head"),

927 + )

928 if config.tie_word_embeddings:

923 - self.lm_head = self.model.decoder.embed_tokens

924 - else:

925 - self.lm_head = ParallelLMHead(

926 - config.vocab_size,

927 - config.hidden_size,

928 - quant_config=quant_config,

929 - prefix=maybe_prefix(prefix, "lm_head"),

930 - )

929 + self.lm_head = self.lm_head.tie_weights(self.model.decoder.embed_tokens)

930 self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)

931

932 def get_language_model(self) -> torch.nn.Module:

vllm/model_executor/models/gemma4_mtp.py

+1 -1

@@ -489,7 +489,7 @@ class Gemma4MTP(nn.Module):

489 prefix=maybe_prefix(prefix, "lm_head"),

490 )

491 if getattr(config, "tie_word_embeddings", True):

492 - self.lm_head.weight = self.model.embed_tokens.weight

492 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

493

494 self.logits_processor = LogitsProcessor(

495 text_config.vocab_size,

vllm/model_executor/models/glm4.py

+7 -8

@@ -252,15 +252,14 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

252 )

253

254 if get_pp_group().is_last_rank:

255 + self.lm_head = ParallelLMHead(

256 + config.vocab_size,

257 + config.hidden_size,

258 + quant_config=quant_config,

259 + prefix=maybe_prefix(prefix, "lm_head"),

260 + )

261 if config.tie_word_embeddings:

256 - self.lm_head = self.model.embed_tokens

257 - else:

258 - self.lm_head = ParallelLMHead(

259 - config.vocab_size,

260 - config.hidden_size,

261 - quant_config=quant_config,

262 - prefix=maybe_prefix(prefix, "lm_head"),

263 - )

262 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

263 else:

264 self.lm_head = PPMissingLayer()

265

vllm/model_executor/models/gpt_neox.py

+1 -1

@@ -285,7 +285,7 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):

285 prefix=maybe_prefix(prefix, "embed_out"),

286 )

287 if self.config.tie_word_embeddings:

288 - self.embed_out.weight = self.gpt_neox.embed_in.weight

288 + self.embed_out = self.embed_out.tie_weights(self.gpt_neox.embed_in)

289 self.logits_processor = LogitsProcessor(config.vocab_size)

290 self.make_empty_intermediate_tensors = (

291 self.gpt_neox.make_empty_intermediate_tensors

vllm/model_executor/models/granite.py

+1 -1

@@ -365,7 +365,7 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

365 prefix=maybe_prefix(prefix, "lm_head"),

366 )

367 if config.tie_word_embeddings:

368 - self.lm_head.weight = self.model.embed_tokens.weight

368 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

369

370 logit_scale = getattr(config, "logit_scale", 1.0)

371 if hasattr(config, "logits_scaling"):

vllm/model_executor/models/granite4_vision.py

+1 -1

@@ -316,7 +316,7 @@ class Granite4VisionLLMForCausalLM(GraniteForCausalLM):

316 prefix=maybe_prefix(prefix, "lm_head"),

317 )

318 if config.tie_word_embeddings:

319 - self.lm_head.weight = self.model.embed_tokens.weight

319 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

320 logit_scale = getattr(config, "logit_scale", 1.0)

321 if hasattr(config, "logits_scaling"):

322 logit_scale /= config.logits_scaling

vllm/model_executor/models/granitemoe.py

+1 -1

@@ -505,7 +505,7 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

505 prefix=maybe_prefix(prefix, "lm_head"),

506 )

507 if config.tie_word_embeddings:

508 - self.lm_head.weight = self.model.embed_tokens.weight

508 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

509

510 self.logits_processor = LogitsProcessor(

511 config.vocab_size,

vllm/model_executor/models/granitemoehybrid.py

+1 -1

@@ -663,7 +663,7 @@ class GraniteMoeHybridForCausalLM(

663 prefix=maybe_prefix(prefix, "lm_head"),

664 )

665 if config.tie_word_embeddings:

666 - self.lm_head.weight = self.model.embed_tokens.weight

666 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

667 self.logits_processor = LogitsProcessor(

668 config.vocab_size,

669 config.vocab_size,

vllm/model_executor/models/granitemoeshared.py

+1 -1

@@ -281,7 +281,7 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

281 prefix=maybe_prefix(prefix, "lm_head"),

282 )

283 if config.tie_word_embeddings:

284 - self.lm_head.weight = self.model.embed_tokens.weight

284 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

285

286 self.logits_processor = LogitsProcessor(

287 config.vocab_size,

vllm/model_executor/models/hrm_text.py

+7 -8

@@ -487,15 +487,14 @@ class HrmTextForCausalLM(nn.Module):

487 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

488 )

489

490 + self.lm_head = ParallelLMHead(

491 + config.vocab_size,

492 + config.hidden_size,

493 + quant_config=quant_config,

494 + prefix=maybe_prefix(prefix, "lm_head"),

495 + )

496 if config.tie_word_embeddings:

491 - self.lm_head = self.model.embed_tokens

492 - else:

493 - self.lm_head = ParallelLMHead(

494 - config.vocab_size,

495 - config.hidden_size,

496 - quant_config=quant_config,

497 - prefix=maybe_prefix(prefix, "lm_head"),

498 - )

497 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

498 self.logits_processor = LogitsProcessor(config.vocab_size)

499

500 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/hunyuan_v1.py

+1 -1

@@ -927,7 +927,7 @@ class HunyuanV1ModelBase(

927 prefix=maybe_prefix(prefix, "lm_head"),

928 )

929 if config.tie_word_embeddings:

930 - self.lm_head.weight = self.model.embed_tokens.weight

930 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

931

932 logit_scale = getattr(config, "logit_scale", 1.0)

933 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/hy_v3.py

+1 -1

@@ -673,7 +673,7 @@ class HYV3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

673 prefix=maybe_prefix(prefix, "lm_head"),

674 )

675 if self.config.tie_word_embeddings:

676 - self.lm_head.weight = self.model.embed_tokens.weight

676 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

677 self.logits_processor = LogitsProcessor(config.vocab_size)

678 self.make_empty_intermediate_tensors = (

679 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/idefics3.py

+1 -1

@@ -539,7 +539,7 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo

539 prefix=maybe_prefix(prefix, "lm_head"),

540 )

541 if self.config.text_config.tie_word_embeddings:

542 - self.lm_head.weight = self.model.text_model.embed_tokens.weight

542 + self.lm_head = self.lm_head.tie_weights(self.model.text_model.embed_tokens)

543 self.logits_processor = LogitsProcessor(config.text_config.vocab_size)

544

545 def _parse_and_validate_image_input(self, **kwargs: object) -> ImageInputs | None:

vllm/model_executor/models/internlm2.py

+1 -1

@@ -352,7 +352,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant):

352 prefix=maybe_prefix(prefix, "output"),

353 )

354 if self.config.tie_word_embeddings:

355 - self.output.weight = self.model.tok_embeddings.weight

355 + self.output = self.output.tie_weights(self.model.tok_embeddings)

356 self.logits_processor = LogitsProcessor(config.vocab_size)

357 self.make_empty_intermediate_tensors = (

358 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/interns1_pro.py

+1 -1

@@ -481,7 +481,7 @@ class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM):

481 prefix=maybe_prefix(prefix, "lm_head"),

482 )

483 if self.config.tie_word_embeddings:

484 - self.lm_head.weight = self.model.embed_tokens.weight

484 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

485 self.logits_processor = LogitsProcessor(self.config.vocab_size)

486 self.make_empty_intermediate_tensors = (

487 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/interns2_mobius.py

+7 -8

@@ -443,15 +443,14 @@ class InternS2MobiusForCausalLM(Qwen3_5ForCausalLMBase):

443 )

444

445 if get_pp_group().is_last_rank:

446 + self.lm_head = ParallelLMHead(

447 + config.vocab_size,

448 + config.hidden_size,

449 + quant_config=self.quant_config,

450 + prefix=maybe_prefix(prefix, "lm_head"),

451 + )

452 if config.tie_word_embeddings:

447 - self.lm_head = self.model.embed_tokens

448 - else:

449 - self.lm_head = ParallelLMHead(

450 - config.vocab_size,

451 - config.hidden_size,

452 - quant_config=self.quant_config,

453 - prefix=maybe_prefix(prefix, "lm_head"),

454 - )

453 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

454 else:

455 self.lm_head = PPMissingLayer()

456

vllm/model_executor/models/iquest_loopcoder.py

+7 -8

@@ -541,15 +541,14 @@ class IQuestLoopCoderForCausalLM(nn.Module):

541 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

542 )

543

544 + self.lm_head = ParallelLMHead(

545 + config.vocab_size,

546 + config.hidden_size,

547 + quant_config=quant_config,

548 + prefix=maybe_prefix(prefix, "lm_head"),

549 + )

550 if config.tie_word_embeddings:

545 - self.lm_head = self.model.embed_tokens

546 - else:

547 - self.lm_head = ParallelLMHead(

548 - config.vocab_size,

549 - config.hidden_size,

550 - quant_config=quant_config,

551 - prefix=maybe_prefix(prefix, "lm_head"),

552 - )

551 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

552

553 self.logits_processor = LogitsProcessor(config.vocab_size)

554

vllm/model_executor/models/mamba.py

+6 -7

@@ -188,14 +188,13 @@ class MambaForCausalLM(

188 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "backbone")

189 )

190

191 + self.lm_head = ParallelLMHead(

192 + config.vocab_size,

193 + config.hidden_size,

194 + prefix=maybe_prefix(prefix, "lm_head"),

195 + )

196 if config.tie_word_embeddings:

192 - self.lm_head = self.backbone.embeddings

193 - else:

194 - self.lm_head = ParallelLMHead(

195 - config.vocab_size,

196 - config.hidden_size,

197 - prefix=maybe_prefix(prefix, "lm_head"),

198 - )

197 + self.lm_head = self.lm_head.tie_weights(self.backbone.embeddings)

198

199 self.logits_processor = LogitsProcessor(config.vocab_size)

200

vllm/model_executor/models/mellum.py

+1 -1

@@ -221,7 +221,7 @@ class MellumForCausalLM(Qwen3MoeForCausalLM):

221 prefix=maybe_prefix(prefix, "lm_head"),

222 )

223 if self.config.tie_word_embeddings:

224 - self.lm_head.weight = self.model.embed_tokens.weight

224 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

225 self.logits_processor = LogitsProcessor(config.vocab_size)

226 self.make_empty_intermediate_tensors = (

227 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/mimo.py

+7 -8

@@ -101,15 +101,14 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):

101 )

102

103 if get_pp_group().is_last_rank:

104 + self.lm_head = ParallelLMHead(

105 + config.vocab_size,

106 + config.hidden_size,

107 + quant_config=quant_config,

108 + prefix=maybe_prefix(prefix, "lm_head"),

109 + )

110 if config.tie_word_embeddings:

105 - self.lm_head = self.model.embed_tokens

106 - else:

107 - self.lm_head = ParallelLMHead(

108 - config.vocab_size,

109 - config.hidden_size,

110 - quant_config=quant_config,

111 - prefix=maybe_prefix(prefix, "lm_head"),

112 - )

111 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

112 else:

113 self.lm_head = PPMissingLayer()

114

vllm/model_executor/models/mixtral.py

+1 -1

@@ -403,7 +403,7 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):

403 prefix=maybe_prefix(prefix, "lm_head"),

404 )

405 if self.config.tie_word_embeddings:

406 - self.lm_head.weight = self.model.embed_tokens.weight

406 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

407 self.logits_processor = LogitsProcessor(config.vocab_size)

408 self.make_empty_intermediate_tensors = (

409 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/molmo.py

+7 -8

@@ -1340,15 +1340,14 @@ class MolmoForCausalLM(

1340

1341 self.img_patch_id = None

1342

1343 + self.lm_head = ParallelLMHead(

1344 + config.embedding_size or config.vocab_size,

1345 + config.hidden_size,

1346 + quant_config=quant_config,

1347 + prefix=maybe_prefix(prefix, "lm_head"),

1348 + )

1349 if self.config.weight_tying:

1344 - self.lm_head = self.model.transformer.wte

1345 - else:

1346 - self.lm_head = ParallelLMHead(

1347 - config.embedding_size or config.vocab_size,

1348 - config.hidden_size,

1349 - quant_config=quant_config,

1350 - prefix=maybe_prefix(prefix, "lm_head"),

1351 - )

1350 + self.lm_head = self.lm_head.tie_weights(self.model.transformer.wte)

1351

1352 self.logits_processor = LogitsProcessor(

1353 config.embedding_size or config.vocab_size

vllm/model_executor/models/moss_audio.py

+7 -8

@@ -835,15 +835,14 @@ class MossQwen3ForCausalLM(Qwen3ForCausalLM):

835 )

836

837 if get_pp_group().is_last_rank:

838 + self.lm_head = ParallelLMHead(

839 + config.vocab_size,

840 + config.hidden_size,

841 + quant_config=quant_config,

842 + prefix=maybe_prefix(prefix, "lm_head"),

843 + )

844 if config.tie_word_embeddings:

839 - self.lm_head = self.model.embed_tokens

840 - else:

841 - self.lm_head = ParallelLMHead(

842 - config.vocab_size,

843 - config.hidden_size,

844 - quant_config=quant_config,

845 - prefix=maybe_prefix(prefix, "lm_head"),

846 - )

845 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

846 else:

847 from .utils import PPMissingLayer

848

vllm/model_executor/models/nemotron.py

+1 -1

@@ -404,7 +404,7 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

404 prefix=maybe_prefix(prefix, "lm_head"),

405 )

406 if config.tie_word_embeddings:

407 - self.lm_head.weight = self.model.embed_tokens.weight

407 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

408

409 logit_scale = getattr(config, "logit_scale", 1.0)

410 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/olmo3.py

+7 -8

@@ -369,15 +369,14 @@ class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

369 self.model = Olmo3Model(

370 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

371 )

372 + self.lm_head = ParallelLMHead(

373 + config.vocab_size,

374 + config.hidden_size,

375 + quant_config=vllm_config.quant_config,

376 + prefix=maybe_prefix(prefix, "lm_head"),

377 + )

378 if config.tie_word_embeddings:

373 - self.lm_head = self.model.embed_tokens

374 - else:

375 - self.lm_head = ParallelLMHead(

376 - config.vocab_size,

377 - config.hidden_size,

378 - quant_config=vllm_config.quant_config,

379 - prefix=maybe_prefix(prefix, "lm_head"),

380 - )

379 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

380 self.logits_processor = LogitsProcessor(config.vocab_size)

381 self.make_empty_intermediate_tensors = (

382 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/olmo_hybrid.py

+7 -8

@@ -397,15 +397,14 @@ class OlmoHybridForCausalLM(

397 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

398 )

399

400 + self.lm_head = ParallelLMHead(

401 + config.vocab_size,

402 + config.hidden_size,

403 + quant_config=vllm_config.quant_config,

404 + prefix=maybe_prefix(prefix, "lm_head"),

405 + )

406 if config.tie_word_embeddings:

401 - self.lm_head = self.model.embed_tokens

402 - else:

403 - self.lm_head = ParallelLMHead(

404 - config.vocab_size,

405 - config.hidden_size,

406 - quant_config=vllm_config.quant_config,

407 - prefix=maybe_prefix(prefix, "lm_head"),

408 - )

407 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

408

409 self.logits_processor = LogitsProcessor(config.vocab_size)

410 self.make_empty_intermediate_tensors = (

vllm/model_executor/models/openpangu.py

+1 -1

@@ -1122,7 +1122,7 @@ class OpenPanguModelBase(nn.Module, SupportsPP, SupportsLoRA):

1122 prefix=maybe_prefix(prefix, "lm_head"),

1123 )

1124 if config.tie_word_embeddings:

1125 - self.lm_head.weight = self.model.embed_tokens.weight

1125 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

1126 else:

1127 self.lm_head = PPMissingLayer()

1128 self.logits_processor = LogitsProcessor(config.vocab_size)

vllm/model_executor/models/opt.py

+6 -7

@@ -349,14 +349,13 @@ class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

349 self.model = OPTModel(

350 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

351 )

352 + self.lm_head = ParallelLMHead(

353 + config.vocab_size,

354 + config.word_embed_proj_dim,

355 + prefix=maybe_prefix(prefix, "lm_head"),

356 + )

357 if self.config.tie_word_embeddings:

353 - self.lm_head = self.model.decoder.embed_tokens

354 - else:

355 - self.lm_head = ParallelLMHead(

356 - config.vocab_size,

357 - config.word_embed_proj_dim,

358 - prefix=maybe_prefix(prefix, "lm_head"),

359 - )

358 + self.lm_head = self.lm_head.tie_weights(self.model.decoder.embed_tokens)

359 self.logits_processor = LogitsProcessor(config.vocab_size)

360 self.make_empty_intermediate_tensors = (

361 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/orion.py

+1 -1

@@ -305,7 +305,7 @@ class OrionForCausalLM(nn.Module, SupportsPP):

305 prefix=maybe_prefix(prefix, "lm_head"),

306 )

307 if self.config.tie_word_embeddings:

308 - self.lm_head.weight = self.model.embed_tokens.weight

308 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

309 self.logits_processor = LogitsProcessor(config.vocab_size)

310 self.make_empty_intermediate_tensors = (

311 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/param2moe.py

+7 -8

@@ -649,15 +649,14 @@ class Param2MoEForCausalLM(

649

650 self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False)

651 if get_pp_group().is_last_rank:

652 + self.lm_head = ParallelLMHead(

653 + config.vocab_size,

654 + config.hidden_size,

655 + quant_config=quant_config,

656 + prefix=maybe_prefix(prefix, "lm_head"),

657 + )

658 if self.tie_word_embeddings:

653 - self.lm_head = self.model.embed_tokens

654 - else:

655 - self.lm_head = ParallelLMHead(

656 - config.vocab_size,

657 - config.hidden_size,

658 - quant_config=quant_config,

659 - prefix=maybe_prefix(prefix, "lm_head"),

660 - )

659 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

660 self.logits_processor = LogitsProcessor(config.vocab_size)

661 else:

662 self.lm_head = PPMissingLayer()

vllm/model_executor/models/qwen2.py

+7 -8

@@ -460,15 +460,14 @@ class Qwen2ForCausalLM(

460 )

461

462 if get_pp_group().is_last_rank:

463 + self.lm_head = ParallelLMHead(

464 + config.vocab_size,

465 + config.hidden_size,

466 + quant_config=quant_config,

467 + prefix=maybe_prefix(prefix, "lm_head"),

468 + )

469 if config.tie_word_embeddings:

464 - self.lm_head = self.model.embed_tokens

465 - else:

466 - self.lm_head = ParallelLMHead(

467 - config.vocab_size,

468 - config.hidden_size,

469 - quant_config=quant_config,

470 - prefix=maybe_prefix(prefix, "lm_head"),

471 - )

470 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

471 else:

472 self.lm_head = PPMissingLayer()

473

vllm/model_executor/models/qwen2_moe.py

+1 -1

@@ -461,7 +461,7 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

461 prefix=maybe_prefix(prefix, "lm_head"),

462 )

463 if self.config.tie_word_embeddings:

464 - self.lm_head.weight = self.model.embed_tokens.weight

464 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

465 self.logits_processor = LogitsProcessor(config.vocab_size)

466 self.make_empty_intermediate_tensors = (

467 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/qwen3.py

+7 -8

@@ -295,15 +295,14 @@ class Qwen3ForCausalLM(

295 )

296

297 if get_pp_group().is_last_rank:

298 + self.lm_head = ParallelLMHead(

299 + config.vocab_size,

300 + config.hidden_size,

301 + quant_config=quant_config,

302 + prefix=maybe_prefix(prefix, "lm_head"),

303 + )

304 if config.tie_word_embeddings:

299 - self.lm_head = self.model.embed_tokens

300 - else:

301 - self.lm_head = ParallelLMHead(

302 - config.vocab_size,

303 - config.hidden_size,

304 - quant_config=quant_config,

305 - prefix=maybe_prefix(prefix, "lm_head"),

306 - )

305 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

306 else:

307 self.lm_head = PPMissingLayer()

308

vllm/model_executor/models/qwen3_5.py

+7 -8

@@ -334,15 +334,14 @@ class Qwen3_5ForCausalLMBase(

334 )

335

336 if get_pp_group().is_last_rank:

337 + self.lm_head = ParallelLMHead(

338 + config.vocab_size,

339 + config.hidden_size,

340 + quant_config=self.quant_config,

341 + prefix=maybe_prefix(prefix, "lm_head"),

342 + )

343 if config.tie_word_embeddings:

338 - self.lm_head = self.model.embed_tokens

339 - else:

340 - self.lm_head = ParallelLMHead(

341 - config.vocab_size,

342 - config.hidden_size,

343 - quant_config=self.quant_config,

344 - prefix=maybe_prefix(prefix, "lm_head"),

345 - )

344 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

345 else:

346 self.lm_head = PPMissingLayer()

347

vllm/model_executor/models/qwen3_5_mtp.py

+7 -8

@@ -238,15 +238,14 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):

238 )

239

240 if get_pp_group().is_last_rank:

241 + self.lm_head = ParallelLMHead(

242 + config.vocab_size,

243 + config.hidden_size,

244 + quant_config=self.quant_config,

245 + prefix=maybe_prefix(prefix, "lm_head"),

246 + )

247 if config.tie_word_embeddings:

242 - self.lm_head = self.model.embed_tokens

243 - else:

244 - self.lm_head = ParallelLMHead(

245 - config.vocab_size,

246 - config.hidden_size,

247 - quant_config=self.quant_config,

248 - prefix=maybe_prefix(prefix, "lm_head"),

249 - )

248 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

249 else:

250 self.lm_head = PPMissingLayer()

251

vllm/model_executor/models/qwen3_moe.py

+1 -1

@@ -575,7 +575,7 @@ class Qwen3MoeForCausalLM(

575 prefix=maybe_prefix(prefix, "lm_head"),

576 )

577 if self.config.tie_word_embeddings:

578 - self.lm_head.weight = self.model.embed_tokens.weight

578 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

579 self.logits_processor = LogitsProcessor(config.vocab_size)

580 self.make_empty_intermediate_tensors = (

581 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+1 -1

@@ -1111,7 +1111,7 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):

1111 config.vocab_size, config.hidden_size, quant_config=quant_config

1112 )

1113 if self.config.tie_word_embeddings:

1114 - self.lm_head.weight = self.model.embed_tokens.weight

1114 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

1115 self.logits_processor = LogitsProcessor(config.vocab_size)

1116 self.make_empty_intermediate_tensors = (

1117 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/qwen3_vl.py

+7 -8

@@ -1707,15 +1707,14 @@ class Qwen3LLMForCausalLM(Qwen3ForCausalLM):

1707 )

1708

1709 if get_pp_group().is_last_rank:

1710 + self.lm_head = ParallelLMHead(

1711 + config.vocab_size,

1712 + config.hidden_size,

1713 + quant_config=quant_config,

1714 + prefix="lm_head",

1715 + )

1716 if config.tie_word_embeddings:

1711 - self.lm_head = self.model.embed_tokens

1712 - else:

1713 - self.lm_head = ParallelLMHead(

1714 - config.vocab_size,

1715 - config.hidden_size,

1716 - quant_config=quant_config,

1717 - prefix="lm_head",

1718 - )

1717 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

1718 else:

1719 self.lm_head = PPMissingLayer()

1720

vllm/model_executor/models/qwen3_vl_moe.py

+1 -1

@@ -168,7 +168,7 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):

168 prefix=maybe_prefix(prefix, "lm_head"),

169 )

170 if self.config.tie_word_embeddings:

171 - self.lm_head.weight = self.model.embed_tokens.weight

171 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

172 self.logits_processor = LogitsProcessor(self.config.vocab_size)

173 self.make_empty_intermediate_tensors = (

174 self.model.make_empty_intermediate_tensors

vllm/model_executor/models/sarvam.py

+7 -8

@@ -615,15 +615,14 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE

615

616 self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)

617 if get_pp_group().is_last_rank:

618 + self.lm_head = ParallelLMHead(

619 + config.vocab_size,

620 + config.hidden_size,

621 + quant_config=quant_config,

622 + prefix=maybe_prefix(prefix, "lm_head"),

623 + )

624 if self.tie_word_embeddings:

619 - self.lm_head = self.model.embed_tokens

620 - else:

621 - self.lm_head = ParallelLMHead(

622 - config.vocab_size,

623 - config.hidden_size,

624 - quant_config=quant_config,

625 - prefix=maybe_prefix(prefix, "lm_head"),

626 - )

625 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

626 self.logits_processor = LogitsProcessor(config.vocab_size)

627 else:

628 self.lm_head = PPMissingLayer()

vllm/model_executor/models/seed_oss.py

+7 -8

@@ -388,15 +388,14 @@ class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

388 )

389

390 if get_pp_group().is_last_rank:

391 + self.lm_head = ParallelLMHead(

392 + config.vocab_size,

393 + config.hidden_size,

394 + quant_config=quant_config,

395 + prefix=maybe_prefix(prefix, "lm_head"),

396 + )

397 if config.tie_word_embeddings:

392 - self.lm_head = self.model.embed_tokens

393 - else:

394 - self.lm_head = ParallelLMHead(

395 - config.vocab_size,

396 - config.hidden_size,

397 - quant_config=quant_config,

398 - prefix=maybe_prefix(prefix, "lm_head"),

399 - )

398 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

399 else:

400 self.lm_head = PPMissingLayer()

401

vllm/model_executor/models/solar.py

+1 -1

@@ -386,7 +386,7 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

386 prefix=maybe_prefix(prefix, "lm_head"),

387 )

388 if config.tie_word_embeddings:

389 - self.lm_head.weight = self.model.embed_tokens.weight

389 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

390

391 logit_scale = getattr(config, "logit_scale", 1.0)

392 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/stablelm.py

+1 -1

@@ -294,7 +294,7 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

294 prefix=f"{prefix}.lm_head",

295 )

296 if self.config.tie_word_embeddings:

297 - self.lm_head.weight = self.model.embed_tokens.weight

297 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

298 self.logits_processor = LogitsProcessor(config.vocab_size)

299 self.make_empty_intermediate_tensors = (

300 self.model.make_empty_intermediate_tensors