diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index ee2134bd1..29b5d3bcd 100644 --- a/transformer_lens/HookedTransformer.py +++ b/transformer_lens/HookedTransformer.py @@ -32,7 +32,6 @@ import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F import tqdm.auto as tqdm from jaxtyping import Float, Int from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase @@ -73,6 +72,7 @@ init_xavier_normal_, init_xavier_uniform_, ) +from transformer_lens.utilities.activation_functions import apply_softcap from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.weight_processing import ProcessWeights @@ -670,9 +670,7 @@ def forward( else: logits = self.unembed(residual) # [batch, pos, d_vocab] if self.cfg.output_logits_soft_cap > 0.0: - logits = self.cfg.output_logits_soft_cap * F.tanh( - logits / self.cfg.output_logits_soft_cap - ) + logits = apply_softcap(logits, self.cfg.output_logits_soft_cap) if return_type == "logits": return logits else: @@ -2186,11 +2184,13 @@ def generate( top_p=top_p, temperature=temperature, freq_penalty=freq_penalty, - tokens=torch.cat( - (input_tokens, torch.cat(sampled_tokens_list, dim=1)), dim=1 - ) - if "sampled_tokens" in locals() - else input_tokens, + tokens=( + torch.cat( + (input_tokens, torch.cat(sampled_tokens_list, dim=1)), dim=1 + ) + if "sampled_tokens" in locals() + else input_tokens + ), ).to(get_device_for_block_index(0, self.cfg)) else: sampled_tokens = utils.sample_logits( diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 01a0646eb..9f0e819ba 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/transformer_lens/components/abstract_attention.py @@ -19,6 +19,7 @@ from transformer_lens.FactoredMatrix import FactoredMatrix from transformer_lens.hook_points import HookPoint from transformer_lens.utilities import get_offset_position_ids +from transformer_lens.utilities.activation_functions import apply_softcap from transformer_lens.utilities.attention import complex_attn_linear, simple_attn_linear if is_bitsandbytes_available(): @@ -522,10 +523,7 @@ def calculate_attention_scores( k, "batch key_pos head_index d_head -> batch head_index d_head key_pos" ) attn_scores = q_ @ k_ / self.attn_scale - if self.cfg.attn_scores_soft_cap > 0: - attn_scores = self.cfg.attn_scores_soft_cap * F.tanh( - attn_scores / self.cfg.attn_scores_soft_cap - ) + attn_scores = apply_softcap(attn_scores, self.cfg.attn_scores_soft_cap) return attn_scores def calculate_z_scores( diff --git a/transformer_lens/config/hooked_transformer_config.py b/transformer_lens/config/hooked_transformer_config.py index 57ed34745..0e4a757ef 100644 --- a/transformer_lens/config/hooked_transformer_config.py +++ b/transformer_lens/config/hooked_transformer_config.py @@ -14,7 +14,10 @@ import numpy as np import torch -from transformer_lens.utilities.activation_functions import SUPPORTED_ACTIVATIONS +from transformer_lens.utilities.activation_functions import ( + SOFTCAP_DISABLED, + SUPPORTED_ACTIVATIONS, +) from transformer_lens.utilities.devices import get_device from .transformer_lens_config import TransformerLensConfig @@ -280,8 +283,8 @@ class HookedTransformerConfig(TransformerLensConfig): decoder_start_token_id: Optional[int] = None tie_word_embeddings: bool = False use_normalization_before_and_after: bool = False - attn_scores_soft_cap: float = -1.0 - output_logits_soft_cap: float = -1.0 + attn_scores_soft_cap: float = SOFTCAP_DISABLED + output_logits_soft_cap: float = SOFTCAP_DISABLED use_NTK_by_parts_rope: bool = False NTK_by_parts_low_freq_factor: float = 1.0 NTK_by_parts_high_freq_factor: float = 4.0 diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index 593407451..9938c6313 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/transformer_lens/config/transformer_bridge_config.py @@ -4,6 +4,8 @@ import torch +from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED + from .transformer_lens_config import TransformerLensConfig @@ -77,8 +79,8 @@ def __init__( decoder_start_token_id: Optional[int] = None, tie_word_embeddings: bool = False, use_normalization_before_and_after: bool = False, - attn_scores_soft_cap: float = -1.0, - output_logits_soft_cap: float = -1.0, + attn_scores_soft_cap: float = SOFTCAP_DISABLED, + output_logits_soft_cap: float = SOFTCAP_DISABLED, use_NTK_by_parts_rope: bool = False, NTK_by_parts_low_freq_factor: float = 1.0, NTK_by_parts_high_freq_factor: float = 4.0, diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 4800b0150..f891f92d0 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -50,6 +50,7 @@ VARIANT_SUBMODULE_NAMES, ) from transformer_lens.model_bridge.get_params_util import get_bridge_params +from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED from transformer_lens.utilities.aliases import resolve_alias from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss @@ -920,7 +921,7 @@ def process_weights( print(f"Processing weights for {self.cfg.model_name}...") # Soft capping (tanh) is not translation-invariant; centering would change output. - if center_unembed and getattr(self.cfg, "output_logits_soft_cap", -1.0) > 0.0: + if center_unembed and getattr(self.cfg, "output_logits_soft_cap", SOFTCAP_DISABLED) > 0.0: import logging logging.warning( diff --git a/transformer_lens/model_bridge/sources/native/model.py b/transformer_lens/model_bridge/sources/native/model.py index 4cae95a2a..5617459e7 100644 --- a/transformer_lens/model_bridge/sources/native/model.py +++ b/transformer_lens/model_bridge/sources/native/model.py @@ -18,6 +18,7 @@ from transformer_lens.config import TransformerBridgeConfig from transformer_lens.utilities import TypedModuleList +from transformer_lens.utilities.activation_functions import apply_softcap # gelu_new = the tanh-approximation HF GPT-2 / HT use; F.gelu(approximate="tanh") # is the exact same formula. @@ -269,9 +270,7 @@ def forward( scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale # Gemma2 soft-cap before the causal mask so masked positions stay -inf. - if self.attn_scores_soft_cap > 0: - c = self.attn_scores_soft_cap - scores = c * torch.tanh(scores / c) + scores = apply_softcap(scores, self.attn_scores_soft_cap) if self.causal: block_mask = self.causal_mask[:seq, :seq] @@ -468,7 +467,5 @@ def forward( ) hidden_states = self.ln_out(hidden_states) logits = self.head(hidden_states) - if self.output_logits_soft_cap > 0: - c = self.output_logits_soft_cap - logits = c * torch.tanh(logits / c) + logits = apply_softcap(logits, self.output_logits_soft_cap) return logits diff --git a/transformer_lens/utilities/activation_functions.py b/transformer_lens/utilities/activation_functions.py index b0a280967..b74f9c20d 100644 --- a/transformer_lens/utilities/activation_functions.py +++ b/transformer_lens/utilities/activation_functions.py @@ -105,6 +105,27 @@ def xielu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, ) +SOFTCAP_DISABLED: float = -1.0 + + +def apply_softcap(x: torch.Tensor, cap: float) -> torch.Tensor: + """Apply soft-cap to a tensor using tanh. + + If cap <= 0, returns x unchanged (softcap disabled). + Otherwise returns cap * tanh(x / cap). + + Args: + x: The tensor to apply soft-cap to. + cap: The soft-cap value. If <= 0, softcap is disabled. + + Returns: + The soft-capped tensor. + """ + if cap <= 0: + return x + return cap * torch.tanh(x / cap) + + # Convenient type for the format of each activation function ActivationFunction = Callable[..., torch.Tensor]