modeling_phi.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
#
# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
# Licensed under the BSD 3-Clause License.

from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Union, Tuple

import math
import torch
import torch.nn as nn
from einops import rearrange, repeat
from transformers import PretrainedConfig, PreTrainedModel
from transformers.activations import ACT2FN
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_moondream import PhiConfig

FusedDense = None


@dataclass
class InferenceParams:
    max_seqlen: int
    max_batch_size: int
    seqlen_offset: int = 0
    batch_size_offset: int = 0
    key_value_memory_dict: Dict[str, Any] = field(default_factory=dict)
    lengths_per_sample: torch.Tensor = None


class Embedding(nn.Module):
    def __init__(self, config: PretrainedConfig):
        super().__init__()
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.drop = nn.Dropout(config.embd_pdrop)

    def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
        return self.drop(self.wte(input_ids.view(-1, input_ids.size(-1))))


def _apply_rotary_emb(x, cos, sin):
    seqlen, rotary_dim = x.size(1), cos.size(1) * 2
    x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
    x1, x2 = x_rot.chunk(2, dim=-1)
    c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1)
    x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1)
    return torch.cat([x_rot.to(x.dtype), x_pass], dim=-1)


def _apply_rotary_emb_kv(
    kv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor
) -> torch.FloatTensor:
    seqlen, rotary_dim = kv.shape[1], cos.shape[-1] * 2
    k_rot = kv[:, :, 0, :, :rotary_dim].chunk(2, dim=-1)
    k_pass = kv[:, :, 0, :, rotary_dim:]
    c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1)
    k_rot = torch.cat(
        [k_rot[0] * c - k_rot[1] * s, k_rot[0] * s + k_rot[1] * c], dim=-1
    )
    return torch.cat(
        [torch.cat([k_rot, k_pass], dim=-1).unsqueeze(2), kv[:, :, 1:2, :, :]], dim=2
    )


