File size: 9,079 Bytes
805e739 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0.
"""
Bidirectional Llama model for embedding tasks.
This module provides a modified LlamaModel that uses bidirectional (non-causal)
attention, suitable for generating embeddings where each token should attend
to all other tokens in the sequence.
Supports transformers version 4.44 and above with a unified forward() implementation.
Version compatibility notes:
- transformers 4.47: Setting _attn_implementation in __init__ had no effect due to
attention initialization order
- transformers 4.48+: Attention refactor (transformers#35235) activated the
_attn_implementation setting, which defaulted to "eager" instead of "sdpa"
- transformers < 4.53: LlamaModel has _update_causal_mask method that can be overridden
- transformers 4.53+: _update_causal_mask removed; masking moved to masking_utils module,
necessitating a full forward() override for custom attention masks
- transformers < 4.54: Decoder layer returns tuple, uses past_key_value (singular)
- transformers 4.54-4.55: Decoder layer returns tensor, uses past_key_value (singular)
- transformers 4.56+: Decoder layer returns tensor, uses past_key_values (plural),
DynamicCache accepts config parameter
- transformers 5.0+: Has native create_bidirectional_mask in masking_utils
"""
import inspect
import torch
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaModel
from transformers.utils import logging
logger = logging.get_logger(__name__)
# Check if native create_bidirectional_mask exists (transformers >= 5.0)
try:
from transformers.masking_utils import create_bidirectional_mask
_HAS_NATIVE_BIDIRECTIONAL_MASK = True
except ImportError:
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
_HAS_NATIVE_BIDIRECTIONAL_MASK = False
# Detect API differences via introspection
_decoder_forward_params = inspect.signature(LlamaDecoderLayer.forward).parameters
_dynamic_cache_init_params = inspect.signature(DynamicCache.__init__).parameters
# past_key_value (singular) in < 4.56, past_key_values (plural) in >= 4.56
_USE_PLURAL_CACHE_PARAM = "past_key_values" in _decoder_forward_params
# DynamicCache accepts config parameter in >= 4.56
_DYNAMIC_CACHE_ACCEPTS_CONFIG = "config" in _dynamic_cache_init_params
class LlamaBidirectionalConfig(LlamaConfig):
"""Configuration for LlamaBidirectionalModel with pooling and temperature settings."""
model_type = "llama_bidirec"
def __init__(
self, pooling: str = "avg", temperature: float = 1.0, **kwargs
) -> None:
"""
Initialize bidirectional Llama configuration.
Args:
pooling: Pooling strategy for embeddings ("avg", "cls", "last", etc.)
temperature: Temperature scaling for embeddings
**kwargs: Additional arguments passed to LlamaConfig
"""
self.pooling = pooling
self.temperature = temperature
super().__init__(**kwargs)
class LlamaBidirectionalModel(LlamaModel):
"""
LlamaModel modified to use bidirectional (non-causal) attention.
In standard Llama, each token can only attend to previous tokens (causal attention).
This model removes that restriction, allowing each token to attend to all tokens
in the sequence, which is useful for embedding tasks.
The key modifications are:
1. Setting is_causal=False on all attention layers
2. Using a bidirectional attention mask instead of causal mask
"""
config_class = LlamaBidirectionalConfig
def __init__(self, config: LlamaConfig) -> None:
super().__init__(config)
for layer in self.layers:
layer.self_attn.is_causal = False
def _create_bidirectional_mask(
self,
input_embeds: torch.Tensor,
attention_mask: torch.Tensor | None,
) -> torch.Tensor | None:
"""
Create bidirectional attention mask.
Args:
input_embeds: Input embeddings tensor of shape (batch_size, seq_len, hidden_size)
attention_mask: Optional 2D attention mask of shape (batch_size, seq_len)
where 1 indicates tokens to attend to and 0 indicates masked tokens
Returns:
4D attention mask suitable for the attention implementation, or None
if no masking is needed
"""
if attention_mask is None:
return None
if _HAS_NATIVE_BIDIRECTIONAL_MASK:
return create_bidirectional_mask(
config=self.config,
input_embeds=input_embeds,
attention_mask=attention_mask,
)
# Fallback for transformers < 5.0 without create_bidirectional_mask
# Flash attention handles 2D masks internally; only pass mask if there
# are actually masked tokens (zeros), otherwise return None for efficiency
if getattr(self.config, "_attn_implementation", None) == "flash_attention_2":
has_masked_tokens = (attention_mask == 0).any()
return attention_mask if has_masked_tokens else None
return _prepare_4d_attention_mask(attention_mask, input_embeds.dtype)
def forward(
self,
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
inputs_embeds: torch.FloatTensor | None = None,
cache_position: torch.LongTensor | None = None,
use_cache: bool | None = None,
**kwargs,
) -> BaseModelOutputWithPast:
"""
Forward pass with bidirectional attention.
Args:
input_ids: Input token IDs of shape (batch_size, seq_len)
attention_mask: Attention mask of shape (batch_size, seq_len)
position_ids: Position IDs for rotary embeddings
past_key_values: Cached key/value states for incremental decoding
inputs_embeds: Pre-computed input embeddings (alternative to input_ids)
cache_position: Position indices for cache updates
use_cache: Whether to return cached key/value states
**kwargs: Additional arguments passed to decoder layers
Returns:
BaseModelOutputWithPast containing last_hidden_state and past_key_values
"""
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError(
"You must specify exactly one of input_ids or inputs_embeds"
)
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
# Initialize cache if needed
if use_cache and past_key_values is None:
if _DYNAMIC_CACHE_ACCEPTS_CONFIG:
past_key_values = DynamicCache(config=self.config)
else:
past_key_values = DynamicCache()
if cache_position is None:
past_seen_tokens = (
past_key_values.get_seq_length() if past_key_values is not None else 0
)
cache_position = torch.arange(
past_seen_tokens,
past_seen_tokens + inputs_embeds.shape[1],
device=inputs_embeds.device,
)
if position_ids is None:
position_ids = cache_position.unsqueeze(0)
bidirectional_mask = self._create_bidirectional_mask(
inputs_embeds, attention_mask
)
hidden_states = inputs_embeds
position_embeddings = self.rotary_emb(hidden_states, position_ids)
# Build decoder layer kwargs with correct cache parameter name
# (past_key_value in < 4.56, past_key_values in >= 4.56)
layer_kwargs = {
"attention_mask": bidirectional_mask,
"position_ids": position_ids,
"use_cache": use_cache,
"cache_position": cache_position,
"position_embeddings": position_embeddings,
}
if _USE_PLURAL_CACHE_PARAM:
layer_kwargs["past_key_values"] = past_key_values
else:
layer_kwargs["past_key_value"] = past_key_values
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
layer_outputs = decoder_layer(hidden_states, **layer_kwargs)
# Decoder returns tuple in < 4.54, tensor in >= 4.54
if isinstance(layer_outputs, tuple):
hidden_states = layer_outputs[0]
else:
hidden_states = layer_outputs
hidden_states = self.norm(hidden_states)
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=past_key_values,
)
|