Hello

#1
by mylesgoose - opened

I tested with the 70B model that was abliterated with the acceptance direction rather than the refusal layer :-). and the script worked.. and the model refuses every request haha. So i wonder if we ask it a bad request it will then accept.
myles@ubuntu11:~/TransformerLens$ /usr/local/bin/python3 /home/myles/TransformerLens/transformer_lens/utilities/test90b.py
Loading checkpoint shards: 100%|
<|begin_of_text|><|start_header_id|>user<|end_header_id|>

<|image|>If I had to write a haiku for this one, it would be: <|eot_id|><|start_header_id|>assistant<|end_header_id|>

I don't think I can help with that. Is there something else I can help with?<|eot_id|>

Funny. it was just a bunny picture... 79 layers to test..
Ill run with the correct layer latter and upload.

import json
import os
import re

import torch
from safetensors import safe_open
from transformers import AutoProcessor, MllamaForConditionalGeneration

total_layers=64
#total_layers=32
cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38, 43, 48, 53, 58, 63, 68, 73, 78, 83, 88, 93, 98] # 90b
#cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38] # 11b

hidden_size = 8192 # From Llama 3.2 90b
#hidden_size = 4096 # From Llama 3.2 11b

b8 = "/home/myles/abliteration/mylesgoose/Llama-3.1-70B-Instruct-abliterated"
print(b8)

#model_id = "/home/myles/Desktop/llama/Llama-3.2-11B-Vision-Instruct"
model_id = "/home/myles/Desktop/llama/Llama-3.2-90B-Vision-Instruct"
print(model_id)

def create_layer_mapping(total_layers=total_layers, cross_attn_layers=cross_attention_layers):
    """
    Creates a mapping from llama-3.1-8b layer indices to llama-3.2-11b layer indices.
    """
    mapping = {}
    shift = 0
    next_cross_attn_idx = 0
    for X in range(total_layers):
        # Check if a cross-attention layer is inserted before this layer
        if next_cross_attn_idx < len(cross_attn_layers) and (X + shift) == cross_attn_layers[next_cross_attn_idx]:
            shift += 1
            next_cross_attn_idx += 1
        Y = X + shift
        mapping[X] = Y
    return mapping

def load_sharded_state_dict(model_dir):
    index_file = os.path.join(model_dir, 'model.safetensors.index.json')
    with open(index_file, 'r') as f:
        index_data = json.load(f)
    weight_map = index_data['weight_map']
    state_dict = {}
    shard_to_params = {}
    for param_name, shard_file in weight_map.items():
        if shard_file not in shard_to_params:
            shard_to_params[shard_file] = []
        shard_to_params[shard_file].append(param_name)
    for shard_file, params_in_shard in shard_to_params.items():
        shard_path = os.path.join(model_dir, shard_file)
        with safe_open(shard_path, framework="pt", device="cpu") as f:
            for name in params_in_shard:
                state_dict[name] = f.get_tensor(name)
    return state_dict

def compare_model_states(model, new_state_dict):
    current_state = model.state_dict()
    unchanged_params = []
    changed_params = []
    missing_params = []

    for name, param in current_state.items():
        if name not in new_state_dict:
            missing_params.append(name)
        elif torch.equal(param.cpu(), new_state_dict[name].cpu()):
            unchanged_params.append(name)
        else:
            changed_params.append(name)

    return {
        'unchanged': unchanged_params,
        'changed': changed_params,
        'missing': missing_params
    }


layer_mapping = create_layer_mapping()

# Load Llama 3.2 state dict
llama_3_2_state_dict = load_sharded_state_dict(model_id)

# Extract the embedding matrix from Llama 3.2
llama_3_2_embeddings = llama_3_2_state_dict['language_model.model.embed_tokens.weight']  # Shape: [128264, 4096] or [128256, 8192]

llama_3_2_state_dict.clear()

b8dict = load_sharded_state_dict(b8)

embed_tokens_weight = b8dict['model.embed_tokens.weight']  # Shape: [128256, 4096] or [128256, 8192]
new_vocab_size = 128264  # From Llama 3.2

new_embed_tokens_weight = torch.zeros((new_vocab_size, hidden_size), dtype=embed_tokens_weight.dtype)

# Copy the existing embeddings
new_embed_tokens_weight[:128256, :] = embed_tokens_weight
# Copy the additional embeddings from Llama 3.2
new_embed_tokens_weight[128256:, :] = llama_3_2_embeddings[128256:, :]

b8dict['model.embed_tokens.weight'] = new_embed_tokens_weight


llama_3_2_embeddings = None

# Adjust Llama 3.1 parameter names to match Llama 3.2 language model
st8dict = {}
for name, param in b8dict.items():
    # Prefix non-layer parameters with 'language_model.'
    if not re.match(r'model\.layers\.\d+\.', name):
        new_name = 'language_model.' + name
    else:
        # Extract the layer index X from 'model.layers.X.'
        match = re.match(r'model\.layers\.(\d+)\.(.+)', name)
        if match:
            X = int(match.group(1))
            suffix = match.group(2)
            # Get the corresponding Y in llama-3.2-11b
            Y = layer_mapping.get(X, X + len(cross_attention_layers))
            new_name = f'language_model.model.layers.{Y}.{suffix}'
        else:
            # If the pattern doesn't match, just prefix with 'language_model.'
            new_name = 'language_model.' + name
    st8dict[new_name] = param

#write st8dict keys to file for verification
with open('st8dict.txt', 'w') as f:
    f.write('\n'.join(st8dict.keys()))


model = MllamaForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)

#original_state = {k: v.clone() for k, v in model.state_dict().items()}

model.load_state_dict(st8dict, strict=False)

b8dict.clear()
st8dict.clear()


'''
result = compare_model_states(model, original_state)

print("Unchanged parameters:", len(result['unchanged']))
print("Changed parameters:", len(result['changed']))
print("Missing parameters:", len(result['missing']))

#write result to file
with open('result.txt', 'w') as f:
    f.write(json.dumps(result, indent=2))
'''


processor = AutoProcessor.from_pretrained(model_id)


model.save_pretrained("Llama-3.2-90B-Vision-Instruct-abliterated")

I was thinking a better process would be to do the reverse. convert the 90b model into the 70b model. and then save the 70b model and abliterate that with transformer lens then run your script. as i have compared the model 70b to the original and they are not exactly the same. for example the 90b model says its text models vocab size is the same as the orignal model in the config file. however its not obviously as it has been trained on the new tokens and hence why you needed to adjsut the size above for the tokenizer. IN my un-expert opinion the 3.2 models from 90b size of the weights is this "metadata": {
"total_size": 141107675136
}, when i ran the script in reverse. so the model size is slightly larger. indicating maybe it had more training? however my script for reversing the model and using the 90b tokenism etc in the 70b model. the models does not output coherent text. so perhaps somethings wrong. can you help me reverse this process so we can actually abliterate the 70b model from the 90b models and then reinsert it?

i added a script for extracting the 8b/70b from 11b/90b. i just tested it with 8b and it seems to be working coherently. weights are different as pointed out here https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct/discussions/32. interestingly he mentions the 70b/90b LM weights are identical.

i got incoherent results when trying to force in the expanded vocab size, maybe that's what you're running into? i assume we're not meant to be touching the image tokens

also, i wasn't having success swapping in abliterated models; they were just as refusal-happy (and sometimes more) as the og model in my experience with vision tasks (though i didn't test it very thoroughly, and not at all for text-only). for example, asking it to identify someone in an image.

but that might just be a problem with the hotswap script, i need to clean it up

Sign up or log in to comment