def _apply_rotary_emb_qkv(
    qkv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor
) -> torch.FloatTensor:
    seqlen, rotary_dim = qkv.shape[1], cos.shape[1] * 2

    c = cos[:seqlen].unsqueeze(1)
    s = sin[:seqlen].unsqueeze(1)

    qkv_rot = torch.stack(
        [
            torch.cat(
                [
                    qkv[:, :, i, :, : rotary_dim // 2] * c
                    - qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * s,
                    qkv[:, :, i, :, : rotary_dim // 2] * s
                    + qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * c,
                ],
                dim=-1,
            ).to(qkv.dtype)
            for i in range(2)
        ],
        dim=2,
    )

    qkv_pass = qkv[:, :, :2, :, rotary_dim:].unsqueeze(2)
    qkv_v = qkv[:, :, 2:3, :, :]
    return torch.cat([qkv_rot, qkv_pass, qkv_v], dim=2)


class RotaryEmbedding(nn.Module):
    # Enhanced Transformer with Rotary Position Embedding (https://arxiv.org/pdf/2104.09864.pdf)
    def __init__(
        self,
        dim: int,
        base: int = 10000,
        scale_base: Optional[float] = None,
        pos_idx_in_fp32: bool = True,
        max_position_embeddings: int = 2048,
        device: Optional[str] = None,
    ) -> None:
        super().__init__()
        # fp32 is preferred since the output of `torch.arange` can be quite large and bf16 would lose a lot of precision
        self.dim, self.base, self.pos_idx_in_fp32, self.device = (
            dim,
            float(base),
            pos_idx_in_fp32,
            device,
        )
        self.max_position_embeddings = max_position_embeddings
        if scale_base is not None:
            raise NotImplementedError

        # Generate and register the non-trainable buffers
        self.register_buffer(
            "inv_freq", self._compute_inv_freq(device), persistent=False
        )
        self.register_buffer(
            "scale", self._calculate_scale(dim, scale_base, device), persistent=False
        )
        self._update_cos_sin_cache(
            max_position_embeddings, device=device, dtype=torch.float32
        )

    def _calculate_scale(self, dim, scale_base, device):
        return (
            (
                (
                    torch.arange(0, dim, 2, device=device, dtype=torch.float32)
                    + 0.4 * dim
                )
                / (1.4 * dim)
            )
            if scale_base is not None
            else None
        )

    def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
        return 1.0 / (
            self.base
            ** (
                torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
                / self.dim
            )
        )

    def _update_cos_sin_cache(
        self,
        seqlen: int,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None:
        self._seq_len_cached = seqlen
        t = torch.arange(
            seqlen,
            device=device,
            dtype=torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype,
        )
        inv_freq = (
            self._compute_inv_freq(device=device)
            if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32
            else self.inv_freq
        )

        freqs = torch.outer(t, inv_freq)

        def apply_scale(freqs, scale, operator, dtype):
            result = operator(freqs)
            return (result / scale).to(dtype) if scale is not None else result.to(dtype)

        if scale := self.scale:
            power = (
                torch.arange(seqlen, dtype=scale.dtype, device=scale.device)
                - seqlen // 2
            ) / self.scale_base
            scale = scale.to(device=power.device) ** power.unsqueeze(1)

        self._cos_cached = apply_scale(
            freqs, 1 / scale if scale is not None else None, torch.cos, dtype
        )
        self._sin_cached = apply_scale(
            freqs, 1 / scale if scale is not None else None, torch.sin, dtype
        )
        if scale is not None:
            self._cos_k_cached = apply_scale(freqs, scale, torch.cos, dtype)
            self._sin_k_cached = apply_scale(freqs, scale, torch.sin, dtype)

    def forward(
        self,
        qkv: torch.Tensor,
        kv: Optional[torch.Tensor] = None,
        seqlen_offset: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        should_update = (
            self._seq_len_cached < qkv.shape[1] + seqlen_offset
            or self._cos_cached.device != qkv.device
            or self._cos_cached.dtype != qkv.dtype
            or (self.training and self._cos_cached.is_inference())
        )

        if should_update:
            self._update_cos_sin_cache(
                qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype
            )

        offset_cos = self._cos_cached[seqlen_offset:]
        offset_sin = self._sin_cached[seqlen_offset:]

        if kv is None:
            return _apply_rotary_emb_qkv(qkv, offset_cos, offset_sin)
        else:
            return _apply_rotary_emb(qkv, offset_cos, offset_sin), _apply_rotary_emb_kv(
                kv, offset_cos, offset_sin
            )


class MLP(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        n_inner: Optional[int] = None,
        act_fn: Optional[str] = None,
    ) -> None:
        super().__init__()
        n_inner = n_inner or getattr(config, "n_inner", None) or 4 * config.n_embd
        act_fn = act_fn or config.activation_function

        self.fc1 = nn.Linear(config.n_embd, n_inner)
        self.fc2 = nn.Linear(n_inner, config.n_embd)
        self.act = ACT2FN[act_fn]

    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        return self.fc2(self.act(self.fc1(hidden_states)))


# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py)
class SelfAttention(nn.Module):
    def __init__(
        self,
        causal: bool = True,
        softmax_scale: Optional[float] = None,
        attention_dropout: float = 0.0,
    ):
        super().__init__()
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.drop = nn.Dropout(attention_dropout)

    @torch.autocast("cpu", enabled=False)
    @torch.autocast("cuda", enabled=False)
    def forward(
        self,
        qkv: torch.FloatTensor,
        causal: Optional[bool] = None,
        key_padding_mask: Optional[torch.BoolTensor] = None,
    ):
        q, k, v = qkv.chunk(3, dim=-1)
        scale = self.softmax_scale or 1.0 / q.size(-1) ** 0.5

        scores = (
            torch.einsum("bthd,bshd->bhts", q.to(torch.float32), k.to(torch.float32))
            * scale
        )
        if causal or self.causal:
            scores.triu_(1).fill_(-10000.0)
        if key_padding_mask is not None:
            scores.masked_fill_(key_padding_mask[:, None, None, :], -10000.0)

        attn = self.drop(torch.softmax(scores, dim=-1).to(v.dtype))
        return torch.einsum("bhts,bshd->bthd", attn, v)


# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py)
class CrossAttention(nn.Module):
    def __init__(self, causal=True, softmax_scale=None, attention_dropout=0.0):
        super().__init__()
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.drop = nn.Dropout(attention_dropout)

    @torch.autocast("cpu", enabled=False)
    @torch.autocast("cuda", enabled=False)
    def forward(
        self,
        q: torch.FloatTensor,
        kv: torch.FloatTensor,
        causal: bool = None,
        key_padding_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.FloatTensor:
        batch_size, seqlen_q = q.shape[0], q.shape[1]
        seqlen_k = kv.shape[1]

        if kv.shape[3] != q.shape[2]:
            kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
        k, v = kv.unbind(dim=2)

        q = q.to(torch.float32)
        k = k.to(torch.float32)

        causal = self.causal if causal is None else causal
        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])

        # Autocast is manually disabled to avoid `torch.einsum` performing the operation using float16, which might lead to overflow
        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)

        if key_padding_mask is not None:
            padding_mask = torch.full(
                (batch_size, seqlen_k),
                -10000.0,
                dtype=scores.dtype,
                device=scores.device,
            )
            padding_mask.masked_fill_(key_padding_mask, 0.0)
            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")

        if causal:
            rows = rearrange(
                torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1"
            )
            cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
            causal_mask = cols > rows + seqlen_k - seqlen_q
            scores = scores.masked_fill(causal_mask, -10000.0)

        attention = torch.softmax(scores, dim=-1).to(v.dtype)
        attention = self.drop(attention)
        output = torch.einsum("bhts,bshd->bthd", attention, v)

        return output


def _find_mha_dims(
    config: PretrainedConfig,
    n_head: Optional[int] = None,
    n_head_kv: Optional[int] = None,
    head_dim: Optional[int] = None,
) -> Tuple[int, int]:
    if n_head is None and head_dim is None:
        head_dim = config.n_embd // config.n_head
        n_head = config.n_head
    elif n_head is None or head_dim is None:
        raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
    if n_head_kv is None:
        n_head_kv = getattr(config, "n_head_kv", None) or n_head
    return n_head, n_head_kv, head_dim


def _update_kv_cache(
    kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int
) -> torch.FloatTensor:
    num_heads, head_dim = kv.shape[-2:]
    layer_memory = inference_params.key_value_memory_dict.setdefault(
        layer_idx,
        torch.empty(
            inference_params.max_batch_size,
            inference_params.max_seqlen,
            2,
            num_heads,
            head_dim,
            dtype=kv.dtype,
            device=kv.device,
        ),
    )

    batch_slice = slice(
        inference_params.batch_size_offset,
        inference_params.batch_size_offset + kv.shape[0],
    )
    seqlen_slice = slice(
        inference_params.seqlen_offset, inference_params.seqlen_offset + kv.shape[1]
    )

    if seqlen_slice.stop >= inference_params.max_seqlen:
        layer_memory = torch.cat((layer_memory, kv), dim=1)
        inference_params.key_value_memory_dict[layer_idx] = layer_memory

    layer_memory[batch_slice, seqlen_slice, ...] = kv
    return layer_memory[batch_slice, : seqlen_slice.stop, ...]


# Multi-head attention layer with rotary embeddings
class MHA(nn.Module):
    def __init__(
        self,
        config,
        dtype=None,
        device=None,
        rotary_dim=None,
        rotary_base=10000.0,
        rotary_scale_base=None,
        n_head=None,
        n_head_kv=None,
        head_dim=None,
        bias=True,
        causal=True,
        softmax_scale=None,
        layer_idx=None,
        return_residual=False,
        checkpointing=False,
    ):
        super().__init__()

        # Set rotary embedding if specified
        self.rotary_dim = rotary_dim or getattr(config, "rotary_dim", 0)
        if self.rotary_dim:
            self.rotary_emb = RotaryEmbedding(
                self.rotary_dim,
                base=rotary_base,
                scale_base=rotary_scale_base,
                device=device,
                max_position_embeddings=config.n_positions,
            )

        # Determine MHA dims from arguments or config
        self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
            config, n_head, n_head_kv, head_dim
        )
        op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
        hidden_size = config.n_embd

        # Choose Linear class based on config, FusedDense is optional
        LinearClass = (
            FusedDense if config.fused_dense and FusedDense is not None else nn.Linear
        )
        self.Wqkv = LinearClass(
            hidden_size, op_size, bias=bias, device=device, dtype=dtype
        )
        self.out_proj = LinearClass(
            hidden_size, hidden_size, bias=bias, device=device, dtype=dtype
        )

        # Initialize attention mechanisms
        attn_kwargs = {
            "causal": causal,
            "softmax_scale": softmax_scale,
            "attention_dropout": config.attn_pdrop,
        }
        self.inner_attn = SelfAttention(**attn_kwargs)
        self.inner_cross_attn = CrossAttention(**attn_kwargs)

        self.layer_idx = layer_idx
        self.return_residual = return_residual
        self.checkpointing = checkpointing

    def _forward_self_attn(
        self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
    ) -> torch.FloatTensor:
        qkv = rearrange(
            self.Wqkv(x), "... (three h d) -> ... three h d", three=3, d=self.head_dim
        )
        if self.rotary_dim > 0:
            qkv = self.rotary_emb(qkv)
        attn_func = (
            torch.utils.checkpoint.checkpoint
            if self.checkpointing
            else lambda f, *args, **kwargs: f(*args, **kwargs)
        )
        return attn_func(self.inner_attn, qkv, key_padding_mask=key_padding_mask)

    def _forward_cross_attn(
        self,
        x: torch.FloatTensor,
        past_key_values: Optional[InferenceParams],
        key_padding_mask: Optional[torch.BoolTensor],
    ) -> torch.FloatTensor:
        qkv = self.Wqkv(x)
        q, kv = (
            qkv[..., : self.n_head * self.head_dim],
            qkv[..., self.n_head * self.head_dim :],
        )
        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)

        seqlen_offset = (
            past_key_values.seqlen_offset if past_key_values is not None else 0
        )
        causal = None if seqlen_offset == 0 else False
        if self.rotary_dim > 0:
            q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)

        if past_key_values is not None:
            kv = _update_kv_cache(kv, past_key_values, self.layer_idx)

        attn_func = (
            torch.utils.checkpoint.checkpoint
            if self.checkpointing
            else lambda fn, *args, **kwargs: fn(*args, **kwargs)
        )

        return attn_func(
            self.inner_cross_attn,
            q,
            kv,
            key_padding_mask=key_padding_mask,
            causal=causal,
        )

    def forward(
        self,
        x: torch.FloatTensor,
        past_key_values: Optional[InferenceParams] = None,
        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
        attention_mask = attention_mask.bool() if attention_mask is not None else None
        use_cross_attn = self.n_head != self.n_head_kv or past_key_values is not None
        attn_output_function = (
            self._forward_cross_attn if use_cross_attn else self._forward_self_attn
        )
        attn_output = (
            attn_output_function(x, past_key_values, attention_mask)
            if use_cross_attn
            else attn_output_function(x, attention_mask)
        )
        output = self.out_proj(rearrange(attn_output, "... h d -> ... (h d)"))
        return (output, x) if self.return_residual else output


# Parallel block. This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
class ParallelBlock(nn.Module):
    def __init__(self, config: PretrainedConfig, block_idx: Optional[int] = None):
        super().__init__()
        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        self.block_idx = block_idx
        self.mixer = MHA(config, layer_idx=block_idx)
        self.mlp = MLP(config)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.FloatTensor:
        residual = hidden_states
        hidden_states = self.ln(hidden_states)

        attn_outputs = self.mixer(
            hidden_states,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
        )
        if isinstance(attn_outputs, tuple):
            attn_outputs = attn_outputs[0]

        attn_outputs = self.resid_dropout(attn_outputs)
        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
        return attn_outputs + feed_forward_hidden_states + residual


class CausalLMHead(nn.Module):
    """Causal Language Modeling head. Simplified version."""

    def __init__(self, config):
        super().__init__()
        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.linear = nn.Linear(config.n_embd, config.vocab_size)

    def forward(self, hidden_states):
        return self.linear(self.ln(hidden_states)).to(torch.float32)


# Improving Language Understanding by Generative Pre-Training
# (https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf)
class CausalLMLoss(nn.Module):
    def __init__(self, shift_labels: bool = True) -> None:
        super().__init__()
        self.shift_labels = shift_labels
        self.loss_fct = nn.CrossEntropyLoss()

    def forward(
        self, logits: torch.FloatTensor, labels: torch.LongTensor
    ) -> torch.FloatTensor:
        if self.shift_labels:
            logits, labels = logits[..., :-1, :], labels[..., 1:]
        return self.loss_fct(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))


