Instructions to use xxparthparekhxx/pluto-124m-base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use xxparthparekhxx/pluto-124m-base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="xxparthparekhxx/pluto-124m-base", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("xxparthparekhxx/pluto-124m-base", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use xxparthparekhxx/pluto-124m-base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "xxparthparekhxx/pluto-124m-base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "xxparthparekhxx/pluto-124m-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/xxparthparekhxx/pluto-124m-base

SGLang

How to use xxparthparekhxx/pluto-124m-base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "xxparthparekhxx/pluto-124m-base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "xxparthparekhxx/pluto-124m-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "xxparthparekhxx/pluto-124m-base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "xxparthparekhxx/pluto-124m-base",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use xxparthparekhxx/pluto-124m-base with Docker Model Runner:
```
docker model run hf.co/xxparthparekhxx/pluto-124m-base
```

pluto-124m-base / modeling_pluto.py

xxparthparekhxx

Pluto 124M base: weights, remote code, tokenizer, model card and evals

fb25867 verified about 17 hours ago

Raw History Blame Contribute Delete

10.7 kB

"""Pluto: a small dense decoder-only LLM (GQA + QK-norm + gated attention + logit soft-capping).
Self-contained: needs only torch and transformers. A port of the dense path of the original
training code; weights are numerically equivalent to it (see test_export_hf.py in the source repo).
    model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
"""

from typing import Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import DynamicCache, GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .configuration_pluto import PlutoConfig


class PlutoRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)


def rope_cos_sin(position_ids: torch.Tensor, head_dim: int, theta: float):
    """(B, T) integer positions -> cos and sin of shape (B, 1, T, head_dim / 2), float32"""
    inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=position_ids.device) / head_dim)
    freqs = position_ids.float()[..., None] * inv_freq  # (B, T, D/2)
    return freqs.cos()[:, None], freqs.sin()[:, None]


def apply_rope(x, cos, sin):
    """Rotate (B, H, T, D) by position, pairing dimension i with i + D/2"""
    x1, x2 = x.float().chunk(2, dim=-1)
    out = torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
    return out.type_as(x)


class PlutoAttention(nn.Module):
    def __init__(self, config: PlutoConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_heads = config.num_attention_heads
        self.n_kv = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.rope_theta = config.rope_theta
        d = config.hidden_size
        self.q_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
        self.g_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False) if config.attn_gate else None
        self.q_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
        self.k_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()

    def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
        b, t, _ = x.shape
        q = self.q_norm(self.q_proj(x).view(b, t, self.n_heads, self.head_dim)).transpose(1, 2)
        k = self.k_norm(self.k_proj(x).view(b, t, self.n_kv, self.head_dim)).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
        cos, sin = rope_cos_sin(position_ids, self.head_dim, self.rope_theta)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if past_key_values is not None:
            k, v = past_key_values.update(k, v, self.layer_idx)  # cached un-repeated, post-RoPE
        if self.n_kv != self.n_heads:
            rep = self.n_heads // self.n_kv
            k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal)
        y = y.transpose(1, 2).reshape(b, t, -1)
        if self.g_proj is not None:
            y = y * torch.sigmoid(self.g_proj(x))
        return self.o_proj(y)


class PlutoMLP(nn.Module):
    def __init__(self, config: PlutoConfig):
        super().__init__()
        self.clamp = config.swiglu_clamp
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        gate, up = self.gate_proj(x), self.up_proj(x)
        if self.clamp is not None:
            gate, up = gate.clamp(max=self.clamp), up.clamp(-self.clamp, self.clamp)
        return self.down_proj(F.silu(gate) * up)


