Deepspeed unable to resume training on Peft
Author: lucasjinreal
Pull request: no

this was my main code:

```python
parser = argparse.ArgumentParser()
parser.add_argument("--wandb", action="store_true", default=False)
parser.add_argument("--prompt_type", type=str, default="chat")
parser.add_argument("--data_path", type=str, default="merge.json")
parser.add_argument("--output_path", type=str, default="out/lora-Vicuna-chat")
parser.add_argument("--model_path", type=str, default="decapoda-research/llama-7b-hf")
parser.add_argument("--num_epoch", type=int, default=3)
parser.add_argument("--micro_batch", type=int, default=4)
parser.add_argument("--total_batch", type=int, default=128)
parser.add_argument("--log_steps", type=int, default=100)
parser.add_argument("--eval_steps", type=int, default=200)
parser.add_argument("--save_steps", type=int, default=200)
parser.add_argument("--warmup_ratio", type=float, default=0.05)
parser.add_argument("--test_size", type=int, default=10)
parser.add_argument("--resume_from_checkpoint", type=str, default=None)
parser.add_argument("--lora_remote_checkpoint", type=str, default=None)
parser.add_argument("--ignore_data_skip", type=bool, default=False)
parser.add_argument("--int8_train", type=bool, default=False)
parser.add_argument("--deepspeed", type=str, default=False)
args = parser.parse_args()
if not args.wandb:
    os.environ["WANDB_MODE"] = "disable"
MICRO_BATCH_SIZE = args.micro_batch  # this could actually be 5 but i like powers of 2
BATCH_SIZE = args.total_batch
MAX_STEPS = None
GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE
EPOCHS = args.num_epoch
LEARNING_RATE = 3e-4  # the Karpathy constant
# CUTOFF_LEN = 2048
CUTOFF_LEN = 512
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
VAL_SET_SIZE = args.test_size  # 2000
TARGET_MODULES = [
    "q_proj",
    "v_proj",
    "k_proj",
    "o_proj",
    "down_proj",
    "gate_proj",
    "up_proj",
]
DATA_PATH = args.data_path
OUTPUT_DIR = args.output_path  # "lora-Vicuna"

device_map = "auto"
world_size = int(os.environ.get("WORLD_SIZE", 1))
ddp = world_size != 1
if ddp:
    device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
    GRADIENT_ACCUMULATION_STEPS = GRADIENT_ACCUMULATION_STEPS // world_size
# we must make sure batch_size and gradient_accumulation_steps not changed for resuming training.
if args.resume_from_checkpoint:
    s_ = utils.check_args_on_resume(args)
    print(f'Resume args check status: {s_}')
    # checkpoint = os.path.join(args.resume_from_checkpoint, 'pytorch_model.bin')

logger = utils.set_file_logger(__name__, OUTPUT_DIR)
# 1. load dataset
logger.info(f">>> processing data from {DATA_PATH}")
logger.info(f">>> using {args}")

train_tokenizer = LlamaTokenizer.from_pretrained(args.model_path, add_eos_token=True)
assert train_tokenizer.eos_token_id == 2, "Tokenizer eos is wrong!!!"
# unk. we want this to be different from the eos token
train_tokenizer.pad_token_id = 0
# cannot use eos in generation!
# tokenizer.padding_side = "left"  # Allow batched inference
test_tokenizer = LlamaTokenizer.from_pretrained(args.model_path)
if args.prompt_type == "instruct":
    PROMPT = prompt.instruct_prompt(train_tokenizer, CUTOFF_LEN)
elif args.prompt_type == "chat":
    PROMPT = prompt.chat_prompt(train_tokenizer, CUTOFF_LEN)
else:
    raise Exception("not support")
data = load_dataset("json", data_files=DATA_PATH)

start = random.randint(1, 100)
examples = Dataset.from_dict(data["train"][start : start + 5]).map(
    PROMPT.preprocess_train
)
for example in examples:
    logger.info(
        f'>>> using prompt {args.prompt_type}, prompt example:\n { train_tokenizer.decode(example["input_ids"]) }'
    )
    logger.info(
        f'>>> tokenizer labels: { train_tokenizer.decode([ 0 if l==-100 else l for l in example["labels"]])}'
    )
    logger.info(
        f'>>> tokenizer example: { example["input_ids"][:10] }...{ example["input_ids"][-10:]}'
    )
num_proc = os.cpu_count()
if VAL_SET_SIZE > 0:
    train_val = data["train"].train_test_split(
        test_size=VAL_SET_SIZE, shuffle=True, seed=42
    )
    train_data = (
        train_val["train"].shuffle().map(PROMPT.preprocess_train, num_proc=num_proc)
    )
    val_data = (
        train_val["test"].shuffle().map(PROMPT.preprocess_train, num_proc=num_proc)
    )
else:
    train_data = data["train"].shuffle().map(PROMPT.preprocess_train, num_proc=num_proc)
    val_data = None
now_max_steps = max((len(data["train"]) - VAL_SET_SIZE) // BATCH_SIZE * EPOCHS, EPOCHS)

logger.info(f">>> load model from {args.model_path}")
model = LlamaForCausalLM.from_pretrained(
    args.model_path,
    load_in_8bit=args.int8_train,
    device_map=device_map,
    torch_dtype=torch.float16,
)
if args.int8_train:
    model = prepare_model_for_int8_training(model)
config = LoraConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    target_modules=TARGET_MODULES,
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

class CustomCallback(TrainerCallback):
    def __init__(self, trainer) -> None:
        super().__init__()
        self.trainer = trainer
        self.generation_config = GenerationConfig(
            temperature=1.0,
            top_p=0.75,
            top_k=40,
            num_beams=2,
            bos_token_id=train_tokenizer.bos_token_id,
            eos_token_id=train_tokenizer.eos_token_id,
            pad_token_id=train_tokenizer.pad_token_id,
            max_new_tokens=1024,  # max_length=max_new_tokens+input_sequence
            min_new_tokens=1,  # min_length=min_new_tokens+input_sequence
            bad_words_ids=test_tokenizer(
                ["\n\nUser:", "\n\nAssistant:"], add_special_tokens=False
            ).input_ids,
        )
        self.repetition_penalty = 1.3
        self.logger = utils.set_file_logger(
            "transformers.trainer", trainer.args.output_dir
        )

    def on_log(self, args, state, control, logs, **kwargs):
        logger.info(logs)


model.print_trainable_parameters()
print(f"peft config of model: {model.peft_config}")
logger.info(f"model.modules_to_save: {model.modules_to_save}")
old_state_dict = model.state_dict
model.state_dict = (
    lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
).__get__(model, type(model))

if torch.__version__ >= "2" and sys.platform != "win32":
    # model = torch.compile(model)
    pass

model.save_pretrained(args.output_path)

# print(f"now FUCK model s: {model.state_dict().keys()}")
# print(f"{torch.load(os.path.join(args.resume_from_checkpoint, 'pytorch_model.bin')).keys()}")

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=MICRO_BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
        warmup_ratio=args.warmup_ratio,
        num_train_epochs=EPOCHS,
        # max_steps=MAX_STEPS,
        learning_rate=LEARNING_RATE,
        fp16=True,
        logging_steps=args.log_steps,
        logging_first_step=True,  # convenient
        evaluation_strategy="steps" if VAL_SET_SIZE > 0 else "no",
        save_strategy="steps",
        save_total_limit=2,
        eval_steps=args.eval_steps if VAL_SET_SIZE > 0 else None,
        save_steps=args.save_steps,
        output_dir=OUTPUT_DIR,
        load_best_model_at_end=True if VAL_SET_SIZE > 0 else False,
        ddp_find_unused_parameters=False if ddp else None,
        report_to="wandb" if args.wandb else [],
        ignore_data_skip=args.ignore_data_skip,
        deepspeed=args.deepspeed,
    ),
    data_collator=PROMPT.data_collator(),
)
trainer.add_callback(CustomCallback(trainer))
model.config.use_cache = False
trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
model.save_pretrained(OUTPUT_DIR)
```

