Reproduction issues on LLaMA3.1 (TOFU): checkpoint vs final‑model evaluation + epoch step‑count mismatch

[ma-kjh]

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task
• My own task or dataset (give details below)

Reproduction

Thank you for your excellent work on OpenUnlearning. I’m currently attempting to reproduce the results using the LLaMA-2-7b-chat-hf model and the GradDiff unlearning method on the 10% forget split of the TOFU benchmark.

Here’s what I’ve done so far:

---

1. Experiment Configuration

I created the following experiment config at:
/configs/experiment/unlearn/tofu/Llama-2-7b-chat-hf.yaml:

# @package _global_

defaults:
  - override /model: Llama-2-7b-chat-hf
  - override /trainer: GradAscent
  - override /data: unlearn
  - override /data/datasets@data.forget: TOFU_QA_forget
  - override /data/datasets@data.retain: TOFU_QA_retain
  - override /eval: tofu

model:
  model_args:
    pretrained_model_name_or_path: open-unlearning/tofu_Llama-2-7b-chat-hf_full

forget_split: forget10
retain_split: retain90
retain_logs_path: null

eval:
  tofu:
    forget_split: ${forget_split}
    retain_logs_path: ${retain_logs_path}
    overwrite: true
    
data:
  anchor: forget
  forget:
    TOFU_QA_forget: 
      args:
        hf_args:
          name: ${forget_split}
  retain:
    TOFU_QA_retain:
      args:
        hf_args:
          name: ${retain_split}

trainer:
  args:
    warmup_epochs: 1.0
    learning_rate: 1e-5
    weight_decay: 0.01
    num_train_epochs: 10

task_name: ???

---

2. Execution Command

I ran the following command:

CUDA_VISIBLE_DEVICES=0,1 python src/train.py \
  --config-name=unlearn.yaml \
  experiment=unlearn/tofu/Llama-2-7b-chat-hf.yaml \
  forget_split=forget10 \
  retain_split=retain90 \
  trainer=GradDiff \
  task_name=Llama-2-7b-chat-hf-GradDiff \
  trainer.args.per_device_train_batch_size=4 \
  trainer.args.gradient_accumulation_steps=4

---

3. Evaluation Results

At the final checkpoints (31 and 60), I obtained the following results:

Checkpoint 31:

{
  "forget_Q_A_ROUGE": 0.5793,
  "retain_Q_A_ROUGE": 0.7175,
  "model_utility": 0.5869,
  "forget_Q_A_Prob": 0.6063,
  "retain_Q_A_Prob": 0.9013
}

Checkpoint 60:

{
  "forget_Q_A_ROUGE": 0.4949,
  "retain_Q_A_ROUGE": 0.6398,
  "model_utility": 0.5714,
  "forget_Q_A_Prob": 0.4290,
  "retain_Q_A_Prob": 0.8337
}

---

Problem

When using a previous version of the codebase, I observed better forgetting and retention behavior under similar conditions.

I’m having trouble identifying what has changed between the previous implementation and the current one. Could you please advise if there are any known issues or differences in:

• Dataset preparation?
• Model loading?
• Trainer logic (e.g., GradDiff)?
• Evaluation configuration?

Any guidance or suggestions for troubleshooting this would be greatly appreciated.

Thank you very much in advance for your help!

Expected behavior

The ROUGE scores shown in the figure could not be reproduced.
Could you please let me know what changes have been made compared to the previous codebase?

Although the above experiment was run for 10 epochs, I also observed similar trends with 5 epochs. Compared to the previous codebase, it seems that the influence of the forget loss is reduced — the ROUGE score on the forget set decreases gradually rather than showing a clear drop.

[molereddy]

Hi, this is due to the different finetuned target models. A target model is the model before the unlearning process starts, post initial finetuning on the TOFU dataset.
While updating to OpenUnlearning, we trained and provided all target models again, while using the same hyperparameters as the original ones from the paper (and got roughly equal initial checkpoint evaluations with the new target models as the original). But upon unlearning, with these models people have been reporting [1] that the same hyperparameters don't work anymore. A similar issue was raised and discussed at that link. You may have to tune your hyperparameters for this new target model to achieve equally good results.

