Module neuralops.attentions.triton_svd_linear_attention
Triton-Accelerated SVD Linear Attention & Chunked Causal Linear Attention.
High-performance GPU implementation using OpenAI Triton language (tl) with: 1. Chunked Causal Linear Attention Kernel (_triton_chunked_causal_linear_attn_fwd): - Keeps the recurrent key-value state S \in R^{D_k x D_v} strictly in GPU SRAM/registers. - Tiled in blocks of BLOCK_CHUNK = 64 to eliminate O(N^2) memory footprint. - Evaluates intra-chunk causal attention and inter-chunk prefix sums in a single fused pass. - Strictly O(N * d^2) FLOPs and O(N * d) memory.
-
Bidirectional Linear Attention Kernel (_triton_bidirectional_linear_attn_fwd):
- Fused parallel reduction of K^T V and K^T 1 for encoder architectures (DistilBERT, RoBERTa).
- Fused query projection with numerical stability epsilon clamping in registers.
-
Non-Zero Identity SVD Initialization (BUG 5 FIX):
- SVDLinearProjection initializes U with identity-like structure (eye for rank dimensions) and V with Kaiming uniform so initial output is non-zero before fine-tuning.
-
Backward Kernel & Clean Autograd (BUG 1 FIX):
- Uses explicit torch.autograd.grad(out, [q_req, k_req, v_req], grad_out) preventing double-backward warnings.
- Bridges _triton_linear_attn_bwd_kernel on CUDA GPUs for bidirectional attention.
- Provides streaming chunked backward propagation for causal decoder attention.
-
Attention Masking & Padding Invariant (BUG 2 FIX):
- Ensures padding tokens on K and V are strictly zeroed out before attention accumulation, preventing contaminated normalization denominator Z.
-
RoPE and GQA Support (BUG 3 & BUG 4 FIX):
- Seamless Rotary Position Embeddings (RoPE) via apply_rotary_pos_emb.
- Full Grouped Query Attention (GQA) support with separate num_kv_heads and repeat_kv.
Functions
def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor]-
Expand source code
def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) elif cos.dim() == 3: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_embed = (q_rot * cos) + (_rotate_half(q_rot) * sin) k_embed = (k_rot * cos) + (_rotate_half(k_rot) * sin) if q_pass.shape[-1] > 0: return torch.cat([q_embed, q_pass], dim=-1), torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed def inject_triton_svd_linear_attention(model: nn.Module = None, rank: int = 64, copy_weights: bool = True, **kwargs: Any) ‑> torch.nn.modules.module.Module | None-
Expand source code
@ATTENTION_REGISTRY.register("inject_triton_svd_linear_attention") @ATTENTION_REGISTRY.register("triton_svd_linear_attention_inject") def inject_triton_svd_linear_attention( model: nn.Module = None, rank: int = 64, copy_weights: bool = True, **kwargs: Any, ) -> nn.Module | None: """Safe factory function to inject TritonSVDLinearAttention into a model.""" from neuralops.core.attention_utils import inject_attention return inject_attention( model, "triton_svd_linear_attention", rank=rank, copy_weights=copy_weights, **kwargs, )Safe factory function to inject TritonSVDLinearAttention into a model.
def triton_linear_attention(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
is_causal: bool = False,
activation_mode: str = 'elu',
eps: float = 1e-06) ‑> torch.Tensor-
Expand source code
def triton_linear_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool = False, activation_mode: str = "elu", eps: float = 1e-6, ) -> torch.Tensor: """Convenience functional wrapper for Triton Linear Attention.""" return _TritonLinearAttentionAutogradFunction.apply(q, k, v, is_causal, activation_mode, eps)Convenience functional wrapper for Triton Linear Attention.
Classes
class SVDLinearProjection (in_features: int, out_features: int, rank: int = 64, bias: bool = True)-
Expand source code
class SVDLinearProjection(nn.Module): r""" Genuine SVD Low-Rank Linear Layer: W \approx U @ V. Factorizes in_features -> out_features through bottleneck rank r: V \in R^{r x in_features} (Down-projection) U \in R^{out_features x r} (Up-projection) y = (x @ V^T) @ U^T + bias """ triton_kernel = "low_rank" def __init__( self, in_features: int, out_features: int, rank: int = 64, bias: bool = True, ): super().__init__() self.in_features = in_features self.out_features = out_features self.rank = min(rank, in_features, out_features) self.V = nn.Parameter(torch.empty(self.rank, in_features)) self.U = nn.Parameter(torch.empty(out_features, self.rank)) if bias: self.bias = nn.Parameter(torch.zeros(out_features)) else: self.register_parameter("bias", None) self.reset_parameters() def reset_parameters(self) -> None: # BUG 5 FIX: Identity-like init for U so initial output is non-zero before fine-tuning nn.init.kaiming_uniform_(self.V, a=math.sqrt(5)) nn.init.zeros_(self.U) with torch.no_grad(): r = min(self.out_features, self.rank) self.U[:r, :r].copy_(torch.eye(r)) if self.bias is not None: nn.init.zeros_(self.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: # Fused low-rank projection: (x @ V^T) @ U^T + bias # Align weights with input dtype (e.g. under fp16/bf16 autocast) so Triton kernel runs directly u = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U v = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V b = ( self.bias.to(x.dtype) if (self.bias is not None and self.bias.dtype != x.dtype) else self.bias ) return triton_dispatch(self.triton_kernel, x, u, v, bias=b) @classmethod def from_pretrained_linear(cls, linear: nn.Linear, rank: int = 64) -> SVDLinearProjection: r""" Initializes low-rank factors directly from pretrained weights via Truncated SVD: W = U_full @ diag(S) @ V_full^T U = U_full[:, :r] @ sqrt(S[:r]) V = sqrt(S[:r]) @ V_full^T[:r, :] """ out_features, in_features = linear.weight.shape proj = cls(in_features, out_features, rank=rank, bias=(linear.bias is not None)) with torch.no_grad(): W = linear.weight.data.float() u, s, vh = torch.linalg.svd(W, full_matrices=False) r = proj.rank s_sqrt = torch.diag(torch.sqrt(s[:r])) proj.U.data.copy_(torch.matmul(u[:, :r], s_sqrt).to(linear.weight.dtype)) proj.V.data.copy_(torch.matmul(s_sqrt, vh[:r, :]).to(linear.weight.dtype)) if linear.bias is not None: proj.bias.data.copy_(linear.bias.data) return projGenuine SVD Low-Rank Linear Layer: W \approx U @ V.
Factorizes in_features -> out_features through bottleneck rank r: V \in R^{r x in_features} (Down-projection) U \in R^{out_features x r} (Up-projection) y = (x @ V^T) @ U^T + bias
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Class variables
var triton_kernel-
The type of the None singleton.
Static methods
def from_pretrained_linear(linear: nn.Linear, rank: int = 64) ‑> SVDLinearProjection-
Initializes low-rank factors directly from pretrained weights via Truncated SVD: W = U_full @ diag(S) @ V_full^T U = U_full[:, :r] @ sqrt(S[:r]) V = sqrt(S[:r]) @ V_full^T[:r, :]
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: # Fused low-rank projection: (x @ V^T) @ U^T + bias # Align weights with input dtype (e.g. under fp16/bf16 autocast) so Triton kernel runs directly u = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U v = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V b = ( self.bias.to(x.dtype) if (self.bias is not None and self.bias.dtype != x.dtype) else self.bias ) return triton_dispatch(self.triton_kernel, x, u, v, bias=b)Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: # BUG 5 FIX: Identity-like init for U so initial output is non-zero before fine-tuning nn.init.kaiming_uniform_(self.V, a=math.sqrt(5)) nn.init.zeros_(self.U) with torch.no_grad(): r = min(self.out_features, self.rank) self.U[:r, :r].copy_(torch.eye(r)) if self.bias is not None: nn.init.zeros_(self.bias)
class TritonSVDLinearAttention (config: Any = None,
embed_dim: int | None = None,
num_heads: int | None = None,
num_kv_heads: int | None = None,
rank: int | None = 64,
is_decoder: bool = False,
activation_mode: str = 'elu',
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register( "triton_svd_linear_attention", supports_causal=True, supports_bidirectional=True ) @ATTENTION_REGISTRY.register( "triton_linear_attention", supports_causal=True, supports_bidirectional=True ) class TritonSVDLinearAttention(BaseAttention): r""" Triton-Accelerated Linear Attention with Genuine SVD Low-Rank Projections, RoPE, GQA, and Masked Attention. """ triton_kernel = "linear_attn" def __init__( self, config: Any = None, embed_dim: int | None = None, num_heads: int | None = None, num_kv_heads: int | None = None, rank: int | None = 64, is_decoder: bool = False, activation_mode: str = "elu", eps: float = 1e-6, pretrained_attn: nn.Module | None = None, **kwargs: Any, ): super().__init__() self.config = config self.is_decoder = bool( is_decoder or getattr(config, "is_decoder", False) or getattr(config, "is_causal", False) ) self.activation_mode = activation_mode self.eps = eps self.triton_kwargs = { "phi": "elu+1" if self.activation_mode == "elu" else self.activation_mode, "eps": self.eps, "cos_weight": None, "sin_weight": None, } # Dimensionality extraction self.embed_dim = ( embed_dim or kwargs.get("hidden_size") or getattr(config, "hidden_size", 768) ) self.num_heads = ( num_heads or kwargs.get("num_attention_heads") or getattr(config, "num_attention_heads", 12) ) # BUG 4 FIX: Support GQA self.num_kv_heads = ( num_kv_heads or kwargs.get("num_kv_heads") or getattr(config, "num_key_value_heads", self.num_heads) ) self.head_dim = kwargs.get("head_dim") or (self.embed_dim // self.num_heads) self.rank = rank or 64 self.kv_dim = self.num_kv_heads * self.head_dim # Genuine SVD Low-Rank Linear Projections if self.rank < self.embed_dim: self.q_proj = SVDLinearProjection( self.embed_dim, self.embed_dim, rank=self.rank, bias=True ) self.k_proj = SVDLinearProjection( self.embed_dim, self.kv_dim, rank=min(self.rank, self.kv_dim), bias=True ) self.v_proj = SVDLinearProjection( self.embed_dim, self.kv_dim, rank=min(self.rank, self.kv_dim), bias=True ) else: self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.k_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) if pretrained_attn is not None: self.copy_weights(pretrained_attn) def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None) k_w, k_b = ( c_w[split_size : 2 * split_size, :], (c_b[split_size : 2 * split_size] if c_b is not None else None), ) v_w, v_b = ( c_w[2 * split_size :, :], (c_b[2 * split_size :] if c_b is not None else None), ) self._copy_single_proj(self.q_proj, q_w, q_b) self._copy_single_proj(self.k_proj, k_w, k_b) self._copy_single_proj(self.v_proj, v_w, v_b) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): w = base_src.weight.data b = getattr(base_src, "bias", None) b_data = b.data if b is not None else None self._copy_single_proj(target_mod, w, b_data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data b = getattr(base_mod, "bias", None) b_data = b.data if b is not None else None if target_name == "out_proj": if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if b_data is not None and target_mod.bias is not None: target_mod.bias.data.copy_(b_data) else: self._copy_single_proj(target_mod, w, b_data) break def _copy_single_proj( self, target_mod: nn.Module, w: torch.Tensor, b: torch.Tensor | None ) -> None: """Decomposes weights into low-rank SVD factors or copies directly.""" if isinstance(target_mod, SVDLinearProjection): with torch.no_grad(): W_f = w.float() u, s, vh = torch.linalg.svd(W_f, full_matrices=False) r = target_mod.rank s_sqrt = torch.diag(torch.sqrt(s[:r])) target_mod.U.data.copy_(torch.matmul(u[:, :r], s_sqrt).to(w.dtype)) target_mod.V.data.copy_(torch.matmul(s_sqrt, vh[:r, :]).to(w.dtype)) if b is not None and target_mod.bias is not None: target_mod.bias.data.copy_(b) else: if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if b is not None and target_mod.bias is not None: target_mod.bias.data.copy_(b) # Alias for API uniformity copy_weights_from_pretrained = copy_weights def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, position_ids: torch.Tensor | None = None, past_key_value: Any | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, Any | None]: B, N, C = hidden_states.shape if torch.is_autocast_enabled() and hidden_states.is_cuda: ac_dtype = ( torch.get_autocast_dtype("cuda") if hasattr(torch, "get_autocast_dtype") else torch.float16 ) hidden_states = hidden_states.to(ac_dtype) if position_embeddings is None: position_embeddings = kwargs.get("position_embeddings") # 1. Project Q, K, V via SVD low-rank factorized projections q = ( self.q_proj(hidden_states) .view(B, N, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) k = ( self.k_proj(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) v = ( self.v_proj(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) # 2. BUG 3 FIX: Apply Rotary Position Embeddings (RoPE) if provided if position_embeddings is not None: cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=position_ids) # 3. BUG 4 FIX: GQA expansion if num_kv_heads < num_heads if self.num_kv_heads != self.num_heads: k = repeat_kv(k, self.num_heads // self.num_kv_heads) v = repeat_kv(v, self.num_heads // self.num_kv_heads) # 4. BUG 2 FIX: Mask handling for padding tokens on K and V if attention_mask is not None: if attention_mask.dim() == 2: mask = attention_mask[:, None, :, None].to(q.dtype) # [B, 1, N, 1] if (mask < 0).any(): mask = (mask > -1000.0).to(q.dtype) k = k * mask v = v * mask elif attention_mask.dim() == 4: if attention_mask.shape[-2] == 1: mask = attention_mask.squeeze(1).squeeze(1)[:, None, :, None].to(q.dtype) else: mask = attention_mask[:, 0, -1, :][:, None, :, None].to(q.dtype) if (mask < 0).any(): mask = (mask > -1000.0).to(q.dtype) k = k * mask v = v * mask # 5. Dispatch Fully Differentiable Triton Linear Attention out = triton_dispatch( self.triton_kernel, q, k, v, is_causal=self.is_decoder, **self.triton_kwargs, ) # 6. Reshape and final output projection out = out.transpose(1, 2).contiguous().view(B, N, C) out = self.out_proj(out) # Numerical safety guard against unbounded linear attention residual accumulation in FP16 if out.dtype == torch.float16: out = torch.clamp(out, min=-65000.0, max=65000.0) return out, NoneTriton-Accelerated Linear Attention with Genuine SVD Low-Rank Projections, RoPE, GQA, and Masked Attention.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var triton_kernel-
The type of the None singleton.
Methods
def copy_weights(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None) k_w, k_b = ( c_w[split_size : 2 * split_size, :], (c_b[split_size : 2 * split_size] if c_b is not None else None), ) v_w, v_b = ( c_w[2 * split_size :, :], (c_b[2 * split_size :] if c_b is not None else None), ) self._copy_single_proj(self.q_proj, q_w, q_b) self._copy_single_proj(self.k_proj, k_w, k_b) self._copy_single_proj(self.v_proj, v_w, v_b) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): w = base_src.weight.data b = getattr(base_src, "bias", None) b_data = b.data if b is not None else None self._copy_single_proj(target_mod, w, b_data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data b = getattr(base_mod, "bias", None) b_data = b.data if b is not None else None if target_name == "out_proj": if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if b_data is not None and target_mod.bias is not None: target_mod.bias.data.copy_(b_data) else: self._copy_single_proj(target_mod, w, b_data) breakTransfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.
def copy_weights_from_pretrained(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None) k_w, k_b = ( c_w[split_size : 2 * split_size, :], (c_b[split_size : 2 * split_size] if c_b is not None else None), ) v_w, v_b = ( c_w[2 * split_size :, :], (c_b[2 * split_size :] if c_b is not None else None), ) self._copy_single_proj(self.q_proj, q_w, q_b) self._copy_single_proj(self.k_proj, k_w, k_b) self._copy_single_proj(self.v_proj, v_w, v_b) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): w = base_src.weight.data b = getattr(base_src, "bias", None) b_data = b.data if b is not None else None self._copy_single_proj(target_mod, w, b_data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data b = getattr(base_mod, "bias", None) b_data = b.data if b is not None else None if target_name == "out_proj": if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if b_data is not None and target_mod.bias is not None: target_mod.bias.data.copy_(b_data) else: self._copy_single_proj(target_mod, w, b_data) breakTransfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.
Inherited members