KeyError: 'tokenizer'

[Ther-nullptr]

When I run PiSSA/scripts/metamath_llama2_7b/run_qlora.sh and to save a checkpoint, I get the following error:

{'loss': 0.3431, 'grad_norm': 0.09538137912750244, 'learning_rate': 1.9506710479124212e-05, 'epoch': 0.13}                                                                                             
 13%|████████████████████▏                                                                                                                                         | 100/781 [33:26<3:49:05, 20.18s/it]Traceback (most recent call last):
  File "/home/yujin-wa20/projects/PiSSA/train.py", line 303, in <module>
    train()
  File "/home/yujin-wa20/projects/PiSSA/train.py", line 292, in train
    trainer.train(resume_from_checkpoint = resume_from_checkpoint_dir)
  File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 2245, in train
    return inner_training_loop(
  File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 2620, in _inner_training_loop
    self._maybe_log_save_evaluate(
  File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 3101, in _maybe_log_save_evaluate
    self.control = self.callback_handler.on_save(self.args, self.state, self.control)
  File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer_callback.py", line 545, in on_save
    return self.call_event("on_save", args, state, control)
  File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer_callback.py", line 556, in call_event
    result = getattr(callback, event)(
  File "/home/yujin-wa20/projects/PiSSA/train.py", line 69, in on_save
    self.save_model(args, state, kwargs)
  File "/home/yujin-wa20/projects/PiSSA/train.py", line 66, in save_model
    kwargs["tokenizer"].save_pretrained(peft_model_path)
KeyError: 'tokenizer'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/yujin-wa20/projects/PiSSA/train.py", line 303, in <module>
[rank0]:     train()
[rank0]:   File "/home/yujin-wa20/projects/PiSSA/train.py", line 292, in train
[rank0]:     trainer.train(resume_from_checkpoint = resume_from_checkpoint_dir)
[rank0]:   File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 2245, in train
[rank0]:     return inner_training_loop(
[rank0]:   File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 2620, in _inner_training_loop
[rank0]:     self._maybe_log_save_evaluate(
[rank0]:   File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer.py", line 3101, in _maybe_log_save_evaluate
[rank0]:     self.control = self.callback_handler.on_save(self.args, self.state, self.control)
[rank0]:   File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer_callback.py", line 545, in on_save
[rank0]:     return self.call_event("on_save", args, state, control)
[rank0]:   File "/home/yujin-wa20/miniconda3/envs/halo/lib/python3.10/site-packages/transformers/trainer_callback.py", line 556, in call_event
[rank0]:     result = getattr(callback, event)(
[rank0]:   File "/home/yujin-wa20/projects/PiSSA/train.py", line 69, in on_save
[rank0]:     self.save_model(args, state, kwargs)
[rank0]:   File "/home/yujin-wa20/projects/PiSSA/train.py", line 66, in save_model
[rank0]:     kwargs["tokenizer"].save_pretrained(peft_model_path)
[rank0]: KeyError: 'tokenizer'

[yuxuan-z19]

@Ther-nullptr A small tweak to train.py should suffice to save the tokenizer along with the PEFT model:

# train.py:L112
class SavePeftModelCallback(transformers.TrainerCallback):
    def __init__(self, tokenzier):
        self.tokenizer = tokenzier

    def save_model(self, args, state, kwargs):
        ...
        tokenizer = kwargs.get("tokenizer", self.tokenizer)
        tokenizer.save_pretrained(peft_model_path)

And then register the callback with the tokenizer:

# train.py:L413
if not script_args.full_finetune:
        trainer.add_callback(SavePeftModelCallback(tokenizer))

This way, the tokenizer will be saved automatically during PEFT training. Hope this helps!