For your research work, you may want to stick with the original paper's target models for Llama-7b. Though ideally, for updating the OpenUnlearning leaderboard, we plan to only use the new versions of target models.

Why target models differ
Compared to the original TOFU repo, we're only aware of using different data preprocessing logic ([2] v/s [3]) in order to standardize the preprocessing and allow for dynamic max batch length padding. Other than that the differences are mostly in the hardware usage, which shouldn't matter significantly.
We are looking into if we've missed anything more.

[ma-kjh]

Thank you for the quick response. I’ll give it a try and leave a comment if I run into any issues.

[ma-kjh]

@molereddy Hello, I have one more question, so I'm leaving a comment here. I might have missed it, but could you please let me know how to measure forget_quality using an already generated TOFU_EVAL.json file? I remember this was included in the original TOFU code, but it seems a bit hard to find in this repository. Thank you so much.

[molereddy]

Our readme includes instructions for downloading this data using setup_data.py. See #73

[ma-kjh]

Thank you so much.

[ma-kjh]

@molereddy

I have a couple of questions:

1. The evaluation results provided for full fine-tuning are slightly different from the evaluation results I get when I load the model and run the first evaluation (eval0) myself. What factors could cause such discrepancies during evaluation?

2. After completing fine-tuning, the final evaluation result also differs slightly from the result I obtain when I load the fine-tuned model and perform evaluation again. Even though I have set do_sample=False during generation, I am unsure why these results are not identical. Shouldn’t they be the same? Did the authors also encounter such issues?

[Dornavineeth]

@ma-kjh Could you please share:

1. The script and configuration files you used for evaluation.
2. The results you're obtaining. Minor differences are fine, especially when running on different machines.

Note: I recommend pulling the latest updates and re-run setup_data.py. In our recent PR, we introduced a fixed seed for evaluation as well. However, using the same seed does not guarantee exactly same results across different machines, due to potential variations at the hardware or system level.

[ma-kjh]

I'm noticing that the evaluation results from checkpoint-last-epoch/TOFU_eval.json or TOFU_summary.json differ from the results I get when I reload the final checkpoint and run evaluation again (e.g., eval0). I'm wondering why this discrepancy occurs — is there something I might be missing?

/.hydra/config.yaml

model:
  model_args:
    pretrained_model_name_or_path: open-unlearning/tofu_Llama-3.2-1B-Instruct_full
    attn_implementation: flash_attention_2
    torch_dtype: bfloat16
  tokenizer_args:
    pretrained_model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
  template_args:
    apply_chat_template: true
    system_prompt: You are a helpful assistant.
    system_prompt_with_special_tokens: '<|begin_of_text|><|start_header_id|>system<|end_header_id|>


      You are a helpful assistant.<|eot_id|>'
    user_start_tag: '<|start_header_id|>user<|end_header_id|>


      '
    user_end_tag: <|eot_id|>
    asst_start_tag: '<|start_header_id|>assistant<|end_header_id|>


      '
    asst_end_tag: <|eot_id|>

trainer:
  handler: GradDiff
  args:
    per_device_train_batch_size: 8
    per_device_eval_batch_size: 16
    gradient_accumulation_steps: 2
    learning_rate: 1.0e-05
    bf16: true
    bf16_full_eval: true
    logging_steps: 5
    output_dir: ${paths.output_dir}
    logging_dir: ${trainer.args.output_dir}/logs
    report_to: tensorboard
    ddp_find_unused_parameters: None
    gradient_checkpointing: false
    optim: paged_adamw_32bit
    save_strategy: 'no'
    save_only_model: true
    weight_decay: 0.01
    do_train: true
    do_eval: true
    eval_on_start: true
    eval_strategy: epoch
    num_train_epochs: 10
    seed: 0
    warmup_epochs: 1.0
    remove_unused_columns: false
  method_args:
    gamma: 0.1
    alpha: 0.9
    retain_loss_type: NLL

