diff --git a/.gitignore b/.gitignore index b526ea4c7..e0254fa5d 100644 --- a/.gitignore +++ b/.gitignore @@ -180,4 +180,4 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ -.vscode/ +.vscode/ \ No newline at end of file diff --git a/README.md b/README.md index 8bde09ce7..f66a325ee 100644 --- a/README.md +++ b/README.md @@ -114,7 +114,7 @@ We provide several variants for each of the components in the unlearning pipelin # Environment setup conda create -n unlearning python=3.11 conda activate unlearning -pip install .[lm_eval] +pip install ".[lm_eval]" pip install --no-build-isolation flash-attn==2.6.3 # Data setup diff --git a/community/leaderboard.md b/community/leaderboard.md index ffb515ab7..ab9aad8f5 100644 --- a/community/leaderboard.md +++ b/community/leaderboard.md @@ -9,7 +9,7 @@ We encourage the community to develop new methods, optimize them for specific be To implement a new method, refer to our [contributing guide](../docs/contributing.md). > [!NOTE] -> The [results.md](../docs/results.md) file is maintained for reproducibility purposes. However, we encourage contributors to update the leaderboard table instead of the reproducibility table. We will continue refining and tuning baseline methods to keep the leaderboard up to date. +> The [results.md](../docs/repro.md) file is maintained for reproducibility purposes. However, we encourage contributors to update the leaderboard table instead of the reproducibility table. We will continue refining and tuning baseline methods to keep the leaderboard up to date. ### TOFU unlearning on the `Llama-2-7b-hf-chat` architecture diff --git a/docs/components.md b/docs/components.md index a0ba85a64..107827634 100644 --- a/docs/components.md +++ b/docs/components.md @@ -38,7 +38,7 @@ class GradDiff(UnlearnTrainer): def __init__(self, gamma, alpha, ...): ... - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): ... ``` diff --git a/requirements.txt b/requirements.txt index 2f39c76e2..5186a418d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ -huggingface-hub==0.29.1 -transformers==4.45.1 +huggingface-hub==0.36.0 +transformers==4.51.3 +hf-xet==1.2.0 numpy==2.2.3 hydra-core==1.3 hydra_colorlog==1.2.0 @@ -12,3 +13,4 @@ scipy==1.14.1 tensorboard==2.18.0 scikit-learn==1.5.2 deepspeed==0.15.4 +wandb==0.21.4 diff --git a/setup.py b/setup.py index 6a5c99c7e..fca95d9d6 100644 --- a/setup.py +++ b/setup.py @@ -18,12 +18,12 @@ install_requires=requirements, # Uses requirements.txt extras_require={ "lm-eval": [ - "lm-eval==0.4.8", - ], # Install using `pip install .[lm-eval]` + "lm-eval==0.4.11", + ], # Install using `pip install ".[lm-eval]"` "dev": [ "pre-commit==4.0.1", "ruff==0.6.9", - ], # Install using `pip install .[dev]` + ], # Install using `pip install ".[dev]"` }, python_requires=">=3.11", ) diff --git a/src/train.py b/src/train.py index a2f81c8d4..9cc6abcb5 100644 --- a/src/train.py +++ b/src/train.py @@ -50,7 +50,7 @@ def main(cfg: DictConfig): model=model, train_dataset=data.get("train", None), eval_dataset=data.get("eval", None), - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=collator, evaluators=evaluators, template_args=template_args, diff --git a/src/trainer/__init__.py b/src/trainer/__init__.py index 0bb7aa3b5..447b2d2dc 100644 --- a/src/trainer/__init__.py +++ b/src/trainer/__init__.py @@ -50,7 +50,7 @@ def load_trainer( model, train_dataset=None, eval_dataset=None, - tokenizer=None, + processing_class=None, data_collator=None, evaluators=None, template_args=None, @@ -70,7 +70,7 @@ def load_trainer( model=model, train_dataset=train_dataset, eval_dataset=eval_dataset, - tokenizer=tokenizer, + processing_class=processing_class, data_collator=data_collator, args=trainer_args, evaluators=evaluators, diff --git a/src/trainer/base.py b/src/trainer/base.py index 05f36a2a4..b10492e2e 100644 --- a/src/trainer/base.py +++ b/src/trainer/base.py @@ -1,21 +1,26 @@ # Modified from https://github.com/huggingface/transformers/blob/v4.45.1/src/transformers/trainer.py -from typing import Dict, List, Optional, Union - -import os import logging -from transformers import Trainer +import os +from typing import Any, Dict, List, Optional, Union + from torch.utils.data import Dataset +from transformers import Trainer from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR -from typing import Any logger = logging.getLogger(__name__) +# When using custom evaluators without an eval dataset, pass a dummy value +# to prevent Trainer from raising on eval_dataset=None when eval_strategy is set +_EVAL_PLACEHOLDER = "_EVAL_PLACEHOLDER" + class FinetuneTrainer(Trainer): def __init__(self, evaluators=None, template_args=None, *args, **kwargs): self.evaluators = evaluators self.template_args = template_args + if kwargs.get("eval_dataset") is None and evaluators: + kwargs["eval_dataset"] = _EVAL_PLACEHOLDER super().__init__(*args, **kwargs) def evaluate( @@ -26,33 +31,30 @@ def evaluate( trial: Dict[str, Any] = None, ) -> Dict[str, float]: # Run a custom evaluator and save results - if self.evaluators: - if self.accelerator.is_local_main_process: - eval_metrics = {} - if self.accelerator.num_processes == 1: - run_dir = self._get_output_dir(trial=trial) - checkpoint_folder = ( - f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" - ) - output_dir = os.path.join(run_dir, checkpoint_folder, "evals") - os.makedirs(output_dir, exist_ok=True) - eval_metrics = {} - for _, evaluator in self.evaluators.items(): - eval_args = { - "output_dir": output_dir, - "template_args": self.template_args, - "model": self.model, - "tokenizer": self.tokenizer, - } - eval_metrics.update(evaluator.evaluate(**eval_args)) - self.log(eval_metrics) - else: - logger.warning( - "Custom evaluator can be run with this Trainer only when a single accelerator process is running." - ) - return eval_metrics - - if eval_dataset is None: + if self.evaluators and self.accelerator.is_local_main_process: + if self.accelerator.num_processes != 1: + logger.warning( + "Custom evaluator can be run with this Trainer only when a single accelerator process is running." + ) + return {} + + run_dir = self._get_output_dir(trial=trial) + checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" + output_dir = os.path.join(run_dir, checkpoint_folder, "evals") + os.makedirs(output_dir, exist_ok=True) + eval_metrics = {} + for _, evaluator in self.evaluators.items(): + eval_args = { + "output_dir": output_dir, + "template_args": self.template_args, + "model": self.model, + "tokenizer": self.processing_class, + } + eval_metrics.update(evaluator.evaluate(**eval_args)) + self.log(eval_metrics) + return eval_metrics + + if eval_dataset is None or eval_dataset == _EVAL_PLACEHOLDER: return {} # Run the default HF Trainer evaluate method when eval dataset is provided return super().evaluate(eval_dataset, ignore_keys, metric_key_prefix) diff --git a/src/trainer/unlearn/base.py b/src/trainer/unlearn/base.py index 683698da2..26ab0ac7f 100644 --- a/src/trainer/unlearn/base.py +++ b/src/trainer/unlearn/base.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Optional, Union import torch from torch import nn @@ -78,10 +78,10 @@ def _prepare_deepspeed(self, model): def prediction_step( self, model: nn.Module, - inputs: Dict[str, Union[torch.Tensor, Any]], + inputs: dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, - ignore_keys: Optional[List[str]] = None, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + ignore_keys: Optional[list[str]] = None, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: """ The only change to this function is calling the Trainer's compute_loss, as it's often overridden by unlearning methods, and we want to maintain the Trainer's evaluation setup. """ @@ -104,7 +104,9 @@ def prediction_step( if ignore_keys is None: if hasattr(self.model, "config"): ignore_keys = getattr( - self.model.config, "keys_to_ignore_at_inference", [] + self.model.config, + "keys_to_ignore_at_inference", + ["past_key_values"], ) else: ignore_keys = [] @@ -146,11 +148,11 @@ def prediction_step( else: if has_labels or loss_without_labels: with self.compute_loss_context_manager(): - ### Call compute_loss of super class since overridden compute_loss is not be applicable to eval_dataset. + ### Call compute_loss of super class since overridden compute_loss is not applicable to eval_dataset. loss, outputs = super().compute_loss( model, inputs, return_outputs=True ) - loss = loss.mean().detach() + loss = loss.detach().mean() if isinstance(outputs, dict): logits = tuple( diff --git a/src/trainer/unlearn/ceu.py b/src/trainer/unlearn/ceu.py index 33da99c3c..90b069d47 100644 --- a/src/trainer/unlearn/ceu.py +++ b/src/trainer/unlearn/ceu.py @@ -86,7 +86,9 @@ def __init__(self, ignore_first_n_answer_tokens=1, *args, **kwargs): super().__init__(*args, **kwargs) self.ignore_first_n_answer_tokens = ignore_first_n_answer_tokens - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] loss, outputs = compute_batch_ceu( model, diff --git a/src/trainer/unlearn/dpo.py b/src/trainer/unlearn/dpo.py index b64b474b4..81ab39c35 100644 --- a/src/trainer/unlearn/dpo.py +++ b/src/trainer/unlearn/dpo.py @@ -9,7 +9,9 @@ def __init__(self, beta=1.0, *args, **kwargs): if self.ref_model is None: self.ref_model = self._prepare_ref_model(self.model) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"]["original"] alternate_inputs = inputs["forget"]["alternate"] diff --git a/src/trainer/unlearn/grad_ascent.py b/src/trainer/unlearn/grad_ascent.py index eda8b4812..ccb7dabab 100644 --- a/src/trainer/unlearn/grad_ascent.py +++ b/src/trainer/unlearn/grad_ascent.py @@ -2,7 +2,9 @@ class GradAscent(UnlearnTrainer): - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"], diff --git a/src/trainer/unlearn/grad_diff.py b/src/trainer/unlearn/grad_diff.py index bfecc19a2..53e3a5b0e 100644 --- a/src/trainer/unlearn/grad_diff.py +++ b/src/trainer/unlearn/grad_diff.py @@ -38,7 +38,9 @@ def compute_retain_loss(self, model, retain_inputs): ) return retain_loss - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"], diff --git a/src/trainer/unlearn/npo.py b/src/trainer/unlearn/npo.py index 7c782d968..3996f8ecf 100644 --- a/src/trainer/unlearn/npo.py +++ b/src/trainer/unlearn/npo.py @@ -9,7 +9,9 @@ def __init__(self, beta=1.0, *args, **kwargs): if self.ref_model is None: self.ref_model = self._prepare_ref_model(self.model) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_loss, forget_outputs = compute_dpo_loss( diff --git a/src/trainer/unlearn/pdu.py b/src/trainer/unlearn/pdu.py index e79bcc58b..773d2159f 100644 --- a/src/trainer/unlearn/pdu.py +++ b/src/trainer/unlearn/pdu.py @@ -1,7 +1,8 @@ import torch -from trainer.unlearn.grad_diff import GradDiff from transformers import TrainerCallback +from trainer.unlearn.grad_diff import GradDiff + class PDU(GradDiff): def __init__( @@ -102,7 +103,9 @@ def post_epoch_dual_param_update(self): ) self.log({"retain_preference": self.preferences[1]}) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"], diff --git a/src/trainer/unlearn/rmu.py b/src/trainer/unlearn/rmu.py index d990d3a38..77bce2897 100644 --- a/src/trainer/unlearn/rmu.py +++ b/src/trainer/unlearn/rmu.py @@ -136,7 +136,9 @@ def compute_retain_loss(self, model, retain_inputs): retain_loss = super().compute_retain_loss(model, retain_inputs) return retain_loss - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"], diff --git a/src/trainer/unlearn/satimp.py b/src/trainer/unlearn/satimp.py index b664390cd..f42d4acbb 100644 --- a/src/trainer/unlearn/satimp.py +++ b/src/trainer/unlearn/satimp.py @@ -14,7 +14,9 @@ def __init__( if self.ref_model is None: self.ref_model = self._prepare_ref_model(self.model) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"], diff --git a/src/trainer/unlearn/simnpo.py b/src/trainer/unlearn/simnpo.py index cb4f7f99c..dc96ee29d 100644 --- a/src/trainer/unlearn/simnpo.py +++ b/src/trainer/unlearn/simnpo.py @@ -10,7 +10,9 @@ def __init__(self, delta=0.0, beta=1.0, *args, **kwargs): self.delta = delta self.beta = beta - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_labels = forget_inputs["labels"] diff --git a/src/trainer/unlearn/undial.py b/src/trainer/unlearn/undial.py index e32147b30..d7c1d77bb 100644 --- a/src/trainer/unlearn/undial.py +++ b/src/trainer/unlearn/undial.py @@ -9,7 +9,9 @@ def __init__(self, beta=1.0, *args, **kwargs): if self.ref_model is None: self.ref_model = self._prepare_ref_model(self.model) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_loss, forget_outputs = compute_undial_loss( model, self.ref_model, forget_inputs, self.beta diff --git a/src/trainer/unlearn/wga.py b/src/trainer/unlearn/wga.py index 08c4bf402..f0b19cbae 100644 --- a/src/trainer/unlearn/wga.py +++ b/src/trainer/unlearn/wga.py @@ -11,7 +11,9 @@ def __init__(self, beta=1.0, gamma=1.0, alpha=1.0, *args, **kwargs): if self.ref_model is None: self.ref_model = self._prepare_ref_model(self.model) - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): forget_inputs = inputs["forget"] forget_inputs = { "input_ids": forget_inputs["input_ids"],