from .lingua_transformer import * def apply_rotary_emb_xattn( xq: torch.Tensor, xk: torch.Tensor, seq_dim: int, freqs_cis_q: torch.Tensor, freqs_cis_k: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: xq_ = xq.reshape(*xq.shape[:-1], -1, 1, 2) # B S H D -> B S H D/2 1 2 xk_ = xk.reshape(*xk.shape[:-1], -1, 1, 2) # B S H D -> B S H D/2 1 2 freqs_cis_q = reshape_for_broadcast( freqs_cis_q, xq_, seq_dim ).float() # S D/2 2 2 -> 1 S 1 D/2 2 2 freqs_cis_k = reshape_for_broadcast( freqs_cis_k, xk_, seq_dim ).float() # S D/2 2 2 -> 1 S 1 D/2 2 2 xq_out = (xq_ * freqs_cis_q).sum(5).flatten(3) xk_out = (xk_ * freqs_cis_k).sum(5).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) class AdaRMSNorm(nn.Module): """ Initialize the RMSNorm normalization layer. Args: dim (int): The dimension of the input tensor. eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. Attributes: eps (float): A small value added to the denominator for numerical stability. weight (nn.Parameter): Learnable scaling parameter. """ def __init__(self, emb_dim, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Linear(emb_dim, dim, bias=True) def _norm(self, x: torch.Tensor): return x * torch.rsqrt((x * x).mean(-1, keepdim=True) + self.eps) def forward(self, x: torch.Tensor, c: torch.Tensor): x = probe.log_stats(x, "resid") output = self._norm(x.float()) return (output * self.weight(c).float()).type_as(x) def reset_parameters(self): # bias to ones, weight to 0s nn.init.ones_(self.weight.bias) nn.init.zeros_(self.weight.weight) class CrossAttention(nn.Module): def __init__( self, dim: int, head_dim: int, n_heads: int, n_kv_heads: int, rope_theta: float, rope_dim: int, ): super().__init__() self.dim = dim self.head_dim = head_dim self.rope_theta = rope_theta self.rope_dim = rope_dim self.n_heads = n_heads self.n_kv_heads = n_kv_heads self.heads_per_group = self.n_heads // self.n_kv_heads self.wq = nn.Linear( dim, n_heads * head_dim, bias=False, ) self.wk = nn.Linear( dim, n_kv_heads * head_dim, bias=False, ) self.wv = nn.Linear( dim, n_kv_heads * head_dim, bias=False, ) self.wo = nn.Linear( n_heads * head_dim, dim, bias=False, ) def forward( self, xq: torch.Tensor, xkv: torch.Tensor, freq_cis: torch.Tensor, tok_idx: Optional[torch.Tensor] = None, cross_tok_idx: Optional[torch.Tensor] = None, mask: Optional[Union[BlockMask, str]] = None, attn_impl: str = "sdpa", ) -> torch.Tensor: # B S D assert attn_impl == "flex_attention", "Only flex_attention is supported for now" bsz, seq_len_q, dim = xq.shape _, seq_len_kv, _ = xkv.shape xq = self.wq(xq.view_as(xq)) xk = self.wk(xkv.view_as(xkv)) xv = self.wv(xkv.view_as(xkv)) output_shape = xq.shape # B S D -> B S H D xq = xq.view(bsz, seq_len_q, self.n_heads, self.head_dim) xk = xk.view(bsz, seq_len_kv, self.n_kv_heads, self.head_dim) xv = xv.view(bsz, seq_len_kv, self.n_kv_heads, self.head_dim) if self.rope_dim==0: pass elif self.rope_dim==1: if tok_idx is not None and cross_tok_idx is not None: xq, xk = apply_rotary_emb_xattn( xq, xk, 1, freq_cis[tok_idx], freq_cis[cross_tok_idx] ) else: xq, xk = apply_rotary_emb_xattn( xq, xk, 1, freq_cis[0:seq_len_q], freq_cis[0:seq_len_kv] ) elif self.rope_dim==4: # Build freqcis_4RoPE by indexing freq_cis with each dimension of tok_idx separately and concatenating # Cat along a new dimension to get [S, head_dim//2, 2, 2] freqcis_parts = [] freqcis_cross_parts = [] for i in range(self.rope_dim): freqcis_parts.append(freq_cis[tok_idx[:, i]]) freqcis_cross_parts.append(freq_cis[cross_tok_idx[:, i]]) freqcis_4RoPE = torch.cat(freqcis_parts, dim=1) freqcis_cross_4RoPE = torch.cat(freqcis_cross_parts, dim=1) xq, xk = apply_rotary_emb_xattn( xq, xk, 1, freqcis_4RoPE, freqcis_cross_4RoPE ) else: print(f"I dont know how to handle {self.rope_dim=} inside xattn.CrossAttention.forward") import IPython; print('\n\nDebug:'); IPython.embed(); import time; time.sleep(0.3) # This condition helps us be easily compatible # with inference by adding a pluggable KVCache if hasattr(self, "kv_cache"): xk, xv = self.kv_cache.update(xk, xv, tok_idx) xk = repeat_kv(xk, self.heads_per_group, dim=2) xv = repeat_kv(xv, self.heads_per_group, dim=2) if attn_impl == "flex_attention": assert mask is None or isinstance(mask, BlockMask) xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv)) if xq.device.type == "mps": # MPS does not support flex_attention; fall back to SDPA with dense mask if mask is not None: S_q, S_kv = xq.shape[2], xk.shape[2] q_idx = torch.arange(S_q, device='cpu') kv_idx = torch.arange(S_kv, device='cpu') dense_bool = mask.mask_mod(0, 0, q_idx.unsqueeze(1), kv_idx.unsqueeze(0)) attn_mask = torch.zeros(1, 1, S_q, S_kv, dtype=xq.dtype, device=xq.device) attn_mask.masked_fill_(~dense_bool.unsqueeze(0).unsqueeze(0).to(xq.device), float("-inf")) else: attn_mask = None output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask) elif xq.device.type == "cuda": output = flex_attention_comp(xq, xk, xv, block_mask=mask) else: output = flex_attention(xq, xk, xv, block_mask=mask) output = output.transpose(1, 2).contiguous() # B H S D -> B S H D elif attn_impl == "sdpa": xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv)) assert mask is None or isinstance(mask, (str, torch.Tensor)) is_causal = (mask == "causal") if isinstance(mask, str) else False mask = mask if isinstance(mask, torch.Tensor) else None output = F.scaled_dot_product_attention( xq, xk, xv, is_causal=is_causal, attn_mask=mask, ) output = output.transpose(1, 2).contiguous() # B H S D -> B S H D else: raise NotImplementedError( f"Attention implementation {attn_impl} not supported" ) output = self.wo(output.reshape(output_shape)) return output def reset_parameters(self, init_std=None, factor=1.0): init_std = init_std or (self.dim ** (-0.5)) for w in [self.wq, self.wk, self.wv]: nn.init.trunc_normal_( w.weight, mean=0.0, std=init_std, a=-3 * init_std, b=3 * init_std, ) nn.init.trunc_normal_( self.wo.weight, mean=0.0, std=init_std / factor, a=-3 * init_std, b=3 * init_std, ) class FourierConditioner(nn.Module): def __init__( self, output_dim: int, input_dim: int = 1, std: float = 0.02, min_val: float = 0.0, max_val: float = 1.0, ): super().__init__() assert input_dim == 1 assert output_dim % 2 == 0 self.output_dim = output_dim self.register_buffer("weight", torch.randn([output_dim // 2, input_dim]) * std) self.min_val, self.max_val = min_val, max_val self.proj = nn.Linear(output_dim, output_dim) def forward(self, x: list[float], device=None): x = (x - self.min_val) / (self.max_val - self.min_val) f = (2 * torch.pi * x.float() @ self.weight.T).type_as(x) return self.proj(torch.cat([f.cos(), f.sin()], dim=-1)) def reset_parameters(self, init_std=None, factor=1.0): init_std = init_std or (self.output_dim ** (-0.5)) self.register_buffer("weight", torch.randn([self.output_dim // 2, 1]).to(self.proj.weight.device) * init_std) nn.init.trunc_normal_( self.proj.weight, mean=0.0, std=init_std / factor, a=-3 * init_std, b=3 * init_std, ) nn.init.zeros_(self.proj.bias) @dataclass class DecoderArgs(BaseTransformerArgs): t_dim: int = 64 n_heads: int = 8 seqlen_t: bool = False class DecoderBlock(nn.Module): def __init__(self, args: DecoderArgs): super().__init__() assert (args.head_dim is not None) or (args.n_heads is not None), ( "Should specify at least head_dim or n_heads" ) self.head_dim = args.head_dim or args.dim // args.n_heads self.n_heads = args.n_heads or args.dim // args.head_dim self.n_kv_heads = args.n_kv_heads or self.n_heads assert args.n_heads % self.n_kv_heads == 0 assert args.dim % args.n_heads == 0 self.cross_attention = CrossAttention( dim=args.dim, head_dim=self.head_dim, n_heads=self.n_heads, n_kv_heads=self.n_kv_heads, rope_theta=args.rope_theta, rope_dim=args.rope_dim, ) self.cross_attention_x_norm = AdaRMSNorm( args.t_dim, args.dim, eps=args.norm_eps ) self.seqlen_t = args.seqlen_t if args.seqlen_t: self.cross_attention_y_norm = RMSNorm( args.dim, eps=args.norm_eps ) else: self.cross_attention_y_norm = AdaRMSNorm( args.t_dim, args.dim, eps=args.norm_eps ) self.attention = Attention( dim=args.dim, head_dim=self.head_dim, n_heads=self.n_heads, n_kv_heads=self.n_kv_heads, rope_theta=args.rope_theta, rope_dim=args.rope_dim, ) self.feed_forward = FeedForward( dim=args.dim, hidden_dim=4 * args.dim, multiple_of=args.multiple_of, ffn_dim_multiplier=args.ffn_dim_multiplier, ) self.attention_norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps) self.ffn_norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps) def forward( self, x: torch.Tensor, y: torch.Tensor, c: torch.Tensor, freq_cis: torch.Tensor, tok_idx: Optional[torch.Tensor] = None, cross_tok_idx: Optional[torch.Tensor] = None, self_attn_mask: Optional[Union[BlockMask, str]] = None, cross_attn_mask: Optional[Union[BlockMask, str]] = None, attn_impl: str = "sdpa", do_idx: Optional[torch.Tensor] = None, print_layerwise_activation_stats: bool = False, ) -> torch.Tensor: if print_layerwise_activation_stats and do_idx is not None: x_normed = self.cross_attention_x_norm(x, c) y_normed = self.cross_attention_y_norm(y, c) if not self.seqlen_t else self.cross_attention_y_norm(y) print(f"\n\tDecoder cross_attn_x_norm: (drop-out) mean={x[:, do_idx, :].mean().item():.6f}, std={x[:, do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={x_normed[:, do_idx, :].mean().item():.6f}, std={x_normed[:, do_idx, :].std().item():.6f}") print(f"\tDecoder cross_attn_x_norm: (non-drop) mean={x[:, ~do_idx, :].mean().item():.6f}, std={x[:, ~do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={x_normed[:, ~do_idx, :].mean().item():.6f}, std={x_normed[:, ~do_idx, :].std().item():.6f}") print(f"\n\tDecoder cross_attn_y_norm: (drop-out) mean={y[:, do_idx, :].mean().item():.6f}, std={y[:, do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={y_normed[:, do_idx, :].mean().item():.6f}, std={y_normed[:, do_idx, :].std().item():.6f}") print(f"\tDecoder cross_attn_y_norm: (non-drop) mean={y[:, ~do_idx, :].mean().item():.6f}, std={y[:, ~do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={y_normed[:, ~do_idx, :].mean().item():.6f}, std={y_normed[:, ~do_idx, :].std().item():.6f}") x = x + self.cross_attention( x_normed, y_normed, freq_cis, tok_idx=tok_idx, cross_tok_idx=cross_tok_idx, mask=cross_attn_mask, attn_impl=attn_impl, ) else: x = x + self.cross_attention( self.cross_attention_x_norm(x, c), self.cross_attention_y_norm(y, c) if not self.seqlen_t else self.cross_attention_y_norm(y), freq_cis, tok_idx=tok_idx, cross_tok_idx=cross_tok_idx, mask=cross_attn_mask, attn_impl=attn_impl, ) if print_layerwise_activation_stats and do_idx is not None: x_normed = self.attention_norm(x, c) print(f"\n\tDecoder self attn_norm: (drop-out) mean={x[:, do_idx, :].mean().item():.6f}, std={x[:, do_idx, :].std().item():.6f}", end=" --> ") print(f" mean={x_normed[:, do_idx, :].mean().item():.6f}, std={x_normed[:, do_idx, :].std().item():.6f}") print(f"\tDecoder self attn_norm: (non-drop) mean={x[:, ~do_idx, :].mean().item():.6f}, std={x[:, ~do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={x_normed[:, ~do_idx, :].mean().item():.6f}, std={x_normed[:, ~do_idx, :].std().item():.6f}") h = x + self.attention( x_normed, freq_cis, tok_idx=tok_idx, mask=self_attn_mask, attn_impl=attn_impl, ) h_normed = self.ffn_norm(h, c) print(f"\n\tDecoder ffn_norm: (drop-out) mean={h[:, do_idx, :].mean().item():.6f}, std={h[:, do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={h_normed[:, do_idx, :].mean().item():.6f}, std={h_normed[:, do_idx, :].std().item():.6f}") print(f"\tDecoder ffn_norm: (non-drop) mean={h[:, ~do_idx, :].mean().item():.6f}, std={h[:, ~do_idx, :].std().item():.6f}", end=" --> ") print(f"mean={h_normed[:, ~do_idx, :].mean().item():.6f}, std={h_normed[:, ~do_idx, :].std().item():.6f}") out = h + self.feed_forward(h_normed) else: h = x + self.attention( self.attention_norm(x, c), freq_cis, tok_idx=tok_idx, mask=self_attn_mask, attn_impl=attn_impl, ) out = h + self.feed_forward(self.ffn_norm(h, c)) return out def init_weights(self, init_std=None, factor=1.0): self.cross_attention.reset_parameters(init_std, factor) self.cross_attention_x_norm.reset_parameters() self.cross_attention_y_norm.reset_parameters() self.attention.reset_parameters(init_std, factor) self.attention_norm.reset_parameters() self.feed_forward.reset_parameters(init_std, factor) self.ffn_norm.reset_parameters()