Instructions to use q-project/Q-U-164M with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use q-project/Q-U-164M with Transformers:

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

pipe = pipeline("text-generation", model="q-project/Q-U-164M", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("q-project/Q-U-164M", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use q-project/Q-U-164M with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "q-project/Q-U-164M"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "q-project/Q-U-164M",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/q-project/Q-U-164M

SGLang

How to use q-project/Q-U-164M 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 "q-project/Q-U-164M" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "q-project/Q-U-164M",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

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 "q-project/Q-U-164M" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "q-project/Q-U-164M",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use q-project/Q-U-164M with Docker Model Runner:
```
docker model run hf.co/q-project/Q-U-164M
```

Q-U-164M / modeling_qagent.py

LakoMoor

Q-U-164M v5: weights, config, tokenizer, architecture code, qpack

1939f28 verified about 1 month ago

Raw History Blame Contribute Delete

10.3 kB

import math

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

from .configuration_qagent import QAgentConfig


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        xf = x.float()
        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
        return (xf * self.weight.float()).to(x.dtype)


class TernaryLinear(nn.Linear):
    """Weight is re-quantised to {-1,0,+1}*alpha (per-output-row absmean alpha) on every forward; gradients pass
    straight through to the latent full-precision master weight (BitNet b1.58-style QAT, trained this way from step 1)."""

    def __init__(self, in_features, out_features, act_quant=False):
        super().__init__(in_features, out_features, bias=False)
        self.act_quant = act_quant

    def ternary(self):
        wf = self.weight.float()
        alpha = wf.abs().mean(dim=1, keepdim=True).clamp(min=1e-5)
        return torch.clamp(torch.round(wf / alpha), -1, 1), alpha

    def forward(self, x):
        wf = self.weight.float()
        alpha = wf.abs().mean(dim=1, keepdim=True).clamp(min=1e-5)
        wq = torch.clamp(torch.round(wf / alpha), -1, 1) * alpha
        w = (wf + (wq - wf).detach()).to(x.dtype)
        if self.act_quant:
            s = x.detach().abs().amax(dim=-1, keepdim=True).clamp(min=1e-5) / 127.0
            x = x + (torch.clamp(torch.round(x / s), -127, 127) * s - x).detach()
        return F.linear(x, w)


def make_linear(cfg, i, o):
    return TernaryLinear(i, o, cfg.act_quant) if cfg.quant_mode == "ternary" else nn.Linear(i, o, bias=False)


def rotate_half(x):
    h = x.shape[-1] // 2
    return torch.cat((-x[..., h:], x[..., :h]), dim=-1)


class RotaryEmbedding(nn.Module):
    """inv_freq is computed on the fly (a non-persistent buffer would be lost/garbled by meta-device init in from_pretrained)."""

    def __init__(self, dim, theta):
        super().__init__()
        self.dim, self.theta = dim, theta

    def forward(self, position_ids, dtype):
        inv = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, device=position_ids.device).float() / self.dim))
        f = position_ids[..., None].float() * inv[None, None, :]
        e = torch.cat((f, f), dim=-1)
        return e.cos().to(dtype), e.sin().to(dtype)


class QAgentAttention(nn.Module):
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.nh, self.nkv, self.hd = cfg.num_attention_heads, cfg.num_key_value_heads, cfg.head_dim
        d = cfg.hidden_size
        self.q_proj = make_linear(cfg, d, self.nh * self.hd)
        self.k_proj = make_linear(cfg, d, self.nkv * self.hd)
        self.v_proj = make_linear(cfg, d, self.nkv * self.hd)
        self.o_proj = make_linear(cfg, self.nh * self.hd, d)
        self.q_norm = RMSNorm(self.hd, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.hd, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
        self.out_gate = nn.Linear(d, self.nh, bias=True) if cfg.attention_output_gate else None  # per-head sigmoid gate

    def forward(self, x, cos, sin, cache=None, mask=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        c, s = cos[:, None], sin[:, None]
        q, k = q * c + rotate_half(q) * s, k * c + rotate_half(k) * s
        if cache is not None:
            k, v = cache.update(k, v, self.layer_idx)
        rep = self.nh // self.nkv
        k, v = k.repeat_interleave(rep, 1), v.repeat_interleave(rep, 1)
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=(mask is None and k.shape[2] == T))
        if self.out_gate is not None:
            o = o * torch.sigmoid(self.out_gate(x)).transpose(1, 2).unsqueeze(-1)
        return self.o_proj(o.transpose(1, 2).reshape(B, T, self.nh * self.hd))


class QAgentMLP(nn.Module):
    """Plain two-matrix SiLU feed-forward (d -> d_ff -> d), both ternary."""

    def __init__(self, cfg):
        super().__init__()
        self.w1 = make_linear(cfg, cfg.hidden_size, cfg.intermediate_size)
        self.w2 = make_linear(cfg, cfg.intermediate_size, cfg.hidden_size)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)))


