compatibility.py

"""
Pyreft <-> Adapters compatibility check.
Script adapted from examples/loreft/train.py.
"""
import sys
sys.path.append("./examples/loreft")

import copy
import os
import torch
import argparse
from tqdm import tqdm, trange
from transformers import (
    AutoConfig,
    AutoTokenizer, 
    AutoModelForCausalLM, 
    AutoModelForSequenceClassification,
    DataCollatorForSeq2Seq,
    DataCollatorWithPadding,
    set_seed,
)
import datetime

from examples.loreft.task_config import task_config
from examples.loreft.dataset import LoReftGLUEDataset, LoReftSupervisedDataset, parse_positions
from examples.loreft.compute_metrics import make_data_collator, make_dataloader

from pyreft import (
    TaskType,
    get_reft_model,
    ReftConfig,
    NoreftIntervention,   # remove ortho.
    LoreftIntervention,
    ConsreftIntervention, # constant bias only
    LobireftIntervention, # low-rank bitfit reft
    DireftIntervention,   # direct edit reft
    NodireftIntervention, # remove ortho + direct edit reft <- this is like LoRA on time-step
    ReftDataCollator
)

import adapters
import adapters.loading
from convert_reft import convert_checkpoint

device = "cuda" if torch.cuda.is_available() else "cpu"
classification_tasks = {"glue"}
residual_stream_component_mapping = {
    "robertaformaskedlm": "roberta.encoder.layer[%s].output"
}
dtype_mapping = {
    "float32": torch.float32,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float8": "float8",
}
intervention_mapping = {
    "NoreftIntervention": NoreftIntervention,
    "LoreftIntervention": LoreftIntervention,
    "ConsreftIntervention": ConsreftIntervention,
    "LobireftIntervention": LobireftIntervention,
    "DireftIntervention": DireftIntervention,
    "NodireftIntervention": NodireftIntervention,
}


