Ethanfel
f7a6f7790d
Modernized MisoTTS integration for ComfyUI with no torchtune/moshi: - vendored plain-torch Llama backbone (csm_llama), parity-verified Δ=0 vs torchtune - transformers.MimiModel codec (bit-identical codes to moshi), drops moshi/bnb/sphn - low-memory loader: streams 32GB fp32 checkpoint to GPU in bf16 (~18GB VRAM) - nodes: Model Loader, Generate (audiobook chunking + voice anchoring), EPUB Loader - pin-free requirements; runs on modern torch / Blackwell GPUs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
309 lines
13 KiB
Python
309 lines
13 KiB
Python
"""
|
|
torchtune-free reimplementation of the Llama3.2 TransformerDecoder used by CSM/MisoTTS.
|
|
|
|
Goal: drop-in replacement for `torchtune.models.llama3_2.llama3_2(...)` followed by
|
|
`_prepare_transformer` (tok_embeddings -> Identity, output -> Identity). The module:
|
|
|
|
* produces an IDENTICAL state_dict key layout, so MisoTTS weights load unchanged
|
|
* computes numerically identical outputs (RoPE scaling, GQA, RMSNorm, SwiGLU, KV-cache)
|
|
* exposes the methods models.py relies on: setup_caches / caches_are_enabled /
|
|
reset_caches / max_seq_len, and forward(h, input_pos=, mask=) returning float32.
|
|
|
|
All math is copied 1:1 from torchtune 0.4.0 so this can be validated by diffing.
|
|
No torchtune / torchao / torch-pin dependency: plain torch.
|
|
"""
|
|
import math
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
# ----------------------------------------------------------------------------- RoPE
|
|
class Llama3ScaledRoPE(nn.Module):
|
|
"""Verbatim port of torchtune.models.llama3_1._position_embeddings.Llama3ScaledRoPE."""
|
|
|
|
def __init__(self, dim, max_seq_len=4096, base=10_000, scale_factor=8,
|
|
low_freq_factor=1, high_freq_factor=4, old_context_len=8192):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.base = base
|
|
self.max_seq_len = max_seq_len
|
|
self.scale_factor = scale_factor
|
|
self.low_freq_factor = low_freq_factor
|
|
self.high_freq_factor = high_freq_factor
|
|
self.old_context_len = old_context_len
|
|
self.is_cache_built = False
|
|
self.rope_init()
|
|
|
|
def rope_init(self):
|
|
freqs = 1.0 / (self.base ** (torch.arange(0, self.dim, 2)[: (self.dim // 2)].float() / self.dim))
|
|
if freqs.is_meta:
|
|
return
|
|
theta = self.apply_scaling(freqs, self.scale_factor, self.low_freq_factor,
|
|
self.high_freq_factor, self.old_context_len)
|
|
self.register_buffer("theta", theta, persistent=False)
|
|
self.build_rope_cache(self.max_seq_len)
|
|
self.is_cache_built = True
|
|
|
|
def build_rope_cache(self, max_seq_len=4096):
|
|
seq_idx = torch.arange(max_seq_len, dtype=self.theta.dtype, device=self.theta.device)
|
|
idx_theta = torch.einsum("i, j -> ij", seq_idx, self.theta).float()
|
|
cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
|
|
self.register_buffer("cache", cache, persistent=False)
|
|
|
|
def apply_scaling(self, freqs, scale_factor, low_freq_factor, high_freq_factor, old_context_len):
|
|
low_freq_wavelen = old_context_len / low_freq_factor
|
|
high_freq_wavelen = old_context_len / high_freq_factor
|
|
new_freqs = []
|
|
for freq in freqs:
|
|
wavelen = 2 * math.pi / freq
|
|
if wavelen < high_freq_wavelen:
|
|
new_freqs.append(freq)
|
|
elif wavelen > low_freq_wavelen:
|
|
new_freqs.append(freq / scale_factor)
|
|
else:
|
|
assert low_freq_wavelen != high_freq_wavelen
|
|
smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
|
|
new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
|
|
return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)
|
|
|
|
def forward(self, x, *, input_pos=None):
|
|
if not self.is_cache_built:
|
|
raise RuntimeError("RoPE cache is not built. Please call rope_init() first.")
|
|
seq_len = x.size(1)
|
|
rope_cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
|
|
xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
|
|
rope_cache = rope_cache.view(-1, xshaped.size(1), 1, xshaped.size(3), 2)
|
|
x_out = torch.stack([
|
|
xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
|
|
xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
|
|
], -1)
|
|
x_out = x_out.flatten(3)
|
|
return x_out.type_as(x)
|
|
|
|
|
|
# ----------------------------------------------------------------------------- RMSNorm
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.scale = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
x_fp32 = x.float()
|
|
x_normed = (x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)
|
|
return x_normed * self.scale
|
|
|
|
|
|
# ----------------------------------------------------------------------------- KV cache
|
|
class KVCache(nn.Module):
|
|
def __init__(self, batch_size, max_seq_len, num_heads, head_dim, dtype):
|
|
super().__init__()
|
|
cache_shape = (batch_size, num_heads, max_seq_len, head_dim)
|
|
self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False)
|
|
self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False)
|
|
self.register_buffer("cache_pos", torch.arange(0, cache_shape[2]), persistent=False)
|
|
self.batch_size = batch_size
|
|
|
|
def reset(self):
|
|
self.k_cache.zero_()
|
|
self.v_cache.zero_()
|
|
self.cache_pos -= self.size
|
|
|
|
@property
|
|
def size(self):
|
|
return self.cache_pos[0].item()
|
|
|
|
def update(self, k_val, v_val):
|
|
bsz, _, seq_len, _ = k_val.shape
|
|
if bsz > self.k_cache.shape[0]:
|
|
raise ValueError("batch size larger than cache")
|
|
assert (self.cache_pos[0] + seq_len) <= self.k_cache.shape[2]
|
|
k_out = self.k_cache
|
|
v_out = self.v_cache
|
|
k_out[:, :, self.cache_pos[:seq_len]] = k_val
|
|
v_out[:, :, self.cache_pos[:seq_len]] = v_val
|
|
self.cache_pos += seq_len
|
|
return k_out, v_out
|
|
|
|
|
|
# ----------------------------------------------------------------------------- Attention
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(self, *, embed_dim, num_heads, num_kv_heads, head_dim,
|
|
pos_embeddings, max_seq_len=4096, attn_dropout=0.0):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.num_kv_heads = num_kv_heads
|
|
self.embed_dim = embed_dim
|
|
self.attn_dropout = attn_dropout
|
|
self.head_dim = head_dim
|
|
self.max_seq_len = max_seq_len
|
|
self.is_causal = True
|
|
self.kv_cache = None
|
|
self.q_proj = nn.Linear(embed_dim, num_heads * head_dim, bias=False)
|
|
self.k_proj = nn.Linear(embed_dim, num_kv_heads * head_dim, bias=False)
|
|
self.v_proj = nn.Linear(embed_dim, num_kv_heads * head_dim, bias=False)
|
|
self.output_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
|
self.pos_embeddings = pos_embeddings
|
|
self.cache_enabled = False
|
|
|
|
def setup_cache(self, batch_size, dtype, max_seq_len):
|
|
if self.kv_cache is not None:
|
|
return
|
|
self.kv_cache = KVCache(batch_size=batch_size, max_seq_len=max_seq_len,
|
|
num_heads=self.num_heads, head_dim=self.head_dim, dtype=dtype)
|
|
self.cache_enabled = True
|
|
|
|
def reset_cache(self):
|
|
self.kv_cache.reset()
|
|
|
|
@staticmethod
|
|
def _sdpa(q, k, v, mask, dropout_p, is_causal):
|
|
if mask is not None:
|
|
mask = mask[:, None, :, :]
|
|
return F.scaled_dot_product_attention(q, k, v, attn_mask=mask,
|
|
dropout_p=dropout_p, is_causal=is_causal)
|
|
|
|
def forward(self, x, y=None, *, mask=None, input_pos=None):
|
|
b, s_x, _ = x.shape
|
|
s_y = y.shape[1] if y is not None else 0
|
|
q = self.q_proj(x)
|
|
q_per_kv = self.num_heads // self.num_kv_heads
|
|
q = q.view(b, s_x, self.num_kv_heads * q_per_kv, self.head_dim)
|
|
if self.pos_embeddings is not None:
|
|
q = self.pos_embeddings(q, input_pos=input_pos)
|
|
q = q.transpose(1, 2)
|
|
|
|
if y is None:
|
|
k = self.kv_cache.k_cache
|
|
v = self.kv_cache.v_cache
|
|
else:
|
|
k = self.k_proj(y)
|
|
v = self.v_proj(y)
|
|
k = k.view(b, s_y, -1, self.head_dim)
|
|
if self.pos_embeddings is not None:
|
|
k = self.pos_embeddings(k, input_pos=input_pos)
|
|
k = k.view(b, s_y, self.num_kv_heads, 1, self.head_dim)
|
|
v = v.view(b, s_y, self.num_kv_heads, 1, self.head_dim)
|
|
if self.num_heads != self.num_kv_heads:
|
|
k = k.expand(b, s_y, self.num_kv_heads, q_per_kv, self.head_dim)
|
|
v = v.expand(b, s_y, self.num_kv_heads, q_per_kv, self.head_dim)
|
|
k = k.reshape(b, s_y, -1, self.head_dim)
|
|
v = v.reshape(b, s_y, -1, self.head_dim)
|
|
k = k.transpose(1, 2)
|
|
v = v.transpose(1, 2)
|
|
if self.kv_cache is not None and self.cache_enabled:
|
|
k, v = self.kv_cache.update(k, v)
|
|
|
|
output = self._sdpa(q, k, v, mask=mask,
|
|
dropout_p=self.attn_dropout if self.training else 0.0,
|
|
is_causal=self.kv_cache is None and mask is None and self.is_causal)
|
|
output = output.transpose(1, 2).contiguous().view(b, s_x, -1)
|
|
return self.output_proj(output)
|
|
|
|
|
|
# ----------------------------------------------------------------------------- MLP
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, hidden_dim):
|
|
super().__init__()
|
|
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
|
|
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
|
|
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
|
|
self.activation = nn.SiLU()
|
|
|
|
def forward(self, x):
|
|
h = self.activation(self.w1(x))
|
|
h = h * self.w3(x)
|
|
return self.w2(h)
|
|
|
|
|
|
# ----------------------------------------------------------------------------- Layer
|
|
class TransformerSelfAttentionLayer(nn.Module):
|
|
def __init__(self, attn, mlp, sa_norm, mlp_norm):
|
|
super().__init__()
|
|
self.attn = attn
|
|
self.mlp = mlp
|
|
self.sa_norm = sa_norm
|
|
self.mlp_norm = mlp_norm
|
|
|
|
def setup_caches(self, batch_size, dtype, decoder_max_seq_len):
|
|
self.attn.setup_cache(batch_size, dtype, max_seq_len=decoder_max_seq_len)
|
|
|
|
def caches_are_enabled(self):
|
|
return self.attn.cache_enabled
|
|
|
|
def reset_cache(self):
|
|
self.attn.reset_cache()
|
|
|
|
def forward(self, x, *, mask=None, input_pos=None):
|
|
h = self.sa_norm(x)
|
|
attn_out = self.attn(h, h, mask=mask, input_pos=input_pos)
|
|
h = attn_out + x
|
|
mlp_out = self.mlp(self.mlp_norm(h))
|
|
return h + mlp_out
|
|
|
|
|
|
# ----------------------------------------------------------------------------- Decoder
|
|
class TransformerDecoder(nn.Module):
|
|
"""Self-attention-only Llama decoder operating on pre-computed embeddings (h).
|
|
|
|
Mirrors torchtune's TransformerDecoder after _prepare_transformer replaced
|
|
tok_embeddings and output with Identity: forward takes embeddings, returns float32.
|
|
"""
|
|
|
|
def __init__(self, *, layers, norm, max_seq_len, num_heads, head_dim):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList(layers)
|
|
self.norm = norm
|
|
self.max_seq_len = max_seq_len
|
|
self.num_heads = num_heads
|
|
self.head_dim = head_dim
|
|
self.embed_dim = num_heads * head_dim
|
|
|
|
def setup_caches(self, batch_size, dtype, *, decoder_max_seq_len=None):
|
|
max_seq_len = decoder_max_seq_len if decoder_max_seq_len is not None else self.max_seq_len
|
|
self.decoder_max_cache_seq_len = max_seq_len
|
|
for layer in self.layers:
|
|
layer.setup_caches(batch_size, dtype, decoder_max_seq_len=max_seq_len)
|
|
|
|
def caches_are_enabled(self):
|
|
return self.layers[0].caches_are_enabled()
|
|
|
|
def reset_caches(self):
|
|
for layer in self.layers:
|
|
layer.reset_cache()
|
|
|
|
def forward(self, h, *, mask=None, input_pos=None):
|
|
for layer in self.layers:
|
|
h = layer(h, mask=mask, input_pos=input_pos)
|
|
h = self.norm(h)
|
|
return h.float()
|
|
|
|
|
|
# ----------------------------------------------------------------------------- builder
|
|
def llama3_2(*, vocab_size, num_layers, num_heads, num_kv_heads, embed_dim,
|
|
max_seq_len, intermediate_dim, attn_dropout=0.0, norm_eps=1e-5,
|
|
rope_base=500_000, scale_factor=32):
|
|
"""Matches torchtune.models.llama3_2.llama3_2(...) + _prepare_transformer.
|
|
|
|
vocab_size is accepted for signature parity but unused (tok_embeddings/output
|
|
are Identity in CSM, so they carry no weights).
|
|
"""
|
|
head_dim = embed_dim // num_heads
|
|
rope = Llama3ScaledRoPE(dim=head_dim, max_seq_len=max_seq_len, base=rope_base, scale_factor=scale_factor)
|
|
layers = []
|
|
for _ in range(num_layers):
|
|
attn = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads, num_kv_heads=num_kv_heads,
|
|
head_dim=head_dim, pos_embeddings=rope, max_seq_len=max_seq_len,
|
|
attn_dropout=attn_dropout)
|
|
mlp = FeedForward(dim=embed_dim, hidden_dim=intermediate_dim)
|
|
layers.append(TransformerSelfAttentionLayer(
|
|
attn=attn, mlp=mlp,
|
|
sa_norm=RMSNorm(embed_dim, eps=norm_eps),
|
|
mlp_norm=RMSNorm(embed_dim, eps=norm_eps),
|
|
))
|
|
return TransformerDecoder(layers=layers, norm=RMSNorm(embed_dim, eps=norm_eps),
|
|
max_seq_len=max_seq_len, num_heads=num_heads, head_dim=head_dim)
|