class QAgentLayer(nn.Module):
    def __init__(self, cfg, i):
        super().__init__()
        self.n1, self.n2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps), RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn, self.mlp = QAgentAttention(cfg, i), QAgentMLP(cfg)

    def forward(self, x, cos, sin, cache=None, mask=None):
        x = x + self.attn(self.n1(x), cos, sin, cache, mask)
        return x + self.mlp(self.n2(x))


class QAgentPreTrainedModel(PreTrainedModel):
    config_class = QAgentConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["QAgentLayer"]

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0.0, self.config.initializer_range)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, RMSNorm):
            nn.init.ones_(m.weight)
        elif isinstance(m, QAgentModel):
            nn.init.zeros_(m.delta_in); nn.init.zeros_(m.delta_out)


class QAgentModel(QAgentPreTrainedModel):
    def __init__(self, cfg):
        super().__init__(cfg)
        # frozen fingerprint table: one +-1 code of cfg.code_bits per token (buffer, NOT a parameter; 0
        # trainable embedding params). Unlike Q-Omni, there is no multimodal id range: every id in the
        # vocabulary (text tokens AND control/dialogue/tool tokens alike) gets the SAME treatment -- one
        # shared frozen fingerprint plus one learned per-id correction, split into separate in/out deltas
        # (Q-Omni's `learned_text_deltas` ablation showed tying embedding and readout roles into a single
        # correction vector measurably hurts both; untying is the validated default here, not an add-on).
        self.register_buffer("codes", torch.sign(torch.randn(cfg.vocab_size, cfg.code_bits)).to(torch.float16), persistent=True)
        self.delta_in = nn.Parameter(torch.zeros(cfg.vocab_size, cfg.code_bits))
        self.delta_out = nn.Parameter(torch.zeros(cfg.vocab_size, cfg.code_bits))
        self.w_in = nn.Linear(cfg.code_bits, cfg.hidden_size, bias=False)   # dense fp io
        self.layers = nn.ModuleList([QAgentLayer(cfg, i) for i in range(cfg.num_hidden_layers)])
        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.rotary = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
        self.post_init()

    def table_in(self, dtype):
        return (self.codes.float() + self.delta_in.float()).to(dtype)

    def table_out(self, dtype):
        return (self.codes.float() + self.delta_out.float()).to(dtype)

    def embed(self, ids):
        dt = torch.float16 if torch.is_autocast_enabled() else self.w_in.weight.dtype
        return self.w_in(F.embedding(ids, self.table_in(dt)))

    def forward(self, input_ids, position_ids=None, past_key_values=None, use_cache=False, attention_mask=None):
        B, T = input_ids.shape
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()
        past = past_key_values.get_seq_length() if past_key_values is not None else 0
        if position_ids is None:
            position_ids = torch.arange(past, past + T, device=input_ids.device)[None].expand(B, -1)
        x = self.embed(input_ids)
        cos, sin = self.rotary(position_ids, x.dtype)
        mask = None
        if past > 0 or attention_mask is not None:
            S = past + T
            qi = torch.arange(past, S, device=x.device)[:, None]
            mask = (torch.arange(S, device=x.device)[None, :] <= qi)[None, None]
            if attention_mask is not None:
                mask = mask & attention_mask[:, None, None, :S].bool()
        for layer in self.layers:
            x = layer(x, cos, sin, past_key_values, mask)
        return self.norm(x), past_key_values