the model training is OK, model save is OK

Got output like this:

```
(base) ➜  checkpoint-1200 git:(main) ll
total 115M
-rw-r--r-- 1 root root  77M May 30 20:37 aa
drwxr-xr-x 2 root root 268M May 30 17:30 global_step1200
-rw-r--r-- 1 root root   15 May 30 17:30 latest
-rw-r--r-- 1 root root  39M May 30 17:30 pytorch_model.bin
-rw-r--r-- 1 root root  16K May 30 17:30 rng_state_0.pth
-rw-r--r-- 1 root root  16K May 30 17:30 rng_state_1.pth
-rw-r--r-- 1 root root 3.1K May 30 17:30 trainer_state.json
-rw-r--r-- 1 root root 5.0K May 30 17:30 training_args.bin
-rwxr--r-- 1 root root  19K May 30 17:30 zero_to_fp32.py
```

but somehow I can't resume the checkpoint, 

From my limited knowledge, resume should send same as output path, and inside output path, we might have checkpiint-800 checkpint-1600 etc.

So I just resume from output_path.

Then it says 

ValueError: Can't find a valid checkpoint at out/lora-Vicuna-chat/

Why?????

I try to send a path like `out/lora-Vicuna-chat/checkpoint-600`, but also failed ,so strange 