def test(
    act_fn: str,
    add_bias: bool,
    model: str,
    layers: str,
    rank: int,
    position: str,
    # epochs: int,
    seed: int,
    intervention_type: str,
    max_n_train_example: int,
    max_n_eval_example: int,
    # is_wandb: bool,
    # wandb_name: str,
    # gradient_accumulation_steps: int,
    batch_size: int,
    output_dir: str,
    task: str,
    # lr: float,
    # schedule: str,
    data_dir: str,
    train_dataset: str,
    eval_dataset: str,
    # save_model: bool,
    # eval_batch_size: int,
    # warmup_ratio: float,
    # weight_decay: float,
    dropout: float,
    test_split: str,
    # train_on_inputs: bool,
    max_length: int,
    # use_normalized_template: bool,
    allow_cls_grad: bool,
    # metric_for_best_model: str,
    dtype: str,
    # logging_steps: int,
    # wandb_dir: str,
    # wandb_proj: str,
    share_weights: bool,
    # greedy_decoding: bool,
    # temperature: float,
    # top_p: float,
    # top_k: float,
    args,
):
    dtype = dtype_mapping[dtype]
    
    # store/log run details
    print(
        f"task: {task}, model: {model}, intervention_type: {intervention_type}, "
        f"layers: {layers}, rank: {rank}, "
        f"position: {position}, "
        f"max_length: {max_length}, allow_cls_grad: {allow_cls_grad}"
    )

    # everything is guarded by a single seed
    set_seed(seed)

    model_name = model
    model_str = model.split("/")[-1]
    train_dataset_str = train_dataset
    now = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")
    if train_dataset is not None:
        run_name = f"{model_str}.{task}.{train_dataset_str}.{test_split}.{now}"
    else:
        run_name = f"{model_str}.{task}.{now}"

    # which layers to intervene on
    if layers != "all":
        layers = [int(l) for l in layers.split(";")]
    else:
        temp_config = AutoConfig.from_pretrained(model)
        layers = [l for l in range(temp_config.num_hidden_layers)]

    # position str takes the following formats:
    # f1 -> first token; f2 -> first two tokens.
    # f1+l1 -> first and last tokens; f2+l2 -> first and last two tokens.
    # fn or ln shares the same intervention.
    if "+" in position and not share_weights:
        layers += layers

    # load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        model_max_length=max_length,
        padding_side="right",
        use_fast=False,
    )
    if tokenizer.unk_token == None and tokenizer.pad_token == None:
        # raw llama3
        print("adding a special padding token...")
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        need_resize = True
    else:
        tokenizer.pad_token = tokenizer.unk_token
        need_resize = False

    # load dataset splits
    assert task in task_config, f"Unrecognized task: {task}"
    train_datasets = task_config[task]["train_datasets"] if train_dataset is None else [train_dataset]
    if task == "glue":
        eval_datasets = [train_dataset]
    else:
        eval_datasets = task_config[task]["eval_datasets"] if eval_dataset is None else [eval_dataset]
        
    ReftDataset = LoReftGLUEDataset if task == "glue" else LoReftSupervisedDataset 
    train_dataset = ReftDataset(
        task, train_datasets[0] if task == "glue" or task == "ultrafeedback_pair" \
            else (os.path.join(data_dir, train_datasets[0]) if data_dir is not None else train_datasets[0]), 
        tokenizer, data_split="train", seed=seed, max_n_example=max_n_train_example,
        **{"num_interventions": len(layers), "position": position, 
           "share_weights": share_weights, "test_split": test_split}
    )
    trigger_tokens = train_dataset.trigger_tokens
    num_labels = train_dataset.num_labels

    all_eval_datasets = {}
    for eval_dataset in eval_datasets:
        test_splits = test_split.split(";")
        all_eval_datasets[eval_dataset] = {}
        for split in test_splits:
            raw_eval = ReftDataset(
                task, eval_dataset if task == "glue" else os.path.join(data_dir, eval_dataset), 
                tokenizer, data_split=split, seed=seed, max_n_example=max_n_eval_example,
                **{"num_interventions": len(layers), "position": position, 
                   "share_weights": share_weights}
            )
            all_eval_datasets[eval_dataset][split] = [raw_eval, raw_eval.raw_dataset]
    eval_datasets = all_eval_datasets

    # load model based on task type.
    if task in classification_tasks:
        config = AutoConfig.from_pretrained(
            model, num_labels=num_labels,
            finetuning_task=train_dataset_str,
            load_in_8bit=True if dtype == "float8" else False,
            device_map=device
        )
        # full precision loading since usually for small models
        model = AutoModelForSequenceClassification.from_pretrained(
            model,
            config=config, # just providing the label
            torch_dtype=dtype if dtype != "float8" else None,
            load_in_8bit=True if dtype == "float8" else False,
            device_map=device
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model,
            torch_dtype=dtype if dtype != "float8" else None,  # save memory
            load_in_8bit=True if dtype == "float8" else False,
            device_map=device
        )
        config = model.config
    if need_resize:
        model.resize_token_embeddings(len(tokenizer))

    intervention_type = intervention_mapping[intervention_type]
        
    # select collator based on the type
    if task in classification_tasks:
        data_collator_fn = DataCollatorWithPadding(
            tokenizer=tokenizer,
            padding="longest"
        )
    else:
        data_collator_fn = DataCollatorForSeq2Seq(
            tokenizer=tokenizer,
            model=model,
            label_pad_token_id=-100,
            padding="longest"
        )
    data_collator = ReftDataCollator(data_collator=data_collator_fn)

    # intervention config based on model type
    intervention_dtype = torch.bfloat16 if isinstance(dtype, str) else dtype
    model_arch = model.config.architectures[0].lower()
    if model_arch in residual_stream_component_mapping:
        representations = [{
            "component": residual_stream_component_mapping[model_arch] % l,
            "low_rank_dimension": rank,
            "intervention": intervention_type(
                embed_dim=config.hidden_size, low_rank_dimension=rank,
                dropout=dropout, dtype=intervention_dtype, act_fn=act_fn, device=device,
                add_bias=add_bias
            )
        } for l in layers]
        task_type=TaskType.SEQ_CLS
    else:
        representations = [{
            "layer": l, "component": "block_output",
            "low_rank_dimension": rank,
            "intervention": intervention_type(
                embed_dim=config.hidden_size, low_rank_dimension=rank,
                dropout=dropout, dtype=intervention_dtype, act_fn=act_fn, device=device,
                add_bias=add_bias
            )
        } for l in layers]
        task_type=TaskType.CAUSAL_LM
    
    reft_config = ReftConfig(representations=representations)
    reft_model = get_reft_model(model, reft_config, set_device=not isinstance(dtype, str))
    reft_model.print_trainable_parameters()

    # for GLUE tasks, we enable gradients on the classifier head.
    # the parameter will be counted as well.
    if task == "glue" and allow_cls_grad:
        for param in reft_model.model.classifier.parameters():
            # reft_model with HF trainer will automatically pick up these params to optimize
            param.requires_grad = True

    n_params = reft_model.count_parameters(include_model=False)

    # save model
    chkpt_folder = f"{output_dir}/{run_name}"
    reft_model.save(chkpt_folder)

    # ensure everything is in eval mode
    reft_model.model.eval()
    for k,v in reft_model.interventions.items():
        _ = v[0].eval()

    print({"n_params": n_params})

    # Copy model
    adapter_model = copy.deepcopy(model)
    adapters.init(adapter_model)

    # Convert pyreft config to Adapters config
    if intervention_type == LoreftIntervention:
        config_class = adapters.LoReftConfig
    elif intervention_type == NoreftIntervention:
        config_class = adapters.NoReftConfig
    elif intervention_type == NodireftIntervention:
        config_class = adapters.DiReftConfig
    else:
        raise ValueError("Unsupported intervention type")
    prefix_positions, suffix_positions = parse_positions(position)
    config = config_class(
        layers=layers,
        r=rank,
        prefix_positions=prefix_positions,
        suffix_positions=suffix_positions,
        tied_weights=share_weights
    )
    adapter_model.add_adapter("my_reft", config=config)

    # Reload pyreft checkpoint to Adapters model
    state_dict = convert_checkpoint(chkpt_folder, "my_reft")
    loader = adapters.loading.AdapterLoader(adapter_model)
    loading_info = {}
    loader.load_from_state_dict(state_dict, "my_reft", loading_info=loading_info)
    assert len(loading_info["unexpected_keys"]) == 0, "Unexpected keys in loading."
    adapter_model = adapter_model.to(device)
    adapter_model.adapter_to("my_reft", dtype=dtype)
    adapter_model.eval()

    # do eval
    for split, (eval_dataset, data_items) in eval_datasets[list(eval_datasets.keys())[0]].items():

        data_collator = data_collator if data_collator is not None else \
            make_data_collator(tokenizer, reft_model.model)
        eval_dataloader = make_dataloader(eval_dataset, batch_size, data_collator, shuffle=False)
        eval_iterator = tqdm(eval_dataloader, position=0, leave=True)

        with torch.inference_mode():
            for step, inputs in enumerate(eval_iterator):
                for k, v in inputs.items():
                    if v is not None and isinstance(v, torch.Tensor):
                        inputs[k] = v.to(device)
                
                # [layers, batch_size, positions]
                intervention_locations = inputs["intervention_locations"].permute(1, 0, 2)
                # intervention_locations -= 1

                # forward pass with pyreft model
                _, cf_outputs = reft_model(
                    {"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"]},
                    unit_locations={"sources->base": (None, intervention_locations.tolist())}
                )

                # forward pass with adapters model
                with adapters.AdapterSetup("my_reft"):
                    adapter_outputs = adapter_model(
                        **{"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"]}
                    )

                if not torch.allclose(cf_outputs.logits, adapter_outputs.logits, atol=1e-3):
                    print(cf_outputs.logits)
                    print(adapter_outputs.logits)
                    raise ValueError("Adapter model is not the same as the reft model.")

        print("All checks successful for {}.".format(split))