class PhiPreTrainedModel(PreTrainedModel):
    config_class = PhiConfig
    base_model_prefix = "transformer"
    supports_gradient_checkpointing = False
    _no_split_modules = ["ParallelBlock"]

    def __init__(self, *inputs, **kwargs) -> None:
        super().__init__(*inputs, **kwargs)

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor = None,
        inputs_embeds: torch.FloatTensor = None,
        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        if input_ids is None and inputs_embeds is None:
            raise ValueError(
                "You have to specify either `input_ids` or `inputs_embeds`."
            )

        max_batch_size = (
            inputs_embeds.shape[0] if inputs_embeds is not None else input_ids.shape[0]
        )
        seqlen_offset = (
            inputs_embeds.shape[1] + input_ids.shape[1] - 2
            if inputs_embeds is not None
            else input_ids.shape[1] - 1
        )

        args = (
            {"inputs_embeds": inputs_embeds}
            if inputs_embeds is not None
            else {"input_ids": input_ids}
        )

        if not isinstance(past_key_values, InferenceParams):
            past_key_values = InferenceParams(
                max_seqlen=self.config.n_positions,
                max_batch_size=max_batch_size,
                seqlen_offset=0,
                batch_size_offset=0,
                key_value_memory_dict={},
                lengths_per_sample=None,
            )
        else:
            past_key_values.seqlen_offset = seqlen_offset
            args = {"input_ids": input_ids[:, -1].unsqueeze(-1)}

        return {
            **args,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
        }


