modules_to_save resulting in empty tensor with deepspeed zero3 LoRA training

[agokrani]

System Info

peft==0.15.0
transformers==4.49.0

Who can help?

@BenjaminBossan I am training a model with LoRA using DeepSpeed zero3, and I have set modules_to_save=["embed_tokens", "lm_head"]. The model trained successfully without any errors. However, while merging the model, I am encountering the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/agokrani/Documents/git/zett/env-zett-311/lib/python3.11/site-packages/transformers/integrations/peft.py", line 239, in load_adapter
    incompatible_keys = set_peft_model_state_dict(
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/agokrani/Documents/git/zett/env-zett-311/lib/python3.11/site-packages/peft/utils/save_and_load.py", line 444, in set_peft_model_state_dict
    load_result = model.load_state_dict(peft_model_state_dict, strict=False)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/agokrani/Documents/git/zett/env-zett-311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2581, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for Phi3ForCausalLM:
	size mismatch for model.embed_tokens.modules_to_save.default.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([151669, 3072]).

Upon inspecting the adapter_model.safetensors file, I noticed that both embed_tokens and lm_head tensors were saved, but they are empty tensors. I have debugged the root cause, and it seems that the following code was removed from the get_peft_model_state_dict function:

if getattr(model, "modules_to_save", None) is not None:
    for key, value in state_dict.items():
        if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save):
            to_return[key.replace("modules_to_save.", "")] = value

Instead, this code was added:

# ADDITIONAL TRAINING MODULES / MODULES_TO_SAVE
for name, module in model.named_modules():
    if isinstance(module, AuxiliaryTrainingWrapper):
        to_return.update({f"{name}.{k}": v for k, v in module.adapter_state_dict(adapter_name).items()})

Although the code above adds the embedding layers to the to_return dict, it only adds empty tensors instead of the actual weights. The easiest way to solve this problem is to re-add the removed code. However, I am not 100% sure if this is the correct approach. Do you have any idea why this might be happening? My guess is that DeepSpeed zero3 might be causing this problem. If the solution is indeed just to restore the removed code, I can create the PR. Please let me know what you think and whether I should create the PR.

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

Just use deepspeed zero3 with latest version of peft and try to save embedding layers.

Expected behavior

The embedding layers shouldn't be empty tensors, but rather the actual modified weights of the embedding layers.

[BenjaminBossan]

The change stems from #2376. This PR should not affect modules_to_save but of course there is always the possibility that we missed something.

The easiest way to solve this problem is to re-add the removed code.

Just to be sure, when you did that, the same script started working correctly again? Is it possible for you to share it?

[agokrani]

Yes, I did check that same script started working correctly again. I am using

https://github.com/agokrani/distillKitPlus/tree/phi4-support-debugging

to start the training. If you want I can share the config file and commands I am using.

[agokrani]

I also see the following in the implementation details of #2376

 In theory this should handle saving modules_to_save as well as trainable_token_indices but that's one point that needs verification and careful review.

[githubnemo]

If you want I can share the config file and commands I am using.

Yes, that would be helpful.

[agokrani]

Following is the config file, I am using:

{
    "project_name": "distil-logits",
    "dataset": {
        "name": "bespoke-stratos-17k/dataset.jsonl",
        "split": "train",
        "logits_file": null,
        "num_samples": 20,
        "select_range": null, 
        "format_function": null
    },
    "models": {
        "teacher": "Qwen/QWQ-32B",
        "student": "agokrani/phi4-mini-instruct-qwq",
        "student_adapter": null,
        "teacher_adapter": null,
        "teacher_vocab_size": null,
        "trust_remote_code": true
    },
    "tokenizer": {
        "max_length": 16384,
        "chat_template": null,
        "pad_token_id": 151643
    },
    "training": {
        "output_dir": "distilled_model",
        "num_train_epochs": 1,
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": 1,
        "save_steps": 100,
        "logging_steps": 1,
        "learning_rate": 2e-5,
        "weight_decay": 0.05,
        "warmup_ratio": 0.1,
        "lr_scheduler_type": "cosine",
        "resume_from_checkpoint": null,
        "fp16": false,
        "bf16": true
    },
    "distillation": {
        "temperature": 2.0,
        "alpha": 0.2
    },
    "model_config": {
        "use_flash_attention": true
    },
    "lora": {
        "enable_training": true,
        "r": 16,
        "alpha": 32,
        "dropout": 0.05,
        "bias": "none",
        "task_type": "CAUSAL_LM",
        "target_modules": "all-linear",
        "modules_to_save": ["embed_tokens", "lm_head"]
    },
    "quantization": {
        "enabled": true
    }, 
    "execution": {
        "use_accelerate": true,
        "accelerate_config": "accelerate_configs/accelerate_zero3.yaml"
    }
}

and following command can be used to launch the training

accelerate launch --config_file config/accelerate_configs/accelerate_zero3.yaml scripts/local/distill_logits.py --config <config_path>

For more details you can look at the README.md file in the repo.