Error for finetuning `/microsoft/analysing_pii_leakage/` for DP-based configuration

Dear community,

I tried to finetune /microsoft/analysing_pii_leakage/ for DP-based configuration using below commmand

python fine_tune.py --config_path ../configs/fine-tune/echr-gpt2-small-dp8.yml

I followed the instructions. However, I am receiving the below error:

-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Saving LM to pii-leakage/microsoft/analysing_pii_leakage/examples/experiments/experiment_00001. Train Size: 118161, Eval Size: 26258
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Train Sample: 72.  On 23 October 2000 the Aliens Office informed the applicant?s lawyer that the time-limit had been extended until such time as the applicant and his child were fully recovered.
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

  0%|                                                                                                                                                                       | 0/7384 [00:00<?, ?it/s]anaconda3/envs/pii-leakage/lib/python3.10/site-packages/torch/nn/modules/module.py:1117: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.
  warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "
anaconda3/envs/pii-leakage/lib/python3.10/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
{'final_epsilon_prv': 0.09999154539803283, 'final_epsilon_rdp': 0, 'epoch': 0}                                                                                                                       
  0%|                                                                                                                                                                       | 0/7384 [00:01<?, ?it/s]Traceback (most recent call last):
  File "pii-leakage/microsoft/analysing_pii_leakage/examples/fine_tune.py", line 81, in <module>
    fine_tune(*parse_args())
  File "pii-leakage/microsoft/analysing_pii_leakage/examples/fine_tune.py", line 73, in fine_tune
    lm.fine_tune(train_dataset, eval_dataset, train_args, privacy_args)
  File "pii-leakage/microsoft/analysing_pii_leakage/src/pii_leakage/models/language_model.py", line 298, in fine_tune
    return self._fine_tune_dp(train_dataset, eval_dataset, train_args, privacy_args)
  File "pii-leakage/microsoft/analysing_pii_leakage/src/pii_leakage/models/language_model.py", line 278, in _fine_tune_dp
    trainer.train()
  File "anaconda3/envs/pii-leakage/lib/python3.10/site-packages/transformers/trainer.py", line 1948, in train
    return inner_training_loop(
  File "anaconda3/envs/pii-leakage/lib/python3.10/site-packages/transformers/trainer.py", line 2289, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "anaconda3/envs/pii-leakage/lib/python3.10/site-packages/dp_transformers/dp_utils.py", line 264, in training_step
    loss.backward()
  File anaconda3/envs/pii-leakage/lib/python3.10/site-packages/torch/_tensor.py", line 488, in backward
    torch.autograd.backward(
  File "anaconda3/envs/pii-leakage/lib/python3.10/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "anaconda3/envs/pii-leakage/lib/python3.10/site-packages/torch/nn/modules/module.py", line 61, in __call__
    raise RuntimeError("You are trying to call the hook of a dead Module!")
RuntimeError: You are trying to call the hook of a dead Module!

I do appreciate your help.

Kind Regards,
Mah

Reproduce the failure with the command in examples/fine_tune.py and follow the call into src/pii_leakage/models/language_model.py, especially _fine_tune_dp. Inspect the traceback's dp_transformers training_step and the hook-related failure first. Done means the DP fine-tuning command completes without the dead Module error.