class PhiModel(PhiPreTrainedModel):
    _keys_to_ignore_on_load_missing = [""]
    _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]

    def __init__(self, config: PhiConfig) -> None:
        super().__init__(config)
        self.embd = Embedding(config)
        self.h = nn.ModuleList(
            [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)]
        )
        self.gradient_checkpointing = config.gradient_checkpointing
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embd.wte

    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
        self.embd.wte = new_embeddings

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        inputs_embeds: torch.FloatTensor = None,
        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.FloatTensor:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of `input_ids` or `inputs_embeds`.")
        hidden_states = self.embd(input_ids) if input_ids is not None else inputs_embeds

        for layer in self.h:
            func = layer.__call__ if self.gradient_checkpointing else layer
            args = (hidden_states, past_key_values, attention_mask)
            hidden_states = (
                torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=True)
                if self.gradient_checkpointing
                else func(*args)
            )

        return hidden_states


class PhiForCausalLM(PhiPreTrainedModel):
    _keys_to_ignore_on_load_missing, _keys_to_ignore_on_load_unexpected = (
        [""],
        [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"],
    )

    def __init__(self, config: PhiConfig) -> None:
        super().__init__(config)
        self.transformer = PhiModel(config)
        self.lm_head = CausalLMHead(config)
        self.loss = CausalLMLoss()
        self.post_init()

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head.linear

    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
        self.lm_head.linear = new_embeddings

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        inputs_embeds: torch.FloatTensor = None,
        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        hidden_states = self.transformer(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
        )
        lm_logits = self.lm_head(hidden_states)
        loss = self.loss(lm_logits, labels) if labels is not None else None

        return CausalLMOutputWithPast(
            loss=loss, logits=lm_logits, past_key_values=past_key_values
        )