class QAgentForCausalLM(QAgentPreTrainedModel, GenerationMixin):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.model = QAgentModel(cfg)
        self.w_out = nn.Linear(cfg.hidden_size, cfg.code_bits, bias=False)  # readout: hidden -> fingerprint space
        self.register_buffer("logit_bias", torch.zeros(cfg.vocab_size), persistent=True)  # frozen unigram log-prior
        self.post_init()

    def logits_of(self, h):
        z = self.w_out(h)
        return (z @ self.model.table_out(z.dtype).t()) / math.sqrt(self.config.code_bits) + self.logit_bias.to(z.dtype)

    def forward(self, input_ids=None, labels=None, position_ids=None, past_key_values=None, use_cache=None,
                attention_mask=None, logits_to_keep=0, **kw):
        h, past_key_values = self.model(input_ids, position_ids, past_key_values, bool(use_cache), attention_mask)
        if labels is None:
            if logits_to_keep:
                h = h[:, -logits_to_keep:]
            return CausalLMOutputWithPast(logits=self.logits_of(h), past_key_values=past_key_values)
        hs, ys = h[:, :-1].reshape(-1, h.size(-1)), labels[:, 1:].reshape(-1)

        def ce(hc, yc):  # chunked + checkpointed: full (tokens x vocab) logits are never materialised
            return F.cross_entropy(self.logits_of(hc).float(), yc, ignore_index=-100, reduction="sum")

        tot = hs.new_zeros((), dtype=torch.float32)
        for i in range(0, hs.size(0), 4096):
            tot = tot + torch.utils.checkpoint.checkpoint(ce, hs[i:i + 4096], ys[i:i + 4096], use_reentrant=False)
        return CausalLMOutputWithPast(loss=tot / (ys != -100).sum().clamp(min=1), past_key_values=past_key_values)

	import math

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

	from .configuration_qagent import QAgentConfig


	class RMSNorm(nn.Module):
	def __init__(self, dim, eps=1e-5):
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(dim))

	def forward(self, x):
	xf = x.float()
	xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
	return (xf * self.weight.float()).to(x.dtype)


	class TernaryLinear(nn.Linear):
	"""Weight is re-quantised to {-1,0,+1}*alpha (per-output-row absmean alpha) on every forward; gradients pass
	straight through to the latent full-precision master weight (BitNet b1.58-style QAT, trained this way from step 1)."""

	def __init__(self, in_features, out_features, act_quant=False):
	super().__init__(in_features, out_features, bias=False)
	self.act_quant = act_quant

	def ternary(self):
	wf = self.weight.float()
	alpha = wf.abs().mean(dim=1, keepdim=True).clamp(min=1e-5)
	return torch.clamp(torch.round(wf / alpha), -1, 1), alpha

	def forward(self, x):
	wf = self.weight.float()
	alpha = wf.abs().mean(dim=1, keepdim=True).clamp(min=1e-5)
	wq = torch.clamp(torch.round(wf / alpha), -1, 1) * alpha
	w = (wf + (wq - wf).detach()).to(x.dtype)
	if self.act_quant:
	s = x.detach().abs().amax(dim=-1, keepdim=True).clamp(min=1e-5) / 127.0
	x = x + (torch.clamp(torch.round(x / s), -127, 127) * s - x).detach()
	return F.linear(x, w)


	def make_linear(cfg, i, o):
	return TernaryLinear(i, o, cfg.act_quant) if cfg.quant_mode == "ternary" else nn.Linear(i, o, bias=False)


	def rotate_half(x):
	h = x.shape[-1] // 2
	return torch.cat((-x[..., h:], x[..., :h]), dim=-1)


	class RotaryEmbedding(nn.Module):
	"""inv_freq is computed on the fly (a non-persistent buffer would be lost/garbled by meta-device init in from_pretrained)."""

	def __init__(self, dim, theta):
	super().__init__()
	self.dim, self.theta = dim, theta

	def forward(self, position_ids, dtype):
	inv = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, device=position_ids.device).float() / self.dim))
	f = position_ids[..., None].float() * inv[None, None, :]
	e = torch.cat((f, f), dim=-1)
	return e.cos().to(dtype), e.sin().to(dtype)


	class QAgentAttention(nn.Module):
	def __init__(self, cfg, layer_idx):
	super().__init__()
	self.layer_idx = layer_idx
	self.nh, self.nkv, self.hd = cfg.num_attention_heads, cfg.num_key_value_heads, cfg.head_dim
	d = cfg.hidden_size
	self.q_proj = make_linear(cfg, d, self.nh * self.hd)
	self.k_proj = make_linear(cfg, d, self.nkv * self.hd)
	self.v_proj = make_linear(cfg, d, self.nkv * self.hd)
	self.o_proj = make_linear(cfg, self.nh * self.hd, d)
	self.q_norm = RMSNorm(self.hd, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
	self.k_norm = RMSNorm(self.hd, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
	self.out_gate = nn.Linear(d, self.nh, bias=True) if cfg.attention_output_gate else None # per-head sigmoid gate

	def forward(self, x, cos, sin, cache=None, mask=None):
	B, T, _ = x.shape
	q = self.q_proj(x).view(B, T, self.nh, self.hd).transpose(1, 2)
	k = self.k_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
	v = self.v_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
	q, k = self.q_norm(q), self.k_norm(k)
	c, s = cos[:, None], sin[:, None]
	q, k = q * c + rotate_half(q) * s, k * c + rotate_half(k) * s
	if cache is not None:
	k, v = cache.update(k, v, self.layer_idx)
	rep = self.nh // self.nkv
	k, v = k.repeat_interleave(rep, 1), v.repeat_interleave(rep, 1)
	o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=(mask is None and k.shape[2] == T))
	if self.out_gate is not None:
	o = o * torch.sigmoid(self.out_gate(x)).transpose(1, 2).unsqueeze(-1)
	return self.o_proj(o.transpose(1, 2).reshape(B, T, self.nh * self.hd))


	class QAgentMLP(nn.Module):
	"""Plain two-matrix SiLU feed-forward (d -> d_ff -> d), both ternary."""

	def __init__(self, cfg):
	super().__init__()
	self.w1 = make_linear(cfg, cfg.hidden_size, cfg.intermediate_size)
	self.w2 = make_linear(cfg, cfg.intermediate_size, cfg.hidden_size)

	def forward(self, x):
	return self.w2(F.silu(self.w1(x)))


	class QAgentLayer(nn.Module):
	def __init__(self, cfg, i):
	super().__init__()
	self.n1, self.n2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps), RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
	self.attn, self.mlp = QAgentAttention(cfg, i), QAgentMLP(cfg)

	def forward(self, x, cos, sin, cache=None, mask=None):
	x = x + self.attn(self.n1(x), cos, sin, cache, mask)
	return x + self.mlp(self.n2(x))


	class QAgentPreTrainedModel(PreTrainedModel):
	config_class = QAgentConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = False
	_no_split_modules = ["QAgentLayer"]

	def _init_weights(self, m):
	if isinstance(m, nn.Linear):
	nn.init.normal_(m.weight, 0.0, self.config.initializer_range)
	if m.bias is not None:
	nn.init.zeros_(m.bias)
	elif isinstance(m, RMSNorm):
	nn.init.ones_(m.weight)
	elif isinstance(m, QAgentModel):
	nn.init.zeros_(m.delta_in); nn.init.zeros_(m.delta_out)


	class QAgentModel(QAgentPreTrainedModel):
	def __init__(self, cfg):
	super().__init__(cfg)
	# frozen fingerprint table: one +-1 code of cfg.code_bits per token (buffer, NOT a parameter; 0
	# trainable embedding params). Unlike Q-Omni, there is no multimodal id range: every id in the
	# vocabulary (text tokens AND control/dialogue/tool tokens alike) gets the SAME treatment -- one
	# shared frozen fingerprint plus one learned per-id correction, split into separate in/out deltas
	# (Q-Omni's `learned_text_deltas` ablation showed tying embedding and readout roles into a single
	# correction vector measurably hurts both; untying is the validated default here, not an add-on).
	self.register_buffer("codes", torch.sign(torch.randn(cfg.vocab_size, cfg.code_bits)).to(torch.float16), persistent=True)
	self.delta_in = nn.Parameter(torch.zeros(cfg.vocab_size, cfg.code_bits))
	self.delta_out = nn.Parameter(torch.zeros(cfg.vocab_size, cfg.code_bits))
	self.w_in = nn.Linear(cfg.code_bits, cfg.hidden_size, bias=False) # dense fp io
	self.layers = nn.ModuleList([QAgentLayer(cfg, i) for i in range(cfg.num_hidden_layers)])
	self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
	self.rotary = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
	self.post_init()

	def table_in(self, dtype):
	return (self.codes.float() + self.delta_in.float()).to(dtype)

	def table_out(self, dtype):
	return (self.codes.float() + self.delta_out.float()).to(dtype)

	def embed(self, ids):
	dt = torch.float16 if torch.is_autocast_enabled() else self.w_in.weight.dtype
	return self.w_in(F.embedding(ids, self.table_in(dt)))

	def forward(self, input_ids, position_ids=None, past_key_values=None, use_cache=False, attention_mask=None):
	B, T = input_ids.shape
	if use_cache and past_key_values is None:
	past_key_values = DynamicCache()
	past = past_key_values.get_seq_length() if past_key_values is not None else 0
	if position_ids is None:
	position_ids = torch.arange(past, past + T, device=input_ids.device)[None].expand(B, -1)
	x = self.embed(input_ids)
	cos, sin = self.rotary(position_ids, x.dtype)
	mask = None
	if past > 0 or attention_mask is not None:
	S = past + T
	qi = torch.arange(past, S, device=x.device)[:, None]
	mask = (torch.arange(S, device=x.device)[None, :] <= qi)[None, None]
	if attention_mask is not None:
	mask = mask & attention_mask[:, None, None, :S].bool()
	for layer in self.layers:
	x = layer(x, cos, sin, past_key_values, mask)
	return self.norm(x), past_key_values


	class QAgentForCausalLM(QAgentPreTrainedModel, GenerationMixin):
	def __init__(self, cfg):
	super().__init__(cfg)
	self.model = QAgentModel(cfg)
	self.w_out = nn.Linear(cfg.hidden_size, cfg.code_bits, bias=False) # readout: hidden -> fingerprint space
	self.register_buffer("logit_bias", torch.zeros(cfg.vocab_size), persistent=True) # frozen unigram log-prior
	self.post_init()

	def logits_of(self, h):
	z = self.w_out(h)
	return (z @ self.model.table_out(z.dtype).t()) / math.sqrt(self.config.code_bits) + self.logit_bias.to(z.dtype)

	def forward(self, input_ids=None, labels=None, position_ids=None, past_key_values=None, use_cache=None,
	attention_mask=None, logits_to_keep=0, **kw):
	h, past_key_values = self.model(input_ids, position_ids, past_key_values, bool(use_cache), attention_mask)
	if labels is None:
	if logits_to_keep:
	h = h[:, -logits_to_keep:]
	return CausalLMOutputWithPast(logits=self.logits_of(h), past_key_values=past_key_values)
	hs, ys = h[:, :-1].reshape(-1, h.size(-1)), labels[:, 1:].reshape(-1)

	def ce(hc, yc): # chunked + checkpointed: full (tokens x vocab) logits are never materialised
	return F.cross_entropy(self.logits_of(hc).float(), yc, ignore_index=-100, reduction="sum")

	tot = hs.new_zeros((), dtype=torch.float32)
	for i in range(0, hs.size(0), 4096):
	tot = tot + torch.utils.checkpoint.checkpoint(ce, hs[i:i + 4096], ys[i:i + 4096], use_reentrant=False)
	return CausalLMOutputWithPast(loss=tot / (ys != -100).sum().clamp(min=1), past_key_values=past_key_values)