def main():
    parser = argparse.ArgumentParser(description="A simple script that takes different arguments.")
    
    parser.add_argument('-task', '--task', type=str, default=None)
    parser.add_argument('-data_dir', '--data_dir', type=str, default="./datasets")
    parser.add_argument('-train_dataset', '--train_dataset', type=str, default=None)
    parser.add_argument('-eval_dataset', '--eval_dataset', type=str, default=None)
    parser.add_argument('-model', '--model', type=str, help='yahma/llama-7b-hf', default='yahma/llama-7b-hf')
    parser.add_argument('-seed', '--seed', type=int, help='42', default=42)
    parser.add_argument('-l', '--layers', type=str, help='2;10;18;26', default='2;10;18;26')
    parser.add_argument('-r', '--rank', type=int, help=8, default=8)
    parser.add_argument('-p', '--position', type=str, help='f1+l1', default='f1+l1')
    parser.add_argument('-max_n_train_example', '--max_n_train_example', type=int, default=None)
    parser.add_argument('-max_n_eval_example', '--max_n_eval_example', type=int, default=None)
    parser.add_argument(
        '-type', '--intervention_type', type=str, 
        help='LoreftIntervention', default="LoreftIntervention")
    parser.add_argument('-batch_size', '--batch_size', type=int, default=4)
    parser.add_argument('-output_dir', '--output_dir', type=str, default="./official_results")
    parser.add_argument('-dropout', '--dropout', type=float, default=0.00)
    parser.add_argument('-act_fn', '--act_fn', type=str, default=None)
    parser.add_argument('-add_bias', '--add_bias', action='store_true')
    parser.add_argument('-test_split', '--test_split', type=str, default="validation")
    parser.add_argument('-max_length', '--max_length', type=int, help=512, default=512)
    parser.add_argument('-allow_cls_grad', '--allow_cls_grad', action='store_true')
    parser.add_argument('-dtype', '--dtype', type=str, default="bfloat16" if device == "cuda" else "float32")
    parser.add_argument('-sw', '--share_weights', action='store_true')

    args = parser.parse_args()

    test(**vars(args), args=args)


if __name__ == "__main__":
    main()