Module neuralops.attentions.cosformer
Cosformer: Rethinking Softmax in Attention with Linear Complexity.
Implements non-negative kernel feature mapping with cosine/sine positional reweighting, cumulative sum causal linear attention for decoders, and associative accumulation for encoders.
Hardened for production across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA architectures: - Robust pretrained weight extraction (GPT-2 fused c_attn, RoBERTa nested, Qwen GQA, DistilBERT) - Strict Frobenius norm verification (>= 1e-4) for all projections - Grouped Query Attention (GQA) with kv_head_size projections and repeat_kv - Rotary Position Embedding (RoPE) support with partial-coordinate slicing - Strict causal masking for decoder models (vectorized prefix sum) - Nested encoder detection to bypass out_proj in inner self-attention - Parameter-efficient fine-tuning via LoRALinear adapters and frozen base weights - Complete value masking (masking both K and V, including last-chunk V masking) - Numerically stable FP32 accumulation with nan_to_num protection - Safe injection factory returning None on extraction failures - Full cross-attention support - Thread-safe injection depth guard - Even head_dim validation for RoPE - Transparent rank-agnostic behavior (Cosformer does not use rank)
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]: r"""Applies rotary position embeddings with partial-RoPE dimension slicing.""" if cos is None or sin is None: return q, k 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] head_dim = q.shape[-1] if rotary_dim < head_dim: q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_rot_embed = (q_rot * cos) + (_rotate_half(q_rot) * sin) k_rot_embed = (k_rot * cos) + (_rotate_half(k_rot) * sin) q = torch.cat([q_rot_embed, q_pass], dim=-1) k = torch.cat([k_rot_embed, k_pass], dim=-1) else: q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) return q, kApplies rotary position embeddings with partial-RoPE dimension slicing.
def inject_cosformer(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: nn.Module | None = None,
**kwargs: Any) ‑> Any | None-
Expand source code
@ATTENTION_REGISTRY.register("cosformer_inject") @ATTENTION_REGISTRY.register("inject_cosformer") def inject_cosformer( layer: Any = None, model: Any = None, model_config: Any = None, comp_config: Any = None, layer_idx: int = 0, total_layers: int = 12, old_attn: nn.Module | None = None, **kwargs: Any, ) -> Any | None: r"""Safe injection factory for CosformerAttention. Fix 3: Thread-safe recursion depth guard. Fix 13: pop pretrained_attn from kwargs to avoid conflict. Fix 14: Removed dead string expression. Fix 16: Clean kwargs before recursive calls. """ depth = _get_inject_depth() if depth >= _MAX_DEPTH: logger.error("Max injection depth reached in inject_cosformer - possible recursion") return None _set_inject_depth(depth + 1) try: target = layer if layer is not None else model if target is None: target = kwargs.get("model", kwargs.get("layer")) if target is None: return None cfg = model_config if model_config is not None else getattr(target, "config", None) # Fix 16: clean kwargs to avoid conflicts in recursive calls kwargs_clean = { k: v for k, v in kwargs.items() if k not in ("model_config", "comp_config", "layer", "old_attn", "model") } # Fix 13: pop pretrained_attn (we use old_attn explicitly) kwargs_clean.pop("pretrained_attn", None) # Whole-model traversal layer_list = None if _is_single_layer(target): pass # handle below elif hasattr(target, "distilbert") and hasattr(target.distilbert, "transformer"): layer_list = target.distilbert.transformer.layer elif hasattr(target, "roberta") and hasattr(target.roberta, "encoder"): layer_list = target.roberta.encoder.layer elif ( hasattr(target, "bert") and hasattr(target.bert, "encoder") and hasattr(target.bert.encoder, "layer") ): layer_list = target.bert.encoder.layer elif hasattr(target, "transformer") and hasattr(target.transformer, "h"): layer_list = target.transformer.h elif hasattr(target, "model") and hasattr(target.model, "layers"): layer_list = target.model.layers elif hasattr(target, "encoder") and hasattr(target.encoder, "layer"): layer_list = target.encoder.layer elif hasattr(target, "layers") and isinstance(target.layers, (torch.nn.ModuleList, list)): layer_list = target.layers if layer_list is not None: num_injected = 0 tot = len(layer_list) for idx, lyr in enumerate(layer_list): res = inject_cosformer( layer=lyr, model_config=cfg, comp_config=comp_config, layer_idx=idx, total_layers=tot, **kwargs_clean, ) if res is not None: num_injected += 1 return target if num_injected > 0 else None layer = target # Resolve source attention module src_attn = old_attn if src_attn is None and layer is not None: if hasattr(layer, "attention"): src_attn = layer.attention elif hasattr(layer, "attn"): src_attn = layer.attn elif hasattr(layer, "self_attn"): src_attn = layer.self_attn elif hasattr(layer, "self"): src_attn = layer.self # Fix 12: consistent hidden_size default (768, not 64) hidden_size = ( getattr(src_attn, "hidden_size", None) or getattr(src_attn, "embed_dim", None) or ( src_attn.q_proj.in_features if hasattr(src_attn, "q_proj") and hasattr(src_attn.q_proj, "in_features") else None ) or (layer.in_features if hasattr(layer, "in_features") else None) or getattr(model_config, "hidden_size", 768) ) num_heads = ( getattr(src_attn, "num_heads", None) or getattr(src_attn, "num_attention_heads", None) or getattr(model_config, "num_attention_heads", 4) ) try: new_attn = CosformerAttention( config=model_config, comp_config=comp_config, layer_id=layer_idx, layer_idx=layer_idx, old_attn=src_attn, hidden_size=hidden_size, num_heads=num_heads, num_attention_heads=num_heads, **kwargs_clean, ) # Fix 5: do NOT override is_nested_encoder here. # copy_weights (called inside CosformerAttention.__init__) # already sets it correctly based on source architecture. if hasattr(layer, "attention"): if hasattr(layer.attention, "self"): layer.attention.self = new_attn else: layer.attention = new_attn elif hasattr(layer, "attn"): layer.attn = new_attn elif hasattr(layer, "self_attn"): layer.self_attn = new_attn elif hasattr(layer, "self"): layer.self = new_attn return new_attn except Exception as e: logger.error(f"Cosformer copy_weights failed: {e}") return None finally: _set_inject_depth(depth)Safe injection factory for CosformerAttention.
Fix 3: Thread-safe recursion depth guard. Fix 13: pop pretrained_attn from kwargs to avoid conflict. Fix 14: Removed dead string expression. Fix 16: Clean kwargs before recursive calls.
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) ‑> torch.Tensor-
Expand source code
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: r"""Expands grouped key/value heads to match query heads.""" if n_rep == 1: return hidden_states batch, num_kv_heads, slen, head_dim = hidden_states.shape hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_kv_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)Expands grouped key/value heads to match query heads.
Classes
class CosformerAttention (config: Any | None = None,
rank: int | None = None,
max_seq_len: int = 512,
act_fun: str = 'relu',
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register( "cosformer", supports_causal=True, supports_bidirectional=True, supports_gqa=True ) @ATTENTION_REGISTRY.register( "cosformer_attention", supports_causal=True, supports_bidirectional=True, supports_gqa=True, ) class CosformerAttention(BaseAttention): r"""Cosformer Attention with Linear Complexity and Exact Multi-Architecture Support.""" is_custom_attention: bool = True triton_kernel = "linear_attn" def __init__( self, config: Any | None = None, rank: int | None = None, max_seq_len: int = 512, act_fun: str = "relu", eps: float = 1e-6, pretrained_attn: nn.Module | None = None, old_attn: nn.Module | None = None, comp_config: Any | None = None, layer_id: int | None = None, layer_idx: int | None = None, **kwargs: Any, ) -> None: source_attn = pretrained_attn if pretrained_attn is not None else old_attn def _cfg_get(k: str, default: Any = None) -> Any: for obj in (config, comp_config): if obj is not None and hasattr(obj, k): val = getattr(obj, k) if val is not None: return val return kwargs.get(k, default) hidden_size = kwargs.get("hidden_size") num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads") if hidden_size is None: hidden_size = _cfg_get("hidden_size", _cfg_get("dim", None)) if hidden_size is None and source_attn is not None: out_layer = getattr( source_attn, "out_lin", getattr( source_attn, "dense", getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)), ), ) hidden_size = ( getattr(out_layer, "out_features", 768) if out_layer is not None else 768 ) elif hidden_size is None: hidden_size = 768 if num_heads is None: num_heads = _cfg_get("num_attention_heads", _cfg_get("num_heads", None)) if num_heads is None and source_attn is not None: num_heads = getattr( source_attn, "n_heads", getattr( source_attn, "num_attention_heads", getattr(source_attn, "num_heads", 12), ), ) elif num_heads is None: num_heads = 12 super().__init__(config, rank) self.dim: int = int(hidden_size) self.hidden_size: int = self.dim self.num_heads: int = int(num_heads) resolved_head_dim = _cfg_get("head_dim", None) if resolved_head_dim is None and source_attn is not None: resolved_head_dim = getattr(source_attn, "head_dim", None) if resolved_head_dim is None: resolved_head_dim = self.dim // max(1, self.num_heads) self.head_dim: int = int(resolved_head_dim) # Fix 11: even head_dim required for RoPE if self.head_dim % 2 != 0: raise ValueError( f"CosformerAttention with RoPE requires even head_dim, got head_dim={self.head_dim}" ) kv_heads = _cfg_get("num_key_value_heads", None) if kv_heads is None and source_attn is not None: kv_heads = getattr( source_attn, "num_key_value_heads", getattr( source_attn, "num_kv_heads", getattr(source_attn, "n_kv_heads", None), ), ) if ( kv_heads is None and hasattr(source_attn, "k_proj") and hasattr(source_attn.k_proj, "out_features") and self.head_dim > 0 ): kv_heads = source_attn.k_proj.out_features // self.head_dim self.num_kv_heads: int = int(kv_heads if kv_heads is not None else self.num_heads) self.all_head_size: int = self.num_heads * self.head_dim self.all_kv_head_size: int = self.num_kv_heads * self.head_dim self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads)) # Fix 9: documented max_seq_len default self.max_seq_len: int = int(_cfg_get("max_position_embeddings", max_seq_len)) self.act_fun: str = act_fun self.eps: float = float(eps) # Static layer index resolved_layer_id = ( layer_id if layer_id is not None else ( layer_idx if layer_idx is not None else kwargs.get("layer_id", kwargs.get("layer_idx", 0)) ) ) self.layer_idx: int = int(resolved_layer_id) # LoRA adaptation configuration self.lora_rank: int = int(_cfg_get("lora_rank", kwargs.get("lora_rank", 16))) self.lora_alpha: float = float(_cfg_get("lora_alpha", kwargs.get("lora_alpha", 32.0))) self.freeze_base: bool = self.lora_rank > 0 # Decoder & nested encoder detection model_type = getattr(config, "model_type", "").lower() if config is not None else "" is_dec = _cfg_get("is_decoder", _cfg_get("is_causal", kwargs.get("is_causal", False))) _decoder_prefixes = ( "gpt", "qwen", "llama", "mistral", "phi", "gemma", "opt", "bloom", "falcon", ) if ( any(model_type.startswith(p) for p in _decoder_prefixes) or hasattr(config, "architectures") and any( "GPT" in str(a) or "CausalLM" in str(a) for a in (getattr(config, "architectures", []) or []) ) ): is_dec = True if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False): is_dec = True self.is_decoder: bool = bool(is_dec) self.is_causal: bool = bool(is_dec) self.is_nested_encoder: bool = bool(kwargs.get("is_nested_encoder", False)) # Multi-head projections equipped with LoRA adapters self.q_proj = LoRALinear( self.dim, self.all_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.k_proj = LoRALinear( self.dim, self.all_kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.v_proj = LoRALinear( self.dim, self.all_kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.out_proj = LoRALinear( self.all_head_size, self.dim, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) dropout_p = ( getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1)) if config else 0.1 ) self.dropout = nn.Dropout(dropout_p) self.rotary_emb = None self.comp_config: Any | None = comp_config self.triton_kwargs = { "phi": "elu+1", "eps": self.eps, "cos_weight": None, "sin_weight": None, } if source_attn is not None: self.copy_weights(source_attn) # ------------------------------------------------------------------ # Compatibility properties # ------------------------------------------------------------------ @property def query(self) -> LoRALinear: return self.q_proj @property def key(self) -> LoRALinear: return self.k_proj @property def value(self) -> LoRALinear: return self.v_proj @property def dense(self) -> LoRALinear: return self.out_proj @property def o_proj(self) -> LoRALinear: return self.out_proj @property def q_lin(self) -> LoRALinear: return self.q_proj @property def k_lin(self) -> LoRALinear: return self.k_proj @property def v_lin(self) -> LoRALinear: return self.v_proj @property def out_lin(self) -> LoRALinear: return self.out_proj @property def W_q(self) -> LoRALinear: return self.q_proj @property def W_k(self) -> LoRALinear: return self.k_proj @property def W_v(self) -> LoRALinear: return self.v_proj @property def c_proj(self) -> LoRALinear: return self.out_proj # ------------------------------------------------------------------ # Weight transfer # ------------------------------------------------------------------ def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights with strict Frobenius norm verification.""" if src is None: return cls_name = getattr(src, "__class__", type(src)).__name__ has_out = ( hasattr(src, "dense") or hasattr(src, "out_proj") or hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") ) if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or ( hasattr(src, "query") and not has_out ): self.is_nested_encoder = True if hasattr(src, "self") and hasattr(src, "output"): self.is_nested_encoder = True if hasattr(src, "rotary_emb"): self.rotary_emb = src.rotary_emb extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2: Fused c_attn + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False): try: c_attn_mod.merge() except Exception as e: logger.debug(f"c_attn_mod.merge failed: {e}") 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 extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) 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() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) base_mod = getattr(mod, "base_layer", mod) if base_mod is not None and hasattr(base_mod, "weight"): extracted[proj_key] = ( base_mod.weight.data.clone(), ( base_mod.bias.data.clone() if getattr(base_mod, "bias", None) is not None else None ), ) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) extracted["o"] = ( base_out.weight.data.clone(), ( base_out.bias.data.clone() if getattr(base_out, "bias", None) is not None else None ), ) # 3. Standard linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q": ["q_lin", "q_proj", "query", "W_q"], "k": ["k_lin", "k_proj", "key", "W_k"], "v": ["v_lin", "v_proj", "value", "W_v"], "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) base_mod = getattr(mod, "base_layer", mod) if hasattr(base_mod, "weight"): w_data = base_mod.weight.data.clone() if ( (base_mod.__class__.__name__ == "Conv1D") or hasattr(base_mod, "nf") or ( w_data.ndim == 2 and proj_key == "o" and w_data.shape[0] < w_data.shape[1] ) ): w_data = w_data.t().contiguous() b_data = getattr(base_mod, "bias", None) extracted[proj_key] = ( w_data, b_data.data.clone() if b_data is not None else None, ) break # Fallback via UniversalAdapter missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug(f"UniversalAdapter fallback: {e}") missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: raise RuntimeError( f"CosformerAttention copy_weights: missing mandatory projections " f"{missing_after} from {src.__class__.__name__}." ) target_map = { "q": self.q_proj, "k": self.k_proj, "v": self.v_proj, "o": self.out_proj, } with torch.no_grad(): for proj_key, target_mod in target_map.items(): w, b = extracted[proj_key] if w is None: if proj_key == "o" and self.is_nested_encoder: nn.init.eye_(target_mod.weight) if target_mod.bias is not None: nn.init.zeros_(target_mod.bias) continue if w.ndim != 2: w = w.reshape(w.size(0), -1) if ( w.shape != target_mod.weight.shape and target_mod.weight.shape == w.t().contiguous().shape ): w = w.t().contiguous() out_d = min(target_mod.weight.size(0), w.size(0)) in_d = min(target_mod.weight.size(1), w.size(1)) target_mod.weight.data[:out_d, :in_d].copy_( w[:out_d, :in_d].to( device=target_mod.weight.device, dtype=target_mod.weight.dtype ) ) if b is not None and target_mod.bias is not None: min_b = min(target_mod.bias.numel(), b.numel()) target_mod.bias.data[:min_b].copy_( b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype) ) elif target_mod.bias is not None: nn.init.zeros_(target_mod.bias) # Frobenius norm verification (>= 1e-4) w_norm = float(torch.norm(target_mod.weight.data).item()) if w_norm < 1e-4: logger.warning( f"CosformerAttention weight copy verification failed " f"for '{proj_key}': norm too small ({w_norm:.2e})" ) if self.freeze_base: target_mod.freeze_base_weights() # ------------------------------------------------------------------ # Positional weighting # ------------------------------------------------------------------ def _get_cos_sin_weights( self, seq_len: int, device: torch.device, dtype: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: r"""Generate positional cosine and sine factors: [1, 1, seq_len, 1]. Fix 6: Uses seq_len directly (simpler and more numerically stable). """ pos = torch.arange(seq_len, device=device, dtype=torch.float32) m = max(1, int(seq_len)) angles = (math.pi / (2.0 * m)) * pos cos_weight = torch.cos(angles).view(1, 1, seq_len, 1).to(dtype=dtype) sin_weight = torch.sin(angles).view(1, 1, seq_len, 1).to(dtype=dtype) return cos_weight, sin_weight # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, past_key_value: Any | None = None, output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, position_ids: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: # ------------------------------------------------------------------ # 0. Input validation # ------------------------------------------------------------------ assert hidden_states.dim() == 3, ( f"CosformerAttention expects 3D hidden_states [B, N, D], " f"got shape {tuple(hidden_states.shape)}" ) B, N, D = hidden_states.shape H, d = self.num_heads, self.head_dim device = hidden_states.device input_dtype = hidden_states.dtype is_cross_attention = encoder_hidden_states is not None kv_input = encoder_hidden_states if is_cross_attention else hidden_states kv_len = kv_input.shape[1] # ------------------------------------------------------------------ # 1. Q/K/V projections # ------------------------------------------------------------------ Q = self.q_proj(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous() K = self.k_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous() V = self.v_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous() # ------------------------------------------------------------------ # 2. RoPE (only for self-attention, not cross-attention) # ------------------------------------------------------------------ pos_emb = ( position_embeddings if position_embeddings is not None else kwargs.get("position_embeddings") ) if pos_emb is not None and not is_cross_attention: if isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2: cos, sin = pos_emb[0], pos_emb[1] Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids) else: logger.warning( "CosformerAttention: unexpected position_embeddings format (%s), skipping RoPE.", type(pos_emb), ) elif self.rotary_emb is not None and not is_cross_attention: rotary_mod = self.rotary_emb pos_ids = position_ids if position_ids is not None else kwargs.get("position_ids") if pos_ids is None: pos_ids = torch.arange(N, device=device).unsqueeze(0).expand(B, -1) cos, sin = None, None try: cos, sin = rotary_mod(Q, seq_len=kv_len) except Exception: try: cos, sin = rotary_mod(Q, pos_ids) except Exception: try: cos, sin = rotary_mod(Q) except Exception: pass if cos is not None and sin is not None: Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=pos_ids) # ------------------------------------------------------------------ # 3. GQA head expansion # ------------------------------------------------------------------ if self.num_kv_heads != self.num_heads: K = repeat_kv(K, self.kv_repeat) V = repeat_kv(V, self.kv_repeat) # ------------------------------------------------------------------ # 4. Kernel feature mapping (ELU+1) — in FP32 for stability # ------------------------------------------------------------------ if self.act_fun == "softplus": Q_k = F.softplus(Q.float()) K_k = F.softplus(K.float()) else: # default: elu + 1 → non-negative kernel feature map Q_k = F.elu(Q.float()) + 1.0 K_k = F.elu(K.float()) + 1.0 V_f = V.float() # ------------------------------------------------------------------ # 5. Token valid mask on both K and V (padding tokens → 0) # ------------------------------------------------------------------ eff_mask = encoder_attention_mask if is_cross_attention else attention_mask valid_mask = None if eff_mask is not None: valid_mask = get_token_valid_mask(eff_mask, B, kv_len, device) # [B, kv_len] bool mask_expanded = valid_mask.unsqueeze(1).unsqueeze(-1).float() # [B, 1, kv_len, 1] K_k = K_k * mask_expanded V_f = V_f * mask_expanded # ------------------------------------------------------------------ # 6. Cos/Sin positional reweighting (cosformer's signature) # ------------------------------------------------------------------ cos_w, sin_w = self._get_cos_sin_weights(N, device, dtype=torch.float32) # cos_w, sin_w: [1, 1, N, 1] Q_cos = Q_k * cos_w Q_sin = Q_k * sin_w K_cos = K_k * cos_w K_sin = K_k * sin_w # ------------------------------------------------------------------ # 7. Linear attention — try Triton, fallback to stable Python path # ------------------------------------------------------------------ is_causal_eval = bool(self.is_decoder and not is_cross_attention and kv_len == N) context = None triton_ok = False if triton_dispatch is not None and not is_cross_attention: try: # Triton path expects FP16 inputs. Force FP16 only if values are bounded. # ELU+1 can produce values up to ~10, times cos/sin up to 1 → safe. Q_t = Q_cos.half() K_t = K_cos.half() K_s_t = K_sin.half() V_t = V_f.half() # Quick sanity check for overflow before calling Triton if ( torch.isnan(Q_t).any() or torch.isnan(K_t).any() or torch.isinf(Q_t).any() or torch.isinf(K_t).any() ): raise RuntimeError("Non-finite values detected before Triton dispatch") dispatch_kwargs = dict(self.triton_kwargs) dispatch_kwargs["cos_weight"] = cos_w.half() dispatch_kwargs["sin_weight"] = sin_w.half() dispatch_kwargs["eps"] = self.eps # Some dispatchers expect (Q_cos, Q_sin, K_cos, K_sin, V) context = triton_dispatch( self.triton_kernel, Q_t, K_t, V_t, Q_sin=Q_k.half() * sin_w.half(), K_sin=K_s_t, is_causal=is_causal_eval, **dispatch_kwargs, ) context = context.float() triton_ok = True except Exception as e: logger.debug("Triton dispatch failed, using Python fallback: %s", e) triton_ok = False # ------------------------------------------------------------------ # 7b. Python fallback: stable causal / bidirectional linear attention # ------------------------------------------------------------------ if not triton_ok or context is None: if is_cross_attention: # Cross-attention: only bidirectional path (no causal masking) context = self._forward_bidirectional( Q_cos, Q_sin, K_cos, K_sin, V_f, B, H, d, kv_len, device, torch.float32, ) elif is_causal_eval: # Causal self-attention for decoders context = self._forward_causal( Q_cos, Q_sin, K_cos, K_sin, V_f, B, H, d, N, device, torch.float32, ) else: # Bidirectional self-attention for encoders context = self._forward_bidirectional( Q_cos, Q_sin, K_cos, K_sin, V_f, B, H, d, N, device, torch.float32, ) # ------------------------------------------------------------------ # 8. Safety: kill any nan/inf before head mask and out_proj # ------------------------------------------------------------------ context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4) # ------------------------------------------------------------------ # 9. Head mask # ------------------------------------------------------------------ if head_mask is not None: if head_mask.dim() == 1: head_mask = head_mask.view(1, -1, 1, 1) elif head_mask.dim() == 2: head_mask = head_mask.view(head_mask.shape[0], head_mask.shape[1], 1, 1) elif head_mask.dim() == 3: head_mask = head_mask.unsqueeze(-1) context = context * head_mask # ------------------------------------------------------------------ # 10. Merge heads → [B, N, H*d] # ------------------------------------------------------------------ context = context.to(input_dtype) context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size) # ------------------------------------------------------------------ # 11. Output projection (or nested encoder bypass) # ------------------------------------------------------------------ if self.is_nested_encoder: final_out = context else: final_out = self.out_proj(context) # ------------------------------------------------------------------ # 12. Final safety clamp — prevents logit explosion downstream # (adjust bounds if you know your task needs larger values) # ------------------------------------------------------------------ final_out = torch.nan_to_num(final_out, nan=0.0, posinf=1e4, neginf=-1e4) final_out = final_out.clamp(min=-1e4, max=1e4) return final_out.to(input_dtype), None # ------------------------------------------------------------------ # Causal forward (Fix 8: V masked in last chunk) # ------------------------------------------------------------------ def _forward_causal( self, Q_cos: torch.Tensor, Q_sin: torch.Tensor, K_cos: torch.Tensor, K_sin: torch.Tensor, V: torch.Tensor, B: int, H: int, d: int, N: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Chunked causal linear attention. Stable in FP32.""" chunk_size = max(16, min(64, N // 8)) if chunk_size >= N: # Single-chunk: explicit causal mask causal_mask = torch.tril(torch.ones(N, N, device=device, dtype=torch.bool)) score_c = torch.matmul(Q_cos, K_cos.transpose(-1, -2).contiguous()) score_s = torch.matmul(Q_sin, K_sin.transpose(-1, -2).contiguous()) scores = (score_c + score_s).masked_fill(~causal_mask, 0.0) num = torch.matmul(scores, V) denom = (score_c + score_s).masked_fill(~causal_mask, 0.0).sum(dim=-1, keepdim=True) return (num / denom.clamp(min=self.eps)).to(dtype) # Multi-chunk pad_len = (chunk_size - (N % chunk_size)) % chunk_size if pad_len > 0: Q_cos = F.pad(Q_cos, (0, 0, 0, pad_len)) Q_sin = F.pad(Q_sin, (0, 0, 0, pad_len)) K_cos = F.pad(K_cos, (0, 0, 0, pad_len)) K_sin = F.pad(K_sin, (0, 0, 0, pad_len)) V = F.pad(V, (0, 0, 0, pad_len)) total_N = N + pad_len num_chunks = total_N // chunk_size qc_c = Q_cos.view(B, H, num_chunks, chunk_size, d) qs_c = Q_sin.view(B, H, num_chunks, chunk_size, d) kc_c = K_cos.view(B, H, num_chunks, chunk_size, d) ks_c = K_sin.view(B, H, num_chunks, chunk_size, d) v_c = V.view(B, H, num_chunks, chunk_size, d) causal_mask = torch.tril( torch.ones(chunk_size, chunk_size, device=device, dtype=torch.bool) ) score_c = torch.matmul(qc_c, kc_c.transpose(-1, -2).contiguous()).masked_fill( ~causal_mask, 0.0 ) score_s = torch.matmul(qs_c, ks_c.transpose(-1, -2).contiguous()).masked_fill( ~causal_mask, 0.0 ) num_intra = torch.matmul(score_c + score_s, v_c) denom_intra = (score_c + score_s).sum(dim=-1, keepdim=True) S_cos = torch.zeros(B, H, d, d, device=device, dtype=torch.float32) S_sin = torch.zeros(B, H, d, d, device=device, dtype=torch.float32) Z_cos = torch.zeros(B, H, d, 1, device=device, dtype=torch.float32) Z_sin = torch.zeros(B, H, d, 1, device=device, dtype=torch.float32) num_list, denom_list = [], [] for c in range(num_chunks): num_inter = torch.matmul(qc_c[:, :, c], S_cos) + torch.matmul(qs_c[:, :, c], S_sin) denom_inter = torch.matmul(qc_c[:, :, c], Z_cos) + torch.matmul(qs_c[:, :, c], Z_sin) num_list.append(num_intra[:, :, c] + num_inter) denom_list.append(denom_intra[:, :, c] + denom_inter) kc_curr = kc_c[:, :, c] ks_curr = ks_c[:, :, c] v_curr = v_c[:, :, c] if c == num_chunks - 1 and pad_len > 0: valid = chunk_size - pad_len kc_curr = kc_curr.clone() kc_curr[:, :, valid:] = 0.0 ks_curr = ks_curr.clone() ks_curr[:, :, valid:] = 0.0 v_curr = v_curr.clone() v_curr[:, :, valid:] = 0.0 S_cos = S_cos + torch.matmul(kc_curr.transpose(-1, -2).contiguous(), v_curr) S_sin = S_sin + torch.matmul(ks_curr.transpose(-1, -2).contiguous(), v_curr) Z_cos = Z_cos + kc_curr.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous() Z_sin = Z_sin + ks_curr.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous() num = torch.cat(num_list, dim=2)[:, :, :N, :] denom = torch.cat(denom_list, dim=2)[:, :, :N, :] return (num / (denom + self.eps).clamp(min=self.eps)).to(dtype) # ------------------------------------------------------------------ # Bidirectional forward # ------------------------------------------------------------------ def _forward_bidirectional( self, Q_cos: torch.Tensor, Q_sin: torch.Tensor, K_cos: torch.Tensor, K_sin: torch.Tensor, V: torch.Tensor, B: int, H: int, d: int, N: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: KV_cos = torch.matmul(K_cos.transpose(-2, -1).contiguous(), V) KV_sin = torch.matmul(K_sin.transpose(-2, -1).contiguous(), V) Z_cos = K_cos.sum(dim=2, keepdim=True) # [B, H, 1, d] Z_sin = K_sin.sum(dim=2, keepdim=True) num = torch.matmul(Q_cos, KV_cos) + torch.matmul(Q_sin, KV_sin) denom = torch.matmul(Q_cos, Z_cos.transpose(-2, -1).contiguous()) + torch.matmul( Q_sin, Z_sin.transpose(-2, -1).contiguous() ) denom = denom.clamp(min=self.eps) context = (num / denom).to(dtype) return torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)Cosformer Attention with Linear Complexity and Exact Multi-Architecture Support.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var is_custom_attention : bool-
The type of the None singleton.
var triton_kernel-
The type of the None singleton.
Instance variables
prop W_k : LoRALinear-
Expand source code
@property def W_k(self) -> LoRALinear: return self.k_proj prop W_q : LoRALinear-
Expand source code
@property def W_q(self) -> LoRALinear: return self.q_proj prop W_v : LoRALinear-
Expand source code
@property def W_v(self) -> LoRALinear: return self.v_proj prop c_proj : LoRALinear-
Expand source code
@property def c_proj(self) -> LoRALinear: return self.out_proj prop dense : LoRALinear-
Expand source code
@property def dense(self) -> LoRALinear: return self.out_proj prop k_lin : LoRALinear-
Expand source code
@property def k_lin(self) -> LoRALinear: return self.k_proj prop key : LoRALinear-
Expand source code
@property def key(self) -> LoRALinear: return self.k_proj prop o_proj : LoRALinear-
Expand source code
@property def o_proj(self) -> LoRALinear: return self.out_proj prop out_lin : LoRALinear-
Expand source code
@property def out_lin(self) -> LoRALinear: return self.out_proj prop q_lin : LoRALinear-
Expand source code
@property def q_lin(self) -> LoRALinear: return self.q_proj prop query : LoRALinear-
Expand source code
@property def query(self) -> LoRALinear: return self.q_proj prop v_lin : LoRALinear-
Expand source code
@property def v_lin(self) -> LoRALinear: return self.v_proj prop value : LoRALinear-
Expand source code
@property def value(self) -> LoRALinear: return self.v_proj
Methods
def copy_weights(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights with strict Frobenius norm verification.""" if src is None: return cls_name = getattr(src, "__class__", type(src)).__name__ has_out = ( hasattr(src, "dense") or hasattr(src, "out_proj") or hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") ) if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or ( hasattr(src, "query") and not has_out ): self.is_nested_encoder = True if hasattr(src, "self") and hasattr(src, "output"): self.is_nested_encoder = True if hasattr(src, "rotary_emb"): self.rotary_emb = src.rotary_emb extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2: Fused c_attn + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False): try: c_attn_mod.merge() except Exception as e: logger.debug(f"c_attn_mod.merge failed: {e}") 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 extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) 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() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) base_mod = getattr(mod, "base_layer", mod) if base_mod is not None and hasattr(base_mod, "weight"): extracted[proj_key] = ( base_mod.weight.data.clone(), ( base_mod.bias.data.clone() if getattr(base_mod, "bias", None) is not None else None ), ) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) extracted["o"] = ( base_out.weight.data.clone(), ( base_out.bias.data.clone() if getattr(base_out, "bias", None) is not None else None ), ) # 3. Standard linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q": ["q_lin", "q_proj", "query", "W_q"], "k": ["k_lin", "k_proj", "key", "W_k"], "v": ["v_lin", "v_proj", "value", "W_v"], "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) base_mod = getattr(mod, "base_layer", mod) if hasattr(base_mod, "weight"): w_data = base_mod.weight.data.clone() if ( (base_mod.__class__.__name__ == "Conv1D") or hasattr(base_mod, "nf") or ( w_data.ndim == 2 and proj_key == "o" and w_data.shape[0] < w_data.shape[1] ) ): w_data = w_data.t().contiguous() b_data = getattr(base_mod, "bias", None) extracted[proj_key] = ( w_data, b_data.data.clone() if b_data is not None else None, ) break # Fallback via UniversalAdapter missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug(f"UniversalAdapter fallback: {e}") missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: raise RuntimeError( f"CosformerAttention copy_weights: missing mandatory projections " f"{missing_after} from {src.__class__.__name__}." ) target_map = { "q": self.q_proj, "k": self.k_proj, "v": self.v_proj, "o": self.out_proj, } with torch.no_grad(): for proj_key, target_mod in target_map.items(): w, b = extracted[proj_key] if w is None: if proj_key == "o" and self.is_nested_encoder: nn.init.eye_(target_mod.weight) if target_mod.bias is not None: nn.init.zeros_(target_mod.bias) continue if w.ndim != 2: w = w.reshape(w.size(0), -1) if ( w.shape != target_mod.weight.shape and target_mod.weight.shape == w.t().contiguous().shape ): w = w.t().contiguous() out_d = min(target_mod.weight.size(0), w.size(0)) in_d = min(target_mod.weight.size(1), w.size(1)) target_mod.weight.data[:out_d, :in_d].copy_( w[:out_d, :in_d].to( device=target_mod.weight.device, dtype=target_mod.weight.dtype ) ) if b is not None and target_mod.bias is not None: min_b = min(target_mod.bias.numel(), b.numel()) target_mod.bias.data[:min_b].copy_( b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype) ) elif target_mod.bias is not None: nn.init.zeros_(target_mod.bias) # Frobenius norm verification (>= 1e-4) w_norm = float(torch.norm(target_mod.weight.data).item()) if w_norm < 1e-4: logger.warning( f"CosformerAttention weight copy verification failed " f"for '{proj_key}': norm too small ({w_norm:.2e})" ) if self.freeze_base: target_mod.freeze_base_weights()Transfer pretrained projection weights with strict Frobenius norm verification.
Inherited members
class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = False,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class LoRALinear(nn.Module): r"""Linear projection layer with frozen base weights and trainable residual LoRA adapters.""" def __init__( self, in_features: int, out_features: int, bias: bool = True, lora_rank: int = 16, lora_alpha: float = 32.0, freeze_base: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() factory_kwargs = {"device": device, "dtype": dtype} self.in_features = in_features self.out_features = out_features self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0 self.freeze_base = freeze_base self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs)) if bias: self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter("bias", None) if self.lora_rank > 0: self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs)) self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs)) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() if freeze_base: self.freeze_base_weights() def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def freeze_base_weights(self) -> None: """Freeze base projection parameters so only LoRA parameters adapt.""" self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0 and self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: self.lora_B.requires_grad = True def unfreeze_base_weights(self) -> None: """Unfreeze base projection parameters.""" self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = True def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling out = out + lora_out return outLinear projection layer with frozen base weights and trainable residual LoRA adapters.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling out = out + lora_out return outDefine 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 freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: """Freeze base projection parameters so only LoRA parameters adapt.""" self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0 and self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: self.lora_B.requires_grad = TrueFreeze base projection parameters so only LoRA parameters adapt.
def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def unfreeze_base_weights(self) ‑> None-
Expand source code
def unfreeze_base_weights(self) -> None: """Unfreeze base projection parameters.""" self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = TrueUnfreeze base projection parameters.