Upload folder using huggingface_hub

README.md

@@ -0,0 +1,84 @@
+---
+language:
+- en
+base_model:
+- talkie-lm/talkie-1930-13b-base
+license: apache-2.0
+pipeline_tag: text-generation
+tags:
+- talkie
+- vllm
+---
+
+# talkie-1930-13b-base (vLLM-servable repackage)
+
+This is a HuggingFace + vLLM-ready repackage of [`talkie-lm/talkie-1930-13b-base`](https://huggingface.co/talkie-lm/talkie-1930-13b-base). The original release ships as a raw torch checkpoint (`final.ckpt`) plus a tiktoken vocab (`vocab.txt`), with no `config.json`, `tokenizer.json`, or HF modeling code, so it can't be loaded by `transformers` or served by vLLM out of the box.
+
+`talkie-1930-13b-base` is a 13B parameter base (completion) language model from the [talkie-lm](https://talkie-lm.com/) project, pretrained on ~260B tokens of pre-1931 English text. This is the **base** model — it is not instruction-tuned and has no chat template; prompt it with raw text and let it continue.
+
+## What this repo adds
+
+| File | What it is |
+|---|---|
+| `model.safetensors` | bf16 weights, ~26 GB. `lm_head_gain` (a learned scalar) is pre-multiplied into `lm_head.weight` so vLLM's transformers backend doesn't need to know about it. |
+| `config.json` | `TalkieConfig` (vocab=65536, hidden=5120, 40 layers × 40 heads, head_dim=128, ctx=2048, RoPE θ=1e6) plus `auto_map` for `AutoConfig`/`AutoModel`/`AutoModelForCausalLM`. |
+| `tokenizer.json`, `tokenizer_config.json` | HF fast BPE built from the original `vocab.txt` (ranks < 65535) with `<|endoftext|>` at id 65535. EOS = pad = `<\|endoftext\|>`. No chat specials (base model). |
+| `generation_config.json` | `eos_token_id=65535`, `pad_token_id=65535`. |
+| `modeling_talkie.py`, `configuration_talkie.py` | HF `PreTrainedModel` implementation with `ALL_ATTENTION_FUNCTIONS` dispatch (vLLM transformers-backend compatible). |
+
+The architecture and tokenizer BPE ranks are identical to the IT repackage; the only differences are `vocab_size` (65536 vs 65540), the absence of the 4 chat special tokens, and no chat template.
+
+## Serving with vLLM
+
+```bash
+vllm serve agaralon/talkie-1930-13b-base-vllm \
+    --model-impl transformers \
+    --trust-remote-code \
+    --dtype bfloat16 \
+    --max-model-len 2048
+```
+
+Then hit it as an OpenAI-style **completions** endpoint (base model — use `/v1/completions`, not chat):
+
+```bash
+curl http://localhost:8000/v1/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "agaralon/talkie-1930-13b-base-vllm",
+        "prompt": "In the year 1925, the great city of London",
+        "temperature": 0.8,
+        "max_tokens": 120
+    }'
+```
+
+### Sampling notes
+
+- **Use `temperature ≥ 0.5`** — greedy decoding (`temperature=0`) can collapse into single-token loops on this architecture.
+- bf16 only — fp8 is broken on this architecture.
+
+## Plain HuggingFace usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+tok = AutoTokenizer.from_pretrained("agaralon/talkie-1930-13b-base-vllm", trust_remote_code=True)
+m = AutoModelForCausalLM.from_pretrained(
+    "agaralon/talkie-1930-13b-base-vllm", trust_remote_code=True, dtype=torch.bfloat16,
+).cuda().eval()
+
+ids = tok(["In the year 1925, the great city of London"], return_tensors="pt").to("cuda")
+out = m.generate(**ids, max_new_tokens=120, do_sample=True, temperature=0.8, top_p=0.9,
+                 pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
+print(tok.decode(out[0, ids.input_ids.shape[1]:], skip_special_tokens=True))
+```
+
+## Provenance
+
+- Weights: `final.ckpt` from `talkie-lm/talkie-1930-13b-base`, with `lm_head_gain.w_g` baked into `lm_head.weight` and cast to bf16.
+- Tokenizer: built from `vocab.txt` from the same release (ranks < 65535, then `<|endoftext|>` at id 65535).
+- Modeling code: adapted from the IT repackage [`awilliamson/talkie-1930-13b-it-vllm`](https://huggingface.co/awilliamson/talkie-1930-13b-it-vllm).
+
+## License
+
+Apache 2.0, matching the upstream `talkie-lm/talkie-1930-13b-base` release.

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "TalkieForCausalLM"
+  ],
+  "model_type": "talkie",
+  "vocab_size": 65536,
+  "hidden_size": 5120,
+  "intermediate_size": 13696,
+  "num_hidden_layers": 40,
+  "num_attention_heads": 40,
+  "head_dim": 128,
+  "max_position_embeddings": 2048,
+  "rope_theta": 1000000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "auto_map": {
+    "AutoConfig": "configuration_talkie.TalkieConfig",
+    "AutoModel": "modeling_talkie.TalkieModel",
+    "AutoModelForCausalLM": "modeling_talkie.TalkieForCausalLM"
+  }
+}

configuration_talkie.py

@@ -0,0 +1,41 @@
+"""Talkie model configuration for HuggingFace Transformers."""
+
+from transformers import PretrainedConfig
+
+
+class TalkieConfig(PretrainedConfig):
+    """Configuration class for the Talkie 13B decoder-only transformer.
+
+    This is a 40-layer, 40-head GPT with RoPE, SwiGLU, RMS normalisation,
+    embedding skip connections, and per-head / per-layer gain parameters.
+    """
+
+    model_type = "talkie"
+
+    def __init__(
+        self,
+        vocab_size: int = 65540,
+        hidden_size: int = 5120,
+        intermediate_size: int = 13696,
+        num_hidden_layers: int = 40,
+        num_attention_heads: int = 40,
+        head_dim: int = 128,
+        max_position_embeddings: int = 2048,
+        rope_theta: float = 1_000_000.0,
+        torch_dtype: str = "bfloat16",
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            torch_dtype=torch_dtype,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 65535,
+  "pad_token_id": 65535
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:675eabef7cc77d8b8fa1d0aa3781959bf514190e66f78099d634ba58e367f9aa
+size 26560483096

modeling_talkie.py

@@ -0,0 +1,469 @@
+"""Talkie 13B transformer — patched for long-context SFT.
+
+Differences vs lewtun/talkie-1930-13b-it-hf upstream:
+1. Liger fused linear cross-entropy in the loss path so the float32 logits
+   tensor (shape S x V) is never materialised in HBM. Roughly 16 GB saved at
+   S=64K, V=65540.
+2. FlashAttention varlen path keyed off `position_ids`. When TRL passes a
+   packed sequence (padding_free=True), tokens from different documents do
+   not attend across boundaries.
+3. Gradient checkpointing on the decoder stack.
+4. RoPE precompute is configurable via config.max_position_embeddings; we set
+   it to 64K at load time.
+"""
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from .configuration_talkie import TalkieConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    _HAS_FA = True
+except ImportError:
+    _HAS_FA = False
+
+try:
+    from liger_kernel.transformers.fused_linear_cross_entropy import (
+        LigerFusedLinearCrossEntropyLoss,
+    )
+    _HAS_LIGER = True
+except ImportError:
+    _HAS_LIGER = False
+
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TalkieCausalLMOutput(CausalLMOutputWithPast):
+    """CausalLMOutputWithPast plus a token_accuracy field expected by TRL when
+    SFTConfig.use_liger_kernel=True."""
+    token_accuracy: Optional[torch.Tensor] = None
+
+
+class TalkieHeadGain(nn.Module):
+    def __init__(self, n_head: int):
+        super().__init__()
+        self.head_g = nn.Parameter(torch.ones(n_head))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.head_g.type_as(x).view(1, 1, -1, 1)
+
+
+class TalkieWeightGain(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w_g = nn.Parameter(torch.ones(1))
+
+    def forward(self, w: torch.Tensor) -> torch.Tensor:
+        return w * self.w_g.type_as(w)
+
+
+class TalkieActGain(nn.Module):
+    def __init__(self, init_value: float):
+        super().__init__()
+        self.a_g = nn.Parameter(torch.ones(1) * init_value)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.a_g.type_as(x)
+
+
+def _apply_rotary_emb(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    assert x.ndim == 4
+    d = x.shape[3] // 2
+    x1 = x[..., :d]
+    x2 = x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], 3).type_as(x)
+
+
+def _precompute_rotary_embeddings(
+    seq_len: int, head_dim: int, base: float, device: torch.device
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+    inv_freq = 1.0 / (base ** (channel_range / head_dim))
+    t = torch.arange(seq_len, dtype=torch.float32, device=device)
+    freqs = torch.outer(t, inv_freq)
+    cos, sin = freqs.cos(), freqs.sin()
+    cos, sin = cos.bfloat16(), sin.bfloat16()
+    cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+    return cos, sin
+
+
+def _gather_rope_per_position(
+    cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Index RoPE tables by position_ids.
+
+    cos/sin: (1, S_table, 1, D_half)
+    position_ids: (B, S)
+    returns (B, S, 1, D_half) bf16
+    """
+    cos_t = cos[0, :, 0, :]  # (S_table, D_half)
+    sin_t = sin[0, :, 0, :]
+    flat = position_ids.reshape(-1)
+    cos_g = cos_t.index_select(0, flat).reshape(*position_ids.shape, 1, cos_t.shape[-1])
+    sin_g = sin_t.index_select(0, flat).reshape(*position_ids.shape, 1, sin_t.shape[-1])
+    return cos_g, sin_g
+
+
+def _cu_seqlens_from_position_ids(position_ids: torch.Tensor) -> torch.Tensor:
+    """Convert per-token position_ids (where each new doc restarts at 0) into
+    cu_seqlens suitable for flash_attn_varlen_func.
+
+    Expects shape (B, S). For B>1 flatten before calling. Returns only cu_seqlens;
+    the caller can pass the total sequence length as an over-approximation of
+    max_seqlen to avoid a forced .item() sync (which torch.compile breaks on).
+    """
+    pos = position_ids.reshape(-1)
+    starts = (pos == 0).nonzero(as_tuple=False).squeeze(-1)
+    cu = torch.cat(
+        [starts, torch.tensor([pos.numel()], device=pos.device, dtype=starts.dtype)]
+    ).to(torch.int32)
+    return cu
+
+
+class TalkieSelfAttention(nn.Module):
+    is_causal = True
+
+    def __init__(self, config: TalkieConfig, layer_idx: int = 0):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.n_head = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.scaling = 1.0 / math.sqrt(self.head_dim)
+        n_state = config.hidden_size
+
+        self.attn_query = nn.Linear(n_state, n_state, bias=False)
+        self.attn_key = nn.Linear(n_state, n_state, bias=False)
+        self.attn_value = nn.Linear(n_state, n_state, bias=False)
+        self.attn_resid = nn.Linear(n_state, n_state, bias=False)
+        self.head_gain = TalkieHeadGain(config.num_attention_heads)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos_sin: Tuple[torch.Tensor, torch.Tensor],
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = x.size()
+        q = self.attn_query(x).view(bsz, seq_len, self.n_head, self.head_dim)
+        k = self.attn_key(x).view(bsz, seq_len, self.n_head, self.head_dim)
+        v = self.attn_value(x).view(bsz, seq_len, self.n_head, self.head_dim)
+
+        cos, sin = cos_sin
+        q, k = _apply_rotary_emb(q, cos, sin), _apply_rotary_emb(k, cos, sin)
+        q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),))
+        q = self.head_gain(q)
+
+        if cu_seqlens is not None and _HAS_FA:
+            assert bsz == 1, "varlen path expects flattened batch"
+            q_f = q.reshape(seq_len, self.n_head, self.head_dim)
+            k_f = k.reshape(seq_len, self.n_head, self.head_dim)
+            v_f = v.reshape(seq_len, self.n_head, self.head_dim)
+            y = flash_attn_varlen_func(
+                q_f,
+                k_f,
+                v_f,
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                causal=True,
+            )
+            y = y.reshape(bsz, seq_len, self.n_head * self.head_dim)
+        else:
+            attn_impl = getattr(self.config, "_attn_implementation", "sdpa")
+            attn_fn = ALL_ATTENTION_FUNCTIONS.get(attn_impl)
+            if attn_fn is None:
+                attn_fn = ALL_ATTENTION_FUNCTIONS["sdpa"]
+            y, _ = attn_fn(
+                self,
+                q.transpose(1, 2),
+                k.transpose(1, 2),
+                v.transpose(1, 2),
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0,
+                is_causal=True,
+                **kwargs,
+            )
+            y = y.contiguous().view(bsz, seq_len, self.n_head * self.head_dim)
+        return self.attn_resid(y)
+
+
+class TalkieMLP(nn.Module):
+    def __init__(self, config: TalkieConfig):
+        super().__init__()
+        n_state = config.hidden_size
+        n_mlp = config.intermediate_size
+
+        self.mlp_gate = nn.Linear(n_state, n_mlp, bias=False)
+        self.mlp_linear = nn.Linear(n_state, n_mlp, bias=False)
+        self.mlp_resid = nn.Linear(n_mlp, n_state, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp_resid(F.silu(self.mlp_gate(x)) * self.mlp_linear(x))
+
+
+class TalkieDecoderLayer(nn.Module):
+    def __init__(self, config: TalkieConfig, layer_idx: int = 0):
+        super().__init__()
+        gain_init = (2 * config.num_hidden_layers) ** -0.5
+
+        self.layer_idx = layer_idx
+        self.attn = TalkieSelfAttention(config, layer_idx=layer_idx)
+        self.attn_gain = TalkieActGain(gain_init)
+        self.mlp = TalkieMLP(config)
+        self.mlp_gain = TalkieActGain(gain_init)
+        self.embed_skip = TalkieActGain(0.0)
+
+    def forward(
+        self,
+        e_x: torch.Tensor,
+        x: torch.Tensor,
+        cos_sin: Tuple[torch.Tensor, torch.Tensor],
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        x = x + self.attn_gain(
+            self.attn(
+                F.rms_norm(x, (x.shape[-1],)),
+                cos_sin,
+                cu_seqlens,
+                max_seqlen,
+                **kwargs,
+            )
+        )
+        x = x + self.mlp_gain(self.mlp(F.rms_norm(x, (x.shape[-1],))))
+        x = x + self.embed_skip(e_x)
+        return x
+
+
+class TalkieModel(PreTrainedModel):
+    """Decoder stack — HF-style forward so vLLM's transformers backend
+    (`AutoModel.from_config(...)`) can host this model."""
+
+    config_class = TalkieConfig
+    _no_split_modules = ["TalkieDecoderLayer"]
+    _supports_gradient_checkpointing = True
+    _supports_attention_backend = True
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+    base_model_prefix = "model"
+    # Empty plan = single-GPU / replicate. Multi-GPU TP would need entries
+    # for q/k/v/o-proj. vLLM tolerates an empty plan when world_size==1.
+    tp_plan = {}
+
+    def __init__(self, config: TalkieConfig):
+        super().__init__(config)
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList(
+            [
+                TalkieDecoderLayer(config, layer_idx=i)
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+        self.gradient_checkpointing = False
+        # Selective activation checkpointing: only checkpoint every Nth layer.
+        # stride=1 => every layer (HF default), stride=2 => half of layers,
+        # stride=N => no layers checkpointed. Set via env at construction time.
+        import os as _os
+        try:
+            self.gc_stride = max(1, int(_os.environ.get("TALKIE_GC_STRIDE", "1")))
+        except ValueError:
+            self.gc_stride = 1
+
+        self._rope_cos: torch.Tensor | None = None
+        self._rope_sin: torch.Tensor | None = None
+
+    def _set_gradient_checkpointing(self, enable: bool = True, gradient_checkpointing_func=None):
+        self.gradient_checkpointing = enable
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def _get_rope(
+        self, seq_len: int, device: torch.device
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        target = max(seq_len, self.config.max_position_embeddings)
+        if (
+            self._rope_cos is None
+            or self._rope_cos.shape[1] < target
+            or self._rope_cos.device != device
+        ):
+            cos, sin = _precompute_rotary_embeddings(
+                target,
+                self.config.head_dim,
+                self.config.rope_theta,
+                device=device,
+            )
+            self._rope_cos = cos
+            self._rope_sin = sin
+        return self._rope_cos[:, :target], self._rope_sin[:, :target]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        if inputs_embeds is None:
+            assert input_ids is not None
+            x = self.embed(input_ids)
+            seq_len = input_ids.shape[1]
+            device = input_ids.device
+        else:
+            x = inputs_embeds
+            seq_len = inputs_embeds.shape[1]
+            device = inputs_embeds.device
+
+        cos_table, sin_table = self._get_rope(seq_len, device)
+        if position_ids is not None:
+            cos_sin = _gather_rope_per_position(cos_table, sin_table, position_ids)
+        else:
+            cos_sin = (cos_table[:, :seq_len], sin_table[:, :seq_len])
+
+        # FlashAttention varlen path is for packed-sequence training only.
+        # During inference (HF generate, vLLM, etc.) we go through
+        # ALL_ATTENTION_FUNCTIONS instead.
+        cu_seqlens, max_seqlen = (None, None)
+        if self.training and position_ids is not None and _HAS_FA:
+            cu_seqlens = _cu_seqlens_from_position_ids(position_ids)
+            max_seqlen = seq_len
+
+        x = F.rms_norm(x, (x.shape[-1],))
+        e_x = x
+        for i, block in enumerate(self.blocks):
+            if (
+                self.gradient_checkpointing
+                and self.training
+                and (i % self.gc_stride == 0)
+            ):
+                x = torch.utils.checkpoint.checkpoint(
+                    block,
+                    e_x,
+                    x,
+                    cos_sin,
+                    cu_seqlens,
+                    max_seqlen,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(e_x, x, cos_sin, cu_seqlens, max_seqlen, **kwargs)
+        x = F.rms_norm(x, (x.shape[-1],))
+
+        if return_dict is False:
+            return (x,)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class TalkieForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = TalkieConfig
+    _no_split_modules = ["TalkieDecoderLayer"]
+    _supports_gradient_checkpointing = True
+    supports_gradient_checkpointing = True
+    _supports_attention_backend = True
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: TalkieConfig):
+        super().__init__(config)
+        self.model = TalkieModel(config)
+        # lm_head is an nn.Linear so the weight key is `lm_head.weight`,
+        # which matches the safetensors layout that vLLM's transformers
+        # backend also expects. The original talkie checkpoint stored a bare
+        # nn.Parameter and a separate `lm_head_gain.w_g` scalar; for serving
+        # we bake the gain into lm_head.weight ahead of time, so no gain
+        # module is needed here.
+        self.lm_head = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False
+        )
+
+        self.post_init()
+
+    def _set_gradient_checkpointing(self, enable: bool = True, gradient_checkpointing_func=None):
+        self.model.gradient_checkpointing = enable
+
+    def _get_rope(self, seq_len: int, device: torch.device):
+        # Backwards-compat shim for inference/fast_generate.py — RoPE tables
+        # now live on the inner TalkieModel.
+        return self.model._get_rope(seq_len, device)
+
+    def get_input_embeddings(self):
+        return self.model.embed
+
+    def set_input_embeddings(self, value):
+        self.model.embed = value
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[CausalLMOutputWithPast, Tuple]:
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            return_dict=False,
+        )
+        hidden_states = outputs[0]
+
+        loss = None
+        if labels is not None and _HAS_LIGER:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fn = LigerFusedLinearCrossEntropyLoss(return_token_accuracy=True)
+            res = loss_fn(
+                self.lm_head.weight,
+                shift_hidden.view(-1, shift_hidden.size(-1)),
+                shift_labels.view(-1),
+            )
+            return TalkieCausalLMOutput(
+                loss=res.loss, logits=None, token_accuracy=res.token_accuracy,
+            )
+
+        logits = self.lm_head(hidden_states)
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous().float()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+        else:
+            logits = logits.float()
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits)

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": null
+}