Currently not working with Gemma 2 models

[DalasNoin]

I tried to run this with Gemma 2 27b it and found that it doesn't quite work. I verified that everything works with qwen/qwen-1_8b-chat.

I get this error message:

Assertion error: All scores have been filtered out

It also seems that KL scores are very large (>10)

I tried to find the reason but could not find a solution so far.

However, I did verify that the chat template worked correctly and it also seemed i could sample text from the model normally, when i placed a breakpoint in the function get_mean_activations which measured the activations.

What seemed odd was that the mean_diff of activations between harmful and harmless prompts was quite large, often between -200 and +200. In comparison, the mean diff of qwen was more like -2 to 2. So possibly there is an issue with the hooks?

The current GemmaModel is designed for Gemma 1 models. It seems the only architectural change is to add a rms norm before and after the MLP. I am not familiar with the details of the Gemma2RMSNorm implementation.

[fblissjr]

@DalasNoin

I made an attempt on a fork here, not very cleanly, using the Hugging Face local-gemma library. Ran out of time to try applying and testing anything, but maybe it will be useful to someone. Made a bunch of small edits to the pipeline code for observability reasons (starting with finding why NaNs and then working my way up).

fork: https://github.com/fblissjr/refusal_direction
gemma2 model: https://github.com/fblissjr/refusal_direction/blob/main/pipeline/model_utils/gemma2_model.py

If anyone gets this working, would love to know what did it, and if it's more generalized than just refusals (since that's the code we have, it's where I started).

[revmag]

Yes, had the same issue.
Mailed the co-author for this:
The issue lies with filtering, the filtering function uses the tokenizer from the chat model, and not from the base model, as the base model's filters all the instructions out ( they score rather poorely).

prompts = chat_model.tokenize_instructions_fn(instructions=instructions[i:i+batch_size]).to('cuda')
logits = chat_model.model(input_ids=prompts.input_ids, attention_mask=prompts.attention_mask).logits
refusal_scores.append(refusal_score(logits, refusal_toks=chat_model.refusal_toks))

You can either first filter the instructions using chat model's tokenzier, and store it and then call it separately for the base model, or just add chat model's tokenizer for the filtering part

[andyrdt]

I just tried with gemma-2-2b-it, and it worked out of the box - I pushed the code and artifacts to the gemma2 branch.

