137 lines
5.5 KiB
Python
137 lines
5.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Inkling logits processor (muP + LoRA aware).
|
|
|
|
Inkling divides the final logits by a muP width multiplier
|
|
(``logits_mup_width_multiplier``). This applies it two ways, depending on
|
|
whether an lm_head LoRA is attached:
|
|
|
|
* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no
|
|
separate elementwise kernel, no extra rounding, no weight mutation.
|
|
* LoRA attached: the LoRA manager wraps this layer in
|
|
``LogitsProcessorWithLoRA``, whose ``forward`` calls
|
|
``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with
|
|
``self`` bound to the wrapper. We detect that via ``base_layer`` and take the
|
|
LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA
|
|
delta), then divide the full logits by the multiplier so the delta is scaled
|
|
too. muP thus composes with the LoRA delta, with the dispatch as the only
|
|
model-side branching.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
|
|
|
|
|
class InklingLogitsProcessor(LogitsProcessor):
|
|
"""``LogitsProcessor`` that applies Inkling's muP logits width multiplier.
|
|
|
|
Args:
|
|
vocab_size: Padded vocabulary size.
|
|
org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``).
|
|
scale: Base logits scale (kept ``1.0`` for the served checkpoint).
|
|
logits_as_input: Whether the input is already logits.
|
|
soft_cap: Optional logit soft cap (``None`` for the served checkpoint).
|
|
logits_mup_width_multiplier: muP width divisor for the final logits;
|
|
``None`` or ``0`` disables it.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
vocab_size: int,
|
|
org_vocab_size: int | None = None,
|
|
scale: float = 1.0,
|
|
logits_as_input: bool = False,
|
|
soft_cap: float | None = None,
|
|
logits_mup_width_multiplier: float | None = None,
|
|
) -> None:
|
|
super().__init__(
|
|
vocab_size=vocab_size,
|
|
org_vocab_size=org_vocab_size,
|
|
scale=scale,
|
|
logits_as_input=logits_as_input,
|
|
soft_cap=soft_cap,
|
|
)
|
|
self.logits_mup_width_multiplier = logits_mup_width_multiplier
|
|
self._logits_zero: torch.Tensor | None = None
|
|
|
|
def forward(
|
|
self,
|
|
lm_head: VocabParallelEmbedding,
|
|
hidden_states: torch.Tensor,
|
|
embedding_bias: torch.Tensor | None = None,
|
|
skip_gather: bool = False,
|
|
) -> torch.Tensor | None:
|
|
# ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper,
|
|
# which calls this forward with ``self`` bound to the wrapper. The
|
|
# wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch
|
|
# ``_lora_forward`` explicitly through the base_layer's class (it
|
|
# provides ``_get_logits``/``logits_as_input``; only ``_lora_forward``
|
|
# lives on this class).
|
|
if hasattr(self, "base_layer"):
|
|
return type(self.base_layer)._lora_forward(
|
|
self, lm_head, hidden_states, embedding_bias, skip_gather
|
|
)
|
|
return self._base_forward(lm_head, hidden_states, embedding_bias, skip_gather)
|
|
|
|
def _lora_forward(
|
|
self,
|
|
lm_head: VocabParallelEmbedding,
|
|
hidden_states: torch.Tensor,
|
|
embedding_bias: torch.Tensor | None = None,
|
|
skip_gather: bool = False,
|
|
) -> torch.Tensor | None:
|
|
# ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits``
|
|
# returns the base logits plus the lm_head LoRA delta. Apply the muP
|
|
# divisor on the full logits so the LoRA delta is scaled too.
|
|
mup_multiplier = self.base_layer.logits_mup_width_multiplier
|
|
mup = 1.0 / mup_multiplier if mup_multiplier else None
|
|
if self.logits_as_input:
|
|
logits = hidden_states
|
|
else:
|
|
logits = self._get_logits(
|
|
hidden_states, lm_head, embedding_bias, skip_gather
|
|
)
|
|
# TODO: fuse this multiplication
|
|
if logits is not None and mup:
|
|
assert self.base_layer.soft_cap is None
|
|
assert self.base_layer.scale == 1.0
|
|
logits = logits * mup
|
|
return logits
|
|
|
|
def _base_forward(
|
|
self,
|
|
lm_head: VocabParallelEmbedding,
|
|
hidden_states: torch.Tensor,
|
|
embedding_bias: torch.Tensor | None = None,
|
|
skip_gather: bool = False,
|
|
) -> torch.Tensor | None:
|
|
mup = self.logits_mup_width_multiplier
|
|
if not mup:
|
|
return super().forward(lm_head, hidden_states, embedding_bias, skip_gather)
|
|
# Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
|
|
# no separate elementwise kernel, no bf16 rounding of scaled logits, and
|
|
# no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
|
|
# soft cap, unit logits scale.
|
|
assert self.soft_cap is None
|
|
assert self.scale == 1.0
|
|
w = lm_head.weight
|
|
if self._logits_zero is None:
|
|
self._logits_zero = w.new_zeros(1)
|
|
logits = torch.addmm(
|
|
self._logits_zero,
|
|
hidden_states,
|
|
w.t(),
|
|
beta=0.0,
|
|
alpha=1.0 / mup,
|
|
)
|
|
if skip_gather:
|
|
return logits
|
|
logits = self._gather_logits(logits)
|
|
if logits is not None:
|
|
logits = logits[..., : self.org_vocab_size]
|
|
return logits
|