LightNovel-Intro-RetNet-400M / modeling_retnet.py

Upload 2 files

4371fee verified 4 months ago

No virus

50.2 kB

	# by syncdoth: https://github.com/syncdoth/RetNet/blob/main/retnet/modeling_retnet.py

	import math
	from dataclasses import dataclass
	from typing import Dict, List, Optional, Tuple, Union

	import numpy as np
	import torch
	import torch.nn.functional as F
	import torch.utils.checkpoint
	from timm.models.layers import drop_path
	from torch import nn
	from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
	from transformers import top_k_top_p_filtering
	from transformers.modeling_outputs import ModelOutput, SequenceClassifierOutputWithPast
	from transformers.modeling_utils import PreTrainedModel
	from transformers.utils import logging

	try:
	from apex.normalization import FusedLayerNorm as LayerNorm
	except ModuleNotFoundError:
	from torch.nn import LayerNorm

	from .configuration_retnet import RetNetConfig

	logger = logging.get_logger(__name__)


	# helper functions
	def split_heads(tensors, bsz, seqlen, num_heads):
	assert isinstance(tensors, (tuple, list))
	return [x.view(bsz, seqlen, num_heads, -1).transpose(1, 2) for x in tensors]


	def rotate_every_two(x):
	x1 = x[:, :, :, ::2]
	x2 = x[:, :, :, 1::2]
	x = torch.stack((-x2, x1), dim=-1)
	return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)')\


	def theta_shift(x, sin, cos):
	return (x * cos) + (rotate_every_two(x) * sin)


	def get_activation_fn(activation):
	if activation == "relu":
	return F.relu
	elif activation == "gelu":
	return F.gelu
	elif activation == "swish":
	return F.silu
	else:
	raise NotImplementedError


	class RMSNorm(nn.Module):
	def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True):
	super().__init__()
	self.normalized_shape = dim
	self.eps = eps
	self.elementwise_affine = elementwise_affine
	if self.elementwise_affine:
	self.weight = nn.Parameter(torch.ones(dim))
	else:
	self.register_parameter("weight", None)

	def _norm(self, x):
	return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

	def forward(self, x):
	output = self._norm(x.float()).type_as(x)
	if self.weight is not None:
	output = output * self.weight
	return output


	class RetNetRelPos(nn.Module):
	def __init__(self, config: RetNetConfig):
	super().__init__()
	self.config = config
	num_heads = config.decoder_retention_heads

	angle = 1.0 / (
	10000 ** torch.linspace(0, 1, config.decoder_embed_dim // num_heads // 2)
	)
	angle = angle.unsqueeze(-1).repeat(1, 2).flatten()
	# decay (gamma)
	if config.use_lm_decay:
	# NOTE: alternative way described in the paper
	s = torch.log(torch.tensor(1 / 32))
	e = torch.log(torch.tensor(1 / 512))
	decay = torch.log(1 - torch.exp(torch.linspace(s, e, num_heads))) # [h,]
	else:
	decay = torch.log(
	1 - 2 ** (-5 - torch.arange(num_heads, dtype=torch.float))
	)
	self.register_buffer("angle", angle)
	self.register_buffer("decay", decay)
	self.recurrent_chunk_size = config.recurrent_chunk_size

	def forward(
	self,
	slen,
	forward_impl="parallel",
	recurrent_chunk_size=None,
	retention_mask=None,
	get_decay_scale=True,
	):
	if forward_impl == "recurrent":
	sin = torch.sin(self.angle * (slen - 1))
	cos = torch.cos(self.angle * (slen - 1))
	retention_rel_pos = ((sin, cos), self.decay.view(1, -1, 1, 1).exp())
	elif forward_impl == "chunkwise":
	if recurrent_chunk_size is None:
	recurrent_chunk_size = self.recurrent_chunk_size
	index = torch.arange(slen).to(self.decay)
	sin = torch.sin(index[:, None] * self.angle[None, :])
	cos = torch.cos(index[:, None] * self.angle[None, :])

	block_index = torch.arange(recurrent_chunk_size).to(self.decay)
	mask = torch.tril(
	torch.ones(recurrent_chunk_size, recurrent_chunk_size)
	).to(self.decay)
	mask = torch.masked_fill(
	block_index[:, None] - block_index[None, :], ~mask.bool(), float("inf")
	)
	mask = torch.exp(mask * self.decay[:, None, None])
	mask = torch.nan_to_num(mask)
	mask = mask.unsqueeze(0) # [1, h, t, t]
	# TODO: need to handle retention_mask
	# scaling
	value_inner_decay = mask[:, :, -1] / mask[:, :, -1].sum(
	dim=-1, keepdim=True
	)
	value_inner_decay = value_inner_decay.unsqueeze(-1)
	scale = mask.sum(dim=-1, keepdim=True).sqrt()
	inner_mask = mask / scale

	cross_decay = torch.exp(self.decay * recurrent_chunk_size)
	query_inner_decay = torch.exp(self.decay[:, None] * (block_index + 1))
	cross_decay = cross_decay[None, :, None, None]
	query_inner_decay = query_inner_decay[None, :, :, None] / (
	scale / mask[:, :, -1].sum(dim=-1)[:, :, None, None]
	)
	# decay_scale (used for kv cache)
	if get_decay_scale:
	decay_scale = self.compute_decay_scale(slen, retention_mask)
	else:
	decay_scale = None
	retention_rel_pos = (
	(sin, cos),
	(
	inner_mask,
	cross_decay,
	query_inner_decay,
	value_inner_decay,
	decay_scale,
	),
	)
	else: # parallel
	index = torch.arange(slen).to(self.decay)
	sin = torch.sin(index[:, None] * self.angle[None, :])
	cos = torch.cos(index[:, None] * self.angle[None, :])
	mask = torch.tril(torch.ones(slen, slen)).to(self.decay)
	mask = torch.masked_fill(
	index[:, None] - index[None, :], ~mask.bool(), float("inf")
	)
	mask = torch.exp(mask * self.decay[:, None, None])
	mask = torch.nan_to_num(mask)
	mask = mask.unsqueeze(0) # [1, h, t, t]
	if retention_mask is not None:
	# this is required for left padding
	mask = mask * retention_mask.float().view(-1, 1, 1, slen).to(mask)

	# scaling
	mask = mask / mask.sum(dim=-1, keepdim=True).sqrt()
	mask = torch.nan_to_num(mask, nan=0.0)
	# decay_scale (used for kv cache)
	if get_decay_scale:
	decay_scale = self.compute_decay_scale(slen, retention_mask)
	else:
	decay_scale = None
	# mask processing for intra decay
	if retention_mask is not None:
	max_non_zero = (
	torch.cumsum(retention_mask, dim=-1).max(dim=-1).indices
	) # [b,]
	intra_decay = mask[range(mask.shape[0]), :, max_non_zero]
	else:
	intra_decay = mask[:, :, -1]

	retention_rel_pos = ((sin, cos), (mask, intra_decay, decay_scale))

	return retention_rel_pos

	def compute_decay_scale(self, slen, retention_mask=None):
	exponent = torch.arange(slen, device=self.decay.device).float()
	decay_scale = self.decay.exp().view(-1, 1) ** exponent.view(1, -1) # [h, t]
	if retention_mask is not None:
	seqlen = retention_mask.sum(dim=-1) # [b,]
	bsz = seqlen.size(0)
	decay_scale = decay_scale.unsqueeze(0).repeat(bsz, 1, 1) # [b, h, t]
	for i, pos in enumerate(seqlen):
	# the formula for decay_scale is `sum(gamma^i) for i in [0, slen).`
	# Since the retention_mask is 0 for padding, we can set the decay_scale
	# to 0 for the padding positions.
	decay_scale[i, :, pos.item() :] = 0
	else:
	bsz = 1
	decay_scale = decay_scale.sum(-1).view(bsz, -1, 1, 1) # [b, h, 1, 1]
	return decay_scale


	class MultiScaleRetention(nn.Module):
	def __init__(
	self,
	config: RetNetConfig,
	gate_fn="swish",
	use_bias=False,
	tensor_parallel=False,
	):
	super().__init__()
	self.config = config
	self.embed_dim = config.decoder_embed_dim
	self.value_dim = config.decoder_value_embed_dim
	self.num_heads = config.decoder_retention_heads
	self.head_dim = self.value_dim // self.num_heads
	self.key_dim = self.embed_dim // self.num_heads
	self.scaling = self.key_dim**-0.5

	self.gate_fn = get_activation_fn(activation=str(gate_fn))

	self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=use_bias)
	self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=use_bias)
	self.v_proj = nn.Linear(self.embed_dim, self.value_dim, bias=use_bias)
	self.g_proj = nn.Linear(self.embed_dim, self.value_dim, bias=use_bias)

	self.out_proj = nn.Linear(self.value_dim, self.embed_dim, bias=use_bias)

	self.group_norm = RMSNorm(
	self.head_dim, eps=config.layernorm_eps, elementwise_affine=False
	)
	self.reset_parameters()

	if tensor_parallel:
	self.decay_proj = nn.Linear(self.num_heads, self.num_heads, bias=False)
	else:
	self.decay_proj = None

	def reset_parameters(self):
	nn.init.xavier_uniform_(self.q_proj.weight, gain=2**-2.5)
	nn.init.xavier_uniform_(self.k_proj.weight, gain=2**-2.5)
	nn.init.xavier_uniform_(self.v_proj.weight, gain=2**-2.5)
	nn.init.xavier_uniform_(self.g_proj.weight, gain=2**-2.5)
	nn.init.xavier_uniform_(self.out_proj.weight)

	def parallel_retention(self, q, k, v, decay_mask):
	"""
	q, # bsz * num_head * len * qk_dim
	k, # bsz * num_head * len * qk_dim
	v, # bsz * num_head * len * v_dim
	decay_mask, # (1 or bsz) * num_head * len * len
	"""
	decay_mask, intra_decay, scale = decay_mask
	# just return retention_rel_pos projected
	# TODO: for shardformer
	if self.decay_proj is not None:
	decay_mask = self.decay_proj(decay_mask.transpose(-1, -3)).transpose(-3, -1)

	# [b, h, t, t]
	retention = q @ k.transpose(-1, -2) # (scaled dot-product)
	retention = retention * decay_mask

	# invariant after normalization
	retention = retention / retention.detach().sum(
	dim=-1, keepdim=True
	).abs().clamp(min=1)

	output = retention @ v # [b, h, t, v_dim / h]
	output = output.transpose(1, 2) # [b, t, h, v_dim / h]

	if self.training: # skip cache
	return output, None, retention

	if self.decay_proj is not None:
	intra_decay = self.decay_proj(intra_decay.transpose(-1, -2)).transpose(
	-2, -1
	)

	# kv cache: [b, h, t, v_dim, qk_dim]
	current_kv = k.unsqueeze(-2) * v.unsqueeze(-1)
	intra_decay = intra_decay[:, :, :, None, None] # [b, h, t, 1, 1]
	current_kv = (current_kv * intra_decay).sum(2) # [b, h, v_dim, qk_dim]

	cache = {"prev_key_value": current_kv, "scale": scale}
	return output, cache, retention

	def recurrent_retention(
	self, q, k, v, decay, past_key_value=None, retention_mask=None
	):
	"""
	q, k, v, # bsz * num_head * 1 * qkv_dim
	past_key_value:
	- "prev_key_value" # bsz * num_head * v_dim * qk_dim
	- "scale" # (1 or bsz) * num_head * 1 * 1
	decay # (1 or bsz) * num_head * 1 * 1
	retention_mask # bsz * 1
	"""
	if retention_mask is not None:
	retention_mask = retention_mask.float().view(-1, 1, 1, 1).to(decay)
	else:
	retention_mask = torch.ones(k.size(0), 1, 1, 1).to(decay)
	# (b, h, v_dim, qk_dim)
	current_kv = k * v.transpose(-1, -2) * retention_mask

	if past_key_value is not None and "prev_key_value" in past_key_value:
	prev_kv = past_key_value["prev_key_value"]
	prev_scale = past_key_value["scale"]
	scale = torch.where(retention_mask == 0, prev_scale, prev_scale * decay + 1)
	# connect prev_kv and current_kv
	# how much to decay prev_kv
	decay_amount = prev_scale.sqrt() * decay / scale.sqrt()
	decay_amount = torch.where(retention_mask == 0, 1, decay_amount)
	prev_kv = prev_kv * decay_amount # decay prev_kv
	current_kv = current_kv / scale.sqrt() # scale current_kv
	current_kv = torch.nan_to_num(
	current_kv, nan=0.0
	) # remove nan, scale might be 0

	current_kv = prev_kv + current_kv
	else:
	scale = torch.ones_like(decay)
	# when retention_mask is 0 at the beginning, setting scale to 1 will
	# make the first retention to use the padding incorrectly. Hence,
	# setting it to 0 here. This is a little ugly, so we might want to
	# change this later. TODO: improve
	scale = torch.where(retention_mask == 0, torch.zeros_like(decay), scale)

	output = torch.sum(q * current_kv, dim=3).unsqueeze(1) # (b, 1, h, d_v)

	cache = {"prev_key_value": current_kv, "scale": scale}
	return output, cache

	def chunkwise_retention(self, q, k, v, decay_mask):
	"""
	q, k, v, # bsz * num_head * seqlen * qkv_dim
	past_key_value:
	- "prev_key_value" # bsz * num_head * v_dim * qk_dim
	- "scale" # (1 or bsz) * num_head * 1 * 1
	decay_mask, # 1 * num_head * chunk_size * chunk_size
	cross_decay, # 1 * num_head * 1 * 1
	inner_decay, # 1 * num_head * chunk_size * 1
	"""
	# TODO: not working properly
	(
	decay_mask,
	cross_decay,
	query_inner_decay,
	value_inner_decay,
	decay_scale,
	) = decay_mask
	bsz, _, tgt_len, _ = v.size()
	chunk_len = decay_mask.size(-1)
	assert tgt_len % chunk_len == 0
	num_chunks = tgt_len // chunk_len

	# [b, n_c, h, t_c, qkv_dim]
	q = q.view(bsz, self.num_heads, num_chunks, chunk_len, self.key_dim).transpose(
	1, 2
	)
	k = k.view(bsz, self.num_heads, num_chunks, chunk_len, self.key_dim).transpose(
	1, 2
	)
	v = v.view(bsz, self.num_heads, num_chunks, chunk_len, self.head_dim).transpose(
	1, 2
	)

	k_t = k.transpose(-1, -2)

	qk_mat = q @ k_t # [b, n_c, h, t_c, t_c]
	qk_mat = qk_mat * decay_mask.unsqueeze(1)
	inner_scale = qk_mat.detach().abs().sum(dim=-1, keepdim=True).clamp(min=1)
	qk_mat = qk_mat / inner_scale
	# [b, n_c, h, t_c, v_dim]
	inner_output = torch.matmul(qk_mat, v)

	# reduce kv in one chunk
	# [b, n_c, h, qk_dim, v_dim]
	kv = k_t @ (v * value_inner_decay)
	# kv = kv.view(bsz, num_chunks, self.num_heads, self.key_dim, self.head_dim)

	kv_recurrent = []
	cross_scale = []
	kv_state = torch.zeros(bsz, self.num_heads, self.key_dim, self.head_dim).to(v)
	kv_scale = torch.ones(bsz, self.num_heads, 1, 1).to(v)

	# accumulate kv by loop
	for i in range(num_chunks):
	kv_recurrent.append(kv_state / kv_scale)
	cross_scale.append(kv_scale)
	kv_state = kv_state * cross_decay + kv[:, i]
	kv_scale = (
	kv_state.detach()
	.abs()
	.sum(dim=-2, keepdim=True)
	.max(dim=-1, keepdim=True)
	.values.clamp(min=1)
	)

	kv_recurrent = torch.stack(kv_recurrent, dim=1)
	cross_scale = torch.stack(cross_scale, dim=1)

	all_scale = torch.maximum(inner_scale, cross_scale)
	align_inner_scale = all_scale / inner_scale
	align_cross_scale = all_scale / cross_scale

	cross_output = (q * query_inner_decay.unsqueeze(1)) @ kv_recurrent
	output = inner_output / align_inner_scale + cross_output / align_cross_scale
	output = output.transpose(2, 3) # [b, n_c, t_c, h, v_dim]

	cache = {"prev_key_value": kv_state.transpose(-2, -1), "scale": decay_scale}
	return output, cache

	def forward(
	self,
	hidden_states: torch.Tensor,
	rel_pos: Tuple[Tuple[torch.Tensor]],
	retention_mask: Optional[torch.Tensor] = None,
	past_key_value: Optional[Tuple[torch.Tensor]] = None,
	forward_impl: str = "parallel",
	output_retentions: Optional[bool] = False,
	) -> Tuple[torch.FloatTensor, torch.FloatTensor, Optional[torch.FloatTensor]]:
	B, T, H = hidden_states.size()
	(sin, cos), decay_mask = rel_pos
	# projections
	q = self.q_proj(hidden_states)
	k = self.k_proj(hidden_states)
	v = self.v_proj(hidden_states)
	g = self.g_proj(hidden_states)
	# multi-head
	q, k, v = split_heads((q, k, v), B, T, self.num_heads)
	k *= self.scaling # for scaled dot product
	# rotate
	# NOTE: theta_shift has bug with mps device.
	qr = theta_shift(q, sin, cos)
	kr = theta_shift(k, sin, cos)

	# retention
	if forward_impl == "parallel":
	retention_out, curr_kv, retention_weights = self.parallel_retention(
	qr, kr, v, decay_mask
	)
	elif forward_impl == "recurrent":
	retention_out, curr_kv = self.recurrent_retention(
	qr,
	kr,
	v,
	decay_mask,
	past_key_value=past_key_value,
	retention_mask=retention_mask,
	)
	elif forward_impl == "chunkwise":
	retention_out, curr_kv = self.chunkwise_retention(qr, kr, v, decay_mask)
	else:
	raise ValueError(f"forward_impl {forward_impl} not supported.")

	# concaat heads
	normed = self.group_norm(retention_out).reshape(B, T, self.value_dim)
	# out gate & proj
	out = self.gate_fn(g) * normed
	out = self.out_proj(out)

	outputs = (out, curr_kv)
	if output_retentions:
	outputs += (retention_weights,) if forward_impl == "parallel" else (None,)
	return outputs


	class FeedForwardNetwork(nn.Module):
	def __init__(
	self,
	embed_dim,
	ffn_dim,
	activation_fn,
	dropout,
	activation_dropout,
	layernorm_eps,
	subln=False,
	use_rms_norm=False,
	):
	super().__init__()
	self.embed_dim = embed_dim
	self.activation_fn = get_activation_fn(activation=str(activation_fn))
	self.activation_dropout_module = torch.nn.Dropout(activation_dropout)
	self.dropout_module = torch.nn.Dropout(dropout)
	self.fc1 = nn.Linear(self.embed_dim, ffn_dim)
	self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
	if subln:
	if use_rms_norm:
	self.ffn_layernorm = RMSNorm(self.embed_dim, eps=layernorm_eps)
	else:
	self.ffn_layernorm = LayerNorm(self.embed_dim, eps=layernorm_eps)
	else:
	self.ffn_layernorm = None

	def reset_parameters(self):
	self.fc1.reset_parameters()
	self.fc2.reset_parameters()
	if self.ffn_layernorm is not None:
	self.ffn_layernorm.reset_parameters()

	def forward(self, x):
	x_shape = x.shape
	x = x.reshape(-1, x.size(-1))
	x = self.fc1(x)
	x = self.activation_fn(x.float()).type_as(x)
	x = self.activation_dropout_module(x)
	if self.ffn_layernorm is not None:
	x = self.ffn_layernorm(x)
	x = self.fc2(x)
	x = x.view(x_shape)
	x = self.dropout_module(x)
	return x


	class GLU(nn.Module):
	def __init__(
	self,
	embed_dim,
	ffn_dim,
	activation_fn,
	dropout,
	activation_dropout,
	):
	super().__init__()
	self.embed_dim = embed_dim
	self.activation_fn = get_activation_fn(activation=str(activation_fn))
	self.activation_dropout_module = torch.nn.Dropout(activation_dropout)
	self.dropout_module = torch.nn.Dropout(dropout)
	self.fc1 = nn.Linear(self.embed_dim, ffn_dim, bias=False)
	self.fc2 = nn.Linear(ffn_dim, self.embed_dim, bias=False)
	self.gate = nn.Linear(self.embed_dim, ffn_dim, bias=False)

	def reset_parameters(self):
	self.fc1.reset_parameters()
	self.fc2.reset_parameters()
	self.gate.reset_parameters()

	def forward(self, x):
	x_shape = x.shape
	x = x.reshape(-1, x.size(-1))
	g = self.gate(x)
	x = self.fc1(x)
	x = self.activation_fn(x.float()).type_as(x) * g
	x = self.activation_dropout_module(x)
	x = self.fc2(x)
	x = x.view(x_shape)
	x = self.dropout_module(x)
	return x


	class DropPath(nn.Module):
	"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""

	def __init__(self, drop_prob=None):
	super(DropPath, self).__init__()
	self.drop_prob = drop_prob

	def forward(self, x):
	return drop_path(x, self.drop_prob, self.training)

	def extra_repr(self):
	return "p={}".format(self.drop_prob)


	class RetNetDecoderLayer(nn.Module):
	def __init__(self, config: RetNetConfig, depth: int, tensor_parallel: bool = False):
	super().__init__()
	self.config = config
	self.embed_dim = config.decoder_embed_dim
	self.dropout_module = torch.nn.Dropout(config.dropout)

	if config.drop_path_rate > 0:
	drop_path_prob = np.linspace(
	0, config.drop_path_rate, config.decoder_layers
	)[depth]
	self.drop_path = DropPath(drop_path_prob)
	else:
	self.drop_path = None

	self.retention = MultiScaleRetention(
	config, use_bias=False, tensor_parallel=tensor_parallel
	)

	self.normalize_before = config.decoder_normalize_before

	self.retention_layer_norm = RMSNorm(self.embed_dim, eps=config.layernorm_eps)

	self.ffn_dim = config.decoder_ffn_embed_dim

	self.ffn = self.build_ffn()

	self.final_layer_norm = RMSNorm(self.embed_dim, eps=config.layernorm_eps)

	if config.deepnorm:
	self.alpha = math.pow(2.0 * config.decoder_layers, 0.25)
	else:
	self.alpha = 1.0

	def build_ffn(self):
	if self.config.use_glu:
	return GLU(
	self.embed_dim,
	self.ffn_dim,
	self.config.activation_fn,
	self.config.dropout,
	self.config.activation_dropout,
	)
	else:
	return FeedForwardNetwork(
	self.embed_dim,
	self.ffn_dim,
	self.config.activation_fn,
	self.config.dropout,
	self.config.activation_dropout,
	self.config.layernorm_eps,
	self.config.subln,
	self.config.use_ffn_rms_norm,
	)

	def residual_connection(self, x, residual):
	return residual * self.alpha + x

	def forward(
	self,
	hidden_states: torch.Tensor,
	retention_rel_pos: Tuple[Tuple[torch.Tensor]],
	retention_mask: Optional[torch.Tensor] = None,
	forward_impl: str = "parallel",
	past_key_value: Optional[Tuple[torch.Tensor]] = None,
	output_retentions: Optional[bool] = False,
	) -> Tuple[torch.FloatTensor, torch.FloatTensor, Optional[torch.FloatTensor]]:
	residual = hidden_states
	if self.normalize_before:
	hidden_states = self.retention_layer_norm(hidden_states)

	msr_outs = self.retention(
	hidden_states,
	retention_rel_pos,
	retention_mask=retention_mask,
	past_key_value=past_key_value,
	forward_impl=forward_impl,
	output_retentions=output_retentions,
	)
	hidden_states = msr_outs[0]
	curr_kv = msr_outs[1]

	hidden_states = self.dropout_module(hidden_states)

	if self.drop_path is not None:
	hidden_states = self.drop_path(hidden_states)

	hidden_states = self.residual_connection(hidden_states, residual)
	if not self.normalize_before:
	hidden_states = self.retention_layer_norm(hidden_states)

	residual = hidden_states
	if self.normalize_before:
	hidden_states = self.final_layer_norm(hidden_states)

	hidden_states = self.ffn(hidden_states)

	if self.drop_path is not None:
	hidden_states = self.drop_path(hidden_states)

	hidden_states = self.residual_connection(hidden_states, residual)
	if not self.normalize_before:
	hidden_states = self.final_layer_norm(hidden_states)

	outputs = (hidden_states, curr_kv)

	if output_retentions:
	outputs += (msr_outs[2],)
	return outputs


	class RetNetPreTrainedModel(PreTrainedModel):
	# copied from LlamaPretrainedModel
	config_class = RetNetConfig
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["RetNetDecoderLayer"]
	_keys_to_ignore_on_load_unexpected = [r"decoder\.version"]

	def _init_weights(self, module):
	"""
	Following original retnet, weights are already initialized in their own
	ways within their own init.
	"""
	pass
	# below is copied from LlamaPretrainedModel
	# std = self.config.initializer_range
	# if isinstance(module, nn.Linear):
	# module.weight.data.normal_(mean=0.0, std=std)
	# if module.bias is not None:
	# module.bias.data.zero_()
	# elif isinstance(module, nn.Embedding):
	# module.weight.data.normal_(mean=0.0, std=std)
	# if module.padding_idx is not None:
	# module.weight.data[module.padding_idx].zero_()


	@dataclass
	class RetNetOutputWithPast(ModelOutput):
	"""
	class for RetNet model's outputs that may also contain a past key/values (to speed up sequential decoding).

	config:
	last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, decoder_embed_dim)`):
	Sequence of hidden-states at the output of the last layer of the model.

	If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
	decoder_embed_dim)` is output.
	past_key_values (`List(Dict(str, torch.FloatTensor))`, optional, returned when `use_cache=True` is passed or when `config.use_cache=True`):
	- "prev_key_value": shape=(bsz * num_head * v_dim * qk_dim)
	- "scale": shape=((1 or bsz) * num_head * 1 * 1)

	Contains pre-computed hidden-states (key and values in the multi-scale retention blocks)
	that can be used (see `past_key_values` input) to speed up sequential decoding.
	hidden_states (`tuple(torch.FloatTensor)`, optional, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
	Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
	one for the output of each layer) of shape `(batch_size, sequence_length, decoder_embed_dim)`.

	Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
	retentions (`tuple(torch.FloatTensor)`, optional, returned when `output_retentions=True` is passed or when `config.output_retentions=True`):
	Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
	sequence_length)`.

	Retentions weights, used for visualization.

	attentions (`tuple(torch.FloatTensor)`, optional, for backward compatibility. Same as retentions.
	"""

	last_hidden_state: torch.FloatTensor = None
	past_key_values: Optional[List[Dict[str, torch.FloatTensor]]] = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	retentions: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None


	class RetNetModel(RetNetPreTrainedModel):
	def __init__(
	self,
	config: RetNetConfig,
	embed_tokens: nn.Embedding = None,
	tensor_parallel: bool = False,
	):
	super().__init__(config)
	self.config = config

	self.dropout_module = torch.nn.Dropout(config.dropout)

	self.embed_dim = config.decoder_embed_dim
	self.embed_scale = (
	1.0 if config.no_scale_embedding else math.sqrt(self.embed_dim)
	)

	if embed_tokens is None:
	embed_tokens = nn.Embedding(
	config.vocab_size, config.decoder_embed_dim, config.pad_token_id
	)
	self.embed_tokens = embed_tokens

	if config.layernorm_embedding:
	self.layernorm_embedding = RMSNorm(self.embed_dim, eps=config.layernorm_eps)
	else:
	self.layernorm_embedding = None

	self.layers = nn.ModuleList([])

	for i in range(config.decoder_layers):
	self.layers.append(
	RetNetDecoderLayer(config, depth=i, tensor_parallel=tensor_parallel)
	)

	self.decoder_layers = len(self.layers)

	if config.decoder_normalize_before:
	self.layer_norm = RMSNorm(self.embed_dim, eps=config.layernorm_eps)
	else:
	self.layer_norm = None

	self.retnet_rel_pos = RetNetRelPos(config)
	self.recurrent_chunk_size = config.recurrent_chunk_size

	if config.deepnorm:
	init_scale = math.pow(8.0 * config.decoder_layers, 0.25)
	for name, p in self.named_parameters():
	if (
	"fc1" in name
	or "fc2" in name
	or "out_proj" in name
	or "v_proj" in name
	):
	p.data.div_(init_scale)

	if config.subln and not config.use_glu:
	init_scale = math.sqrt(math.log(config.decoder_layers * 2))
	for name, p in self.named_parameters():
	if (
	"fc1" in name
	or "fc2" in name
	or "out_proj" in name
	or "v_proj" in name
	):
	p.data.mul_(init_scale)

	self.gradient_checkpointing = False
	self.post_init()

	def get_input_embeddings(self):
	return self.embed_tokens

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

	def forward_embedding(
	self,
	input_ids,
	forward_impl,
	inputs_embeds=None,
	past_key_values=None,
	):
	# if past_key_values is not None:
	if forward_impl == "recurrent":
	input_ids = input_ids[:, -1:]

	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)

	embed = self.embed_scale * inputs_embeds

	if self.layernorm_embedding is not None:
	embed = self.layernorm_embedding(embed)

	embed = self.dropout_module(embed)

	return embed

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	retention_mask: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[Dict[str, torch.FloatTensor]]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	output_retentions: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	use_cache: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	forward_impl: Optional[str] = "parallel",
	recurrent_chunk_size: Optional[int] = None,
	retention_rel_pos: Optional[Tuple[torch.Tensor]] = None,
	) -> Union[Tuple, RetNetOutputWithPast]:
	if output_retentions is None and output_attentions is not None:
	output_retentions = output_attentions
	output_retentions = (
	output_retentions
	if output_retentions is not None
	else self.config.output_retentions
	)
	output_hidden_states = (
	output_hidden_states
	if output_hidden_states is not None
	else self.config.output_hidden_states
	)
	use_cache = use_cache if use_cache is not None else self.config.use_cache

	return_dict = (
	return_dict if return_dict is not None else self.config.use_return_dict
	)

	# retrieve input_ids and inputs_embeds
	if input_ids is not None and inputs_embeds is not None:
	raise ValueError(
	"You cannot specify both input_ids and inputs_embeds at the same time"
	)
	elif input_ids is not None:
	batch_size, seq_length = input_ids.shape
	elif inputs_embeds is not None:
	batch_size, seq_length, _ = inputs_embeds.shape
	else:
	raise ValueError("You have to specify either input_ids or inputs_embeds")

	# embed tokens
	if inputs_embeds is None:
	inputs_embeds = self.forward_embedding(
	input_ids, forward_impl, inputs_embeds, past_key_values
	)

	if retention_mask is None and attention_mask is not None:
	retention_mask = attention_mask
	if retention_mask is not None and forward_impl == "recurrent":
	retention_mask = retention_mask[:, -1:]

	hidden_states = inputs_embeds

	# handling chunking here
	if recurrent_chunk_size is None:
	recurrent_chunk_size = self.recurrent_chunk_size
	need_pad_for_chunkwise = (
	forward_impl == "chunkwise" and seq_length % recurrent_chunk_size != 0
	)
	if need_pad_for_chunkwise:
	padding_len = recurrent_chunk_size - seq_length % recurrent_chunk_size
	slen = seq_length + padding_len
	hidden_states = F.pad(hidden_states, (0, 0, 0, padding_len))
	else:
	slen = seq_length
	# relative position
	if retention_rel_pos is None:
	retention_rel_pos = self.retnet_rel_pos(
	slen,
	forward_impl=forward_impl,
	recurrent_chunk_size=recurrent_chunk_size,
	retention_mask=retention_mask,
	get_decay_scale=not self.training,
	)

	# start running through the decoder layers
	all_hidden_states = () if output_hidden_states else None
	all_retentions = () if output_retentions else None
	# layers * [bsz, num_head, qk_dim, decoder_embed_dim]
	next_decoder_cache = () if use_cache else None

	for idx, layer in enumerate(self.layers):
	if output_hidden_states:
	all_hidden_states += (hidden_states,)
	past_key_value = (
	past_key_values[idx] if past_key_values is not None else None
	)

	if self.gradient_checkpointing and self.training:

	def create_custom_forward(module):
	def custom_forward(*inputs):
	return module(*inputs, output_retentions)

	return custom_forward

	layer_outputs = torch.utils.checkpoint.checkpoint(
	create_custom_forward(layer),
	hidden_states,
	retention_rel_pos,
	retention_mask,
	forward_impl,
	past_key_value,
	)
	else:
	layer_outputs = layer(
	hidden_states,
	retention_rel_pos,
	retention_mask=retention_mask,
	forward_impl=forward_impl,
	past_key_value=past_key_value,
	output_retentions=output_retentions,
	)

	hidden_states = layer_outputs[0]

	if use_cache:
	next_decoder_cache += (layer_outputs[1],)

	if output_retentions:
	all_retentions += (layer_outputs[2],)

	next_cache = next_decoder_cache if use_cache else None

	if need_pad_for_chunkwise:
	hidden_states = hidden_states[:, :seq_length, :]

	if self.layer_norm is not None:
	hidden_states = self.layer_norm(hidden_states)

	# add hidden states from the last decoder layer
	if output_hidden_states:
	all_hidden_states += (hidden_states,)

	if not return_dict:
	return tuple(
	v
	for v in [hidden_states, next_cache, all_hidden_states, all_retentions]
	if v is not None
	)
	return RetNetOutputWithPast(
	last_hidden_state=hidden_states,
	past_key_values=next_cache,
	hidden_states=all_hidden_states,
	retentions=all_retentions,
	attentions=all_retentions,
	)


	@dataclass
	class RetNetCausalLMOutputWithPast(ModelOutput):
	"""
	class for RetNet causal language model (or autoregressive) outputs.

	config:
	loss (`torch.FloatTensor` of shape `(1,)`, optional, returned when `labels` is provided):
	Language modeling loss (for next-token prediction).
	logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
	Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
	past_key_values (`List(Dict(str, torch.FloatTensor))`, optional, returned when `use_cache=True` is passed or when `config.use_cache=True`):
	- "prev_key_value": shape=(bsz * num_head * v_dim * qk_dim)
	- "scale": shape=((1 or bsz) * num_head * 1 * 1)

	Contains pre-computed hidden-states (key and values in the multi-scale retention blocks)
	that can be used (see `past_key_values` input) to speed up sequential decoding.
	hidden_states (`tuple(torch.FloatTensor)`, optional, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
	Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
	one for the output of each layer) of shape `(batch_size, sequence_length, decoder_embed_dim)`.

	Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
	retentions (`tuple(torch.FloatTensor)`, optional, returned when `output_retentions=True` is passed or when `config.output_retentions=True`):
	Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
	sequence_length)`.

	Retentions weights, used for visualization.

	attentions (`tuple(torch.FloatTensor)`, optional, for backward compatibility. Same as retentions.
	"""

	loss: Optional[torch.FloatTensor] = None
	logits: torch.FloatTensor = None
	past_key_values: Optional[List[Dict[str, torch.FloatTensor]]] = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	retentions: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None


	class RetNetForCausalLM(RetNetPreTrainedModel):
	def __init__(
	self,
	config: RetNetConfig,
	embed_tokens: nn.Embedding = None,
	tensor_parallel: bool = False,
	) -> None:
	super().__init__(config)
	self.model = RetNetModel(
	config, embed_tokens=embed_tokens, tensor_parallel=tensor_parallel
	)
	self.lm_head = nn.Linear(
	config.decoder_embed_dim, config.vocab_size, bias=False
	)
	# init here
	torch.nn.init.normal_(
	self.lm_head.weight, mean=0, std=config.decoder_embed_dim**-0.5
	)

	self.post_init()

	def get_input_embeddings(self):
	return self.model.embed_tokens

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

	def get_output_embeddings(self):
	return self.lm_head

	def set_output_embeddings(self, new_embeddings):
	self.lm_head = new_embeddings

	def set_decoder(self, decoder):
	self.model = decoder

	def get_decoder(self):
	return self.model

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	retention_mask: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	past_key_values: Optional[List[torch.FloatTensor]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_retentions: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	forward_impl: Optional[str] = None,
	recurrent_chunk_size: Optional[int] = None,
	retention_rel_pos: Optional[Tuple[torch.Tensor]] = None,
	) -> Union[Tuple, RetNetCausalLMOutputWithPast]:
	if output_retentions is None and output_attentions is not None:
	output_retentions = output_attentions
	output_retentions = (
	output_retentions
	if output_retentions is not None
	else self.config.output_retentions
	)
	output_hidden_states = (
	output_hidden_states
	if output_hidden_states is not None
	else self.config.output_hidden_states
	)
	return_dict = (
	return_dict if return_dict is not None else self.config.use_return_dict
	)
	forward_impl = (
	forward_impl if forward_impl is not None else self.config.forward_impl
	)
	recurrent_chunk_size = (
	recurrent_chunk_size
	if recurrent_chunk_size is not None
	else self.config.recurrent_chunk_size
	)

	if retention_mask is None and attention_mask is not None:
	retention_mask = attention_mask

	outputs = self.model(
	input_ids,
	retention_mask=retention_mask,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	output_retentions=output_retentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	forward_impl=forward_impl,
	use_cache=use_cache,
	recurrent_chunk_size=recurrent_chunk_size,
	retention_rel_pos=retention_rel_pos,
	)

	hidden_states = outputs[0]
	logits = self.lm_head(hidden_states)

	loss = None
	if labels is not None:
	# Shift so that tokens < n predict n
	shift_logits = logits[..., :-1, :].contiguous()
	shift_labels = labels[..., 1:].contiguous()
	# Flatten the tokens
	loss_fct = nn.CrossEntropyLoss()
	shift_logits = shift_logits.view(-1, self.config.vocab_size)
	shift_labels = shift_labels.view(-1)
	# Enable model parallelism
	shift_labels = shift_labels.to(shift_logits.device)
	loss = loss_fct(shift_logits, shift_labels)

	if self.config.z_loss_coeff > 0:
	# z_loss from PaLM paper
	# z_loss = 1e-4 * log(log(z)), where z = sum(exp(logits))
	z_loss = torch.logsumexp(shift_logits, dim=-1).log().mean()
	loss += self.config.z_loss_coeff * z_loss

	if not return_dict:
	output = (logits,) + outputs[1:]
	return (loss,) + output if loss is not None else output

	return RetNetCausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	retentions=outputs.retentions,
	attentions=outputs.retentions,
	)

	def _crop_past_key_values(model, past_key_values, maximum_length):
	"""Since retnet's kv do not have length, no need to crop. Just return"""
	return past_key_values

	def prepare_inputs_for_generation(
	self,
	input_ids,
	past_key_values=None,
	attention_mask=None,
	inputs_embeds=None,
	**kwargs,
	):
	# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
	if inputs_embeds is not None and past_key_values is None:
	model_inputs = {"inputs_embeds": inputs_embeds}
	else:
	model_inputs = {"input_ids": input_ids}

	forward_impl = kwargs.get("forward_impl", "parallel")
	if past_key_values is not None:
	forward_impl = "recurrent"

	model_inputs.update(
	{
	"past_key_values": past_key_values,
	"use_cache": kwargs.get("use_cache"),
	"attention_mask": attention_mask,
	"forward_impl": forward_impl,
	}
	)
	return model_inputs

	@staticmethod
	def _reorder_cache(past_key_values, beam_idx):
	reordered_past = ()
	for layer_past in past_key_values: # dict
	layer_past_kv = layer_past["prev_key_value"] # [b, h, v_dim / h, qk_dim]
	layer_past_scale = layer_past["scale"] # [b, h, 1, 1]
	if layer_past_scale.size(0) > 1:
	# this means that retention_mask is not None, so the scale for
	# each batch is different. We need to select the correct scale then.
	# NOTE: during huggingface generate, it will generate attention_mask
	# if it is None, so this linke will always be true. Still, having
	# this line here for safety.
	layer_past_scale = layer_past_scale.index_select(0, beam_idx)
	reordered_past += (
	{
	"prev_key_value": layer_past_kv.index_select(0, beam_idx),
	"scale": layer_past_scale,
	},
	)
	return reordered_past

	def sample_token(self, logit, do_sample=False, top_k=1, top_p=1.0, temperature=1.0):
	if not do_sample:
	return torch.argmax(logit, dim=-1, keepdim=True)
	filtered = top_k_top_p_filtering(logit / temperature, top_k=top_k, top_p=top_p)
	return torch.multinomial(torch.softmax(filtered, dim=-1), num_samples=1)

	@torch.inference_mode()
	def custom_generate(
	self,
	input_ids: torch.LongTensor = None,
	retention_mask: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	parallel_compute_prompt=True,
	max_new_tokens=20,
	bos_token_id=0,
	eos_token_id=0,
	do_sample=False,
	top_k=0,
	top_p=1.0,
	temperature=1.0,
	early_stopping=True,
	):
	if retention_mask is None and attention_mask is not None:
	retention_mask = attention_mask

	if input_ids is not None:
	if input_ids.shape[1] == 1:
	past_key_values = None
	elif parallel_compute_prompt:
	ret_mask = (
	retention_mask[:, :-1] if retention_mask is not None else None
	)
	outputs = self(
	input_ids[:, :-1],
	retention_mask=ret_mask,
	forward_impl="parallel",
	return_dict=True,
	use_cache=True,
	)
	past_key_values = outputs.past_key_values
	else:
	past_key_values = None
	for p_i in range(input_ids.shape[1] - 1):
	ret_mask = (
	retention_mask[:, : p_i + 1]
	if retention_mask is not None
	else None
	)
	outputs = self(
	input_ids[:, : p_i + 1],
	retention_mask=ret_mask,
	forward_impl="recurrent",
	past_key_values=past_key_values,
	return_dict=True,
	use_cache=True,
	)
	past_key_values = outputs.past_key_values

	generated = input_ids
	else:
	generated = torch.tensor([[bos_token_id]]).to(self.lm_head.weight.device)
	past_key_values = None

	for i in range(max_new_tokens):
	outputs = self(
	generated,
	retention_mask=retention_mask,
	forward_impl="recurrent",
	past_key_values=past_key_values,
	use_cache=True,
	return_dict=True,
	)
	logit = outputs.logits[:, -1, :] # [batch_size, vocab_size]
	past_key_values = outputs.past_key_values
	token = self.sample_token(
	logit,
	do_sample=do_sample,
	top_k=top_k,
	top_p=top_p,
	temperature=temperature,
	)
	generated = torch.cat([generated, token], dim=-1)
	if retention_mask is not None:
	retention_mask = torch.cat(
	[retention_mask, torch.ones_like(token)], dim=-1
	)
	if early_stopping and (token == eos_token_id).all():
	break
	return generated