Instructions to use Ayaka157/RetNet-0.5B-v21 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use Ayaka157/RetNet-0.5B-v21 with Transformers:

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

pipe = pipeline("text-generation", model="Ayaka157/RetNet-0.5B-v21", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Ayaka157/RetNet-0.5B-v21", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use Ayaka157/RetNet-0.5B-v21 with vLLM:

Install from pip and serve model

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

Use Docker

docker model run hf.co/Ayaka157/RetNet-0.5B-v21

SGLang

How to use Ayaka157/RetNet-0.5B-v21 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 "Ayaka157/RetNet-0.5B-v21" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "Ayaka157/RetNet-0.5B-v21",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

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

Docker Model Runner
How to use Ayaka157/RetNet-0.5B-v21 with Docker Model Runner:
```
docker model run hf.co/Ayaka157/RetNet-0.5B-v21
```

RetNet-0.5B-v21 / modeling_retnet.py

Ayaka157

Upload folder using huggingface_hub

a9e39bc verified 22 days ago

Raw History Blame Contribute Delete

32.2 kB

"""HuggingFace-style RetNet implementation — algorithm identical to the
handwritten retnet/model.py, wrapped for the mature ecosystem (Trainer,
Accelerate, safetensors, PreTrainedModel).
Parameter names are kept byte-for-byte identical to the handwritten model so
that ``load_state_dict(old_checkpoint["model_state"], strict=True)`` works
without any renaming.
    model = RetNetForCausalLM(config)
    model.load_state_dict(strip_optimized_prefix(old["model_state"]))
"""
from __future__ import annotations

import os
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig

#: Opt-in batched-chunk chunkwise retention (see ``retention_fast.py``).
#: Mathematically identical to the reference per-chunk loop but replaces the
#: Python loop's per-chunk slicing + small matmuls with 3 batched matmuls and a
#: tiny state scan.  Verified equivalent by ``scripts/test_fast_chunkwise.py``.
FAST_CHUNKWISE = os.environ.get("RETNET_FAST_CHUNKWISE", "0") == "1"

# NOTE(2026-10-10): 试过给 ``MultiScaleRetention._rotations`` 加一层
# ``(length, offset, device, dtype)`` 键的 sin/cos 表缓存，理由是 profile 里
# ``triton_poi_fused__to_copy*`` 合计 42.9 ms（占一个优化步 5.9%），而这 20 层
# 每次前向算的都是同一张表。**实测没有收益**：缓存开 1080 ms / 关 1084 ms
# （同一供电状态下背靠背），逐字符还原成原始实现也仍是 1078 ms。
# 结论是 inductor 已经把这部分处理掉了，Python 侧的缓存只是在编译区域里
# 多插一层间接。**不要再加这个缓存。**


class RetNetHFConfig(PretrainedConfig):
    model_type = "retnet_hf"

    def __init__(self, vocab_size=8192, d_model=512, num_layers=6, num_heads=8,
                 value_dim=None, ffn_dim=None, dropout=0.1, norm_eps=1e-5,
                 tie_weights=True, embed_scale=None, prompt_loss_weight=0.1,
                 use_gla=False, use_swiglu=False, use_growth_ln=False,
                 growth_old_dim=None, checkpoint_layers=-1, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.value_dim = value_dim if value_dim is not None else 2 * d_model
        self.ffn_dim = ffn_dim if ffn_dim is not None else 2 * d_model
        self.dropout = dropout
        self.norm_eps = norm_eps
        self.tie_weights = tie_weights
        self.embed_scale = embed_scale
        self.prompt_loss_weight = prompt_loss_weight
        self.use_gla = use_gla
        self.use_swiglu = use_swiglu
        self.use_growth_ln = use_growth_ln
        self.growth_old_dim = growth_old_dim
        #: 只对前 N 层做梯度检查点；-1 = 全部（沿用 HF 的全局开关）。
        #: 分层取舍的依据：重算开销与显存占用都随被检查点的层数线性变化，
        #: 而激活值总量也是。取中间值可以在同样的 8 GiB 里换到更大的物理 batch，
        #: 而更大的 batch 又能摊薄内核启动数（见 docs/15 §0.2）。
        self.checkpoint_layers = checkpoint_layers

    @property
    def key_dim(self) -> int:
        return self.d_model // self.num_heads


def rotate_every_two(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    x_rot = torch.stack((-x2, x1), dim=-1)
    return x_rot.flatten(-2)


def theta_shift(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
    return x * cos + rotate_every_two(x) * sin


class MultiScaleRetention(nn.Module):
    def __init__(self, config: RetNetHFConfig):
        super().__init__()
        d_model, num_heads, value_dim = config.d_model, config.num_heads, config.value_dim
        if d_model % num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        if value_dim % num_heads != 0:
            raise ValueError("value_dim must be divisible by num_heads")
        self.d_model = d_model
        self.num_heads = num_heads
        self.key_dim = d_model // num_heads
        self.value_dim = value_dim
        self.value_head_dim = value_dim // num_heads
        self.scale = self.key_dim ** -0.5

        gamma = 1.0 - torch.pow(2.0, -5.0 - torch.arange(num_heads, dtype=torch.float32))
        if config.use_gla:
            # GLA: 可学习逐头门控（初始值 = RetNet 固定公式，训练中自适应）
            self.gamma = nn.Parameter(gamma.clone())
        else:
            self.register_buffer("gamma", gamma, persistent=True)
        inv_freq = 1.0 / (10000.0 ** torch.linspace(0.0, 1.0, steps=self.key_dim // 2,
                                                     dtype=torch.float32))
        self.register_buffer("inv_freq", inv_freq.repeat_interleave(2), persistent=False)

        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, value_dim, bias=False)
        self.g_proj = nn.Linear(d_model, value_dim, bias=False)
        self.out_proj = nn.Linear(value_dim, d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)
        self.group_norm = nn.GroupNorm(num_heads, value_dim, eps=config.norm_eps)

    def _init_weights(self):
        for name, p in self.named_parameters():
            if name.endswith("_proj.weight"):
                if name.split(".")[-1].startswith(("q_", "k_", "v_", "g_")):
                    nn.init.xavier_uniform_(p, gain=2 ** -2.5)
                else:
                    nn.init.xavier_uniform_(p, gain=1.0)
            elif name == "gamma":
                # GLA 的逐头门控，初值必须是 __init__ 里按 RetNet 固定公式算出的
                # 1 - 2^(-5-h)，训练中再自适应。**绝不能置零**：
                # 之前这里写的是「1 维 Parameter 一律 zeros_」兜底 bias，可本模块
                # 所有 _proj 都是 bias=False，于是唯一的 1 维 Parameter 就是 gamma，
                # 门控被置零 → log(0) = -inf → 掩码对角线 exp(-inf*0) = NaN，
                # 批量 chunkwise 前向整层 NaN（use_gla=True 时每一条路径都中招）。
                # 参考实现用 gamma ** delta，0**0 = 1 良定义，所以两条路径在
                # gamma=0 处还**不等价**。详见 tests/test_gla_gamma_init.py。
                continue
            elif name.endswith("bias"):
                nn.init.zeros_(p)

    def _rotations(self, length, offset, device, dtype):
        position = torch.arange(offset, offset + length, device=device, dtype=torch.float32)
        freqs = torch.outer(position, self.inv_freq.to(device=device, dtype=torch.float32))
        return torch.sin(freqs).to(dtype=dtype), torch.cos(freqs).to(dtype=dtype)

    def _project_and_rotate(self, x, offset):
        bsz, tgt_len, _ = x.shape
        q = self.q_proj(x).view(bsz, tgt_len, self.num_heads, self.key_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, tgt_len, self.num_heads, self.key_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, tgt_len, self.num_heads, self.value_head_dim).transpose(1, 2)
        g = self.g_proj(x)
        k = k * self.scale
        sin, cos = self._rotations(tgt_len, offset, x.device, q.dtype)
        q = theta_shift(q, sin, cos)
        k = theta_shift(k, sin, cos)
        return q, k, v, g

    @staticmethod
    def _head_group_norm(x, num_heads, group_norm, bsz, tgt_len):
        """GroupNorm over ``[B,H,L,vhd]``，输出 ``[B,L,H*vhd]``。
        NOTE(2026-10-10): 试过改写成「先在 ``[B,H,L,vhd]`` 上手工算 mean/var
        归一化、最后再 permute 一次」，理由是省掉参考实现里那次物化拷贝。
        **实测反而慢 9.2%**（190 vs 174 ms/微步，真实 Trainer 路径，两轮一致）：
        参考实现把 permute 放在**归一化之前**，那次拷贝随即被逐元素归一化核吸收；
        而放到**之后**，它直接成为 ``out_proj`` 这个外部 cuBLAS 调用的输入，
        再也融不进任何东西。等价性没问题（fp32 相对差 1.6e-7，bf16 5.1e-3 均
        在 1 ULP 量级，见 ``tests/test_head_group_norm.py``），但收益是负的。
        **不要再动这个布局。**
        """
        x = x.transpose(1, 2).reshape(bsz * tgt_len, num_heads * x.shape[-1])
        x = group_norm(x.unsqueeze(-1)).squeeze(-1)
        return x.reshape(bsz, tgt_len, -1)

    def _decay_mask(self, length, device, dtype):
        idx = torch.arange(length, device=device)
        delta = idx[:, None] - idx[None, :]
        causal = delta >= 0
        exp = delta.masked_fill(~causal, 0)
        mask = self.gamma.to(device=device, dtype=torch.float32)[:, None, None] ** exp
        mask = mask.masked_fill(~causal, 0.0)
        mask = mask / mask.sum(dim=-1, keepdim=True).sqrt().clamp_min(1e-8)
        return mask.to(dtype=dtype)

    def _parallel_retention(self, q, k, v):
        bsz, _, tgt_len, _ = q.shape
        mask = self._decay_mask(tgt_len, q.device, q.dtype)
        scores = torch.matmul(q, k.transpose(-1, -2)) * mask
        scores = scores / scores.detach().abs().sum(dim=-1, keepdim=True).clamp_min(1.0)
        return torch.matmul(scores, v)

    def _parallel_final_state(self, k, v, state):
        bsz, _, tgt_len, key_dim = k.shape
        idx = torch.arange(tgt_len, device=k.device)
        weights = self.gamma.to(device=k.device, dtype=torch.float32)[:, None] ** (tgt_len - 1 - idx)
        new_state = torch.einsum("bhtk,bhtv,ht->bhkv", k, v, weights.to(k.dtype))
        if state is not None:
            gamma_pow = self.gamma.to(device=k.device, dtype=k.dtype) ** tgt_len
            new_state = gamma_pow.view(1, -1, 1, 1) * state + new_state
        return new_state

    def _recurrent_retention(self, q, k, v, state):
        bsz, _, tgt_len, key_dim = q.shape
        value_head_dim = v.shape[-1]
        if state is None:
            state = torch.zeros(bsz, self.num_heads, key_dim, value_head_dim,
                                device=q.device, dtype=q.dtype)
        outputs = []
        gamma = self.gamma.to(device=q.device, dtype=q.dtype)
        for t in range(tgt_len):
            state = gamma.view(1, -1, 1, 1) * state + \
                k[:, :, t, :].unsqueeze(-1) * v[:, :, t, :].unsqueeze(-2)
            out_t = torch.sum(q[:, :, t, :].unsqueeze(-1) * state, dim=-2)
            outputs.append(out_t)
        output = torch.stack(outputs, dim=2)
        return output, state

    def _chunkwise_retention(self, q, k, v, state, chunk_size):
        """Equation (7): inner-chunk parallel + cross-chunk recurrent state."""
        if FAST_CHUNKWISE:
            from .retention_fast import chunkwise_batched
            return chunkwise_batched(q, k, v, self.gamma, state, chunk_size)
        bsz, _, tgt_len, key_dim = q.shape
        if tgt_len % chunk_size != 0:
            raise ValueError(f"chunkwise mode requires seq_len ({tgt_len}) to be a "
                             f"multiple of chunk_size ({chunk_size})")
        if state is None:
            state = torch.zeros(bsz, self.num_heads, key_dim, v.shape[-1],
                                device=q.device, dtype=q.dtype)
        gamma = self.gamma.to(device=q.device, dtype=q.dtype)
        outputs = []
        rel = torch.arange(chunk_size, device=q.device)
        delta = rel[:, None] - rel[None, :]
        causal = delta >= 0
        inner_mask = torch.where(causal, gamma[:, None, None] ** delta.masked_fill(~causal, 0),
                                 torch.zeros_like(delta, dtype=torch.float32))
        cross_decay = gamma[:, None] ** (rel + 1)[None, :]
        value_decay = gamma[:, None] ** (chunk_size - 1 - rel)[None, :]
        chunk_decay = gamma ** chunk_size

        for i in range(tgt_len // chunk_size):
            sl = slice(i * chunk_size, (i + 1) * chunk_size)
            qi, ki, vi = q[:, :, sl], k[:, :, sl], v[:, :, sl]
            inner_scores = torch.matmul(qi, ki.transpose(-1, -2)) * inner_mask
            inner = torch.matmul(inner_scores, vi)
            cross = torch.matmul(qi * cross_decay[None, :, :, None], state)
            outputs.append(inner + cross)
            weighted_v = vi * value_decay[None, :, :, None]
            state = chunk_decay.view(1, -1, 1, 1) * state + \
                torch.matmul(ki.transpose(-1, -2), weighted_v)

        output = torch.cat(outputs, dim=2)
        return output, state

    def forward(self, x, mode="parallel", state=None, offset=0, return_state=False,
                chunk_size=64):
        bsz, tgt_len, _ = x.shape
        q, k, v, g = self._project_and_rotate(x, offset)

        if mode == "parallel":
            output, new_state = self._parallel_retention(q, k, v), None
            if return_state:
                new_state = self._parallel_final_state(k, v, state)
        elif mode == "recurrent":
            output, new_state = self._recurrent_retention(q, k, v, state)
        elif mode == "chunkwise":
            output, new_state = self._chunkwise_retention(q, k, v, state, chunk_size)
        else:
            raise ValueError(f"unknown mode: {mode}")

        output = self._head_group_norm(output, self.num_heads, self.group_norm, bsz, tgt_len)
        output = F.silu(g) * output
        output = self.dropout(output)
        output = self.out_proj(output)
        if return_state:
            return output, new_state
        return output


class GrowthLayerNorm(nn.Module):
    """Two-segment LayerNorm for function-preserving width growth.
    The old ``old_dim`` channels are normalized exactly as before (statistics
    unaffected by new channels); the new channels are scaled independently with
    ``gamma_new``/``beta_new`` and stay a zero side branch while
    ``gamma_new=0, beta_new=0``.  This is what makes pre-LN width expansion
    strictly function-preserving (BERT2BERT-style FPI for LayerNorm).
    state_dict keys: ``ln_old.weight``, ``ln_old.bias``, ``gamma_new``, ``beta_new``.
    """

    def __init__(self, old_dim: int, new_dim: int, eps: float = 1e-5,
                 gamma_new: float = 1.0, beta_new: float = 0.0):
        super().__init__()
        self.ln_old = nn.LayerNorm(old_dim, eps=eps)
        self.gamma_new = nn.Parameter(torch.full((new_dim - old_dim,), gamma_new))
        self.beta_new = nn.Parameter(torch.full((new_dim - old_dim,), beta_new))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        old = self.ln_old(x[..., : self.ln_old.normalized_shape[0]])
        new = x[..., self.ln_old.normalized_shape[0]:] * self.gamma_new + self.beta_new
        return torch.cat([old, new], dim=-1)


class RetNetBlock(nn.Module):
    def _norm(self, config: RetNetHFConfig, dim: int):
        if config.use_growth_ln:
            return GrowthLayerNorm(config.growth_old_dim, dim, eps=config.norm_eps,
                                   gamma_new=1.0, beta_new=0.0)
        return nn.LayerNorm(dim, eps=config.norm_eps)

    def __init__(self, config: RetNetHFConfig, layer_idx: int = 0):
        super().__init__()
        self.layer_idx = layer_idx
        #: 该层是否做梯度检查点。-1 表示「跟随 HF 的全局开关」。
        self.checkpoint_layers = getattr(config, "checkpoint_layers", -1)
        self.msr = MultiScaleRetention(config)
        self.norm1 = self._norm(config, config.d_model)
        self.gradient_checkpointing = False
        if config.use_swiglu:
            # SwiGLU FFN：y = silu(xW_gate) * (xW_up) -> W_down
            self.ffn = nn.Sequential()  # placeholder; real modules below
            self.gate_proj = nn.Linear(config.d_model, config.ffn_dim, bias=False)
            self.up_proj = nn.Linear(config.d_model, config.ffn_dim, bias=False)
            self.down_proj = nn.Linear(config.ffn_dim, config.d_model, bias=False)
            self.swiglu = True
        else:
            self.ffn = nn.Sequential(
                nn.Linear(config.d_model, config.ffn_dim, bias=False),
                nn.GELU(),
                nn.Linear(config.ffn_dim, config.d_model, bias=False),
            )
            self.swiglu = False
        self.norm2 = self._norm(config, config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def _block_forward(self, x, mode, state, offset, return_state, chunk_size=64):
        msr_out = self.msr(self.norm1(x), mode=mode, state=state, offset=offset,
                           return_state=return_state, chunk_size=chunk_size)
        if return_state:
            y, new_state = msr_out
        else:
            y, new_state = msr_out, None
        y = x + y
        if self.swiglu:
            ffn_out = F.silu(self.gate_proj(self.norm2(y))) * self.up_proj(self.norm2(y))
            ffn_out = self.down_proj(ffn_out)
        else:
            ffn_out = self.ffn(self.norm2(y))
        y = y + self.dropout(ffn_out)
        if return_state:
            return y, new_state
        return y

    def forward(self, x, mode="parallel", state=None, offset=0, return_state=False,
                chunk_size=64):
        do_ckpt = self.gradient_checkpointing and self.training
        if do_ckpt and 0 <= self.checkpoint_layers <= self.layer_idx:
            do_ckpt = False          # 分层策略：只检查点前 checkpoint_layers 层
        if do_ckpt:
            return torch.utils.checkpoint.checkpoint(
                self._block_forward, x, mode, state, offset, return_state, chunk_size,
                use_reentrant=False)
        return self._block_forward(x, mode, state, offset, return_state, chunk_size)


class RetNetForCausalLM(PreTrainedModel):
    """RetNet language model with a from-scratch token embedding.
    Forward signature is Trainer-compatible:
        forward(input_ids, labels=None, wcodes=None, ...) -> CausalLMOutput
    When labels+wcodes are given the loss is the wcode-weighted cross entropy
    (wcode: 0=pad, 1=reply, 2=prompt, 3=plain LM) exactly like train.py.
    """
    config_class = RetNetHFConfig
    base_model_prefix = "retnet"
    supports_gradient_checkpointing = True

    #: 本模型**不使用** ``num_items_in_batch`` 做损失归一化（loss 已按权重码
    #: ``w.sum()`` 归一），但 ``forward`` 签名末尾的 ``**_kwargs`` 会让 HF 的启发式
    #: 判定「模型自己会处理归一化」，从而**跳过** ``loss = loss / grad_accum_steps``。
    #: 后果是梯度被放大 accum 倍、等效学习率成倍上升 —— 静默，且只在 accum>1 时发作。
    #: HF 按 **``accepts_loss_kwargs``**（不是 ``model_accepts_loss_kwargs``）查找；
    #: 找不到时它会去猜 forward 签名，本项目会被误判为 True。
    accepts_loss_kwargs = False
    model_accepts_loss_kwargs = False

    def __init__(self, config: RetNetHFConfig):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.embed_dropout = nn.Dropout(config.dropout)
        self.embed_scale = config.embed_scale if config.embed_scale is not None \
            else config.d_model ** 0.5
        self.blocks = nn.ModuleList([RetNetBlock(config, i)
                                     for i in range(config.num_layers)])
        if config.use_growth_ln:
            self.final_norm = GrowthLayerNorm(config.growth_old_dim, config.d_model,
                                              eps=config.norm_eps,
                                              gamma_new=1.0, beta_new=0.0)
        else:
            self.final_norm = nn.LayerNorm(config.d_model, eps=config.norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        if config.tie_weights:
            self.lm_head.weight = self.embedding.weight
            self._tied_weights_keys = ["lm_head.weight"]
        self._init_weights_local()

    def _init_weights_local(self):
        nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
        for block in self.blocks:
            if block.swiglu:
                nn.init.xavier_uniform_(block.gate_proj.weight)
                nn.init.xavier_uniform_(block.up_proj.weight)
                nn.init.xavier_uniform_(block.down_proj.weight)
            else:
                for name, p in block.ffn.named_parameters():
                    if name.endswith("weight"):
                        nn.init.xavier_uniform_(p)
            block.msr._init_weights()
        if not self.config.tie_weights:
            nn.init.xavier_uniform_(self.lm_head.weight)

    def tie_weights(self):
        """共享 lm_head 与 embedding（from_pretrained 后恢复 tie）。"""
        if self.config.tie_weights:
            self.lm_head.weight = self.embedding.weight

    def set_checkpoint_layers(self, n: int):
        """只对前 ``n`` 层做梯度检查点（n<0 = 全部）。
        必须在 HF 启用梯度检查点**之后**调用，因为它会把所有层的
        ``gradient_checkpointing`` 置 True。
        """
        self.config.checkpoint_layers = n
        for b in self.blocks:
            b.checkpoint_layers = n
        return n

    def _forward_hidden(self, input_ids, mode="parallel", states=None, offset=0,
                        return_state=False, chunk_size=64, return_hidden=False):
        if input_ids.ndim != 2:
            raise ValueError("input_ids must be [batch, seq]")
        x = self.embedding(input_ids) * self.embed_scale
        x = self.embed_dropout(x)
        if states is None:
            states = [None] * len(self.blocks)
        if len(states) != len(self.blocks):
            raise ValueError("states must contain one entry per block")
        new_states: List[Optional[torch.Tensor]] = []
        for block, block_state in zip(self.blocks, states):
            if return_state or (mode == "recurrent" and block_state is not None):
                x, ns = block(x, mode=mode, state=block_state, offset=offset,
                              return_state=True, chunk_size=chunk_size)
                new_states.append(ns)
            else:
                x = block(x, mode=mode, state=block_state, offset=offset,
                          chunk_size=chunk_size)
                new_states.append(block_state)
        x = self.final_norm(x)
        if return_hidden:
            # 融合交叉熵需要的是 LM head 的**输入**（hidden），不是输出。
            # 这里提前返回可以完全不物化 [B,L,V]。
            if return_state or mode == "recurrent":
                return x, new_states
            return x
        logits = self.lm_head(x)
        if return_state or mode == "recurrent":
            return logits, new_states
        return logits

    def forward(self, input_ids, labels=None, wcodes=None, mode=None,
                states=None, offset=0, return_state=False, chunk_size=None, **_kwargs):
        # NOTE(2026-10-10): 试过在这里调 torch.compiler.cudagraph_mark_step_begin()
        # 以启用 torch.compile(mode="reduce-overhead") 的 CUDA Graphs。**无效**——
        # 错误从损失移到 _forward_hidden，仍是 "accessing tensor output of
        # CUDAGraphs that has been overwritten / Gradient addition node due to
        # multiple use of tensor"：这是 CUDA Graphs 与梯度累积（同一张图跑 4 次、
        # 各自反向）的硬冲突，不是漏调一个 API 能解决的。已回滚。
        # NOTE(2026-10-10): 上面这个默认值以前写的是 64，而 HF Trainer 调用模型时
        # **不会**传这个关键字，于是 ``config.chunk_size`` 和任务 JSON 里的
        # ``chunk_size`` 一直是**死字段**——项目以为在跑 128，实际永远 64。
        # 实测：``config.chunk_size`` 设 64/128/256 的 loss **逐位相同**，而显式
        # 传 ``chunk_size=`` 会改变 loss。``bench_step.py`` 的
        # ``--sweep chunk_size=...`` 因此扫的是空变量，``docs/15`` §6 的
        # 「chunk 512/1024 反而更慢，最优是 128」不成立。
        # 现改成 None 再回落到 config，任务配置才真正生效。
        # ⚠️ 对已经在跑的任务，这会让 chunk 从 64 跳到配置值 —— 数值会跳变，
        # 只能在**新开一轮**时启用，不要热切。
        if chunk_size is None:
            chunk_size = int(getattr(self.config, "chunk_size", 64) or 64)
        if mode is None:
            mode = getattr(self.config, "mode", "parallel")
        _fused = (os.environ.get("RETNET_FUSED_CE") == "1"
                  and labels is not None and wcodes is not None)
        logits = None
        # `loss` 必须在所有分支之前初始化：不传 labels 时（普通推理、prefill、
        # 以及 HF 的 generate 路径）上面几个分支都不会赋值，line 536 的
        # CausalLMOutput(loss=loss, ...) 会直接 UnboundLocalError。
        loss = None
        if _fused:
            # 融合路径：LM head 的输入直接进 fused CE，完全不物化 [B,L,V]
            hidden, new_states = self._forward_hidden(
                input_ids, mode=mode, states=states, offset=offset,
                return_state=True, chunk_size=chunk_size, return_hidden=True)
            B_, L_, D_ = hidden.shape
            w = torch.zeros_like(wcodes, dtype=torch.float32)
            w = torch.where(wcodes == 1, torch.ones_like(w), w)
            w = torch.where(wcodes == 2,
                            torch.full_like(w, self.config.prompt_loss_weight), w)
            w = torch.where(wcodes == 3, torch.ones_like(w), w)
            from .fused_ce import fused_linear_cross_entropy
            loss = fused_linear_cross_entropy(
                hidden.reshape(B_ * L_, D_), self.lm_head.weight,
                labels.reshape(-1), w.reshape(-1))
        elif return_state or mode == "recurrent" or mode == "chunkwise":
            logits, new_states = self._forward_hidden(input_ids, mode=mode, states=states,
                                                      offset=offset, return_state=True,
                                                      chunk_size=chunk_size)
        else:
            logits = self._forward_hidden(input_ids, mode=mode, states=states, offset=offset)
            new_states = None

        if not _fused:
            # 非融合路径：沿用原来的显式 fp32 上采样。
            # NOTE(2026-10-10): 试过去掉 ``logits.float()``（让 ``F.cross_entropy``
            # 直接吃 bf16）。loss 逐位一致但**步时 0% 变化**——profile 把这两个
            # 融合核记成 3.9%，是归因假象。故保留原状。
            if labels is not None and wcodes is not None:
                logits_f = logits.float()
                token_loss = F.cross_entropy(
                    logits_f.reshape(-1, logits_f.shape[-1]), labels.reshape(-1),
                    reduction="none")
                w = torch.zeros_like(wcodes, dtype=torch.float32)
                w = torch.where(wcodes == 1, torch.ones_like(w), w)
                w = torch.where(wcodes == 2,
                                torch.full_like(w, self.config.prompt_loss_weight), w)
                w = torch.where(wcodes == 3, torch.ones_like(w), w)
                loss = (token_loss * w.reshape(-1)).sum() / w.sum().clamp_min(1.0)
            elif labels is not None:
                logits_f = logits.float()
                loss = F.cross_entropy(logits_f.reshape(-1, logits_f.shape[-1]),
                                       labels.reshape(-1))

        from transformers.modeling_outputs import CausalLMOutput
        return CausalLMOutput(loss=loss, logits=logits)

    @torch.inference_mode()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 32,
                 temperature: float = 1.0, top_k: Optional[int] = None,
                 top_p: Optional[float] = None, repetition_penalty: float = 1.0,
                 eos_id: Optional[int] = None, device: Optional[torch.device] = None,
                 penalize_prompt: bool = False, no_repeat_ngram_size: int = 0,
                 on_token=None) -> torch.Tensor:
        """Autoregressive generation (same algorithm/sampling as handwritten).
        ``on_token(token_id)`` is called for every accepted token, in order, as
        soon as it is sampled.  Returning a truthy value stops generation
        immediately (the token is still part of the returned sequence).  Used by
        the interactive play script to stream output and to cut a reply at the
        first ``用户：`` / ``助手：`` turn boundary.
        """
        device = device or next(self.parameters()).device
        ids = input_ids.to(device)
        states = [None] * len(self.blocks)
        if ids.shape[1] > 1:
            logits, states = self._forward_hidden(ids, mode="parallel", return_state=True)
            next_logits = logits[:, -1]
        else:
            next_logits, states = self._forward_hidden(ids, mode="recurrent", states=states,
                                                       offset=0, return_state=True)
            next_logits = next_logits[:, -1]
        generated = list(ids[0].tolist())
        prompt_len = ids.shape[1]
        offset = ids.shape[1]

        for _ in range(max_new_tokens):
            penalty_source = generated if penalize_prompt else generated[prompt_len:]
            banned = []
            if no_repeat_ngram_size and len(generated) >= no_repeat_ngram_size:
                n = no_repeat_ngram_size
                suffix = tuple(generated[-(n - 1):]) if n > 1 else ()
                for i in range(len(generated) - n + 1):
                    if tuple(generated[i:i + n - 1]) == suffix:
                        banned.append(generated[i + n - 1])
            if temperature > 0:
                logits = next_logits / temperature
                if repetition_penalty and repetition_penalty != 1.0:
                    for token_id in set(penalty_source):
                        score = logits[0, token_id]
                        logits[0, token_id] = score / repetition_penalty if score >= 0 \
                            else score * repetition_penalty
                for token_id in banned:
                    logits[0, token_id] = -torch.inf
                if top_k is not None:
                    k = min(top_k, logits.shape[-1])
                    top_values, _ = torch.topk(logits, k)
                    logits[logits < top_values[:, -1:]] = -torch.inf
                probs = F.softmax(logits, dim=-1)
                if top_p is not None and top_p < 1.0:
                    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
                    cumsum = sorted_probs.cumsum(dim=-1)
                    remove = cumsum > top_p
                    remove[..., 1:] = remove[..., :-1].clone()
                    remove[..., 0] = False
                    filtered = torch.zeros_like(probs)
                    filtered.scatter_(1, sorted_indices[~remove].unsqueeze(0),
                                      sorted_probs[~remove].unsqueeze(0))
                    probs = filtered / filtered.sum(dim=-1, keepdim=True)
                next_id = torch.multinomial(probs, num_samples=1).item()
            else:
                logits = next_logits.clone()
                for token_id in banned:
                    logits[0, token_id] = -torch.inf
                next_id = logits.argmax(dim=-1).item()
            if eos_id is not None and next_id == eos_id:
                break
            generated.append(next_id)
            if on_token is not None and on_token(next_id):
                break
            next_token = torch.tensor([[next_id]], device=device, dtype=torch.long)
            next_logits, states = self._forward_hidden(next_token, mode="recurrent",
                                                       states=states, offset=offset,
                                                       return_state=True)
            next_logits = next_logits[:, -1]
            offset += 1
        return torch.tensor([generated], device=device, dtype=torch.long)