⚙️ 记录一下transformers中的Llama是怎么写attention的，并介绍如何无侵入式修改Llama为双向Llama，并取消RoPE旋转位置编码
class LlamaAttention(nn.Module):
"""Multi-headed attention from 'Attention Is All You Need' paper"""
def __init__(self, config: LlamaConfig, layer_idx: int):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
self.scaling = self.head_dim ** -0.5
self.attention_dropout = config.attention_dropout
self.is_causal = True
self.q_proj = nn.Linear(
config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
)
self.k_proj = nn.Linear(
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
)
self.v_proj = nn.Linear(
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
)
self.o_proj = nn.Linear(
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
)
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
def forward(
self,
hidden_states: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor],
attention_mask: Optional[torch.Tensor],
past_key_values: Optional[Cache] = None,
cache_position: Optional[torch.LongTensor] = None,
**kwargs: Unpack[TransformersKwargs],
) -> tuple[torch.Tensor, torch.Tensor]:
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, self.head_dim)
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
if past_key_values is not None:
# sin and cos are specific to RoPE models; cache_position needed for the static cache
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
attention_interface: Callable = eager_attention_forward
if self.config._attn_implementation != "eager":
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
attn_output, attn_weights = attention_interface(
self,
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
**kwargs,
)
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
attn_output = self.o_proj(attn_output)
return attn_output, attn_weights
def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: Optional[torch.Tensor],
scaling: float,
dropout: float = 0.0,
**kwargs: Unpack[TransformersKwargs],
):
key_states = repeat_kv(key, module.num_key_value_groups)
value_states = repeat_kv(value, module.num_key_value_groups)
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
if attention_mask is not None:
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
attn_weights = attn_weights + causal_mask
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()
return attn_output, attn_weights
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
"""
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
"""
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
- 很多其他模型的attention几乎都是在llama attention基础上改的（例如Qwen3，添加了滑动窗口），transformers的好处是不同实现方法的attention（eager，sdpa，flash等）都写了统一的函数接口，调用起来比较方便
🥺 生物上，有很多模型其实是不需要位置编码的，注意力也是双向的，有时候输入也可能是多个序列，比如：以单细胞基因扰动预测任务为例，模型的输入不仅仅是control细胞的基因表达，同时也要融入被敲除基因的embedding，如果想直接调用transformers写的Llama模型，可能不太方便，因为有旋转位置编码要取消，模型的输入也要添加，attention_mask也要改成双向的。最近的虚拟细胞挑战赛Arc Institute对STATE的写法确实很精彩，也让我学习到了很多，官方以一种无侵入式的方式将原始的Llama改成了双向Llama，还添加了额外的输入。
- refer to：https://github.com/ArcInstitute/state/blob/main/src/state/tx/models/utils.py
- 这种方式虽然写起来很优雅，但是如果transformers库对Llama做更改了的话，维护起来就比较麻烦了，还是更推荐从头自己写
import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.models.llama.modeling_llama import LlamaModel
from transformers.models.llama.configuration_llama import LlamaConfig
class NoRoPE(nn.Module):
"""
A drop-in replacement for LlamaRotaryEmbedding that always returns:
cos = all ones, sin = all zeros
of shape (batch_size, seq_len, head_dim), so rotary has no effect.
"""
def __init__(self, head_dim: int):
super().__init__()
self.head_dim = head_dim
def forward(self, hidden_states: torch.Tensor, position_ids: torch.LongTensor):
# hidden_states: (batch_size, seq_len, hidden_dim)
batch_size, seq_len, _hidden_dim = hidden_states.shape
# Create cos = ones, sin = zeros
# shape --> (batch_size, seq_len, head_dim)
cos = hidden_states.new_ones(batch_size, seq_len, self.head_dim)
sin = hidden_states.new_zeros(batch_size, seq_len, self.head_dim)
return cos, sin
class LlamaBidirectionalModel(LlamaModel):
"""
A drop-in replacement for LlamaModel with bidirectional attention.
By overriding _update_causal_mask to return None, all tokens attend to each other.
"""
def __init__(self, config: LlamaConfig):
super().__init__(config)
self.rotary_emb = NoRoPE(
head_dim=config.head_dim,
)
# Explicitly disable causal attention
self.config.is_causal = False
# force every layer to be non-causal
for layer in self.layers:
if hasattr(layer, "self_attn"):
layer.self_attn.is_causal = False
def _update_causal_mask(
self,
attention_mask: torch.Tensor,
input_tensor: torch.Tensor,
cache_position: torch.Tensor,
past_key_values,
output_attentions: bool = False,
):
# By returning None, we disable any causal‐(look‐ahead) masking.
# The only mask that remains is whatever "attention_mask" the user has passed
# (e.g. padding‐mask), which will be handled by Flash/SDPA internally as non‐causal.
return None
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor = None,
past_key_values=None,
inputs_embeds: torch.FloatTensor = None,
use_cache: bool = None,
output_attentions: bool = None,
output_hidden_states: bool = None,
cache_position: torch.LongTensor = None,
**flash_attn_kwargs,
):
flash_attn_kwargs["is_causal"] = False
# If no attention_mask is provided, create an all-ones mask (no masking)
# This ensures bidirectional attention with correct device/dtype
if attention_mask is None:
# Get batch size (B) and sequence length (S) from input_embeds if available, else from input_ids.
# If neither is available, fall back to attention_mask=None and log a warning.
B = None
S = None
if inputs_embeds is not None:
B, S = inputs_embeds.size(0), inputs_embeds.size(1)
if B and S:
attention_mask = torch.ones((B, 1, S, S), dtype=torch.float, device=inputs_embeds.device)
return super().forward(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
cache_position=cache_position,
**flash_attn_kwargs,
)
def get_bidirectional_llama_backbone(config: LlamaConfig) -> PreTrainedModel:
model = LlamaBidirectionalModel(config)
model.embed_tokens.weight.requires_grad = False
model.embed_tokens.weight.zero_()
return model
- 我们可以只用Llama的模型骨架，把input_embs传给模型骨架，而不是从input_ids开始，前面的数据编码头我们就可以随意替换了