The only issue I ran into was needing to specify attn_implementation=eager to avoid NaN issues. (see huggingface/transformers#32390).

[I do not endorse the previous comment from @revmag about base/chat tokenizers - the code only works with chat models, and does not use base models anywhere.]

[andyrdt]

I also got gemma-2-9b-it working, although the direction selection algorithm as implemented doesn't arrive at a very good direction - I found best results by using the candidate direction from pos=-1, layer = 23.

For gemma-2-27b-it I reproduced @DalasNoin's initial observation of very high KL for basically all ablations. I find it weird that things are fairly smooth for 2b and 9b, but break with 27b - I think this requires deeper investigation.

[dribnet]

I found your gemma-2-2b-it results useful - should we merge the gemma2 branch into main to make it an offical part of this repo?

[nikihasrati]

I have the same problem trying to get the code to work with lmsys/vicuna-13b-v1.5. I get the error assert len(filtered_scores) > 0, "All scores have been filtered out!" and all the refusal scores in direction_evaluations.json are NaN. The problem persists even when setting attn_implementation=eager. Here's my code:

import torch
import functools

from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import List
from torch import Tensor
from jaxtyping import Float

from pipeline.utils.utils import get_orthogonalized_matrix
from pipeline.model_utils.model_base import ModelBase


# Chat templates from: https://ollama.com/library/vicuna:13b-v1.5-16k-fp16
VICUNA_DEFAULT_SYSTEM_PROMPT = """A chat between a curious user and an artificial intelligence assistant. 
The assistant gives helpful, detailed, and polite answers to the user's questions."""

VICUNA_CHAT_TEMPLATE = "USER: {instruction}\nASSISTANT: "

VICUNA_CHAT_TEMPLATE_WITH_SYSTEM = "{system_prompt}\nUSER: {instruction}\nASSISTANT: "

VICUNA_REFUSAL_TOKS = [306, 1094] # ['I', 'As']

def format_instruction_vicuna_chat(
    instruction: str,
    system: str=None,
    output: str=None,
    include_trailing_whitespace=True
):
    if system is not None:
        if system == "default":
            system = VICUNA_DEFAULT_SYSTEM_PROMPT
        formatted_instruction = VICUNA_CHAT_TEMPLATE_WITH_SYSTEM.format(instruction=instruction, system_prompt=system)
    else:
        formatted_instruction = VICUNA_CHAT_TEMPLATE.format(instruction=instruction)

    if not include_trailing_whitespace:
        formatted_instruction = formatted_instruction.rstrip()

    if output is not None:
        formatted_instruction += output

    return formatted_instruction

def tokenize_instructions_vicuna_chat(
    tokenizer: AutoTokenizer,
    instructions: List[str],
    outputs: List[str]=None,
    system: str=None,
    include_trailing_whitespace=True
):
    if outputs is not None:
        prompts = [
            format_instruction_vicuna_chat(instruction=instruction, output=output, system=system, include_trailing_whitespace=include_trailing_whitespace)
            for instruction, output in zip(instructions, outputs)
        ]
    else:
        prompts = [
            format_instruction_vicuna_chat(instruction=instruction, system=system, include_trailing_whitespace=include_trailing_whitespace)
            for instruction in instructions
        ]

    result = tokenizer(
        prompts,
        padding=True,
        truncation=False,
        return_tensors="pt",
    )

    return result

def orthogonalize_vicuna_weights(model, direction: Float[Tensor, "d_model"]):
    model.model.embed_tokens.weight.data = get_orthogonalized_matrix(model.model.embed_tokens.weight.data, direction)

    for block in model.model.layers:
        block.self_attn.o_proj.weight.data = get_orthogonalized_matrix(block.self_attn.o_proj.weight.data.T, direction).T
        block.mlp.down_proj.weight.data = get_orthogonalized_matrix(block.mlp.down_proj.weight.data.T, direction).T

def act_add_vicuna_weights(model, direction: Float[Tensor, "d_model"], coeff, layer):
    dtype = model.model.layers[layer-1].mlp.down_proj.weight.dtype
    device = model.model.layers[layer-1].mlp.down_proj.weight.device

    bias = (coeff * direction).to(dtype=dtype, device=device)

    model.model.layers[layer-1].mlp.down_proj.bias = torch.nn.Parameter(bias)

class VicunaModel(ModelBase):
    def _load_model(self, model_path, dtype=torch.float16):
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=dtype,
            device_map="auto",
            attn_implementation="eager"
        )
        return model

    def _load_tokenizer(self, model_path):
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
        tokenizer.pad_token = tokenizer.eos_token
        return tokenizer
    
    def _get_tokenize_instructions_fn(self):
        return functools.partial(tokenize_instructions_vicuna_chat, tokenizer=self.tokenizer)
    
    def _get_model_block_modules(self):
        return self.model.model.layers

    def _get_eoi_toks(self):
        return self.tokenizer.encode(VICUNA_CHAT_TEMPLATE.split("{instruction}")[-1], add_special_tokens=False)
    
    def _get_refusal_toks(self):
        return VICUNA_REFUSAL_TOKS

    def _get_attn_modules(self):
        return torch.nn.ModuleList([block_module.self_attn for block_module in self.model_block_modules])
    
    def _get_mlp_modules(self):
        return torch.nn.ModuleList([block_module.mlp for block_module in self.model_block_modules])
    
    def _get_orthogonalization_mod_fn(self, direction):
        return functools.partial(orthogonalize_vicuna_weights, direction=direction)
    
    def _get_act_add_mod_fn(self, direction, coeff, layer):
        return functools.partial(act_add_vicuna_weights, direction=direction, coeff=coeff, layer=layer)

[nikihasrati]

I fixed the problem (ran into another error downstream but at least I'm able to extract the refusal direction now). The problem: the performance of the model was so bad that none of the prompts in harmful_val were being refused, leading to an empty list. I ended up using the following chat template and the performance of the model was much better. Another problem was that I was just using the tokens for "I" and "As" that the tokenizer spit out, but digging into it a bit more I noticed that the model was using different tokens for "I" and "As" (no idea why).

# Chat templates from: https://huggingface.co/junelee/wizard-vicuna-13b/discussions/1
VICUNA_DEFAULT_SYSTEM_PROMPT = """A chat between a curious user and an artificial intelligence assistant. 
The assistant gives helpful, detailed, and polite answers to the user's questions."""

VICUNA_CHAT_TEMPLATE = "### Human:\n{instruction}\n### Assistant:\n"

VICUNA_CHAT_TEMPLATE_WITH_SYSTEM = "{system_prompt}\n\n### Human:\n{instruction}\n### Assistant:\n"

VICUNA_REFUSAL_TOKS = [29902, 2887] # ['I', 'As']

[NamburiSrinath]

Hi @andyrdt,

Is there a fix for this problem in codebase? I am facing the same issue when running on a compressed version of Llama-2-7B-Chat