Tokle-3M / modeling_tokle.py

Upload 10 files

ff59975 verified 25 days ago

13 kB

"""Tokle: decoder-only transformer with RMSNorm, RoPE, GQA and SwiGLU.
Module and parameter names follow the usual Transformers layout.
Works on Transformers 4.x and 5.x; the v5 differences are marked inline.
"""
import math
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import __version__ as _transformers_version
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
from transformers.modeling_utils import PreTrainedModel

try:  # as a Transformers dynamic module (trust_remote_code)
    from .configuration_tokle import TokleConfig
except ImportError:  # running the files straight out of a checkout
    from configuration_tokle import TokleConfig


# ----------------------------------------------------------------------------- norms and rope
class TokleRMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        # Always reduce in fp32: in fp16 the mean of squares underflows on the
        # long tail and the norm comes back slightly wrong.
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dt)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


def build_rope(head_dim, max_seq_len, theta, device=None):
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
                                        device=device) / head_dim))
    t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
    f = torch.outer(t, inv)
    return torch.cos(f), torch.sin(f)


def apply_rope(x, cos, sin):
    """Half-split rotation (first half against second), not the interleaved
    variant. The two are not interchangeable -- swapping them degrades output
    without raising anything."""
    b, h, t, d = x.shape
    x1, x2 = x.float().chunk(2, dim=-1)
    c = cos[:t].view(1, 1, t, d // 2).to(x1.device)
    s = sin[:t].view(1, 1, t, d // 2).to(x1.device)
    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1).to(x.dtype)


# ----------------------------------------------------------------------------- blocks
class TokleAttention(nn.Module):
    def __init__(self, config: TokleConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.head_dim = config.head_dim
        self.rep = config.n_head // config.n_kv_head
        self.q_proj = nn.Linear(config.d_model, config.n_head * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_head * config.head_dim, config.d_model, bias=False)
        self.q_norm = TokleRMSNorm(config.head_dim, config.norm_eps)
        self.k_norm = TokleRMSNorm(config.head_dim, config.norm_eps)

    def forward(self, x, cos, sin, attn_mask):
        b, t, _ = x.shape
        q = self.q_proj(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
        # QK-norm is applied before rope.
        q = apply_rope(self.q_norm(q), cos, sin)
        k = apply_rope(self.k_norm(k), cos, sin)
        k = k.repeat_interleave(self.rep, dim=1)
        v = v.repeat_interleave(self.rep, dim=1)
        if attn_mask is None:
            # Nothing padded: let SDPA build the causal mask itself, which is
            # the fast kernel path and what an unpadded forward always hits.
            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.dtype))
        return self.o_proj(o.transpose(1, 2).reshape(b, t, -1))


class TokleMLP(nn.Module):
    def __init__(self, config: TokleConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
        self.up_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
        self.down_proj = nn.Linear(config.ffn_hidden, config.d_model, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TokleDecoderLayer(nn.Module):
    def __init__(self, config: TokleConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
        self.self_attn = TokleAttention(config, layer_idx)
        self.post_attention_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
        self.mlp = TokleMLP(config)

    def forward(self, x, cos, sin, attn_mask):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask)
        return x + self.mlp(self.post_attention_layernorm(x))


# ----------------------------------------------------------------------------- model
class ToklePreTrainedModel(PreTrainedModel):
    config_class = TokleConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["TokleDecoderLayer"]
    _supports_sdpa = True

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, 0.0, 0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, 0.0, 0.02)
        elif isinstance(module, TokleRMSNorm):
            nn.init.ones_(module.weight)
        # Residual-output projections get scaled down by depth so the residual
        # stream does not blow up at init (GPT-2 trick).
        if isinstance(module, (TokleAttention, TokleMLP)):
            out = module.o_proj if isinstance(module, TokleAttention) else module.down_proj
            nn.init.normal_(out.weight, 0.0, 0.02 / math.sqrt(2 * self.config.n_layer))


class TokleModel(ToklePreTrainedModel):
    def __init__(self, config: TokleConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList(
            [TokleDecoderLayer(config, i) for i in range(config.n_layer)])
        self.norm = TokleRMSNorm(config.d_model, config.norm_eps)
        # A plain dict, not a buffer. Transformers 5 builds models on the meta
        # device and fills only tensors the weights file contains, so a
        # non-persistent buffer would come back as uninitialised memory.
        # Keyed by device, so .to()/.cuda() need not move anything.
        self._rope_cache = {}
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def rope(self, device):
        key = str(device)
        if key not in self._rope_cache:
            self._rope_cache[key] = build_rope(
                self.config.head_dim, self.config.max_seq_len,
                self.config.rope_theta, device=device)
        return self._rope_cache[key]

    def _causal_pad_mask(self, attention_mask, t, dtype, device):
        """Additive (b, 1, 1, t) padding mask folded together with the causal
        mask, or None when nothing is padded so the caller can take the
        is_causal fast path."""
        if attention_mask is None:
            return None
        attention_mask = attention_mask.to(torch.bool)
        if bool(attention_mask.all()):
            return None
        m = torch.zeros(attention_mask.shape, dtype=dtype, device=device)
        m = m.masked_fill(~attention_mask, torch.finfo(dtype).min)[:, None, None, :]
        # Under causal attention with right padding, position 0 always sees
        # itself, so no row ends up fully masked (which would give NaNs).
        causal = torch.ones(t, t, dtype=torch.bool, device=device).triu(1)
        return m.masked_fill(causal, torch.finfo(dtype).min)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutput]:
        # getattr, not config.use_return_dict: that property is deprecated in
        # Transformers 5 and warns on every call.
        return_dict = (return_dict if return_dict is not None
                       else getattr(self.config, "return_dict", True))
        output_hidden_states = (output_hidden_states if output_hidden_states is not None
                                else self.config.output_hidden_states)
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")

        x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        t = x.shape[1]
        if t > self.config.max_seq_len:
            raise ValueError(
                f"Sequence length {t} exceeds max_seq_len {self.config.max_seq_len}.")

        cos, sin = self.rope(x.device)
        mask = self._causal_pad_mask(attention_mask, t, x.dtype, x.device)

        hidden_states = () if output_hidden_states else None
        for layer in self.layers:
            if output_hidden_states:
                hidden_states += (x,)
            if self.gradient_checkpointing and self.training:
                x = self._gradient_checkpointing_func(layer.__call__, x, cos, sin, mask)
            else:
                x = layer(x, cos, sin, mask)
        x = self.norm(x)
        if output_hidden_states:
            hidden_states += (x,)

        if not return_dict:
            return tuple(v for v in (x, hidden_states) if v is not None)
        return BaseModelOutput(last_hidden_state=x, hidden_states=hidden_states)


# Transformers 5 turned _tied_weights_keys from a list of names into a
# {tied: source} mapping and calls .keys() on it, so the old form raises.
_TRANSFORMERS_V5 = int(_transformers_version.split(".")[0]) >= 5


class TokleForCausalLM(ToklePreTrainedModel, GenerationMixin):
    _tied_weights_keys = ({"lm_head.weight": "model.embed_tokens.weight"}
                          if _TRANSFORMERS_V5 else ["lm_head.weight"])

    def __init__(self, config: TokleConfig):
        super().__init__(config)
        self.model = TokleModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def get_decoder(self):
        return self.model

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutput]:
        return_dict = (return_dict if return_dict is not None
                       else getattr(self.config, "return_dict", True))
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        logits = self.lm_head(out.last_hidden_state)
        # Rows past real_vocab_size only pad the matrix to a multiple of 64 and
        # carry no meaning, so keep them out of sampling and out of any
        # log-likelihood an evaluation harness sums.
        rv = self.config.real_vocab_size
        if self.config.mask_padded_vocab_logits and rv is not None and rv < self.config.vocab_size:
            logits = logits.clone()
            logits[..., rv:] = float("-inf")

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels,
                                      vocab_size=self.config.vocab_size, **kwargs)

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutput(loss=loss, logits=logits, hidden_states=out.hidden_states)

	"""Tokle: decoder-only transformer with RMSNorm, RoPE, GQA and SwiGLU.

	Module and parameter names follow the usual Transformers layout.

	Works on Transformers 4.x and 5.x; the v5 differences are marked inline.
	"""
	import math
	from typing import Optional, Tuple, Union

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import __version__ as _transformers_version
	from transformers.generation import GenerationMixin
	from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
	from transformers.modeling_utils import PreTrainedModel

	try: # as a Transformers dynamic module (trust_remote_code)
	from .configuration_tokle import TokleConfig
	except ImportError: # running the files straight out of a checkout
	from configuration_tokle import TokleConfig


	# ----------------------------------------------------------------------------- norms and rope
	class TokleRMSNorm(nn.Module):
	def __init__(self, d, eps=1e-5):
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(d))

	def forward(self, x):
	# Always reduce in fp32: in fp16 the mean of squares underflows on the
	# long tail and the norm comes back slightly wrong.
	dt = x.dtype
	x = x.float()
	x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
	return (x * self.weight.float()).to(dt)

	def extra_repr(self):
	return f"{tuple(self.weight.shape)}, eps={self.eps}"


	def build_rope(head_dim, max_seq_len, theta, device=None):
	inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
	device=device) / head_dim))
	t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
	f = torch.outer(t, inv)
	return torch.cos(f), torch.sin(f)


	def apply_rope(x, cos, sin):
	"""Half-split rotation (first half against second), not the interleaved
	variant. The two are not interchangeable -- swapping them degrades output
	without raising anything."""
	b, h, t, d = x.shape
	x1, x2 = x.float().chunk(2, dim=-1)
	c = cos[:t].view(1, 1, t, d // 2).to(x1.device)
	s = sin[:t].view(1, 1, t, d // 2).to(x1.device)
	return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1).to(x.dtype)


	# ----------------------------------------------------------------------------- blocks
	class TokleAttention(nn.Module):
	def __init__(self, config: TokleConfig, layer_idx: int):
	super().__init__()
	self.layer_idx = layer_idx
	self.n_head = config.n_head
	self.n_kv_head = config.n_kv_head
	self.head_dim = config.head_dim
	self.rep = config.n_head // config.n_kv_head
	self.q_proj = nn.Linear(config.d_model, config.n_head * config.head_dim, bias=False)
	self.k_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
	self.v_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
	self.o_proj = nn.Linear(config.n_head * config.head_dim, config.d_model, bias=False)
	self.q_norm = TokleRMSNorm(config.head_dim, config.norm_eps)
	self.k_norm = TokleRMSNorm(config.head_dim, config.norm_eps)

	def forward(self, x, cos, sin, attn_mask):
	b, t, _ = x.shape
	q = self.q_proj(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2)
	k = self.k_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
	v = self.v_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
	# QK-norm is applied before rope.
	q = apply_rope(self.q_norm(q), cos, sin)
	k = apply_rope(self.k_norm(k), cos, sin)
	k = k.repeat_interleave(self.rep, dim=1)
	v = v.repeat_interleave(self.rep, dim=1)
	if attn_mask is None:
	# Nothing padded: let SDPA build the causal mask itself, which is
	# the fast kernel path and what an unpadded forward always hits.
	o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	else:
	o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.dtype))
	return self.o_proj(o.transpose(1, 2).reshape(b, t, -1))


	class TokleMLP(nn.Module):
	def __init__(self, config: TokleConfig):
	super().__init__()
	self.gate_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
	self.up_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
	self.down_proj = nn.Linear(config.ffn_hidden, config.d_model, bias=False)

	def forward(self, x):
	return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


	class TokleDecoderLayer(nn.Module):
	def __init__(self, config: TokleConfig, layer_idx: int):
	super().__init__()
	self.input_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
	self.self_attn = TokleAttention(config, layer_idx)
	self.post_attention_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
	self.mlp = TokleMLP(config)

	def forward(self, x, cos, sin, attn_mask):
	x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask)
	return x + self.mlp(self.post_attention_layernorm(x))


	# ----------------------------------------------------------------------------- model
	class ToklePreTrainedModel(PreTrainedModel):
	config_class = TokleConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["TokleDecoderLayer"]
	_supports_sdpa = True

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	nn.init.normal_(module.weight, 0.0, 0.02)
	if module.bias is not None:
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	nn.init.normal_(module.weight, 0.0, 0.02)
	elif isinstance(module, TokleRMSNorm):
	nn.init.ones_(module.weight)
	# Residual-output projections get scaled down by depth so the residual
	# stream does not blow up at init (GPT-2 trick).
	if isinstance(module, (TokleAttention, TokleMLP)):
	out = module.o_proj if isinstance(module, TokleAttention) else module.down_proj
	nn.init.normal_(out.weight, 0.0, 0.02 / math.sqrt(2 * self.config.n_layer))


	class TokleModel(ToklePreTrainedModel):
	def __init__(self, config: TokleConfig):
	super().__init__(config)
	self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
	self.layers = nn.ModuleList(
	[TokleDecoderLayer(config, i) for i in range(config.n_layer)])
	self.norm = TokleRMSNorm(config.d_model, config.norm_eps)
	# A plain dict, not a buffer. Transformers 5 builds models on the meta
	# device and fills only tensors the weights file contains, so a
	# non-persistent buffer would come back as uninitialised memory.
	# Keyed by device, so .to()/.cuda() need not move anything.
	self._rope_cache = {}
	self.gradient_checkpointing = False
	self.post_init()

	def get_input_embeddings(self):
	return self.embed_tokens

	def set_input_embeddings(self, value):
	self.embed_tokens = value

	def rope(self, device):
	key = str(device)
	if key not in self._rope_cache:
	self._rope_cache[key] = build_rope(
	self.config.head_dim, self.config.max_seq_len,
	self.config.rope_theta, device=device)
	return self._rope_cache[key]

	def _causal_pad_mask(self, attention_mask, t, dtype, device):
	"""Additive (b, 1, 1, t) padding mask folded together with the causal
	mask, or None when nothing is padded so the caller can take the
	is_causal fast path."""
	if attention_mask is None:
	return None
	attention_mask = attention_mask.to(torch.bool)
	if bool(attention_mask.all()):
	return None
	m = torch.zeros(attention_mask.shape, dtype=dtype, device=device)
	m = m.masked_fill(~attention_mask, torch.finfo(dtype).min)[:, None, None, :]
	# Under causal attention with right padding, position 0 always sees
	# itself, so no row ends up fully masked (which would give NaNs).
	causal = torch.ones(t, t, dtype=torch.bool, device=device).triu(1)
	return m.masked_fill(causal, torch.finfo(dtype).min)

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	**kwargs,
	) -> Union[Tuple, BaseModelOutput]:
	# getattr, not config.use_return_dict: that property is deprecated in
	# Transformers 5 and warns on every call.
	return_dict = (return_dict if return_dict is not None
	else getattr(self.config, "return_dict", True))
	output_hidden_states = (output_hidden_states if output_hidden_states is not None
	else self.config.output_hidden_states)
	if (input_ids is None) == (inputs_embeds is None):
	raise ValueError("Pass exactly one of input_ids or inputs_embeds.")

	x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
	t = x.shape[1]
	if t > self.config.max_seq_len:
	raise ValueError(
	f"Sequence length {t} exceeds max_seq_len {self.config.max_seq_len}.")

	cos, sin = self.rope(x.device)
	mask = self._causal_pad_mask(attention_mask, t, x.dtype, x.device)

	hidden_states = () if output_hidden_states else None
	for layer in self.layers:
	if output_hidden_states:
	hidden_states += (x,)
	if self.gradient_checkpointing and self.training:
	x = self._gradient_checkpointing_func(layer.__call__, x, cos, sin, mask)
	else:
	x = layer(x, cos, sin, mask)
	x = self.norm(x)
	if output_hidden_states:
	hidden_states += (x,)

	if not return_dict:
	return tuple(v for v in (x, hidden_states) if v is not None)
	return BaseModelOutput(last_hidden_state=x, hidden_states=hidden_states)


	# Transformers 5 turned _tied_weights_keys from a list of names into a
	# {tied: source} mapping and calls .keys() on it, so the old form raises.
	_TRANSFORMERS_V5 = int(_transformers_version.split(".")[0]) >= 5


	class TokleForCausalLM(ToklePreTrainedModel, GenerationMixin):
	_tied_weights_keys = ({"lm_head.weight": "model.embed_tokens.weight"}
	if _TRANSFORMERS_V5 else ["lm_head.weight"])

	def __init__(self, config: TokleConfig):
	super().__init__(config)
	self.model = TokleModel(config)
	self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
	self.post_init()

	def get_input_embeddings(self):
	return self.model.embed_tokens

	def set_input_embeddings(self, value):
	self.model.embed_tokens = value

	def get_output_embeddings(self):
	return self.lm_head

	def set_output_embeddings(self, new_embeddings):
	self.lm_head = new_embeddings

	def get_decoder(self):
	return self.model

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	**kwargs,
	) -> Union[Tuple, CausalLMOutput]:
	return_dict = (return_dict if return_dict is not None
	else getattr(self.config, "return_dict", True))
	out = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	inputs_embeds=inputs_embeds,
	output_hidden_states=output_hidden_states,
	return_dict=True,
	)
	logits = self.lm_head(out.last_hidden_state)
	# Rows past real_vocab_size only pad the matrix to a multiple of 64 and
	# carry no meaning, so keep them out of sampling and out of any
	# log-likelihood an evaluation harness sums.
	rv = self.config.real_vocab_size
	if self.config.mask_padded_vocab_logits and rv is not None and rv < self.config.vocab_size:
	logits = logits.clone()
	logits[..., rv:] = float("-inf")

	loss = None
	if labels is not None:
	loss = self.loss_function(logits=logits, labels=labels,
	vocab_size=self.config.vocab_size, **kwargs)

	if not return_dict:
	return (loss, logits) if loss is not None else (logits,)
	return CausalLMOutput(loss=loss, logits=logits, hidden_states=out.hidden_states)