Unlearning dynamics curve on multi-gpu

[sh-qiangchen]

Thank you very much for maintaining the repository. This is an absolutely fantastic job.
I would like to seek help for the following three questions:

1. Which checkpoint is being evaluated in the file open-unlearning/eval.py? trainer.save_model(output_dir) would save the last checkpoint, rather than the best model.
2. For the forget10 split, per_device_train_batch_size=8 and 4 A800 GPUs, why was the following trainer_state.json obtained? （400/8/4=12.5 steps for one epoch, why doesn't it match here?）

{
 "log_history": [
    {
      "epoch": 1.6,
      "grad_norm": 3.0877197396864764,
      "learning_rate": 8.620689655172414e-06,
      "loss": -0.0265,
      "step": 5
    },
    {
      "epoch": 3.2,
      "grad_norm": 31.033850214532777,
      "learning_rate": 6.896551724137932e-06,
      "loss": -0.7045,
      "step": 10
    },
    {
      "epoch": 4.8,
      "grad_norm": 60.4976325117357,
      "learning_rate": 5.172413793103449e-06,
      "loss": -3.8988,
      "step": 15
    },
    {
      "epoch": 6.4,
      "grad_norm": 98.01466701483473,
      "learning_rate": 3.448275862068966e-06,
      "loss": -10.4987,
      "step": 20
    },
    {
      "epoch": 8.0,
      "grad_norm": 812.0205372731976,
      "learning_rate": 1.724137931034483e-06,
      "loss": -24.3783,
      "step": 25
    },
    {
      "epoch": 9.6,
      "grad_norm": 609.3995261258575,
      "learning_rate": 0.0,
      "loss": -67.0836,
      "step": 30
    },
    {
      "epoch": 9.6,
      "step": 30,
      "total_flos": 0.0,
      "train_loss": -17.765076827506224,
      "train_runtime": 322.3017,
      "train_samples_per_second": 12.411,
      "train_steps_per_second": 0.093
    }
  ],
  "logging_steps": 5,
  "max_steps": 30,
  "num_input_tokens_seen": 0,
  "num_train_epochs": 10
}

3. How can I use FinetuneTrainer.evaluate in multi-GPU setting (currently it only works on a single GPU)? Because I see that most papers present the unlearning dynamics curve of evaluation. In other words, I want to obtain evaluation metrics during the training process using multiple GPUs.

if self.evaluator:
    if self.accelerator.is_local_main_process:
        eval_metrics = {}
        if self.accelerator.num_processes == 1:
            run_dir = self._get_output_dir(trial=trial)

[molereddy]

We're investigating and will elaborate soon but for 1, 2 you may want to read the last few messages of this discussion from yesterday: #88 (comment)

[molereddy]

For 3: we suggest you save the checkpoints and evaluate them after to get the trajectories, since we currently don't support train-time evals on multiple GPUs
Our code is designed to do evaluations on a single GPU for the below reasons
There are few reasons why we think it won't work if we do it on multiple GPUs, for example: 1) Our metrics implementation assumes data and model are on same device and simplifies things for people to contribute metrics 2) Our current batched generation and evaluation code would need to be reworked to aggregate results from devices (if evaluated with deepspeed). In general we're not totally sure what can go wrong
One solution to this can be to avoid deepspeed/accelerate and just set device_map="auto", then self.accelerator.num_processes will be 1 and eval during train will run. But we've not tested if this will work in all cases. On methods like NPO, it might fail because it involves deepcopy of a model already loaded on multiple GPUs.

[molereddy]

For 1: there is no accepted definition of a best validation metric during unlearning. So we leave it to the users to decide their best model.
For the rest of the discussion let us finish it on #88

[molereddy]

We apologize, this earlier response wasn't entirely accurate in the below point.

One solution to this can be to avoid deepspeed/accelerate and just set device_map="auto", then self.accelerator.num_processes will be 1 and eval during train will run.

You actually don't even need to set device_map="auto" and use model parallel. Simply running the multi-gpu command without deepspeed (as @\ma-kjh did) will set the process to use distributed data parallel, which has num_processes as 1. Since there is only a single process, the eval code runs.

A example command that will give you multi-gpu training with evaluations so you can get you dynamics curve

CUDA_VISIBLE_DEVICES=0,1 python src/train.py --config-name=unlearn.yaml 
experiment=unlearn/tofu/default task_name=some_name trainer=GradDiff \
trainer.args.num_train_epochs=10 trainer.args.eval_strategy=epoch

We did not anticipate using DDP to train on multiple GPUs without deepspeed, thus our documentation in this regard is inaccurate as we assumed deepspeed usage. We will clarify this in the next update, but you can simply get your required training dynamics on multiple GPUs using DDP as above.

Of course, another option is our earlier solution of saving the model at each epoch and evaluating separately after.

[molereddy]

Let us know if this resolves things, and close the issue if it does!

[sh-qiangchen]

Let us know if this resolves things, and close the issue if it does!

Thank you for your timely response, which has resolved a question that has been troubling me for a long time.
Best regards!