Checkpoint from Multi-Token Prediction Done Right (arXiv:2610.38331)

README.md

@@ -0,0 +1,83 @@
+---
+license: mit
+language:
+- en
+datasets:
+- HuggingFaceFW/fineweb-edu
+tags:
+- multi-token-prediction
+- speculative-decoding
+- soft-mtp
+library_name: transformers
+pipeline_tag: text-generation
+inference: false
+---
+
+# Soft-MTP — 332m
+
+Multi-token prediction against the model's own next-token verifier. One of twelve checkpoints released with
+[**Multi-Token Prediction Done Right**](https://arxiv.org/abs/2610.38331) (arXiv:2610.38331):
+three objectives x four sizes, trained under an identical recipe so they are directly comparable.
+
+Code: [JaydenTeoh/Soft-MTP](https://github.com/JaydenTeoh/Soft-MTP) — `configs/332m_fineweb.yaml`.
+
+## Objective
+
+Each MTP head is trained against the model's own next-token distribution `q` rather than the realized token `e_y`, via a straight-through KL-recovery coupling that costs one backward pass and no vocabulary softmax. See the paper for the gradient decomposition this removes.
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-332m-model", trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained("gpt2")          # the corpus tokenizer
+
+ids = tok("The capital of France is", return_tensors="pt").input_ids
+print(tok.decode(model.generate(ids, max_new_tokens=20)[0]))
+```
+
+`trust_remote_code=True` is required: the architecture (`model_type: hf_mtp`) is defined by the
+modelling files in this repo, not by `transformers`.
+
+The tokenizer is **stock GPT-2** — this repo ships no tokenizer files. `vocab_size` is padded from
+50257 to 50304 for alignment; the extra rows are unused.
+
+## Model
+
+| | |
+|---|---|
+| hidden size x layers | 768 x 36 |
+| attention heads | 12 |
+| context length | 2048 |
+| MTP heads | 3 (DeepSeek-style cascade, 1 block each) |
+| parameters | 357M total |
+| **deployed (NTP path)** | **357M** |
+
+MTP heads are auxiliary: they are discarded at inference, so what you deploy is an ordinary next-token transformer. They can also serve as a draft model for self-speculative decoding.
+
+## Training
+
+| | |
+|---|---|
+| data | FineWeb-Edu, GPT-2 BPE |
+| tokens | 31.5B (60,000 steps x 256 x 2048) |
+| optimizer | Muon on 2D hidden matrices + AdamW on embeddings/head/norms |
+| peak LR | 0.0015, cosine to 1e-05, 2,000 warmup steps |
+| weight decay | 0.1 |
+| precision | bf16 autocast over fp32 master weights |
+| seed | 13 |
+
+Single run per cell, one seed. The twelve checkpoints share seed, data order and schedule, so
+differences between them are attributable to the objective.
+
+## Citation
+
+```bibtex
+@article{teoh2026multitokenpredictionright,
+  title={Multi-Token Prediction Done Right},
+  author={Jayden Teoh},
+  journal={arXiv preprint arXiv:2610.38331},
+  year={2026},
+}
+```

blocks.py

@@ -0,0 +1,167 @@
+"""Pure-PyTorch Llama-style transformer block (no FLA/triton; runs on CPU + GPU).
+
+RMSNorm (fp32-upcast, matches Llama) + RoPE (rotate_half convention) + GQA attention via
+`F.scaled_dot_product_attention` (uses FlashAttention on CUDA) + SwiGLU. Verified equal to
+HF `LlamaForCausalLM` to 1e-4. One `Block` is reused
+for trunk layers and MTP heads, and `torch.compile`s cleanly (regional compile).
+"""
+
+from contextlib import nullcontext
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        in_dtype = x.dtype  # upcast LOW precision to fp32 (matches Llama; stable under bf16),
+        h = x.float() if in_dtype in (torch.float16, torch.bfloat16) else x  # but never downcast fp64
+        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (self.weight * h).to(in_dtype)
+
+
+def build_rope_cache(seq_len, head_dim, theta, device, dtype=torch.float32):
+    """cos/sin table [T, head_dim]. Trig computed in fp32, cast to `dtype` last."""
+    dev = device if isinstance(device, torch.device) else torch.device(device)
+    # Disable any outside autocast context if any, to really force fp32 (autocast has no policy for
+    # `meta`, which the param-count probes build on, so the context is guarded rather than assumed).
+    ctx = (torch.autocast(device_type=dev.type, enabled=False)
+           if dev.type in ("cpu", "cuda", "xpu") else nullcontext())
+    with ctx:
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+        freqs = torch.outer(torch.arange(seq_len, device=device).float(), inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)  # [T, head_dim]
+        return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def _rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(q, k, cos, sin):
+    """cos/sin are [T, D] (batch-shared positions) or [B, T, D] (per-row `position_ids`, the HF
+    convention — `generate()` derives them as `attention_mask.cumsum(-1) - 1`, so left-padded rows
+    get 0..n-1 on their REAL tokens). Both broadcast to q/k's [B, H, T, D]."""
+    if cos.dim() == 2:
+        cos, sin = cos[None, None], sin[None, None]  # [1, 1, T, D]
+    else:
+        cos, sin = cos[:, None], sin[:, None]        # [B, 1, T, D]
+    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.n_heads, self.n_kv = cfg.num_heads, cfg.num_kv_heads
+        self.hd = cfg.hidden_size // cfg.num_heads
+        self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.hd, bias=False)
+        self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.hd, cfg.hidden_size, bias=False)
+
+    def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
+                cache_readonly=False, attention_mask=None):
+        """`attention_mask` is the pre-built ADDITIVE 4D mask [B, 1, T, k_len] (causal + padding), or
+        None for the training/cache-free path which uses the compile-clean `is_causal` kernel."""
+        B, T, _ = x.shape
+        q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
+        q, k = apply_rope(q, k, cos, sin)
+        if past_key_values is not None:
+            if cache_readonly:
+                # READ-ONLY: attend against the cache without writing to it. transformers'
+                # `update()` is unconditionally append-in-place, so this is hand-rolled. The sequential
+                # drafter needs it — its last positions consume unconfirmed DRAFT tokens, which must
+                # attend but must never enter the cache, or a rejected draft poisons later drafts.
+                layer = past_key_values.layers[layer_idx]
+                pk, pv = getattr(layer, "keys", None), getattr(layer, "values", None)
+                if pk is not None and pk.numel():
+                    k = torch.cat([pk, k], dim=-2)
+                    v = torch.cat([pv, v], dim=-2)
+            else:                                        # inference KV cache (eager): append + read full K/V
+                k, v = past_key_values.update(k, v, layer_idx)
+        if self.n_kv != self.n_heads:  # GQA: expand kv heads
+            rep = self.n_heads // self.n_kv
+            k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
+        if attention_mask is None and past_key_values is not None and cache_position is not None:
+            # SAFETY NET: a cached call whose caller forgot the pre-built mask would otherwise run
+            # BIDIRECTIONAL attention with no error — exactly the f0816ed regression that left
+            # `_dsmtp_layer` non-causal on multi-position chunks while every decode gate kept
+            # passing (the chunk's LAST position is causal by accident, and that is the only one
+            # sampled from). Model-level callers still pass the shared mask, so this never fires on
+            # the trunk/sequential paths and their build-once economics are unchanged.
+            attention_mask = build_causal_mask(cache_position, k.size(2), None, x.device, q.dtype)
+        if past_key_values is None and attention_mask is None:
+            # training / cache-free without padding: clean compiled is_causal kernel. This mirrors
+            # HF's contract exactly — `create_causal_mask` returns None when the sdpa backend can
+            # use `is_causal` (no padding, no cache), and the attention interface falls back to it.
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:                                            # cached, and/or padded cache-free: additive mask
+            # `attention_mask` here is the FINAL additive mask, already combining causal + padding,
+            # built ONCE per forward by `build_causal_mask` and shared by every layer — the same split
+            # HF uses (DeepseekV3Model calls `create_causal_mask` and DeepseekV3Attention only passes
+            # the result to the attention interface). Building it here instead would reallocate the
+            # same [B,1,T,k_len] tensor once per layer.
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+        return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))
+
+
+def build_causal_mask(cache_position, k_len, attention_mask, device, dtype):
+    """The additive [B, 1, T, k_len] mask for the CACHED path: causal AND padding, built once per
+    forward and shared by every layer (HF's `create_causal_mask` plays this role for DeepSeek-V3 /
+    Llama; their attention modules never construct a mask themselves).
+
+    causal : key position <= query's absolute `cache_position`.
+    padding: `attention_mask` is the HF convention [B, k_len], 1 = real token, 0 = pad. Without this
+             term a query attends PAD keys, which silently degrades ANY batched generation with mixed
+             prompt lengths — lm-eval left-pads its generate_until batches.
+
+    Returns a 2D [T, k_len] mask when there is no padding to apply, so the common single-sequence case
+    does not pay for a [B,1,T,k_len] broadcast.
+    """
+    T = cache_position.numel()
+    kpos = torch.arange(k_len, device=device)
+    allow = kpos[None, :] <= cache_position[:, None]
+    mask = torch.zeros(T, k_len, device=device, dtype=dtype).masked_fill(~allow, float("-inf"))
+    if attention_mask is None:
+        return mask
+    am = attention_mask[:, :k_len]
+    mask = mask[None, None].expand(am.size(0), 1, T, k_len).clone()
+    return mask.masked_fill_(am[:, None, None, :] == 0, float("-inf"))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Block(nn.Module):
+    """Pre-norm transformer block; reused for trunk layers and MTP heads."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
+        self.attn = Attention(cfg)
+        self.mlp_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+
+    def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
+                cache_readonly=False, attention_mask=None):
+        x = x + self.attn(self.attn_norm(x), cos, sin, past_key_values, layer_idx, cache_position,
+                          cache_readonly, attention_mask)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "HFMTPForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.HFMTPConfig",
+    "AutoModelForCausalLM": "model.HFMTPForCausalLM"
+  },
+  "ce_kernel": "naive",
+  "dtype": "float32",
+  "hidden_size": 768,
+  "intermediate_size": 2048,
+  "kl_recovery": "ghost",
+  "lambda_kl": 1.0,
+  "max_position_embeddings": 2048,
+  "model_type": "hf_mtp",
+  "mtp_arch": "deepseek",
+  "mtp_detach_lmhead": false,
+  "mtp_head_layers": 1,
+  "mtp_per_head_norm": true,
+  "mtp_share_weights": false,
+  "mtp_weight": 1.0,
+  "mtp_weight_decay": null,
+  "n_mtp": 3,
+  "norm_eps": 1e-05,
+  "num_heads": 12,
+  "num_hidden_layers": 36,
+  "num_kv_heads": 12,
+  "objective": "soft_mtp",
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "vocab_size": 50304
+}

config.py

@@ -0,0 +1,71 @@
+"""HF config for the MTP model.
+
+Standalone `PretrainedConfig` (no FLA). `num_hidden_layers` is the NTP trunk depth; `n_mtp`
+heads are added on top (Meta) or sequential (DeepSeek), per `mtp_arch`. `objective` lives on the
+config so the model's `forward` can return the right loss for the HF `Trainer` (model-returns-
+loss pattern); the training script sets it.
+"""
+
+from transformers import PretrainedConfig
+
+# `soft_mtp` = KL / soft-target MTP (the mirror of `hard_mtp`'s one-hot targets). Do NOT confuse with
+# metric names that contain `kl` (e.g. `mtp_kl_verifier` = KL-to-verifier) — those are unrelated.
+OBJECTIVES = ("ntp", "hard_mtp", "soft_mtp", "direct_kl")
+
+
+class HFMTPConfig(PretrainedConfig):
+    model_type = "hf_mtp"
+
+    def __init__(
+        self,
+        vocab_size=50304,
+        hidden_size=768,
+        num_hidden_layers=12,        # NTP trunk depth
+        num_heads=12,
+        num_kv_heads=None,           # None -> = num_heads (no GQA)
+        n_mtp=3,                     # number of MTP heads / sequential depths
+        mtp_head_layers=1,           # transformer blocks per MTP head/depth
+        mtp_arch="deepseek",         # "deepseek" (sequential, DeepSeek-V3 style) | "meta" (parallel)
+                                     # The two differ in what a head READS, not just its shape: a
+                                     # sequential depth consumes the teacher-forced token x_{t+k} and
+                                     # is cropped to T-k, a meta head ignores input_ids entirely.
+                                     # Code built from this default rather than from a shipped
+                                     # config therefore gets a structurally different model.
+        mtp_per_head_norm=False,     # DeepSeek-faithful: each MTP head gets its OWN final RMSNorm before
+                                     # the SHARED lm_head (official `shared_head.norm` is per-module; only
+                                     # the head weight is shared). False (default) = one shared final norm
+                                     # for NTP + all heads.
+
+        max_position_embeddings=1024,
+        rope_theta=10000.0,
+        norm_eps=1e-5,
+        intermediate_size=None,      # None -> ~8/3 * hidden, rounded to 128
+        tie_word_embeddings=False,   # untied by default
+        objective="soft_mtp",        # ntp | hard_mtp | soft_mtp | direct_kl
+        mtp_weight=1.0,              # auxiliary weight on the MTP loss term (the STABLE-phase value if scheduled)
+        ce_kernel="naive",           # MTP-head CE: "naive" (materialize logits) | "liger" (fused linear CE, low mem)
+        **kwargs,
+    ):
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.n_mtp = n_mtp
+        self.mtp_head_layers = mtp_head_layers
+        self.mtp_arch = mtp_arch
+        self.mtp_per_head_norm = mtp_per_head_norm
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        if intermediate_size is None:
+            inter = int(8 * hidden_size / 3)
+            intermediate_size = (inter + 127) // 128 * 128
+        self.intermediate_size = intermediate_size
+        self.objective = objective
+        self.mtp_weight = mtp_weight
+        self.ce_kernel = ce_kernel
+        super().__init__(vocab_size=vocab_size, tie_word_embeddings=tie_word_embeddings, **kwargs)
+        assert hidden_size % num_heads == 0
+        assert num_heads % self.num_kv_heads == 0
+        assert mtp_arch in ("meta", "deepseek")
+        assert self.objective in OBJECTIVES, f"unknown objective {objective!r}"

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.6"
+}

loss.py

@@ -0,0 +1,180 @@
+"""MTP losses: ntp | hard_mtp | soft_mtp | direct_kl.
+
+soft_mtp trains each MTP head against the model's own next-token distribution instead of the
+realized token, using a straight-through KL-recovery coupling: adding `s_k * (sg(z_k) - z_k)` into the
+aligned slice of the NTP hidden leaves the forward value exactly unchanged, while the backward
+routes the NTP CE's own hidden gradient into head k. The head nets `coeff_k * (p_k - q)W`, the KL
+gradient, in one backward — no teacher cache, no second softmax.
+
+`labels` is PRE-SHIFTED and the same length as `input_ids`: `labels[:, t]` is the target of
+position t. Head k is supervised on its first `T - k` positions with target `labels[:, k:]`; NTP is
+head 0 with target `labels`. This is NOT HF's convention (`ForCausalLMLoss` shifts internally); we
+shift in the collator, so every input position is supervised.
+
+⚠ Teacher forcing and supervision read DIFFERENT sequences: the sequential forces depth k from
+`input_ids[:, k:]`, supervision comes from `labels`. They coincide in pretraining, which is the only
+reason one tensor ever sufficed — under SFT `labels` is IGNORE_INDEX over the prompt and embedding
+-100 would crash.
+"""
+
+import importlib
+
+import torch
+import torch.nn.functional as F
+
+IGNORE_INDEX = -100
+
+
+def direct_kl_loss_sum(student_logits, teacher_logits_aligned, labels, lam=1.0, ignore_index=IGNORE_INDEX):
+    """Forward KL D(q‖p) against the detached verifier, summed over non-ignored positions."""
+    # torch's KL has no ignore_index (only cross_entropy does), so the mask is hand-rolled.
+    mask = (labels != ignore_index).unsqueeze(-1)
+    hi = lambda t: t.float() if t.dtype in (torch.float16, torch.bfloat16) else t   # never DOWNcast fp64
+    logq = F.log_softmax(hi(teacher_logits_aligned), dim=-1).detach()
+    q = logq.exp()
+    logp = F.log_softmax(hi(student_logits), dim=-1)
+    kl = q * (logq - logp)                      # [..., V]
+    return lam * (kl * mask).sum()
+
+
+def _ce_mean(logits, target, ignore_index=IGNORE_INDEX):
+    if logits.dtype in (torch.float16, torch.bfloat16):
+        logits = logits.float()
+    return F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1), ignore_index=ignore_index)
+
+
+_LIGER_FLCE = None
+
+
+@torch.compiler.disable
+def _liger_flce(W, z, y, ignore_index):
+    # The disable is required: dynamo cannot trace Liger's autograd Function under bf16 autocast
+    # (its fused addmm(out_dtype=fp32) fails fake-tensor propagation). Costs a graph break per CE
+    # call, not the kernel.
+    global _LIGER_FLCE
+    if _LIGER_FLCE is None:
+        # importlib, not a literal `from liger_kernel...` line: liger is OPTIONAL, but HF's
+        # trust_remote_code loader statically regexes every import in a vendored modelling file
+        # and refuses to load unless each one is installed. A real import statement here would
+        # make an optional dependency mandatory for anyone loading this model from the Hub.
+        mod = importlib.import_module("liger_kernel.transformers.fused_linear_cross_entropy")
+        _LIGER_FLCE = mod.LigerFusedLinearCrossEntropyLoss(ignore_index=ignore_index, reduction="mean")
+    return _LIGER_FLCE(W, z.reshape(-1, z.size(-1)), y.reshape(-1))
+
+
+def ce_loss(z, W, y, ce_kernel, ignore_index=IGNORE_INDEX):
+    """Mean per-token CE against `W`, for any head. `naive` materializes `[*, V]` logits; `liger`
+    streams over the vocabulary and never does, which is the dominant memory term at large V."""
+    if ce_kernel == "liger":
+        return _liger_flce(W, z, y, ignore_index)
+    return _ce_mean(F.linear(z, W), y, ignore_index)
+
+
+def _direct_kl_loss(model, hiddens, h_ntp, labels, W, mtp_weight, ce_kernel):
+    """Full-KD ablation: the KL gradient reaches the trunk AND the shared `W`, where soft_mtp
+    recovers it into the trunk only and leaves `W` the hard push. Materializes `[B, T, V]` softmaxes.
+
+    Logging units stay cross-arm comparable: `mtp_heads` is the counterfactual hard CE (what the
+    other objectives report), `mtp_heads_kl` is the KL that actually enters `total`.
+    """
+    ntp_hidden = h_ntp
+    ntp_loss = ce_loss(ntp_hidden, W, labels, ce_kernel)
+    parts = {"ntp": ntp_loss.detach(), "mtp_heads": [], "mtp_heads_kl": []}
+    hi = lambda t: t.float() if t.dtype in (torch.float16, torch.bfloat16) else t
+    verifier_logits = (hi(ntp_hidden) @ hi(W).t()).detach()
+    head_terms = []
+    for head in range(1, model.n_heads_total):
+        head_targets = labels[:, head:]
+        head_span = head_targets.size(1)
+        head_hidden = model.final_norm(hiddens[head], head)[:, :head_span]
+        student_logits = hi(head_hidden) @ hi(W).t()
+        # The teacher is a plain crop: head `head` at t and NTP at t+head are supervised under the
+        # SAME condition, so every student position has a real teacher and there is no tail to mask.
+        teacher_logits = verifier_logits[:, head:head + head_span]
+        n_head_supervised = (head_targets != IGNORE_INDEX).sum().clamp(min=1)
+        head_terms.append(direct_kl_loss_sum(student_logits, teacher_logits, head_targets)
+                          / n_head_supervised)
+        parts["mtp_heads"].append(_ce_mean(student_logits.detach(), head_targets).detach())
+        parts["mtp_heads_kl"].append(head_terms[-1].detach())
+    mtp_mean = torch.stack(head_terms).mean()
+    total = ntp_loss + mtp_weight * mtp_mean
+    parts["mtp"] = torch.stack(parts["mtp_heads"]).mean()
+    parts["total"] = total.detach()
+    return total, parts
+
+
+def combined_mtp_loss(model, H, input_ids, labels, objective=None, mtp_weight=None):
+    """One scalar loss for ntp | hard_mtp | soft_mtp | direct_kl -> (total, parts).
+
+    hard_mtp and soft_mtp share one head loop: identical per-head CE against the same undetached
+    `W`, differing only in that soft_mtp also accumulates the KL-recovery term and runs the NTP CE on
+    the coupled hidden. The loss VALUE is the hard CE either way — the KL enters through the gradient.
+
+    `parts` (detached, for logging): {'ntp', 'mtp', 'mtp_heads', 'total'}; direct_kl adds
+    'mtp_heads_kl', the terms that actually enter `total`.
+    """
+    cfg = model.config
+    objective = cfg.objective if objective is None else objective
+    mtp_weight = cfg.mtp_weight if mtp_weight is None else mtp_weight
+    ce_kernel = getattr(cfg, "ce_kernel", "naive")
+    W = model.lm_head.weight
+
+    hiddens = model.all_head_hiddens(H, input_ids)       # [h0, ..., h_n_mtp] pre-norm
+    h_ntp = model.final_norm(hiddens[0], 0)
+
+    assert labels.shape[:2] == input_ids.shape[:2], (
+        f"labels must be PRE-SHIFTED and the same shape as input_ids (labels[:, t] = target of "
+        f"position t); got {tuple(labels.shape)} vs {tuple(input_ids.shape)}. HF's convention is the "
+        f"other one (unshifted labels, shifted inside the loss) -- collate with mtp_collate_dict."
+    )
+    # Guard, not a branch: with an empty head slice `mtp_mean` would average over fewer heads than
+    # `kl_recovery_scale` divides by, silently turning soft_mtp into a partial-removal arm.
+    assert labels.size(1) > model.n_heads_total - 1, (
+        f"seq_len ({labels.size(1)}) must exceed n_mtp ({model.n_heads_total - 1}): the deepest head "
+        f"would have no target."
+    )
+
+    if objective == "ntp" or cfg.n_mtp == 0:
+        ntp_loss = ce_loss(h_ntp, W, labels, ce_kernel)
+        return ntp_loss, {"ntp": ntp_loss.detach(), "mtp_heads": [],
+                          "mtp": torch.zeros((), device=H.device), "total": ntp_loss.detach()}
+
+    if objective == "direct_kl":
+        return _direct_kl_loss(model, hiddens, h_ntp, labels, W, mtp_weight, ce_kernel)
+
+    # ---- shared head loop: hard_mtp | soft_mtp ---------------------------------------------------
+    is_soft = objective == "soft_mtp"
+    ntp_targets = labels
+    ntp_hidden = h_ntp
+    n_heads = model.n_heads_total - 1
+
+    # float64 BEFORE dividing: an int64/int64 tensor divide is float32 in torch, and a trailing
+    # .to() only re-rounds what was already lost.
+    n_ntp_supervised = ((ntp_targets != IGNORE_INDEX).sum().clamp(min=1).to(torch.float64)
+                        if is_soft else None)
+
+    head_losses = []
+    for head in range(1, model.n_heads_total):
+        # Head k predicts the token k ahead, so head has seq_len-k supervised positions.
+        head_targets = labels[:, head:]
+        head_hidden = model.final_norm(hiddens[head], head)[:, :head_targets.size(1)]
+        head_losses.append(ce_loss(head_hidden, W, head_targets, ce_kernel))
+        if not is_soft:
+            continue
+
+        # straight-through gradient routing
+        # s_k undoes the NTP CE's 1/T mean and re-applies k-th MTP head's own
+        # 1/(T-k) mean and the 1/K head average, so the routed gradient
+        # lands at exactly the scale the MTP head's own CE gradient uses.
+        kl_recovery_scale = (mtp_weight * n_ntp_supervised
+                       / (n_heads * (head_targets != IGNORE_INDEX).sum().clamp(min=1).to(torch.float64))
+                       ).to(head_hidden.dtype)
+
+        # ntp_hidden is numerically unchanged; only its backward changes
+        ntp_hidden[:, head:] = ntp_hidden[:, head:] + kl_recovery_scale * (head_hidden.detach() - head_hidden)
+
+    ntp_loss = ce_loss(ntp_hidden, W, ntp_targets, ce_kernel)
+    mtp_mean = torch.stack(head_losses).mean()
+    total = ntp_loss + mtp_weight * mtp_mean
+    return total, {"ntp": ntp_loss.detach(), "mtp_heads": [m.detach() for m in head_losses],
+                   "mtp": mtp_mean.detach(), "total": total.detach()}

model.py

@@ -0,0 +1,359 @@
+"""MTP model — parallel (Meta-style) and sequential (DeepSeek-style) heads in one class.
+
+HF `PreTrainedModel` over the verified pure-PyTorch `Block` (no FLA/triton), so it runs under
+the HF `Trainer` with a single combined backward and also imports/runs locally. `forward`
+returns the combined MTP loss (model-returns-loss pattern) so `Trainer` picks it up directly.
+
+  trunk: num_hidden_layers Blocks -> H (pre-norm); NTP reads H directly.
+  meta:     head j  = mtp_head_layers Block(s) on H (parallel)        -> predicts x_{t+1+j}
+  deepseek: depth k = block_k(proj_k[RMSNorm(h^{k-1}); RMSNorm(emb(x_{t+k}))])  (sequential)
+  shared final RMSNorm + lm_head (so every head goes through the same W -> recovery identity).
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, DynamicCache, GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .blocks import RMSNorm, Block, build_causal_mask, build_rope_cache
+from .config import HFMTPConfig
+from .loss import combined_mtp_loss
+
+
+@dataclass
+class MTPCausalLMOutput(CausalLMOutputWithPast):
+    """`CausalLMOutputWithPast` + the NAMED per-step train-loss breakdown (detached tensors), so the
+    Trainer logs by field name instead of unpacking a positional tensor. All fields are plain tensors
+    (pytree-safe under whole-model torch.compile / DDP). Units are cross-arm-consistent: `mtp_loss`
+    is hard CE on EVERY arm (for direct_kl the counterfactual CE); `kl_loss` is the KL training loss,
+    present only on direct_kl (there `loss = ntp_loss + w*mean(kl_loss)`, else `+ w*mean(mtp_loss)`).
+
+    Logging convention (implemented in `MTPTrainer.log`): a scalar field logs under its own name; a
+    per-head vector field `<x>_loss` ([n_mtp]) logs its MEAN under `<x>_loss` plus `<x>_h{j}_loss`
+    per head (1-based)."""
+    ntp_loss: Optional[torch.FloatTensor] = None  # scalar: NTP CE
+    mtp_loss: Optional[torch.FloatTensor] = None  # [n_mtp]: per-head hard CE (MTP arms)
+    kl_loss: Optional[torch.FloatTensor] = None   # [n_mtp]: per-head KL (direct_kl only)
+
+    @classmethod
+    def breakdown_fields(cls):
+        """The breakdown = exactly the fields this class ADDS over `CausalLMOutputWithPast` —
+        derived, not hand-listed, so adding a field here is the only change ever needed."""
+        from dataclasses import fields
+        base = {f.name for f in fields(CausalLMOutputWithPast)}
+        return tuple(f.name for f in fields(cls) if f.name not in base)
+
+
+class HFMTPPreTrainedModel(PreTrainedModel):
+    config_class = HFMTPConfig
+    _no_split_modules = ["Block"]
+
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=std)
+
+
+class HFMTPForCausalLM(HFMTPPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: HFMTPConfig):
+        super().__init__(config)
+        Blk = Block
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.trunk = nn.ModuleList(Blk(config) for _ in range(config.num_hidden_layers))
+        D, base = config.n_mtp, config.num_hidden_layers
+        if config.mtp_arch == "meta":
+            self.mtp_heads = nn.ModuleList(
+                nn.ModuleList(Blk(config) for _ in range(config.mtp_head_layers)) for _ in range(D)
+            )
+        else:  # deepseek sequential
+            self.proj = nn.ModuleList(nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) for _ in range(D))
+            self.norms_h = nn.ModuleList(RMSNorm(config.hidden_size, config.norm_eps) for _ in range(D))
+            self.norms_e = nn.ModuleList(RMSNorm(config.hidden_size, config.norm_eps) for _ in range(D))
+            self.mtp_blocks = nn.ModuleList(
+                nn.ModuleList(Blk(config) for _ in range(config.mtp_head_layers)) for _ in range(D)
+            )
+        self.norm = RMSNorm(config.hidden_size, config.norm_eps)   # NTP (head 0) final norm; also the
+        # shared final norm for the MTP heads UNLESS mtp_per_head_norm (then each head has its own, below —
+        # the DeepSeek-faithful `shared_head.norm`; the lm_head weight stays shared either way).
+        if getattr(config, "mtp_per_head_norm", False):
+            self.mtp_final_norms = nn.ModuleList(
+                RMSNorm(config.hidden_size, config.norm_eps) for _ in range(D))
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self._rope = None
+        self.post_init()
+
+    # --- HF plumbing ---
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    # --- core ---
+    @property
+    def n_heads_total(self):
+        return 1 + self.config.n_mtp
+
+    def _rope_cache(self, T, device, dtype):
+        """Cache the fp32 table and cast on return, as HF does (it buffers only the dtype-stable
+        `inv_freq` and ends with `cos.to(dtype=x.dtype)`). Caching the CAST result went stale across a
+        dtype change -- the check below sees length and device but not dtype -- so an fp32 forward then
+        `.to(bfloat16)` fed fp32 cos/sin to a bf16 model and SDPA raised on mixed q/k/v dtypes."""
+        if self._rope is None or self._rope[0].size(0) < T or self._rope[0].device != device:
+            hd = self.config.hidden_size // self.config.num_heads
+            self._rope = build_rope_cache(T, hd, self.config.rope_theta, device, torch.float32)
+        cos, sin = self._rope
+        return cos[:T].to(dtype), sin[:T].to(dtype)
+
+    def trunk_forward(self, input_ids, past_key_values=None, use_cache=False, cache_position=None,
+                      attention_mask=None, position_ids=None, inputs_embeds=None):
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
+        T = x.size(1)
+        if past_key_values is None and cache_position is None:
+            if attention_mask is not None and not bool(attention_mask.all()):
+                # padded cache-free batch (e.g. lm-eval generate_until with use_cache=False): the
+                # additive causal+padding mask replaces `is_causal`. The `.all()` skip mirrors HF's
+                # `allow_is_causal_skip` — an all-ones mask keeps the fast kernel. Eval-only branch
+                # (training never passes a mask), so the CPU sync is fine and compile never sees it.
+                cp = torch.arange(T, device=x.device)
+                attn_mask = build_causal_mask(cp, T, attention_mask, x.device, x.dtype)
+                if position_ids is None:
+                    cos, sin = self._rope_cache(T, x.device, x.dtype)
+                else:
+                    cos_all, sin_all = self._rope_cache(int(position_ids.max()) + 1, x.device, x.dtype)
+                    cos, sin = cos_all[position_ids], sin_all[position_ids]
+                for block in self.trunk:
+                    x = block(x, cos, sin, attention_mask=attn_mask)
+                return x
+            # training / cache-free fast path (compile-clean: contiguous [0:T] rope + is_causal SDPA)
+            cos, sin = self._rope_cache(T, x.device, x.dtype)
+            for block in self.trunk:
+                x = block(x, cos, sin)
+            return x  # pre-norm trunk hidden H (NTP reads this)
+        # inference with a KV cache (eager): used by the cached self-spec decoders + self-spec eval.
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        if cache_position is None:
+            past = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past, past + T, device=x.device)
+        # RoPE: per-row `position_ids` when given (HF convention — generate() derives them from the
+        # attention mask, so left-padded rows count 0..n-1 on real tokens), else the batch-shared
+        # absolute `cache_position` (equivalent up to a uniform shift, which relative RoPE cancels).
+        if position_ids is not None:
+            cos_all, sin_all = self._rope_cache(int(position_ids.max()) + 1, x.device, x.dtype)
+            cos, sin = cos_all[position_ids], sin_all[position_ids]
+        else:
+            cos_all, sin_all = self._rope_cache(int(cache_position.max()) + 1, x.device, x.dtype)
+            cos, sin = cos_all[cache_position], sin_all[cache_position]
+        # Build the additive causal+padding mask ONCE and share it across layers — HF's split
+        # (DeepseekV3Model calls create_causal_mask; the attention modules only consume it). Per-layer
+        # construction would reallocate the same [B,1,T,k_len] tensor num_hidden_layers times.
+        k_len = int(cache_position.max()) + 1
+        attn_mask = build_causal_mask(cache_position, k_len, attention_mask, x.device, x.dtype)
+        for i, block in enumerate(self.trunk):
+            x = block(x, cos, sin, past_key_values, i, cache_position, attention_mask=attn_mask)
+        return x
+
+    def _meta_head(self, H, j):
+        if j == 0:
+            return H
+        cos, sin = self._rope_cache(H.size(1), H.device, H.dtype)
+        x = H
+        for block in self.mtp_heads[j - 1]:
+            x = block(x, cos, sin)
+        return x
+
+    def _dsmtp_layer(self, h_prev, token_ids, depth,
+                     past_key_values=None, cache_position=None, cache_readonly=False):
+        """Sequential depth `depth`. `h_prev` is h^{depth-1}; `token_ids[t]` is the teacher-forced token
+        x_{t+depth}.
+
+        INFERENCE KV CACHE (optional; the training path is untouched and bit-identical). With
+        `past_key_values`/`cache_position` the depth's blocks attend over the cached prefix, so a
+        drafter can extend the sequential by a few positions per round instead of re-running it over the
+        whole sequence. Mirrors `trunk_forward`'s cached branch: rope at ABSOLUTE positions, one
+        `layer_idx` per block within THIS depth's own cache.
+
+        `cache_readonly=True` attends against the cache without writing. The caller owns the split:
+        depth d at position t consumes x_{t+d}, so positions with t+d <= (last confirmed index) are
+        CONFIRMED and may be written (append-only, monotone, no crop -- the cache stays contiguous),
+        while the trailing positions consume unconfirmed DRAFTS and must be read-only or a rejected
+        draft poisons later draft distributions. See
+        `decoding/sequential.py`.
+        """
+        i = depth - 1
+        e = self.embed(token_ids.clamp(min=0))
+        h = self.proj[i](torch.cat((self.norms_h[i](h_prev), self.norms_e[i](e)), dim=-1))
+        # Carrier-precision parity: under bf16 autocast the proj emits bf16 while the trunk carrier
+        # h_prev is fp32, an asymmetry no production stack has. Match the carrier rather than forcing
+        # .float() — under a weight-cast bf16 model a forced fp32 would crash the block GEMMs.
+        h = h.to(h_prev.dtype)
+        if past_key_values is None and cache_position is None:
+            # training / cache-free fast path — UNCHANGED (contiguous [0:T] rope + is_causal SDPA)
+            cos, sin = self._rope_cache(h.size(1), h.device, h.dtype)
+            for block in self.mtp_blocks[i]:
+                h = block(h, cos, sin)
+        else:
+            past = past_key_values.get_seq_length() if past_key_values is not None else 0
+            if cache_position is None:
+                cache_position = torch.arange(past, past + h.size(1), device=h.device)
+            cos_all, sin_all = self._rope_cache(int(cache_position.max()) + 1, h.device, h.dtype)
+            cos, sin = cos_all[cache_position], sin_all[cache_position]   # rope at ABSOLUTE positions
+            # Causal mask built once and shared by this depth's blocks. k_len is the post-update key
+            # width `past + T`, NOT cache_position.max()+1 — the latter under-counts for a read-only
+            # overlap chunk, whose new keys duplicate already-cached positions.
+            attn_mask = build_causal_mask(cache_position, past + h.size(1), None, h.device, h.dtype)
+            for bi, block in enumerate(self.mtp_blocks[i]):
+                h = block(h, cos, sin, past_key_values, bi, cache_position, cache_readonly, attn_mask)
+        return h
+
+    # Public alias (the cache-free sequential self-spec decoder calls `model.dsmtp_layer(...)`, matching
+    # the FLA `dsmtp_transformer` interface).
+    def dsmtp_layer(self, h_prev, token_ids, depth,
+                    past_key_values=None, cache_position=None, cache_readonly=False):
+        return self._dsmtp_layer(h_prev, token_ids, depth,
+                                 past_key_values=past_key_values, cache_position=cache_position,
+                                 cache_readonly=cache_readonly)
+
+    def head_forward(self, H, j):
+        """Pre-norm hidden of head j (0 = NTP trunk). Matches the FLA model interface so
+        the shared `metrics.distill.evaluate_distill` / `trunk_grad_agreement` run on HF too.
+        DeepSeek depth j needs the sequential context -> use `all_head_hiddens` instead there."""
+        if self.config.mtp_arch == "meta":
+            return self._meta_head(H, j)
+        raise NotImplementedError("DeepSeek head_forward: use all_head_hiddens(H, labels)")
+
+    def all_head_hiddens(self, H, input_ids=None):
+        """Pre-norm hidden per head [h^0, ..., h^{n_mtp}]. Meta: parallel (input_ids unused).
+
+        DeepSeek: teacher-forced sequential. Depth k consumes the INPUT TOKEN x_{t+k}, i.e.
+        `input_ids[:, k:]`, so it is defined for t in [0, T-k) and h^k is CROPPED to T-k positions --
+        exactly the span on which head k is supervised (it predicts x_{t+1+k}).
+
+        ⚠ TEACHER FORCING READS INPUTS, NOT LABELS. These are the same sequence only in
+        pretraining. Under supervised fine-tuning `labels` is IGNORE_INDEX over the prompt, and
+        embedding -100 would crash -- so the sequential must never be fed labels. Nothing is read past the
+        input, so there is no padded target matrix and no mask.
+
+        The crop is per-depth and the loop is unrolled at trace time, so every slice length is a
+        compile-time constant -- no dynamic shapes, no recompilation.
+        """
+        if self.config.mtp_arch == "meta":
+            return [self._meta_head(H, j) for j in range(self.n_heads_total)]
+        T = input_ids.size(1)
+        outs, h_prev = [H], H
+        for k in range(1, self.n_heads_total):
+            L = T - k                                  # positions whose teacher token is in-window
+            if L <= 0:
+                break
+            h_prev = self._dsmtp_layer(h_prev[:, :L], input_ids[:, k:k + L], k)
+            outs.append(h_prev)
+        return outs
+
+    def final_norm(self, hidden, head_idx=0):
+        """Final RMSNorm before the shared lm_head. head_idx 0 = NTP -> self.norm. With
+        mtp_per_head_norm, head_idx k>=1 uses the head's OWN norm (DeepSeek `shared_head.norm`);
+        otherwise every head reuses self.norm (a NO-OP threading when the flag is off -> callers can
+        always pass the index and stay bit-identical)."""
+        if head_idx == 0 or not getattr(self.config, "mtp_per_head_norm", False):
+            return self.norm(hidden)
+        return self.mtp_final_norms[head_idx - 1](hidden)
+
+    def logits(self, hidden, head_idx=0):
+        return F.linear(self.final_norm(hidden, head_idx), self.lm_head.weight)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: "int | torch.Tensor" = 0,
+        return_all_heads: Optional[bool] = None,
+        **kwargs,
+    ) -> MTPCausalLMOutput:
+        """Signature follows `DeepseekV3ForCausalLM.forward` (HF `modeling_deepseek_v3.py`): HF
+        argument order and names (`position_ids`, `inputs_embeds`, int-or-tensor `logits_to_keep`),
+        and `past_key_values` returned on the output.
+
+        ⚠ `position_ids` in this signature is BEHAVIOUR, not decoration.
+        `GenerationMixin.prepare_inputs_for_generation` inspects the signature and, finding the name,
+        derives per-row positions as `attention_mask.cumsum(-1) - 1`, so a left-padded row gets RoPE
+        0..n-1 on its real tokens. That makes batched and unbatched generation exactly equal.
+
+        ⚠ The cache args are load-bearing. If they fall into `**kwargs` and get dropped while
+        `trunk_forward` runs cache-free, `generate()` (which defaults to `use_cache=True`) feeds only
+        the newest token from step 2 on, believing the prefix is cached — the model then sees one
+        token with no context and emits degenerate repetition after a plausible first token.
+
+        The cache is created here, not in `trunk_forward`, only so we hold a reference to return;
+        `Attention.forward` is what appends to it, in place, so returning the same object is correct.
+
+        Training is untouched: with no cache or positions, `trunk_forward` takes its compile-clean
+        `is_causal` fast path.
+        """
+        if (labels is not None or return_all_heads) and input_ids is None:
+            raise ValueError("MTP paths teacher-force from token ids: `labels`/`return_all_heads` "
+                             "require `input_ids`, not `inputs_embeds`")
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        H = self.trunk_forward(input_ids, past_key_values=past_key_values, use_cache=use_cache,
+                               cache_position=cache_position, attention_mask=attention_mask,
+                               position_ids=position_ids, inputs_embeds=inputs_embeds)
+        if return_all_heads:  # eval/distill: teacher-forced all-head logits [B, T-n_mtp, 1+n_mtp, V]
+            # Head j is supervised on T-j positions, so the sequential returns RAGGED hiddens. Stacking
+            # needs a rectangle -> take the common prefix (the deepest head's span, T-n_mtp). Nothing
+            # meaningful is dropped: the tail positions this omits are exactly the ones no deep head
+            # can predict, and every consumer compares heads AT THE SAME position.
+            hs = self.all_head_hiddens(H, input_ids)
+            n_pos = min(h.size(1) for h in hs)
+            stacked = torch.stack([self.logits(h[:, :n_pos], j) for j, h in enumerate(hs)], dim=2)
+            return MTPCausalLMOutput(logits=stacked)
+        if labels is not None:  # training: combined MTP loss, one backward (full logits omitted to save mem)
+            total, parts = combined_mtp_loss(self, H, input_ids, labels)
+            # NAMED detached breakdown on the output (see MTPCausalLMOutput): the Trainer logs by field
+            # name — no positional packing. `logits` stays None during training (full logits are
+            # deliberately not materialized to save memory; eval uses return_all_heads).
+            heads = parts["mtp_heads"]
+            kl_heads = parts.get("mtp_heads_kl")               # direct_kl only
+            return MTPCausalLMOutput(
+                loss=total,
+                ntp_loss=parts["ntp"],
+                mtp_loss=torch.stack(heads) if heads else None,
+                kl_loss=torch.stack(kl_heads) if kl_heads else None,
+            )
+        # `logits_to_keep` (DeepSeek-V3 / Llama convention): generation needs only the last position,
+        # and at V=50304 a 900-token prefill would otherwise materialise ~181 MB of unused fp32 logits.
+        # HF's exact form: int -> trailing slice (0 -> slice(0, None) == everything); Tensor -> fancy
+        # index of explicit positions (`DeepseekV3ForCausalLM.slice_indices`).
+        sl = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        return MTPCausalLMOutput(logits=self.logits(H[:, sl, :]), past_key_values=past_key_values)
+
+
+try:
+    AutoConfig.register(HFMTPConfig.model_type, HFMTPConfig)
+    AutoModelForCausalLM.register(HFMTPConfig, HFMTPForCausalLM)
+except ValueError:
+    pass

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd2151603c308bc0b7f10b3d1c695e3a96a6799783592fa26376ff32f6d2ac2
+size 1427681336

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}