Skip to content

Upgrading transformers version to 4.51.3 to support recent models #175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Dornavineeth merged 12 commits into from
Mar 7, 2026
Merged

Upgrading transformers version to 4.51.3 to support recent models #175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
56b662e
bump transformers to 4.51.3
filyp Jan 26, 2026
31546d3
simplify installation
filyp Jan 26, 2026
9987ee3
make FinetuneTrainer more readable
filyp Jan 26, 2026
4b43848
use trainer.processing_class instead of trainer.tokenizer, to support…
filyp Jan 29, 2026
e180d7f
fix ruff linter errors
filyp Feb 11, 2026
1a3b7c8
format for ruff==0.6.6
filyp Feb 18, 2026
367fee9
preserve compute_loss backwards compatibility
filyp Feb 22, 2026
d6f1f1a
revert gitignore
filyp Feb 22, 2026
f1d2572
hide the trainer fix that allows for evaluating when eval_dataset=None
filyp Feb 22, 2026
d3dfad6
apply changes from the review
filyp Mar 6, 2026
faf942f
typo fix
filyp Mar 6, 2026
2ff5029
fix: lint
Dornavineeth Mar 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .gitignore
View file
Open in desktop
Copy link
Copy Markdown
Collaborator

@Dornavineeth Dornavineeth Feb 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like these additions in gitignore are specific to your usecases.
can you revert this?

Copy link
Copy Markdown
Contributor Author

@filyp filyp Feb 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup, done

Original file line number Diff line number Diff line change
Expand Up @@ -180,4 +180,4 @@ cython_debug/
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/

.vscode/
.vscode/
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion community/leaderboard.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/components.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
...
```

Expand Down
6 changes: 4 additions & 2 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
6 changes: 3 additions & 3 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,12 @@
install_requires=requirements, # Uses requirements.txt
extras_require={
Copy link
Copy Markdown
Collaborator

@Dornavineeth Dornavineeth Feb 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see https://github.com/locuslab/open-unlearning/pull/175/changes#r2837147182

"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",
)
2 changes: 1 addition & 1 deletion src/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/trainer/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
66 changes: 34 additions & 32 deletions src/trainer/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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(
Expand All @@ -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)
16 changes: 9 additions & 7 deletions src/trainer/unlearn/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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.
Copy link
Copy Markdown
Collaborator

@Dornavineeth Dornavineeth Feb 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you also include these docstrings in the updated code including prediction_step

"""
Expand All @@ -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 = []
Expand Down Expand Up @@ -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.
Copy link
Copy Markdown
Collaborator

@Dornavineeth Dornavineeth Feb 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the line I am referring here:

https://github.com/locuslab/open-unlearning/pull/175/changes#r2837139317

### 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(
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/ceu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/dpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/grad_ascent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/grad_diff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/npo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
7 changes: 5 additions & 2 deletions src/trainer/unlearn/pdu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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__(
Expand Down Expand Up @@ -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"],
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/rmu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/satimp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/simnpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/undial.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion src/trainer/unlearn/wga.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down