data:
  forget:
    TOFU_QA_forget:
      handler: QADataset
      args:
        hf_args:
          name: ${forget_split}
          split: train
          path: locuslab/TOFU
        question_key: question
        answer_key: answer
        max_length: 512
  retain:
    TOFU_QA_retain:
      handler: QADataset
      args:
        hf_args:
          name: ${retain_split}
          split: train
          path: locuslab/TOFU
        question_key: question
        answer_key: answer
        max_length: 512
  anchor: forget
collator:
  DataCollatorForSupervisedDataset:
    handler: DataCollatorForSupervisedDataset
    args:
      padding_side: right
eval:
  tofu:
    metrics:
      forget_quality:
        pre_compute:
          forget_truth_ratio:
            pre_compute:
              forget_Q_A_PARA_Prob:
                datasets:
                  TOFU_QA_forget_para:
                    handler: QADataset
                    args:
                      hf_args:
                        name: ${eval.tofu.forget_split}_perturbed
                        split: train
                        path: locuslab/TOFU
                      question_key: question
                      answer_key: paraphrased_answer
                      max_length: 512
                collators:
                  DataCollatorForSupervisedDataset:
                    handler: DataCollatorForSupervisedDataset
                    args:
                      padding_side: right
                      index: index
                handler: probability
                batch_size: 32
                access_key: correct
              forget_Q_A_PERT_Prob:
                datasets:
                  TOFU_QA_forget_pert:
                    handler: QADataset
                    args:
                      hf_args:
                        name: ${eval.tofu.forget_split}_perturbed
                        split: train
                        path: locuslab/TOFU
                      question_key: question
                      answer_key: perturbed_answer
                      max_length: 512
                collators:
                  DataCollatorForSupervisedDataset:
                    handler: DataCollatorForSupervisedDataset
                    args:
                      padding_side: right
                      index: index
                handler: probability
                batch_size: 32
                access_key: wrong
            handler: truth_ratio
            aggregator: closer_to_1_better
            access_key: forget
        reference_logs:
          retain_model_logs:
            path: ${eval.tofu.retain_logs_path}
            include:
              forget_truth_ratio:
                access_key: retain
        handler: forget_quality
      forget_Q_A_Prob:
        datasets:
          TOFU_QA_forget:
            handler: QADataset
            args:
              hf_args:
                name: ${eval.tofu.forget_split}
                split: train
                path: locuslab/TOFU
              question_key: question
              answer_key: answer
              max_length: 512
        collators:
          DataCollatorForSupervisedDataset:
            handler: DataCollatorForSupervisedDataset
            args:
              padding_side: right
              index: index
        handler: probability
        batch_size: 32
      forget_Q_A_ROUGE:
        datasets:
          TOFU_QA_forget:
            handler: QADataset
            args:
              hf_args:
                name: ${eval.tofu.forget_split}
                split: train
                path: locuslab/TOFU
              question_key: question
              answer_key: answer
              max_length: 512
              predict_with_generate: true
        collators:
          DataCollatorForSupervisedDataset:
            handler: DataCollatorForSupervisedDataset
            args:
              padding_side: left
              index: index
        generation_args:
          do_sample: false
          top_p: null
          temperature: null
          max_new_tokens: 200
          use_cache: true
        handler: rouge
        rouge_type: rougeL_recall
        batch_size: 32

[Dornavineeth]

Can you share the bash command you used for

1. Unlearning
2. Evaluating it after unlearning
   Would be great if you can share the exact numbers too.

Also what do you mean by eval0?

[ma-kjh]

CUDA_VISIBLE_DEVICES=0,1 python src/train.py \
--config-name=unlearn.yaml experiment=unlearn/tofu/default   \ 
forget_split=forget10 retain_split=retain90 \
trainer=GradDiff_01_09 task_name=Llama-3.2-1B-Instruct-GradDiff-gamma-0.1-aplha-0.9-ep10 \
trainer.args.per_device_train_batch_size=8 \ 
trainer.args.gradient_accumulation_steps=2 \ 
trainer.args.num_train_epochs=10

