File size: 6,064 Bytes
685e018 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Pre-LN Transformer encoder that routes self-attention through ``flex_attention``.
The block-diffusion masks are structured (block-causal over thought slots plus optional
key padding), so expressing them as a ``flex_attention`` block mask lets the fused kernel
skip fully-masked blocks instead of materializing a dense ``[batch * heads, L, L]`` score
tensor. This keeps attention memory and time near-linear in the trace length, which the
default ``nn.TransformerEncoder`` math path does not.
The layer geometry mirrors ``nn.TransformerEncoderLayer(norm_first=True, activation='gelu')``
so behaviour matches the default path up to floating-point error; the attention core is
verified against a dense masked-softmax reference before use.
"""
from __future__ import annotations
import math
import os
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
# Compiled dynamic=False is fastest for fixed-shape training. Variable-length generation
# recompiles per new sequence length, so serving sets MDLM_FLEX_EAGER=1 to run eager and
# trade steady-state speed for the absence of per-length recompilation stalls.
_FLEX_EAGER = os.environ.get('MDLM_FLEX_EAGER') == '1'
_flex_compiled = flex_attention if _FLEX_EAGER else torch.compile(flex_attention, dynamic=False)
def build_block_mask(
attn_mask: Tensor | None,
padding_mask: Tensor | None,
batch_size: int,
seq_len: int,
device: torch.device,
) -> BlockMask | None:
"""Build a broadcast-over-heads ``BlockMask`` from project boolean masks.
``attn_mask`` follows the src_mask convention (``True`` blocks a key), shaped ``[L, L]``
or ``[batch, L, L]``. ``padding_mask`` follows the src_key_padding_mask convention
(``True`` marks padding). Returns ``None`` when neither constrains attention.
"""
if attn_mask is None and padding_mask is None:
return None
shared = attn_mask is not None and attn_mask.dim() == 2
def mask_mod(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
keep = torch.ones_like(q_idx, dtype=torch.bool)
if attn_mask is not None:
blocked = attn_mask[q_idx, kv_idx] if shared else attn_mask[b, q_idx, kv_idx]
keep = keep & ~blocked
if padding_mask is not None:
keep = keep & ~padding_mask[b, kv_idx]
return keep
return create_block_mask(
mask_mod, batch_size, None, seq_len, seq_len, device=device, _compile=not _FLEX_EAGER
)
def flex_self_attention(
query: Tensor, key: Tensor, value: Tensor, block_mask: BlockMask | None
) -> Tensor:
"""Multi-head self-attention over ``[batch, heads, L, head_dim]`` tensors."""
if block_mask is None:
return flex_attention(query, key, value)
return _flex_compiled(query, key, value, block_mask=block_mask)
class FlexEncoderLayer(nn.Module):
"""Pre-LN Transformer block with a ``flex_attention`` self-attention core."""
def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float) -> None:
super().__init__()
if d_model % n_heads != 0:
raise ValueError("d_model must be divisible by n_heads")
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def _split_heads(self, projected: Tensor) -> Tensor:
batch_size, seq_len, _ = projected.shape
return projected.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
def forward(self, hidden: Tensor, block_mask: BlockMask | None) -> Tensor:
normed = self.norm1(hidden)
query = self._split_heads(self.q_proj(normed))
key = self._split_heads(self.k_proj(normed))
value = self._split_heads(self.v_proj(normed))
attended = flex_self_attention(query, key, value, block_mask)
batch_size, _, seq_len, _ = attended.shape
attended = attended.transpose(1, 2).reshape(batch_size, seq_len, -1)
hidden = hidden + self.dropout(self.out_proj(attended))
normed = self.norm2(hidden)
feed_forward = self.linear2(self.dropout(F.gelu(self.linear1(normed))))
return hidden + self.dropout(feed_forward)
class FlexEncoder(nn.Module):
"""Stack of :class:`FlexEncoderLayer` blocks with a final layer norm."""
def __init__(
self,
d_model: int,
n_heads: int,
d_ff: int,
dropout: float,
n_layers: int,
activation_checkpointing: bool,
) -> None:
super().__init__()
self.layers = nn.ModuleList(
FlexEncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
)
self.norm = nn.LayerNorm(d_model)
self.activation_checkpointing = activation_checkpointing
def init_residual_outputs(self, n_layers: int) -> None:
residual_std = 0.02 / math.sqrt(2 * n_layers)
for layer in self.layers:
nn.init.normal_(layer.out_proj.weight, mean=0.0, std=residual_std)
nn.init.normal_(layer.linear2.weight, mean=0.0, std=residual_std)
def forward(self, hidden: Tensor, block_mask: BlockMask | None) -> Tensor:
use_checkpoint = (
self.activation_checkpointing and self.training and torch.is_grad_enabled()
)
for layer in self.layers:
if use_checkpoint:
hidden = torch.utils.checkpoint.checkpoint(
layer, hidden, block_mask, use_reentrant=False
)
else:
hidden = layer(hidden, block_mask)
return self.norm(hidden)
|