Comments:
Hi, there, I got same error with deepspeed. training noramlly and resume got `Can't find a valid checkpoint`

first I try resume_from_checkpoint with `out/lora-Vicuna-chat` (output_path) got  `Can't find a valid checkpoint`
then I send `out/lora-Vicuna-chat/checkpoint-6000` I can not load the lora weights........

```
"base_model.model.model.layers.31.self_attn.k_proj.lora_B.default.weight", "base_model.model.model.layers.31.self_attn.v_proj.weight", "base_model.model.model.layers.31.self_attn.v_proj.lora_A.default.weight", 
"base_model.model.model.layers.31.self_attn.v_proj.lora_B.default.weight", "base_model.model.model.layers.31.self_attn.o_proj.weight", "base_model.model.model.layers.31.self_attn.o_proj.lora_A.default.weight", 
"base_model.model.model.layers.31.self_attn.o_proj.lora_B.default.weight", "base_model.model.model.layers.31.self_attn.rotary_emb.inv_freq", "base_model.model.model.layers.31.mlp.gate_proj.weight", 
"base_model.model.model.layers.31.mlp.gate_proj.lora_A.default.weight", "base_model.model.model.layers.31.mlp.gate_proj.lora_B.default.weight", "base_model.model.model.layers.31.mlp.down_proj.weight", 
"base_model.model.model.layers.31.mlp.down_proj.lora_A.default.weight", "base_model.model.model.layers.31.mlp.down_proj.lora_B.default.weight", "base_model.model.model.layers.31.mlp.up_proj.weight", 
"base_model.model.model.layers.31.mlp.up_proj.lora_A.default.weight", "base_model.model.model.layers.31.mlp.up_proj.lora_B.default.weight", "base_model.model.model.layers.31.input_layernorm.weight", 
  Unexpected key(s) in state_dict:"base_model.model.model.layers.31.self_attn.q_proj.lora_A.weight", "base_model.model.model.layers.31.self_attn.q_proj.lora_B.weight", "base_model.model.model.layers.31.self_attn.k_proj.lora_A.weight", 
"base_model.model.model.layers.31.self_attn.k_proj.lora_B.weight", "base_model.model.model.layers.31.self_attn.v_proj.lora_A.weight", "base_model.model.model.layers.31.self_attn.v_proj.lora_B.weight", 
"base_model.model.model.layers.31.self_attn.o_proj.lora_A.weight", "base_model.model.model.layers.31.self_attn.o_proj.lora_B.weight", "base_model.model.model.layers.31.mlp.gate_proj.lora_A.weight", 
"base_model.model.model.layers.31.mlp.gate_proj.lora_B.weight", "base_model.model.model.layers.31.mlp.down_proj.lora_A.weight", "base_model.model.model.layers.31.mlp.down_proj.lora_B.weight", 
```

the model with some suffix `default`, but samed model didn't have......

I am confused so much

cc @pacman100 

I am sorry but I found it might related about this:

```
model.state_dict = (
    lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
).__get__(model, type(model))

if torch.__version__ >= "2" and sys.platform != "win32":
    model = torch.compile(model)

print("\n If there's a warning about missing keys above, please disregard :)")

trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
```

I am replace model_stateduct after create Trainer, digging into code, found that `get_peft_model_state_dict` will replace the peft model sate keynname with some {adapter_name} as suffix.

Does this line of code must before create trainer? There is really lack documentation mentationed about this.

If so, then why must need users to do this manually when resume or same? 

And when using PeftModel.from_pretrained, it actually can set_peft_model_statedict automatically.....

These behaviour really makes me very confused.

@sgugger Sorry for pin again, but this problem obstacles me and make me very confused, please help me clarify, I made a clear code analysis to adress this problem: https://github.com/huggingface/peft/issues/746

Hello, looking into this and https://github.com/huggingface/peft/issues/746

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.