class PlutoDecoderLayer(nn.Module):
    def __init__(self, config: PlutoConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = PlutoAttention(config, layer_idx)
        self.post_attention_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = PlutoMLP(config)

    def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
        x = x + self.self_attn(self.input_layernorm(x), position_ids, attn_mask, is_causal, past_key_values)
        return x + self.mlp(self.post_attention_layernorm(x))


class PlutoPreTrainedModel(PreTrainedModel):
    config_class = PlutoConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["PlutoDecoderLayer"]
    _supports_sdpa = True

    @torch.no_grad()
    def _init_weights(self, module):
        std = 0.02
        if isinstance(module, nn.Linear):
            module.weight.normal_(mean=0.0, std=std)
        elif isinstance(module, nn.Embedding):
            module.weight.normal_(mean=0.0, std=std)
        elif isinstance(module, PlutoRMSNorm):
            module.weight.fill_(1.0)


class PlutoModel(PlutoPreTrainedModel):
    def __init__(self, config: PlutoConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(PlutoDecoderLayer(config, i) for i in range(config.num_hidden_layers))
        self.norm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids and inputs_embeds")
        use_cache = self.config.use_cache if use_cache is None else use_cache
        x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        b, t, _ = x.shape

        if use_cache and past_key_values is None:
            try:
                past_key_values = DynamicCache(config=self.config)
            except TypeError:  # older transformers
                past_key_values = DynamicCache()
        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
        cache_position = torch.arange(past_len, past_len + t, device=x.device)
        if position_ids is None:
            position_ids = cache_position[None].expand(b, -1)

        # Boolean mask, True = may attend; None lets SDPA take its causal fast paths.
        total = past_len + t
        padded = attention_mask is not None and attention_mask.shape[-1] == total and not bool(attention_mask.all())
        if not padded and t == 1:
            attn_mask, is_causal = None, False
        elif not padded and past_len == 0:
            attn_mask, is_causal = None, True
        else:
            keys = torch.arange(total, device=x.device)
            attn_mask = (keys[None, :] <= cache_position[:, None])[None, None]  # (1, 1, T, total)
            if padded:
                attn_mask = attn_mask & attention_mask.bool()[:, None, None, :]
                # Query rows that are all padding would softmax over nothing (NaN); let them see everything
                attn_mask = attn_mask | ~attn_mask.any(dim=-1, keepdim=True)
            is_causal = False

        for layer in self.layers:
            x = layer(x, position_ids, attn_mask, is_causal, past_key_values)
        return BaseModelOutputWithPast(
            last_hidden_state=self.norm(x), past_key_values=past_key_values if use_cache else None
        )


class PlutoForCausalLM(PlutoPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: PlutoConfig):
        super().__init__(config)
        self.model = PlutoModel(config)
        self.lm_head = nn.Linear(config.hidden_size, 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 forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
        )
        h = out.last_hidden_state
        if isinstance(logits_to_keep, int):
            keep = slice(-logits_to_keep, None) if logits_to_keep else slice(None)
        else:
            keep = logits_to_keep  # tensor of positions
        logits = self.lm_head(h[:, keep, :]).float()
        cap = self.config.logit_softcap
        if cap:
            logits = cap * torch.tanh(logits / cap)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].reshape(-1, logits.size(-1))
            loss = F.cross_entropy(shift_logits, labels[:, 1:].reshape(-1).to(logits.device), ignore_index=-100)
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values)

	"""Pluto: a small dense decoder-only LLM (GQA + QK-norm + gated attention + logit soft-capping).

	Self-contained: needs only torch and transformers. A port of the dense path of the original
	training code; weights are numerically equivalent to it (see test_export_hf.py in the source repo).

	model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
	"""

	from typing import Optional, Union

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import DynamicCache, GenerationMixin, PreTrainedModel
	from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

	from .configuration_pluto import PlutoConfig


	class PlutoRMSNorm(nn.Module):
	def __init__(self, dim: int, eps: float = 1e-6):
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(dim))

	def forward(self, x):
	return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)


	def rope_cos_sin(position_ids: torch.Tensor, head_dim: int, theta: float):
	"""(B, T) integer positions -> cos and sin of shape (B, 1, T, head_dim / 2), float32"""
	inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=position_ids.device) / head_dim)
	freqs = position_ids.float()[..., None] * inv_freq # (B, T, D/2)
	return freqs.cos()[:, None], freqs.sin()[:, None]


	def apply_rope(x, cos, sin):
	"""Rotate (B, H, T, D) by position, pairing dimension i with i + D/2"""
	x1, x2 = x.float().chunk(2, dim=-1)
	out = torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
	return out.type_as(x)


	class PlutoAttention(nn.Module):
	def __init__(self, config: PlutoConfig, layer_idx: int):
	super().__init__()
	self.layer_idx = layer_idx
	self.n_heads = config.num_attention_heads
	self.n_kv = config.num_key_value_heads
	self.head_dim = config.head_dim
	self.rope_theta = config.rope_theta
	d = config.hidden_size
	self.q_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
	self.k_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
	self.v_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
	self.o_proj = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
	self.g_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False) if config.attn_gate else None
	self.q_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
	self.k_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()

	def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
	b, t, _ = x.shape
	q = self.q_norm(self.q_proj(x).view(b, t, self.n_heads, self.head_dim)).transpose(1, 2)
	k = self.k_norm(self.k_proj(x).view(b, t, self.n_kv, self.head_dim)).transpose(1, 2)
	v = self.v_proj(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
	cos, sin = rope_cos_sin(position_ids, self.head_dim, self.rope_theta)
	q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
	if past_key_values is not None:
	k, v = past_key_values.update(k, v, self.layer_idx) # cached un-repeated, post-RoPE
	if self.n_kv != self.n_heads:
	rep = self.n_heads // self.n_kv
	k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
	y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal)
	y = y.transpose(1, 2).reshape(b, t, -1)
	if self.g_proj is not None:
	y = y * torch.sigmoid(self.g_proj(x))
	return self.o_proj(y)


	class PlutoMLP(nn.Module):
	def __init__(self, config: PlutoConfig):
	super().__init__()
	self.clamp = config.swiglu_clamp
	self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
	self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
	self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

	def forward(self, x):
	gate, up = self.gate_proj(x), self.up_proj(x)
	if self.clamp is not None:
	gate, up = gate.clamp(max=self.clamp), up.clamp(-self.clamp, self.clamp)
	return self.down_proj(F.silu(gate) * up)


	class PlutoDecoderLayer(nn.Module):
	def __init__(self, config: PlutoConfig, layer_idx: int):
	super().__init__()
	self.input_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
	self.self_attn = PlutoAttention(config, layer_idx)
	self.post_attention_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
	self.mlp = PlutoMLP(config)

	def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
	x = x + self.self_attn(self.input_layernorm(x), position_ids, attn_mask, is_causal, past_key_values)
	return x + self.mlp(self.post_attention_layernorm(x))


	class PlutoPreTrainedModel(PreTrainedModel):
	config_class = PlutoConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = False
	_no_split_modules = ["PlutoDecoderLayer"]
	_supports_sdpa = True

	@torch.no_grad()
	def _init_weights(self, module):
	std = 0.02
	if isinstance(module, nn.Linear):
	module.weight.normal_(mean=0.0, std=std)
	elif isinstance(module, nn.Embedding):
	module.weight.normal_(mean=0.0, std=std)
	elif isinstance(module, PlutoRMSNorm):
	module.weight.fill_(1.0)


	class PlutoModel(PlutoPreTrainedModel):
	def __init__(self, config: PlutoConfig):
	super().__init__(config)
	self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
	self.layers = nn.ModuleList(PlutoDecoderLayer(config, i) for i in range(config.num_hidden_layers))
	self.norm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
	self.post_init()

	def get_input_embeddings(self):
	return self.embed_tokens

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

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[DynamicCache] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	use_cache: Optional[bool] = None,
	**kwargs,
	) -> BaseModelOutputWithPast:
	if (input_ids is None) == (inputs_embeds is None):
	raise ValueError("Specify exactly one of input_ids and inputs_embeds")
	use_cache = self.config.use_cache if use_cache is None else use_cache
	x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
	b, t, _ = x.shape

	if use_cache and past_key_values is None:
	try:
	past_key_values = DynamicCache(config=self.config)
	except TypeError: # older transformers
	past_key_values = DynamicCache()
	past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
	cache_position = torch.arange(past_len, past_len + t, device=x.device)
	if position_ids is None:
	position_ids = cache_position[None].expand(b, -1)

	# Boolean mask, True = may attend; None lets SDPA take its causal fast paths.
	total = past_len + t
	padded = attention_mask is not None and attention_mask.shape[-1] == total and not bool(attention_mask.all())
	if not padded and t == 1:
	attn_mask, is_causal = None, False
	elif not padded and past_len == 0:
	attn_mask, is_causal = None, True
	else:
	keys = torch.arange(total, device=x.device)
	attn_mask = (keys[None, :] <= cache_position[:, None])[None, None] # (1, 1, T, total)
	if padded:
	attn_mask = attn_mask & attention_mask.bool()[:, None, None, :]
	# Query rows that are all padding would softmax over nothing (NaN); let them see everything
	attn_mask = attn_mask \| ~attn_mask.any(dim=-1, keepdim=True)
	is_causal = False

	for layer in self.layers:
	x = layer(x, position_ids, attn_mask, is_causal, past_key_values)
	return BaseModelOutputWithPast(
	last_hidden_state=self.norm(x), past_key_values=past_key_values if use_cache else None
	)


	class PlutoForCausalLM(PlutoPreTrainedModel, GenerationMixin):
	_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

	def __init__(self, config: PlutoConfig):
	super().__init__(config)
	self.model = PlutoModel(config)
	self.lm_head = nn.Linear(config.hidden_size, 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 forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[DynamicCache] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	logits_to_keep: Union[int, torch.Tensor] = 0,
	**kwargs,
	) -> CausalLMOutputWithPast:
	out = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	)
	h = out.last_hidden_state
	if isinstance(logits_to_keep, int):
	keep = slice(-logits_to_keep, None) if logits_to_keep else slice(None)
	else:
	keep = logits_to_keep # tensor of positions
	logits = self.lm_head(h[:, keep, :]).float()
	cap = self.config.logit_softcap
	if cap:
	logits = cap * torch.tanh(logits / cap)

	loss = None
	if labels is not None:
	shift_logits = logits[:, :-1].reshape(-1, logits.size(-1))
	loss = F.cross_entropy(shift_logits, labels[:, 1:].reshape(-1).to(logits.device), ignore_index=-100)
	return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values)