""" 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)