axolotl: deepseed multiGPU resume from checkpoint fails

Please check that this issue hasn’t been reported before.

• I searched previous Bug Reports didn’t find any similar reports.

Expected Behavior

should work

Current behaviour

en-hi-spot, pid=16745)     trainer.train(resume_from_checkpoint=resume_from_checkpoint)
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/transformers/trainer.py", line 1543, in train
(en-hi-spot, pid=16745)     return inner_training_loop(
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/transformers/trainer.py", line 1699, in _inner_training_loop
(en-hi-spot, pid=16745)     deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 402, in deepspeed_load_checkpoint
(en-hi-spot, pid=16745)     load_path, _ = deepspeed_engine.load_checkpoint(
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2724, in load_checkpoint
(en-hi-spot, pid=16745)     load_path, client_states = self._load_checkpoint(load_dir,
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2794, in _load_checkpoint
(en-hi-spot, pid=16745)     self.load_module_state_dict(checkpoint=checkpoint,
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2587, in load_module_state_dict
(en-hi-spot, pid=16745)     self.module.load_state_dict(
(en-hi-spot, pid=16745)   File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
(en-hi-spot, pid=16745)     raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
(en-hi-spot, pid=16745) RuntimeError: Error(s) in loading state_dict for PeftModelForCausalLM:
(en-hi-spot, pid=16745) 	Missing key(s) in state_dict: "base_model.model.model.embed_tokens.weight", "base_model.model.model.layers.0.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.0.self_attn.k_proj.base_layer.weight", "base_model.model.model.layers.0.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.0.self_attn.o_proj.base_layer.weight", "base_model.model.model.layers.0.mlp.gate_proj.base_layer.weight", "base_model.model.model.layers.0.mlp.up_proj.base_layer.weight", "base_model.model.model.layers.0.mlp.down_proj.base_layer.weight", "base_model.model.model.layers.0.input_layernorm.weight", "base_model.model.model.layers.0.post_attention_layernorm.weight", "base_model.model.model.layers.1.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.1.self_attn.k_proj.base_layer.weight", "base_model.model.model.layers.1.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.1.self_attn.o_proj.base_layer.weight", "base_model.model.model.layers.1.mlp.gate_proj.base_layer.weight", "base_model.model.model.layers.1.mlp.up_proj.base_layer.weight", "base_model.model.model.layers.1.mlp.down_proj.base_layer.weight", "base_model.model.model.layers.1.input_layernorm.weight", "base_model.model.model.layers.1.post_attention_layernorm.weight", "base_model.model.model.layers.2.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.2.self_attn.k_proj.base_layer.weight", "base_model.model.model.layers.2.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.2.self_attn.o_proj.base_layer.weight", "base_model.model.model.layers.2.mlp.gate_proj.base_layer.weight", "base_model.model.model.layers.2.mlp.up_proj.base_layer.weight", "base_model.model.model.layers.2.mlp.down_proj.base_layer.weight", "base_model.model.model.layers.2.input_layernorm.weight", "base_model.model.model.layers.2.post_attention_layernorm.weight", "base_model.model.model.layers.3.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.3.self_attn.k_proj.base_layer.weight", "base_model.model.model.layers.3.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.3.self_attn.o_proj.base_layer.weight", "base_model.model.model.layers.3.mlp.gate_proj.base_layer.weight", "base_model.model.model.layers.3.mlp.up_proj.base_layer.weight", "base_model.model.model.layers.3.mlp.down_proj.base_layer.weight", "base_model.model.model.layers.3.input_layernorm.weight", "base_model.model.model.layers.3.post_attention_layernorm.weight", "base_model.model.model.layers.4.self_attn.q_proj.base_layer.weight", "base_model.model.model.layers.4.self_attn.k_proj.base_layer.weight", "base_model.model.model.layers.4.self_attn.v_proj.base_layer.weight", "base_model.model.model.layers.4.self_attn.o_proj.base_layer.weight", "base_model.model.model.layers.4.mlp.gate_proj.base_layer.weight", "base_model.model.model.layers.4.mlp.up_proj.base_layer.weight", "base_model.model.model.layers.4.mlp.down_proj.base_layer.weight", "base_model.model.model.layers.4.input_layernorm.weight", "base_model.model.model.layers.4.post_attention_layernorm.weight", "base_model.model.model.laye

Steps to reproduce

when resuming training from checkpoint

Config yaml

base_model: unsloth/tinyllama

model_type: LlamaForCausalLM tokenizer_type: LlamaTokenizer is_llama_derived_model: true

load_in_8bit: false load_in_4bit: true strict: false

chat_template: chatml datasets:

• path: manishiitg/aditi-chat-instruct-hi-v1-dedupe type: completion

wandb_project: tiny-aditi
hub_model_id: manishiitg/tinyllama-chat-instruct-hi-v1 hf_use_auth_token: true

dataset_prepared_path: val_set_size: 0 output_dir: /sky-notebook/manishiitg/tinyllama-chat-instruct-hi-v1

sequence_len: 4096 sample_packing: true eval_sample_packing: false pad_to_sequence_len: true

adapter: qlora lora_model_dir: lora_r: 32 lora_alpha: 16 lora_dropout: 0.05 lora_target_linear: true lora_fan_in_fan_out:

wandb_entity: wandb_watch: wandb_name: wandb_log_model:

gradient_accumulation_steps: 4 micro_batch_size: 14 num_epochs: 4 optimizer: paged_adamw_32bit lr_scheduler: cosine learning_rate: 0.0002

train_on_inputs: false group_by_length: false bf16: true fp16: false tf32: false

gradient_checkpointing: true early_stopping_patience: resume_from_checkpoint: auto_resume_from_checkpoints: true ## manage check point resume from here local_rank: logging_steps: 1 xformers_attention: flash_attention: true

warmup_steps: 10 eval_steps: 0 eval_table_size: eval_table_max_new_tokens: 128 save_steps: 100 ## increase based on your dataset save_strategy: steps debug: deepspeed: weight_decay: 0.0 fsdp: fsdp_config: special_tokens: bos_token: “” eos_token: “” unk_token: “<unk>”

Possible solution

No response

Which Operating Systems are you using?

• Linux
• macOS
• Windows

Python Version

3.10

axolotl branch-commit

main

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this bug has not been reported yet.
• I am using the latest version of axolotl.
• I have provided enough information for the maintainers to reproduce and diagnose the issue.