vllm.model_executor.models.transformers.causal

Transformers modeling backend mixin for causal language models.

Classes:

• CausalMixin –

CausalMixin

Bases: VllmModelForTextGeneration, Base

Methods:

• compute_logits_local –

  Logits for this TP rank's vocab shard, for batch-sharded sampling.

• load_weights –

  A thin wrapper around Base.load_weights to handle the lm_head bias.

Source code in vllm/model_executor/models/transformers/causal.py

class CausalMixin(VllmModelForTextGeneration, Base):
    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
        # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO
        super(VllmModelForTextGeneration, self).__init__(
            vllm_config=vllm_config, prefix=prefix
        )

        tie_word_embeddings = self._get_tie_word_embeddings()

        if self.pp_group.is_last_rank:
            self.lm_head = ParallelLMHead(
                self.text_config.vocab_size,
                self.text_config.hidden_size,
                quant_config=self.quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
            if tie_word_embeddings:
                for module in self.model.get_input_embeddings().modules():
                    if isinstance(module, VocabParallelEmbedding):
                        self.lm_head = self.lm_head.tie_weights(module)
                        break

            self.logits_processor = LogitsProcessor(
                self.text_config.vocab_size,
                scale=getattr(self.text_config, "logit_scale", 1.0),
                soft_cap=getattr(self.text_config, "final_logit_softcapping", None),
            )
        else:
            self.lm_head = PPMissingLayer()

    def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
        """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
        lm_head_bias = set()
        load_device = self.vllm_config.load_config.device or self.device_config.device

        def auto_load_lm_head_bias(weights):
            for name, weight in weights:
                if name.endswith("lm_head.bias") and self.pp_group.is_last_rank:
                    # load_weights is called outside vLLM's device context,
                    # so we must enter it again when registering the bias.
                    with torch.device(load_device):
                        self.lm_head._register_bias()
                    self.lm_head.bias.weight_loader(self.lm_head.bias, weight)
                    lm_head_bias.add(name)
                else:
                    yield name, weight

        return super().load_weights(auto_load_lm_head_bias(weights)) | lm_head_bias

    def compute_logits_local(self, hidden_states: "torch.Tensor") -> "torch.Tensor":
        """Logits for this TP rank's vocab shard, for batch-sharded sampling."""
        bias = self.lm_head.bias
        return self.logits_processor(
            self.lm_head, hidden_states, bias, skip_gather=True
        )

    def compute_logits(self, hidden_states: "torch.Tensor") -> "torch.Tensor | None":
        logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)
        return logits

compute_logits_local(hidden_states)

Logits for this TP rank's vocab shard, for batch-sharded sampling.

Source code in vllm/model_executor/models/transformers/causal.py

def compute_logits_local(self, hidden_states: "torch.Tensor") -> "torch.Tensor":
    """Logits for this TP rank's vocab shard, for batch-sharded sampling."""
    bias = self.lm_head.bias
    return self.logits_processor(
        self.lm_head, hidden_states, bias, skip_gather=True
    )

load_weights(weights)

A thin wrapper around Base.load_weights to handle the lm_head bias.

Source code in vllm/model_executor/models/transformers/causal.py

def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
    """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
    lm_head_bias = set()
    load_device = self.vllm_config.load_config.device or self.device_config.device

    def auto_load_lm_head_bias(weights):
        for name, weight in weights:
            if name.endswith("lm_head.bias") and self.pp_group.is_last_rank:
                # load_weights is called outside vLLM's device context,
                # so we must enter it again when registering the bias.
                with torch.device(load_device):
                    self.lm_head._register_bias()
                self.lm_head.bias.weight_loader(self.lm_head.bias, weight)
                lm_head_bias.add(name)
            else:
                yield name, weight

    return super().load_weights(auto_load_lm_head_bias(weights)) | lm_head_bias