GradDiff_01_09

defaults:
  - finetune

handler: GradDiff
method_args:
  gamma: 0.1
  alpha: 0.9
  retain_loss_type: NLL

/open-unlearning/saves/unlearn/Llama-3.2-1B-Instruct-GradDiff-gamma-0.1-aplha-0.9-ep10/checkpoint-0/evals

Also what do you mean by eval0?

Sorry, I meant the evaluation result of checkpoint-0 (eval0).

[molereddy]

To be clear, is the discrepancy that when you run GradDiff unlearning, you see different results when you check the final checkpoints eval logs (that were created during unlearning) compared to when you run evaluation on the final checkpoint separately?
Can you share these differences in the results?

[ma-kjh]

To be clear, is the discrepancy that when you run GradDiff unlearning, you see different results when you check the final checkpoints eval logs (that were created during unlearning) compared to when you run evaluation on the final checkpoint separately?

Yes, that's correct.

The differences are very minor, but I was expecting identical performance if the generation process was truly the same. So it seems like there might be slight variations introduced during evaluation — possibly due to factors like a specific seed or something similar. Here's a subset of the results for reference:

{
    "forget_Q_A_PARA_Prob": 0.08903286061016842,
    "forget_Q_A_PERT_Prob": 0.04107259407039237,
    "forget_Q_A_Prob": 0.6961516639590264,
    "forget_Q_A_ROUGE": 0.5595364443170646,
    "forget_truth_ratio": 0.5076989560333867,
    "model_utility": 0.5843485661144942,
    "ra_Q_A_PERT_Prob": 0.004951504494462521,
    "ra_Q_A_Prob": 0.012509137191664195,
    "ra_Q_A_Prob_normalised": 0.4029763631866845,
    "ra_Q_A_ROUGE": 0.7895,
    "ra_Truth_Ratio": 0.5204860241064739,
    "retain_Q_A_PARA_Prob": 0.07882677946239709,
    "retain_Q_A_PERT_Prob": 0.03392576826864133,
    "retain_Q_A_Prob": 0.8609627432376147,
    "retain_Q_A_ROUGE": 0.7620648336468085,
    "retain_Truth_Ratio": 0.5138321446751211,
    "wf_Q_A_PERT_Prob": 0.0013983926368962888,
    "wf_Q_A_Prob": 0.0036429426018894963,
    "wf_Q_A_Prob_normalised": 0.4145705999070623,
    "wf_Q_A_ROUGE": 0.8321937321937323,
    "wf_Truth_Ratio": 0.5886950979743497
}

{
    "forget_Q_A_PARA_Prob": 0.09749096598126926,
    "forget_Q_A_PERT_Prob": 0.04551721187017392,
    "forget_Q_A_Prob": 0.698550868332386,
    "forget_Q_A_ROUGE": 0.548337072063773,
    "forget_truth_ratio": 0.5121633673643747,
    "model_utility": 0.5851525699014697,
    "ra_Q_A_PERT_Prob": 0.005680906512739058,
    "ra_Q_A_Prob": 0.014331501303968253,
    "ra_Q_A_Prob_normalised": 0.4059826104709217,
    "ra_Q_A_ROUGE": 0.7795000000000001,
    "ra_Truth_Ratio": 0.5242252289253063,
    "retain_Q_A_PARA_Prob": 0.0861762935621664,
    "retain_Q_A_PERT_Prob": 0.03761727208955787,
    "retain_Q_A_Prob": 0.8563998416811228,
    "retain_Q_A_ROUGE": 0.7536560653964418,
    "retain_Truth_Ratio": 0.5117162442860267,
    "wf_Q_A_PERT_Prob": 0.001642537959798236,
    "wf_Q_A_Prob": 0.00472154998595545,
    "wf_Q_A_Prob_normalised": 0.41980525401627733,
    "wf_Q_A_ROUGE": 0.8307692307692308,
    "wf_Truth_Ratio": 0.5908485741526334
}