Module neuralops.attentions.unified_rope
Universal Rotary Position Embedding Attention (UnifiedRoPE). Hardened for production use across diverse transformer architectures: - DistilBERT (Flat MHA) - RoBERTa / BERT (Nested Encoder Self-Attention) - GPT-2 (Causal Decoder with Conv1D) - Qwen / LLaMA (GQA with Rotary Position Embeddings)
Functions
def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None,
unsqueeze_dim: int = 1) ‑> 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, unsqueeze_dim: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: """Applies Rotary Position Embedding (RoPE) with partial dimension and independent seq-len support.""" if cos is None or sin is None: return q, k if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(1) sin = sin.unsqueeze(0).unsqueeze(1) elif cos.dim() == 3: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) while cos.dim() < q.dim(): cos = cos.unsqueeze(0) sin = sin.unsqueeze(0) head_dim = q.shape[-1] rotary_dim = 2 * (min(cos.shape[-1], head_dim, k.shape[-1]) // 2) if rotary_dim == 0: return q, k q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] cos_q = cos[..., : q.shape[-2], :rotary_dim].to(q.device, q.dtype) sin_q = sin[..., : q.shape[-2], :rotary_dim].to(q.device, q.dtype) cos_k = cos[..., : k.shape[-2], :rotary_dim].to(k.device, k.dtype) sin_k = sin[..., : k.shape[-2], :rotary_dim].to(k.device, k.dtype) q_rot_embed = (q_rot * cos_q) + (_rotate_half(q_rot) * sin_q) k_rot_embed = (k_rot * cos_k) + (_rotate_half(k_rot) * sin_k) q = torch.cat([q_rot_embed, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot_embed k = torch.cat([k_rot_embed, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot_embed return q, kApplies Rotary Position Embedding (RoPE) with partial dimension and independent seq-len support.
def inject_unified_rope(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: torch.nn.modules.module.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("unified_rope_inject") @ATTENTION_REGISTRY.register("compression_unified_rope_inject") @ATTENTION_REGISTRY.register("inject_unified_rope") def inject_unified_rope( 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, ): 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) # Whole-model traversal layer_list = None if 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, "transformer") and hasattr(target.transformer, "h"): layer_list = target.transformer.h elif hasattr(target, "h") and isinstance(target.h, (torch.nn.ModuleList, list)): layer_list = target.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, "bert") and hasattr(target.bert, "encoder") and hasattr(target.bert.encoder, "layer") ): layer_list = target.bert.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_unified_rope( layer=lyr, model_config=cfg, comp_config=( comp_config if "comp_config" in locals() or "comp_config" in kwargs else None ), layer_idx=idx, total_layers=tot, **kwargs, ) if res is not None: num_injected += 1 return target if num_injected > 0 else None layer = target """Safe injection hook for UnifiedRoPE.""" try: src = old_attn if src is None: for attr in ("attention", "self_attn", "attn", "self"): if hasattr(layer, attr): src = getattr(layer, attr) break if src is None: return None if model_config is None and hasattr(layer, "config"): model_config = layer.config if model_config is None and hasattr(src, "config"): model_config = src.config lora_rank = kwargs.pop("lora_rank", _cfg_get(comp_config, "lora_rank", 16)) rank = kwargs.pop("rank", _cfg_get(comp_config, "rank", 32)) new_attn = UnifiedRoPE( config=model_config, rank=rank, lora_rank=lora_rank, pretrained_attn=src, **kwargs, ) for attr in ("self_attn", "attn", "self"): if hasattr(layer, attr): setattr(layer, attr, new_attn) return new_attn if hasattr(layer, "attention"): if hasattr(layer.attention, "self"): layer.attention.self = new_attn else: layer.attention = new_attn return new_attn layer.attention = new_attn return new_attn except Exception as e: logger.error("inject_unified_rope failed at layer %d: %s", layer_idx, e) return None
Classes
class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 0,
lora_alpha: float = 1.0,
freeze_base: bool = True)-
Expand source code
class LoRALinear(nn.Module): """Linear projection with frozen base weights and trainable LoRA adapters.""" def __init__( self, in_features: int, out_features: int, bias: bool = True, lora_rank: int = 0, lora_alpha: float = 1.0, freeze_base: bool = True, ): super().__init__() self.in_features = in_features self.out_features = out_features self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.scaling = lora_alpha / max(1, lora_rank) self.weight = nn.Parameter(torch.empty(out_features, in_features)) if bias: self.bias = nn.Parameter(torch.empty(out_features)) else: self.register_parameter("bias", 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 lora_rank > 0: self.lora_A = nn.Parameter(torch.empty(lora_rank, in_features)) self.lora_B = nn.Parameter(torch.empty(out_features, lora_rank)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) self.lora_A.requires_grad = True self.lora_B.requires_grad = True else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) if freeze_base: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False def forward(self, x: torch.Tensor) -> torch.Tensor: orig_dtype = x.dtype w = self.weight.to(orig_dtype) if self.weight.dtype != orig_dtype else self.weight b = ( self.bias.to(orig_dtype) if (self.bias is not None and self.bias.dtype != orig_dtype) else self.bias ) out = F.linear(x, w, b) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: a = self.lora_A.to(orig_dtype) if self.lora_A.dtype != orig_dtype else self.lora_A bb = self.lora_B.to(orig_dtype) if self.lora_B.dtype != orig_dtype else self.lora_B lora_out = F.linear(F.linear(x, a), bb) * self.scaling out = out + lora_out return out.to(orig_dtype)Linear projection with frozen base weights and trainable 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: orig_dtype = x.dtype w = self.weight.to(orig_dtype) if self.weight.dtype != orig_dtype else self.weight b = ( self.bias.to(orig_dtype) if (self.bias is not None and self.bias.dtype != orig_dtype) else self.bias ) out = F.linear(x, w, b) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: a = self.lora_A.to(orig_dtype) if self.lora_A.dtype != orig_dtype else self.lora_A bb = self.lora_B.to(orig_dtype) if self.lora_B.dtype != orig_dtype else self.lora_B lora_out = F.linear(F.linear(x, a), bb) * self.scaling out = out + lora_out return out.to(orig_dtype)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.
class UnifiedRoPE (config: Any | None = None,
rank: int | None = None,
dim: int | None = None,
num_heads: int | None = None,
max_seq_len: int = 4096,
base: int = 10000,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("compression_unified_rope", supports_causal=True, supports_gqa=True) @ATTENTION_REGISTRY.register("unified_rope", supports_causal=True, supports_gqa=True) @ATTENTION_REGISTRY.register("unified_rope_attention", supports_causal=True, supports_gqa=True) class UnifiedRoPE(BaseAttention): r"""Universal Rotary Position Embedding Attention module.""" supports_gqa: bool = True supports_causal: bool = True is_custom_attention: bool = True def __init__( self, config: Any | None = None, rank: int | None = None, dim: int | None = None, num_heads: int | None = None, max_seq_len: int = 4096, base: int = 10000, dropout: float = 0.1, pretrained_attn: nn.Module | None = None, lora_rank: int = 16, lora_alpha: float = 32.0, is_decoder: bool | None = None, **kwargs: Any, ) -> None: super().__init__(config=config, rank=rank) src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn") if dim is None: if config is not None and not isinstance(config, int): dim = _cfg_get( config, "hidden_size", _cfg_get(config, "dim", _cfg_get(config, "d_model", None)), ) if dim is None and src is not None: dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None)) if dim is None and hasattr(src, "q_proj") and hasattr(src.q_proj, "in_features"): dim = src.q_proj.in_features elif dim is None and hasattr(src, "q_lin") and hasattr(src.q_lin, "in_features"): dim = src.q_lin.in_features elif dim is None and hasattr(src, "query") and hasattr(src.query, "in_features"): dim = src.query.in_features elif dim is None and hasattr(src, "c_attn") and hasattr(src.c_attn, "weight"): dim = ( src.c_attn.weight.shape[0] if src.c_attn.weight.shape[0] != 3 * src.c_attn.weight.shape[1] else src.c_attn.weight.shape[1] ) if dim is None: dim = kwargs.get("hidden_size", 768) if num_heads is None: if config is not None and not isinstance(config, int): num_heads = _cfg_get( config, "num_attention_heads", _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)), ) if num_heads is None and src is not None: num_heads = getattr( src, "num_heads", getattr(src, "num_attention_heads", getattr(src, "n_head", None)), ) if num_heads is None: num_heads = 12 self.hidden_size: int = int(dim) self.num_heads: int = max(1, int(num_heads)) self.head_dim: int = self.hidden_size // self.num_heads self.max_seq_len: int = max_seq_len self.base: int = base self.scale: float = 1.0 / math.sqrt(self.head_dim) self.lora_rank: int = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank)) self.lora_alpha: float = kwargs.get( "lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha) ) cfg_kv = ( _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None)) if config is not None else None ) if cfg_kv is None and src is not None: cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None)) self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads self.num_key_value_groups: int = max(1, self.num_heads // max(1, self.num_kv_heads)) self.q_head_size: int = self.num_heads * self.head_dim self.kv_head_size: int = self.num_kv_heads * self.head_dim # Decoder / Causal flag if is_decoder is not None: self.is_decoder: bool = is_decoder elif config is not None: self.is_decoder = bool(_cfg_get(config, "is_decoder", False)) else: self.is_decoder = False self.is_nested_encoder: bool = False self.q_proj = LoRALinear( self.hidden_size, self.q_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.k_proj = LoRALinear( self.hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.v_proj = LoRALinear( self.hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.out_proj = LoRALinear( self.q_head_size, self.hidden_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) dropout_p = ( getattr( config, "dropout", getattr(config, "attention_probs_dropout_prob", dropout), ) if config else dropout ) self.dropout = nn.Dropout(dropout_p) rotary_dim = 2 * (self.head_dim // 2) self.rotary_dim = rotary_dim if rotary_dim > 0: inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim)) else: inv_freq = torch.empty(0) self.register_buffer("inv_freq", inv_freq, persistent=False) self._build_cache(max_seq_len) if pretrained_attn is not None: self.copy_weights(pretrained_attn) def _build_cache(self, seq_len: int) -> None: """Precomputes cosine and sine frequency tables up to seq_len.""" if self.rotary_dim == 0 or self.inv_freq.numel() == 0: self.register_buffer( "cos_cached", torch.ones(seq_len, 0, device=self.inv_freq.device), persistent=False, ) self.register_buffer( "sin_cached", torch.zeros(seq_len, 0, device=self.inv_freq.device), persistent=False, ) return t = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos(), persistent=False) self.register_buffer("sin_cached", emb.sin(), persistent=False) def _apply_rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor: """Applies RoPE rotation: [B, H, L, D_head].""" if self.rotary_dim == 0 or self.cos_cached.shape[-1] == 0: return x if seq_len > self.max_seq_len: self._build_cache(seq_len) self.max_seq_len = seq_len rotary_dim = min(self.rotary_dim, self.cos_cached.shape[-1], x.shape[-1]) rotary_dim = 2 * (rotary_dim // 2) if rotary_dim == 0: return x cos = self.cos_cached[:seq_len, :rotary_dim].unsqueeze(0).unsqueeze(0).to(x.device, x.dtype) sin = self.sin_cached[:seq_len, :rotary_dim].unsqueeze(0).unsqueeze(0).to(x.device, x.dtype) if rotary_dim < x.shape[-1]: x_rot = x[..., :rotary_dim] x_pass = x[..., rotary_dim:] rotated = (x_rot * cos) + (_rotate_half(x_rot) * sin) return torch.cat((rotated, x_pass), dim=-1) else: return (x * cos) + (_rotate_half(x) * sin) def copy_weights(self, old_attn: Any) -> None: """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.""" with torch.no_grad(): w_q, b_q = None, None w_k, b_k = None, None w_v, b_v = None, None w_o, b_o = None, None # Check nested encoder container src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__ if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or ( hasattr(old_attn, "query") and not hasattr(old_attn, "dense") and not hasattr(old_attn, "out_proj") ): self.is_nested_encoder = True if hasattr(old_attn, "self") and hasattr(old_attn, "output"): self.is_nested_encoder = True src = old_attn.self if hasattr(old_attn.output, "dense"): w_o = old_attn.output.dense.weight b_o = old_attn.output.dense.bias else: src = old_attn # Pattern 1: GPT-2 Fused c_attn if hasattr(src, "c_attn") and hasattr(src, "c_proj"): c_attn_w = src.c_attn.weight.detach() c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size: c_attn_w = c_attn_w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size w_q = c_attn_w[:q_dim, :] w_k = c_attn_w[q_dim : q_dim + k_dim, :] w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :] if c_attn_b is not None: b_q = c_attn_b[:q_dim] b_k = c_attn_b[q_dim : q_dim + k_dim] b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim] out_w = src.c_proj.weight.detach() if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size: out_w = out_w.t().contiguous() w_o = out_w b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None # Pattern 2: RoBERTa / BERT separate query, key, value elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"): w_q = src.query.weight.detach() b_q = src.query.bias.detach() if src.query.bias is not None else None w_k = src.key.weight.detach() b_k = src.key.bias.detach() if src.key.bias is not None else None w_v = src.value.weight.detach() b_v = src.value.bias.detach() if src.value.bias is not None else None if w_o is None: for attr in ("dense", "out_proj", "o_proj"): if hasattr(src, attr): proj = getattr(src, attr) w_o = proj.weight.detach() b_o = proj.bias.detach() if proj.bias is not None else None break # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj) elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"): w_q = src.q_proj.weight.detach() b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None w_k = src.k_proj.weight.detach() b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None w_v = src.v_proj.weight.detach() b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None if hasattr(src, "o_proj"): w_o = src.o_proj.weight.detach() b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None elif w_o is None: for attr in ( "out_proj", "out_lin", "dense", "c_proj", "wo", "proj", "output", ): if hasattr(src, attr): proj_layer = getattr(src, attr) if hasattr(proj_layer, "dense"): proj_layer = proj_layer.dense if hasattr(proj_layer, "weight"): w_o = proj_layer.weight.detach() b_o = ( proj_layer.bias.detach() if proj_layer.bias is not None else None ) break # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin) elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"): w_q = src.q_lin.weight.detach() b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None w_k = src.k_lin.weight.detach() b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None w_v = src.v_lin.weight.detach() b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None if hasattr(src, "out_lin"): w_o = src.out_lin.weight.detach() b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None elif w_o is None: for attr in ( "out_proj", "o_proj", "dense", "c_proj", "wo", "proj", "output", ): if hasattr(src, attr): proj_layer = getattr(src, attr) if hasattr(proj_layer, "dense"): proj_layer = proj_layer.dense if hasattr(proj_layer, "weight"): w_o = proj_layer.weight.detach() b_o = ( proj_layer.bias.detach() if proj_layer.bias is not None else None ) break # Pattern 5: UniversalAdapter fallback if w_q is None or w_k is None or w_v is None or w_o is None: try: adapter = UniversalAdapter() if w_q is None: w_q, b_q = adapter.get_weight_bias(src, "q") if w_k is None: w_k, b_k = adapter.get_weight_bias(src, "k") if w_v is None: w_v, b_v = adapter.get_weight_bias(src, "v") if w_o is None: w_o, b_o = adapter.get_weight_bias(src, "o") except Exception as e: logger.debug("UniversalAdapter fallback failed: %s", e) # Strict validation missing = [] if w_q is None: missing.append("Q projection") if w_k is None: missing.append("K projection") if w_v is None: missing.append("V projection") if w_o is None and not self.is_nested_encoder: missing.append("Out projection") if missing: raise RuntimeError( f"UnifiedRoPE: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}" ) def _copy_layer(target_layer, w, b, src_sub=None): if w is not None: if w.shape == target_layer.weight.shape: target_layer.weight.copy_(w.to(target_layer.weight.dtype)) elif w.t().contiguous().shape == target_layer.weight.shape: target_layer.weight.copy_(w.t().contiguous().to(target_layer.weight.dtype)) else: min_0 = min(target_layer.weight.shape[0], w.shape[0]) min_1 = min(target_layer.weight.shape[1], w.shape[1]) target_layer.weight.zero_() target_layer.weight[:min_0, :min_1].copy_( w[:min_0, :min_1].to(target_layer.weight.dtype) ) if b is not None and target_layer.bias is not None: if b.shape == target_layer.bias.shape: target_layer.bias.copy_(b.to(target_layer.bias.dtype)) else: min_b = min(target_layer.bias.shape[0], b.shape[0]) target_layer.bias.zero_() target_layer.bias[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype)) elif target_layer.bias is not None: target_layer.bias.zero_() if ( src_sub is not None and hasattr(src_sub, "lora_A") and hasattr(src_sub, "lora_B") ): if ( src_sub.lora_A is not None and src_sub.lora_B is not None and target_layer.lora_A is not None and target_layer.lora_B is not None ): src_lA = getattr(src_sub.lora_A, "weight", src_sub.lora_A) src_lB = getattr(src_sub.lora_B, "weight", src_sub.lora_B) if hasattr(src_lA, "data") and hasattr(src_lB, "data"): r_min = min(target_layer.lora_A.shape[0], src_lA.shape[0]) in_min = min(target_layer.lora_A.shape[1], src_lA.shape[1]) target_layer.lora_A.data[:r_min, :in_min].copy_( src_lA.data[:r_min, :in_min] ) target_layer.lora_A.requires_grad = True out_min = min(target_layer.lora_B.shape[0], src_lB.shape[0]) r_min_b = min(target_layer.lora_B.shape[1], src_lB.shape[1]) target_layer.lora_B.data[:out_min, :r_min_b].copy_( src_lB.data[:out_min, :r_min_b] ) target_layer.lora_B.requires_grad = True if hasattr(src_sub, "lora_scaling"): target_layer.scaling = src_sub.lora_scaling sub_q = getattr(src, "q_proj", getattr(src, "q_lin", getattr(src, "query", None))) sub_k = getattr(src, "k_proj", getattr(src, "k_lin", getattr(src, "key", None))) sub_v = getattr(src, "v_proj", getattr(src, "v_lin", getattr(src, "value", None))) sub_o = getattr( src, "o_proj", getattr( src, "out_lin", getattr(src, "dense", getattr(src, "out_proj", None)), ), ) _copy_layer(self.q_proj, w_q, b_q, sub_q) _copy_layer(self.k_proj, w_k, b_k, sub_k) _copy_layer(self.v_proj, w_v, b_v, sub_v) if w_o is not None: _copy_layer(self.out_proj, w_o, b_o, sub_o) f_norm_o = float(torch.linalg.norm(self.out_proj.weight)) if f_norm_o < 1e-4: logger.warning(f"UnifiedRoPE: Out projection norm is near zero: {f_norm_o}") 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, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor: """Forward pass computing RoPE attention, returning (output, attention_weights).""" B, L, _ = hidden_states.shape q = ( self.q_proj(hidden_states) .view(B, L, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) k = ( self.k_proj(hidden_states) .view(B, L, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) v = ( self.v_proj(hidden_states) .view(B, L, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) # Apply RoPE pos_emb = kwargs.get("position_embeddings") if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2: cos, sin = pos_emb q, k = apply_rotary_pos_emb(q, k, cos, sin) else: q = self._apply_rope(q, L) k = self._apply_rope(k, L) # GQA expansion if self.num_kv_heads != self.num_heads: k = repeat_kv(k, self.num_key_value_groups) v = repeat_kv(v, self.num_key_value_groups) scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q.float(), k.float().transpose(-1, -2).contiguous()) * scale # Causal mask for decoders if self.is_decoder and L > 1: causal_mask = torch.triu( torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype), diagonal=1, ) scores = scores + causal_mask.unsqueeze(0).unsqueeze(0) mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask") if mask is not None: scores = apply_attention_mask(scores, mask) attn_weights = F.softmax(scores, dim=-1) attn_weights = torch.nan_to_num(attn_weights, nan=0.0) attn_weights = self.dropout(attn_weights.to(v.dtype)) 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.size(0), head_mask.size(1), 1, 1) elif head_mask.dim() == 3: head_mask = head_mask.unsqueeze(-1) elif head_mask.dim() == 4: head_mask = head_mask attn_weights = attn_weights * head_mask.to(attn_weights.dtype) context = torch.matmul(attn_weights, v) context = context.transpose(1, 2).contiguous().view(B, L, self.num_heads * self.head_dim) if self.is_nested_encoder: if context.shape[-1] != self.hidden_size: output = self.out_proj(context) else: output = context else: output = self.out_proj(context) if output.shape[-1] != hidden_states.shape[-1]: output = ( F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1])) if output.shape[-1] < hidden_states.shape[-1] else output[..., : hidden_states.shape[-1]] ) return ( output.to(hidden_states.dtype), (attn_weights if output_attentions else None), )Universal Rotary Position Embedding Attention module.
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 supports_causal : bool-
The type of the None singleton.
var supports_gqa : bool-
The type of the None singleton.
Methods
def copy_weights(self, old_attn: Any) ‑> None-
Expand source code
def copy_weights(self, old_attn: Any) -> None: """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.""" with torch.no_grad(): w_q, b_q = None, None w_k, b_k = None, None w_v, b_v = None, None w_o, b_o = None, None # Check nested encoder container src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__ if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or ( hasattr(old_attn, "query") and not hasattr(old_attn, "dense") and not hasattr(old_attn, "out_proj") ): self.is_nested_encoder = True if hasattr(old_attn, "self") and hasattr(old_attn, "output"): self.is_nested_encoder = True src = old_attn.self if hasattr(old_attn.output, "dense"): w_o = old_attn.output.dense.weight b_o = old_attn.output.dense.bias else: src = old_attn # Pattern 1: GPT-2 Fused c_attn if hasattr(src, "c_attn") and hasattr(src, "c_proj"): c_attn_w = src.c_attn.weight.detach() c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size: c_attn_w = c_attn_w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size w_q = c_attn_w[:q_dim, :] w_k = c_attn_w[q_dim : q_dim + k_dim, :] w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :] if c_attn_b is not None: b_q = c_attn_b[:q_dim] b_k = c_attn_b[q_dim : q_dim + k_dim] b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim] out_w = src.c_proj.weight.detach() if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size: out_w = out_w.t().contiguous() w_o = out_w b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None # Pattern 2: RoBERTa / BERT separate query, key, value elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"): w_q = src.query.weight.detach() b_q = src.query.bias.detach() if src.query.bias is not None else None w_k = src.key.weight.detach() b_k = src.key.bias.detach() if src.key.bias is not None else None w_v = src.value.weight.detach() b_v = src.value.bias.detach() if src.value.bias is not None else None if w_o is None: for attr in ("dense", "out_proj", "o_proj"): if hasattr(src, attr): proj = getattr(src, attr) w_o = proj.weight.detach() b_o = proj.bias.detach() if proj.bias is not None else None break # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj) elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"): w_q = src.q_proj.weight.detach() b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None w_k = src.k_proj.weight.detach() b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None w_v = src.v_proj.weight.detach() b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None if hasattr(src, "o_proj"): w_o = src.o_proj.weight.detach() b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None elif w_o is None: for attr in ( "out_proj", "out_lin", "dense", "c_proj", "wo", "proj", "output", ): if hasattr(src, attr): proj_layer = getattr(src, attr) if hasattr(proj_layer, "dense"): proj_layer = proj_layer.dense if hasattr(proj_layer, "weight"): w_o = proj_layer.weight.detach() b_o = ( proj_layer.bias.detach() if proj_layer.bias is not None else None ) break # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin) elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"): w_q = src.q_lin.weight.detach() b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None w_k = src.k_lin.weight.detach() b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None w_v = src.v_lin.weight.detach() b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None if hasattr(src, "out_lin"): w_o = src.out_lin.weight.detach() b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None elif w_o is None: for attr in ( "out_proj", "o_proj", "dense", "c_proj", "wo", "proj", "output", ): if hasattr(src, attr): proj_layer = getattr(src, attr) if hasattr(proj_layer, "dense"): proj_layer = proj_layer.dense if hasattr(proj_layer, "weight"): w_o = proj_layer.weight.detach() b_o = ( proj_layer.bias.detach() if proj_layer.bias is not None else None ) break # Pattern 5: UniversalAdapter fallback if w_q is None or w_k is None or w_v is None or w_o is None: try: adapter = UniversalAdapter() if w_q is None: w_q, b_q = adapter.get_weight_bias(src, "q") if w_k is None: w_k, b_k = adapter.get_weight_bias(src, "k") if w_v is None: w_v, b_v = adapter.get_weight_bias(src, "v") if w_o is None: w_o, b_o = adapter.get_weight_bias(src, "o") except Exception as e: logger.debug("UniversalAdapter fallback failed: %s", e) # Strict validation missing = [] if w_q is None: missing.append("Q projection") if w_k is None: missing.append("K projection") if w_v is None: missing.append("V projection") if w_o is None and not self.is_nested_encoder: missing.append("Out projection") if missing: raise RuntimeError( f"UnifiedRoPE: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}" ) def _copy_layer(target_layer, w, b, src_sub=None): if w is not None: if w.shape == target_layer.weight.shape: target_layer.weight.copy_(w.to(target_layer.weight.dtype)) elif w.t().contiguous().shape == target_layer.weight.shape: target_layer.weight.copy_(w.t().contiguous().to(target_layer.weight.dtype)) else: min_0 = min(target_layer.weight.shape[0], w.shape[0]) min_1 = min(target_layer.weight.shape[1], w.shape[1]) target_layer.weight.zero_() target_layer.weight[:min_0, :min_1].copy_( w[:min_0, :min_1].to(target_layer.weight.dtype) ) if b is not None and target_layer.bias is not None: if b.shape == target_layer.bias.shape: target_layer.bias.copy_(b.to(target_layer.bias.dtype)) else: min_b = min(target_layer.bias.shape[0], b.shape[0]) target_layer.bias.zero_() target_layer.bias[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype)) elif target_layer.bias is not None: target_layer.bias.zero_() if ( src_sub is not None and hasattr(src_sub, "lora_A") and hasattr(src_sub, "lora_B") ): if ( src_sub.lora_A is not None and src_sub.lora_B is not None and target_layer.lora_A is not None and target_layer.lora_B is not None ): src_lA = getattr(src_sub.lora_A, "weight", src_sub.lora_A) src_lB = getattr(src_sub.lora_B, "weight", src_sub.lora_B) if hasattr(src_lA, "data") and hasattr(src_lB, "data"): r_min = min(target_layer.lora_A.shape[0], src_lA.shape[0]) in_min = min(target_layer.lora_A.shape[1], src_lA.shape[1]) target_layer.lora_A.data[:r_min, :in_min].copy_( src_lA.data[:r_min, :in_min] ) target_layer.lora_A.requires_grad = True out_min = min(target_layer.lora_B.shape[0], src_lB.shape[0]) r_min_b = min(target_layer.lora_B.shape[1], src_lB.shape[1]) target_layer.lora_B.data[:out_min, :r_min_b].copy_( src_lB.data[:out_min, :r_min_b] ) target_layer.lora_B.requires_grad = True if hasattr(src_sub, "lora_scaling"): target_layer.scaling = src_sub.lora_scaling sub_q = getattr(src, "q_proj", getattr(src, "q_lin", getattr(src, "query", None))) sub_k = getattr(src, "k_proj", getattr(src, "k_lin", getattr(src, "key", None))) sub_v = getattr(src, "v_proj", getattr(src, "v_lin", getattr(src, "value", None))) sub_o = getattr( src, "o_proj", getattr( src, "out_lin", getattr(src, "dense", getattr(src, "out_proj", None)), ), ) _copy_layer(self.q_proj, w_q, b_q, sub_q) _copy_layer(self.k_proj, w_k, b_k, sub_k) _copy_layer(self.v_proj, w_v, b_v, sub_v) if w_o is not None: _copy_layer(self.out_proj, w_o, b_o, sub_o) f_norm_o = float(torch.linalg.norm(self.out_proj.weight)) if f_norm_o < 1e-4: logger.warning(f"UnifiedRoPE: Out projection norm is near zero: {f_norm_o}")Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.
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,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor-
Expand source code
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, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor: """Forward pass computing RoPE attention, returning (output, attention_weights).""" B, L, _ = hidden_states.shape q = ( self.q_proj(hidden_states) .view(B, L, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) k = ( self.k_proj(hidden_states) .view(B, L, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) v = ( self.v_proj(hidden_states) .view(B, L, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) # Apply RoPE pos_emb = kwargs.get("position_embeddings") if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2: cos, sin = pos_emb q, k = apply_rotary_pos_emb(q, k, cos, sin) else: q = self._apply_rope(q, L) k = self._apply_rope(k, L) # GQA expansion if self.num_kv_heads != self.num_heads: k = repeat_kv(k, self.num_key_value_groups) v = repeat_kv(v, self.num_key_value_groups) scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q.float(), k.float().transpose(-1, -2).contiguous()) * scale # Causal mask for decoders if self.is_decoder and L > 1: causal_mask = torch.triu( torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype), diagonal=1, ) scores = scores + causal_mask.unsqueeze(0).unsqueeze(0) mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask") if mask is not None: scores = apply_attention_mask(scores, mask) attn_weights = F.softmax(scores, dim=-1) attn_weights = torch.nan_to_num(attn_weights, nan=0.0) attn_weights = self.dropout(attn_weights.to(v.dtype)) 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.size(0), head_mask.size(1), 1, 1) elif head_mask.dim() == 3: head_mask = head_mask.unsqueeze(-1) elif head_mask.dim() == 4: head_mask = head_mask attn_weights = attn_weights * head_mask.to(attn_weights.dtype) context = torch.matmul(attn_weights, v) context = context.transpose(1, 2).contiguous().view(B, L, self.num_heads * self.head_dim) if self.is_nested_encoder: if context.shape[-1] != self.hidden_size: output = self.out_proj(context) else: output = context else: output = self.out_proj(context) if output.shape[-1] != hidden_states.shape[-1]: output = ( F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1])) if output.shape[-1] < hidden_states.shape[-1] else output[..., : hidden_states.shape[-1]] ) return ( output.to(hidden_states.dtype), (attn_weights if output_attentions else None), )Forward pass computing RoPE attention, returning (output, attention_weights).
Inherited members