Hugging Face's logo

Angshul
/
TransWave

Text Generation
Transformers
English
transwave_meyer
causal-lm
byte-level
raw-byte
attention-free
meyer
wavelet
custom_code
Model card Files
xet
 Community
TransWave / transwave_core.py
Angshul's picture
Angshul
Release TransWave 96.13M C4 checkpoint
ea51a3d verified
Raw History Blame Contribute Delete
47.9 kB
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as torch_checkpoint
ARCH = "causal_transwave_blocklifting_meyer_byte_lm_g4_v2"
BYTE_VOCAB = 256
BYTE_DIM = 96
BYTES_PER_BLOCK = 8
N_SCALES = 8
MODEL_DIM = 768
MODEL_CONTEXT = 4096
RAW_BYTES_PER_SEQUENCE = MODEL_CONTEXT * BYTES_PER_BLOCK
N_LAYERS = 12
N_GROUPS = 12
CHANNELS_PER_GROUP = BYTE_DIM // N_GROUPS
N_MEYER_BANDS = 12
N_BAND_CONTROLS = 8
MEYER_KERNEL = 512
LOCAL_KERNEL = 7
FFN_DIM = 2816
BYTE_DECODER_DIM = 256
HEAD_CHUNK = 32768
assert BYTE_DIM * BYTES_PER_BLOCK == MODEL_DIM
assert RAW_BYTES_PER_SEQUENCE == 32768
assert BYTE_DIM % N_GROUPS == 0
class CTWLocalBlockDelta(
nn.Module
):
def __init__(
self,
dim=BYTE_DIM,
groups=12,
expansion=2,
):
super().__init__()
if (
dim
%
groups
!=
0
):
raise ValueError(
"dim must be divisible by groups"
)
hidden = (
dim
*
expansion
)
if (
hidden
%
groups
!=
0
):
raise ValueError(
"expanded dimension must be divisible by groups"
)
self.depthwise = nn.Conv1d(
dim,
dim,
kernel_size=3,
padding=1,
groups=dim,
bias=True,
)
self.mix_in = nn.Conv1d(
dim,
hidden,
kernel_size=1,
groups=groups,
bias=True,
)
self.mix_out = nn.Conv1d(
hidden,
dim,
kernel_size=1,
groups=groups,
bias=True,
)
# Haar-like initialization:
# residual branch starts exactly at zero.
nn.init.zeros_(
self.mix_out.weight
)
nn.init.zeros_(
self.mix_out.bias
)
def forward(
self,
x,
):
y = x.transpose(
1,
2,
)
y = self.depthwise(
y
)
y = F.silu(
y
)
y = self.mix_in(
y
)
y = F.silu(
y
)
y = self.mix_out(
y
)
return y.transpose(
1,
2,
)
class CTWLiftingStage(
nn.Module
):
def __init__(
self,
dim=BYTE_DIM,
groups=12,
expansion=2,
):
super().__init__()
self.predict_delta = (
CTWLocalBlockDelta(
dim=dim,
groups=groups,
expansion=expansion,
)
)
self.update_delta = (
CTWLocalBlockDelta(
dim=dim,
groups=groups,
expansion=expansion,
)
)
def predict(
self,
even,
):
return (
even
+
self.predict_delta(
even
)
)
def update(
self,
detail,
):
return (
0.5
*
detail
+
self.update_delta(
detail
)
)
def forward(
self,
x,
):
if (
x.ndim != 3
or
x.shape[1] % 2 != 0
):
raise ValueError(
"Expected [N,T,C] with even T."
)
even = x[
:,
0::2,
:
]
odd = x[
:,
1::2,
:
]
detail = (
odd
-
self.predict(
even
)
)
coarse = (
even
+
self.update(
detail
)
)
return (
coarse,
detail,
)
class CTWCrossScaleMixer(
nn.Module
):
def __init__(
self,
scales=N_SCALES,
dim=BYTE_DIM,
groups=N_GROUPS,
):
super().__init__()
if (
dim
%
groups
!=
0
):
raise ValueError(
"dim must be divisible by groups"
)
self.scales = int(
scales
)
self.dim = int(
dim
)
self.groups = int(
groups
)
self.channels_per_group = (
self.dim
//
self.groups
)
# M_g = I + Delta_g
self.scale_delta = nn.Parameter(
torch.zeros(
self.groups,
self.scales,
self.scales,
)
)
# Content-dependent gain.
self.gate_weight = nn.Parameter(
torch.zeros(
self.groups,
self.scales,
self.channels_per_group,
)
)
self.gate_bias = nn.Parameter(
torch.zeros(
self.groups,
self.scales,
)
)
def forward(
self,
z,
):
if (
z.ndim != 4
or
z.shape[2]
!=
self.scales
or
z.shape[3]
!=
self.dim
):
raise ValueError(
f"Expected "
f"[B,L,{self.scales},{self.dim}]"
)
b, l, s, _ = z.shape
zg = z.reshape(
b,
l,
s,
self.groups,
self.channels_per_group,
)
gate_logits = torch.einsum(
"blsgc,gsc->blsg",
zg,
self.gate_weight,
)
gate_logits = (
gate_logits
/
math.sqrt(
self.channels_per_group
)
+
self.gate_bias.transpose(
0,
1,
)
)
# Zero logits -> gain exactly 1.
gain = (
2.0
*
torch.sigmoid(
gate_logits
)
)
zg = (
zg
*
gain
.to(
zg.dtype
)
.unsqueeze(
-1
)
)
eye = torch.eye(
self.scales,
device=z.device,
dtype=z.dtype,
).unsqueeze(
0
)
mix = (
eye
+
self.scale_delta.to(
z.dtype
)
)
mixed = torch.einsum(
"blsgc,gts->bltgc",
zg,
mix,
)
return mixed.reshape(
b,
l,
s,
self.dim,
)
class CTWBlockLiftingTokenizer(
nn.Module
):
def __init__(
self,
):
super().__init__()
self.byte_embedding = nn.Embedding(
BYTE_VOCAB,
BYTE_DIM,
)
nn.init.normal_(
self.byte_embedding.weight,
mean=0.0,
std=(
BYTE_DIM
**
-0.5
),
)
self.stages = nn.ModuleList(
[
CTWLiftingStage(
dim=BYTE_DIM,
groups=N_GROUPS,
expansion=2,
)
for _ in range(
3
)
]
)
self.cross_scale = (
CTWCrossScaleMixer(
scales=N_SCALES,
dim=BYTE_DIM,
groups=N_GROUPS,
)
)
def forward(
self,
byte_ids,
):
if (
byte_ids.ndim
!=
2
):
raise ValueError(
"Expected byte IDs [B,8L]."
)
b, t = byte_ids.shape
if (
t
%
BYTES_PER_BLOCK
!=
0
):
raise ValueError(
"Byte length must be divisible by 8."
)
l = (
t
//
BYTES_PER_BLOCK
)
# -------------------------------------------------------------
# CRITICAL CAUSAL CHANGE:
#
# each completed eight-byte block becomes an independent item
# before ANY lifting convolution is applied.
# -------------------------------------------------------------
blocks = (
byte_ids
.long()
.reshape(
b,
l,
BYTES_PER_BLOCK,
)
)
x = self.byte_embedding(
blocks.reshape(
b
*
l,
BYTES_PER_BLOCK,
)
)
details = []
coarse = x
for stage in self.stages:
(
coarse,
detail,
) = stage(
coarse
)
details.append(
detail
)
# D0=4, D1=2, D2=1, S3=1.
packed = torch.cat(
[
details[0],
details[1],
details[2],
coarse,
],
dim=1,
)
if (
packed.shape[1]
!=
N_SCALES
):
raise RuntimeError(
"Lifting scale packing mismatch."
)
z = packed.reshape(
b,
l,
N_SCALES,
BYTE_DIM,
)
return self.cross_scale(
z
)
# =============================================================================
# RMS NORM / SWIGLU
# =============================================================================
class CTWRMSNorm(
nn.Module
):
def __init__(
self,
dim,
eps=1e-6,
):
super().__init__()
self.weight = nn.Parameter(
torch.ones(
dim
)
)
self.eps = eps
def forward(
self,
x,
):
xf = x.float()
y = (
xf
*
torch.rsqrt(
xf.pow(2).mean(
dim=-1,
keepdim=True,
)
+
self.eps
)
)
return (
y.to(
x.dtype
)
*
self.weight.to(
x.dtype
)
)
class CTWSwiGLU(
nn.Module
):
def __init__(
self,
dim,
hidden,
):
super().__init__()
self.gate = nn.Linear(
dim,
hidden,
)
self.value = nn.Linear(
dim,
hidden,
)
self.out = nn.Linear(
hidden,
dim,
)
for m in (
self.gate,
self.value,
self.out,
):
nn.init.xavier_uniform_(
m.weight
)
nn.init.zeros_(
m.bias
)
def forward(
self,
x,
):
return self.out(
F.silu(
self.gate(
x
)
)
*
self.value(
x
)
)
# =============================================================================
# MEYER BASIS
# =============================================================================
def ctw_meyer_nu(
x,
):
x = x.clamp(
0.0,
1.0,
)
return (
x**4
*
(
35.0
-
84.0 * x
+
70.0 * x**2
-
20.0 * x**3
)
)
def ctw_build_meyer_basis(
kernel_size,
bands,
controls,
):
nf = (
kernel_size
//
2
+
1
)
f = torch.linspace(
0.0,
1.0,
nf,
dtype=torch.float32,
)
edges = torch.linspace(
0.0,
1.0,
bands + 1,
dtype=torch.float32,
)
windows = []
for j in range(
bands
):
left = float(
edges[j]
)
right = float(
edges[j + 1]
)
delta = (
0.35
*
(
right
-
left
)
)
if j == 0:
left_gate = torch.ones_like(
f
)
else:
a = (
left
-
delta
)
b = (
left
+
delta
)
u = (
f
-
a
) / max(
b - a,
1e-8,
)
left_gate = torch.sin(
0.5
*
math.pi
*
ctw_meyer_nu(
u
)
)
left_gate = torch.where(
f <= a,
torch.zeros_like(
f
),
left_gate,
)
left_gate = torch.where(
f >= b,
torch.ones_like(
f
),
left_gate,
)
if j == bands - 1:
right_gate = torch.ones_like(
f
)
else:
a = (
right
-
delta
)
b = (
right
+
delta
)
u = (
f
-
a
) / max(
b - a,
1e-8,
)
right_gate = torch.cos(
0.5
*
math.pi
*
ctw_meyer_nu(
u
)
)
right_gate = torch.where(
f <= a,
torch.ones_like(
f
),
right_gate,
)
right_gate = torch.where(
f >= b,
torch.zeros_like(
f
),
right_gate,
)
windows.append(
left_gate
*
right_gate
)
windows = torch.stack(
windows,
dim=0,
)
windows = (
windows
/
torch.sqrt(
windows.pow(2).sum(
dim=0,
keepdim=True,
)
+
1e-12
)
)
degree = (
controls
-
1
)
basis = []
for j in range(
bands
):
left = edges[j]
right = edges[j + 1]
u = (
(
f
-
left
)
/
(
right
-
left
+
1e-12
)
).clamp(
0.0,
1.0,
)
row = []
for m in range(
controls
):
bern = (
math.comb(
degree,
m,
)
*
u.pow(
m
)
*
(
1.0
-
u
).pow(
degree
-
m
)
)
row.append(
windows[j]
*
bern
)
basis.append(
torch.stack(
row,
dim=0,
)
)
return torch.stack(
basis,
dim=0,
)
# =============================================================================
# SAME-TIME FULL 8x8 SCALE MIXING
# =============================================================================
class CTWScaleMixer(
nn.Module
):
def __init__(
self,
):
super().__init__()
eye = torch.eye(
N_SCALES
)
self.weight = nn.Parameter(
eye
.unsqueeze(
0
)
.repeat(
N_GROUPS,
1,
1,
)
)
def forward(
self,
x,
):
b, l, s, c = x.shape
xg = x.reshape(
b,
l,
s,
N_GROUPS,
CHANNELS_PER_GROUP,
)
y = torch.einsum(
"blsgc,gts->bltgc",
xg,
self.weight,
)
return y.reshape(
b,
l,
s,
c,
)
# =============================================================================
# STRICTLY CAUSAL MEYER FIR
#
# h[k] is a nonnegative lag:
#
# y[t] = SUM_{k=0}^{K-1} h[k] x[t-k]
#
# FFT only implements this zero-padded LINEAR convolution efficiently.
# =============================================================================
class CTWCausalMeyerFIR(
nn.Module
):
def __init__(
self,
):
super().__init__()
basis = ctw_build_meyer_basis(
MEYER_KERNEL,
N_MEYER_BANDS,
N_BAND_CONTROLS,
)
self.register_buffer(
"basis",
basis,
persistent=False,
)
imag_mask = torch.ones(
MEYER_KERNEL
//
2
+
1,
dtype=torch.float32,
)
imag_mask[0] = 0.0
if (
MEYER_KERNEL
%
2
==
0
):
imag_mask[-1] = 0.0
self.register_buffer(
"imag_mask",
imag_mask,
persistent=False,
)
shape = (
N_GROUPS,
N_SCALES,
N_MEYER_BANDS,
N_BAND_CONTROLS,
)
self.coeff_re = nn.Parameter(
torch.empty(
shape
)
)
self.coeff_im = nn.Parameter(
torch.empty(
shape
)
)
nn.init.normal_(
self.coeff_re,
mean=0.0,
std=1e-3,
)
nn.init.normal_(
self.coeff_im,
mean=0.0,
std=1e-3,
)
def causal_kernel(
self,
):
with torch.autocast(
device_type=self.coeff_re.device.type,
enabled=False,
):
basis = self.basis.float()
re = torch.einsum(
"gsbm,bmf->gsf",
self.coeff_re.float(),
basis,
)
im = torch.einsum(
"gsbm,bmf->gsf",
self.coeff_im.float(),
basis,
)
im = (
im
*
self.imag_mask[
None,
None,
:
]
)
spectrum = torch.complex(
re,
im,
)
h = torch.fft.irfft(
spectrum,
n=MEYER_KERNEL,
dim=-1,
)
return h
def forward(
self,
x,
):
# [B,L,S,96]
b, l, s, c = x.shape
# [B,G,S,Cg,L]
xg = (
x
.reshape(
b,
l,
s,
N_GROUPS,
CHANNELS_PER_GROUP,
)
.permute(
0,
3,
2,
4,
1,
)
.contiguous()
)
linear_length = (
l
+
MEYER_KERNEL
-
1
)
nfft = (
1
<<
(
linear_length
-
1
).bit_length()
)
with torch.autocast(
device_type=self.coeff_re.device.type,
enabled=False,
):
xf = xg.float()
h = self.causal_kernel()
X = torch.fft.rfft(
xf,
n=nfft,
dim=-1,
)
H = torch.fft.rfft(
h,
n=nfft,
dim=-1,
)
Y = torch.fft.irfft(
X
*
H[
None,
:,
:,
None,
:
],
n=nfft,
dim=-1,
)
# First L outputs of LINEAR convolution.
Y = Y[
...,
:l
]
return (
Y
.permute(
0,
4,
2,
1,
3,
)
.contiguous()
.reshape(
b,
l,
s,
c,
)
.to(
x.dtype
)
)
# =============================================================================
# STRICTLY LEFT-CAUSAL LOCAL CONVOLUTION
# =============================================================================
class CTWCausalLocalConv(
nn.Module
):
def __init__(
self,
):
super().__init__()
self.conv = nn.Conv1d(
MODEL_DIM,
MODEL_DIM,
kernel_size=LOCAL_KERNEL,
groups=MODEL_DIM,
bias=True,
)
nn.init.normal_(
self.conv.weight,
mean=0.0,
std=0.02,
)
nn.init.zeros_(
self.conv.bias
)
def forward(
self,
x,
):
xt = x.transpose(
1,
2,
)
xt = F.pad(
xt,
(
LOCAL_KERNEL
-
1,
0,
),
)
return (
self.conv(
xt
)
.transpose(
1,
2,
)
)
# =============================================================================
# ONE CAUSAL TRANSWAVE LAYER
# =============================================================================
class CTWLayer(
nn.Module
):
def __init__(
self,
):
super().__init__()
self.norm1 = CTWRMSNorm(
MODEL_DIM
)
self.input_gate = nn.Linear(
MODEL_DIM,
MODEL_DIM,
)
self.output_gate = nn.Linear(
MODEL_DIM,
MODEL_DIM,
)
nn.init.zeros_(
self.input_gate.weight
)
nn.init.constant_(
self.input_gate.bias,
2.0,
)
nn.init.zeros_(
self.output_gate.weight
)
nn.init.constant_(
self.output_gate.bias,
2.0,
)
self.scale_in = (
CTWScaleMixer()
)
self.meyer = (
CTWCausalMeyerFIR()
)
self.scale_out = (
CTWScaleMixer()
)
self.local = (
CTWCausalLocalConv()
)
self.alpha_meyer = nn.Parameter(
torch.tensor(
0.10,
dtype=torch.float32,
)
)
self.alpha_local = nn.Parameter(
torch.tensor(
0.10,
dtype=torch.float32,
)
)
self.norm2 = CTWRMSNorm(
MODEL_DIM
)
self.ffn = CTWSwiGLU(
MODEL_DIM,
FFN_DIM,
)
def forward(
self,
x,
):
b, l, s, c = x.shape
flat = x.reshape(
b,
l,
MODEL_DIM,
)
n = self.norm1(
flat
)
gate_in = torch.sigmoid(
self.input_gate(
n
)
).reshape(
b,
l,
s,
c,
)
gate_out = torch.sigmoid(
self.output_gate(
n
)
).reshape(
b,
l,
s,
c,
)
spectral = self.meyer(
self.scale_in(
x
*
gate_in
)
)
spectral = self.scale_out(
spectral
)
local = self.local(
n
).reshape(
b,
l,
s,
c,
)
x = (
x
+
self.alpha_meyer.to(
x.dtype
)
*
spectral
*
gate_out
+
self.alpha_local.to(
x.dtype
)
*
local
)
flat = x.reshape(
b,
l,
MODEL_DIM,
)
flat = (
flat
+
self.ffn(
self.norm2(
flat
)
)
)
return flat.reshape(
b,
l,
s,
c,
)
# =============================================================================
# EXACT EIGHT-BYTE AUTOREGRESSIVE DECODER
# =============================================================================
class CTWByteDecoder(
nn.Module
):
def __init__(
self,
shared_embedding,
):
super().__init__()
self.shared_embedding = (
shared_embedding
)
self.context_proj = nn.Linear(
MODEL_DIM,
BYTE_DECODER_DIM,
)
self.gru = nn.GRUCell(
BYTE_DIM,
BYTE_DECODER_DIM,
)
self.norm = CTWRMSNorm(
BYTE_DECODER_DIM
)
self.to_byte_space = nn.Linear(
BYTE_DECODER_DIM,
BYTE_DIM,
)
self.bos = nn.Parameter(
torch.zeros(
BYTE_DIM
)
)
self.byte_position = nn.Parameter(
torch.zeros(
BYTES_PER_BLOCK,
BYTE_DIM,
)
)
nn.init.normal_(
self.bos,
mean=0.0,
std=0.02,
)
nn.init.normal_(
self.byte_position,
mean=0.0,
std=0.01,
)
def nll(
self,
context,
target_blocks,
):
b, l, _ = context.shape
n = (
b
*
l
)
hidden = torch.tanh(
self.context_proj(
context
)
).reshape(
n,
BYTE_DECODER_DIM,
)
previous = (
self.bos[
None,
:
]
.expand(
n,
-1,
)
)
total_loss = None
for r in range(
BYTES_PER_BLOCK
):
decoder_input = (
previous
+
self.byte_position[
r
][
None,
:
]
)
hidden = self.gru(
decoder_input,
hidden,
)
z = self.to_byte_space(
self.norm(
hidden
)
)
target = (
target_blocks[
:,
:,
r
]
.reshape(
-1
)
.long()
)
position_loss = None
for start in range(
0,
n,
HEAD_CHUNK,
):
end = min(
start
+
HEAD_CHUNK,
n,
)
logits = F.linear(
z[
start:end
],
self.shared_embedding.weight,
)
piece = F.cross_entropy(
logits.float(),
target[
start:end
],
reduction="sum",
)
position_loss = (
piece
if position_loss is None
else
position_loss
+
piece
)
total_loss = (
position_loss
if total_loss is None
else
total_loss
+
position_loss
)
# Exact teacher forcing:
# byte r+1 may condition on true byte r.
previous = self.shared_embedding(
target
)
return (
total_loss
/
float(
n
*
BYTES_PER_BLOCK
)
)
# =============================================================================
# COMPLETE CAUSAL MODEL
# =============================================================================
class CTWModel(
nn.Module
):
def __init__(
self,
):
super().__init__()
self.tokenizer = (
CTWBlockLiftingTokenizer()
)
self.position = nn.Parameter(
torch.zeros(
1,
MODEL_CONTEXT,
MODEL_DIM,
)
)
nn.init.normal_(
self.position,
mean=0.0,
std=0.01,
)
self.layers = nn.ModuleList(
[
CTWLayer()
for _ in range(
N_LAYERS
)
]
)
self.final_norm = CTWRMSNorm(
MODEL_DIM
)
self.bos_context = nn.Parameter(
torch.zeros(
MODEL_DIM
)
)
nn.init.normal_(
self.bos_context,
mean=0.0,
std=0.02,
)
self.decoder = CTWByteDecoder(
self.tokenizer.byte_embedding
)
def encode_completed_blocks(
self,
raw_bytes,
checkpoint_layers,
):
if (
raw_bytes.ndim
!=
2
):
raise ValueError(
"Expected raw byte IDs [B,T]."
)
b, t = raw_bytes.shape
if (
t
%
BYTES_PER_BLOCK
!=
0
):
raise ValueError(
"Raw byte length must be divisible by 8."
)
l = (
t
//
BYTES_PER_BLOCK
)
self.alpha_meyer = nn.Parameter(
torch.tensor(
0.10,
dtype=torch.float32,
)
)
self.alpha_local = nn.Parameter(
torch.tensor(
0.10,
dtype=torch.float32,
)
)
self.norm2 = CTWRMSNorm(
MODEL_DIM
)
self.ffn = CTWSwiGLU(
MODEL_DIM,
FFN_DIM,
)
def forward(
self,
x,
):
b, l, s, c = x.shape
flat = x.reshape(
b,
l,
MODEL_DIM,
)
n = self.norm1(
flat
)
gate_in = torch.sigmoid(
self.input_gate(
n
)
).reshape(
b,
l,
s,
c,
)
gate_out = torch.sigmoid(
self.output_gate(
n
)
).reshape(
b,
l,
s,
c,
)
spectral = self.meyer(
self.scale_in(
x
*
gate_in
)
)
spectral = self.scale_out(
spectral
)
local = self.local(
n
).reshape(
b,
l,
s,
c,
)
x = (
x
+
self.alpha_meyer.to(
x.dtype
)
*
spectral
*
gate_out
+
self.alpha_local.to(
x.dtype
)
*
local
)
flat = x.reshape(
b,
l,
MODEL_DIM,
)
flat = (
flat
+
self.ffn(
self.norm2(
flat
)
)
)
return flat.reshape(
b,
l,
s,
c,
)
# =============================================================================
# EXACT EIGHT-BYTE AUTOREGRESSIVE DECODER
# =============================================================================
class CTWByteDecoder(
nn.Module
):
def __init__(
self,
shared_embedding,
):
super().__init__()
self.shared_embedding = (
shared_embedding
)
self.context_proj = nn.Linear(
MODEL_DIM,
BYTE_DECODER_DIM,
)
self.gru = nn.GRUCell(
BYTE_DIM,
BYTE_DECODER_DIM,
)
self.norm = CTWRMSNorm(
BYTE_DECODER_DIM
)
self.to_byte_space = nn.Linear(
BYTE_DECODER_DIM,
BYTE_DIM,
)
self.bos = nn.Parameter(
torch.zeros(
BYTE_DIM
)
)
self.byte_position = nn.Parameter(
torch.zeros(
BYTES_PER_BLOCK,
BYTE_DIM,
)
)
nn.init.normal_(
self.bos,
mean=0.0,
std=0.02,
)
nn.init.normal_(
self.byte_position,
mean=0.0,
std=0.01,
)
def nll(
self,
context,
target_blocks,
):
b, l, _ = context.shape
n = (
b
*
l
)
hidden = torch.tanh(
self.context_proj(
context
)
).reshape(
n,
BYTE_DECODER_DIM,
)
previous = (
self.bos[
None,
:
]
.expand(
n,
-1,
)
)
total_loss = None
for r in range(
BYTES_PER_BLOCK
):
decoder_input = (
previous
+
self.byte_position[
r
][
None,
:
]
)
hidden = self.gru(
decoder_input,
hidden,
)
z = self.to_byte_space(
self.norm(
hidden
)
)
target = (
target_blocks[
:,
:,
r
]
.reshape(
-1
)
.long()
)
position_loss = None
for start in range(
0,
n,
HEAD_CHUNK,
):
end = min(
start
+
HEAD_CHUNK,
n,
)
logits = F.linear(
z[
start:end
],
self.shared_embedding.weight,
)
piece = F.cross_entropy(
logits.float(),
target[
start:end
],
reduction="sum",
)
position_loss = (
piece
if position_loss is None
else
position_loss
+
piece
)
total_loss = (
position_loss
if total_loss is None
else
total_loss
+
position_loss
)
# Exact teacher forcing:
# byte r+1 may condition on true byte r.
previous = self.shared_embedding(
target
)
return (
total_loss
/
float(
n
*
BYTES_PER_BLOCK
)
)
# =============================================================================
# COMPLETE CAUSAL MODEL
# =============================================================================
class CTWModel(
nn.Module
):
def __init__(
self,
):
super().__init__()
self.tokenizer = (
CTWBlockLiftingTokenizer()
)
self.position = nn.Parameter(
torch.zeros(
1,
MODEL_CONTEXT,
MODEL_DIM,
)
)
nn.init.normal_(
self.position,
mean=0.0,
std=0.01,
)
self.layers = nn.ModuleList(
[
CTWLayer()
for _ in range(
N_LAYERS
)
]
)
self.final_norm = CTWRMSNorm(
MODEL_DIM
)
self.bos_context = nn.Parameter(
torch.zeros(
MODEL_DIM
)
)
nn.init.normal_(
self.bos_context,
mean=0.0,
std=0.02,
)
self.decoder = CTWByteDecoder(
self.tokenizer.byte_embedding
)
def encode_completed_blocks(
self,
raw_bytes,
checkpoint_layers,
):
if (
raw_bytes.ndim
!=
2
):
raise ValueError(
"Expected raw byte IDs [B,T]."
)
b, t = raw_bytes.shape
if (
t
%
BYTES_PER_BLOCK
!=
0
):
raise ValueError(
"Raw byte length must be divisible by 8."
)
l = (
t
//
BYTES_PER_BLOCK
)
if (
l
>
MODEL_CONTEXT
):
raise ValueError(
f"Context {l} exceeds "
f"MODEL_CONTEXT={MODEL_CONTEXT}."
)
x = self.tokenizer(
raw_bytes.long()
)
pos = (
self.position[
:,
:l,
:
]
.reshape(
1,
l,
N_SCALES,
BYTE_DIM,
)
)
x = (
x
+
pos.to(
x.dtype
)
)
for layer in self.layers:
if (
checkpoint_layers
and
self.training
):
x = torch_checkpoint(
layer,
x,
use_reentrant=False,
preserve_rng_state=False,
)
else:
x = layer(
x
)
h = self.final_norm(
x.reshape(
b,
l,
MODEL_DIM,
)
)
blocks = (
raw_bytes
.long()
.reshape(
b,
l,
BYTES_PER_BLOCK,
)
)
return (
h,
blocks,
)
def contexts(
self,
raw_bytes,
checkpoint_layers=False,
):
(
h,
blocks,
) = self.encode_completed_blocks(
raw_bytes,
checkpoint_layers,
)
b, l, _ = h.shape
context = torch.empty_like(
h
)
# First block has no preceding block.
context[
:,
0,
:
] = self.bos_context.to(
h.dtype
)
# -------------------------------------------------------------
# CRITICAL CAUSAL SHIFT
#
# prediction for block b receives h[b-1], never h[b].
# -------------------------------------------------------------
if l > 1:
context[
:,
1:,
:
] = h[
:,
:-1,
:
]
return (
context,
blocks,
)
def nll(
self,
raw_bytes,
checkpoint_layers=True,
):
(
context,
blocks,
) = self.contexts(
raw_bytes,
checkpoint_layers,
)
return self.decoder.nll(
context,
blocks,
)