PEFT LoRA with gradient checkpointing

[rockerBOO]

Getting the following error when using gradient checkpointing with PEFT LoRA training.

NotImplementedError self.get_input_embeddings()

Traceback (most recent call last):
  File "/home/rockerboo/code/caption-train/moondream.py", line 388, in <module>
    main(args)
  File "/home/rockerboo/code/caption-train/moondream.py", line 116, in main
    model = get_peft_model(model, config)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/peft/mapping.py", line 136, in get_peft_model
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config, adapter_name=adapter_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/peft/peft_model.py", line 1094, in __init__
    super().__init__(model, peft_config, adapter_name)
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/peft/peft_model.py", line 133, in __init__
    model = self._prepare_model_for_gradient_checkpointing(model)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/peft/peft_model.py", line 421, in _prepare_model_for_gradient_checkpointing
    model.enable_input_require_grads()
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1649, in enable_input_require_grads
    self._require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1668, in get_input_embeddings
    raise NotImplementedError
NotImplementedError

Basically using the same script as the Finetuning notebook in this repo but adding the Lora PEFT to it.

    target_modules = ["mixer.Wqkv"]
    config = LoraConfig(
        r=1,
        lora_alpha=1,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=target_modules,
    )

    # Ideally this also helps to start up the grads for LoRA training to work with gradient checkpointing
    # model.enable_input_require_grads()

    # enabling gradient checkpointing before the PEFT model
    model.text_model.transformer.gradient_checkpointing_enable()

    model = get_peft_model(model, config)

With model.enable_input_require_grads() has similar NotImplementedError

Traceback (most recent call last):
  File "/home/rockerboo/code/caption-train/moondream.py", line 388, in <module>
    main(args)
  File "/home/rockerboo/code/caption-train/moondream.py", line 112, in main
    model.enable_input_require_grads()
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1649, in enable_input_require_grads
    self._require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1668, in get_input_embeddings
    raise NotImplementedError
NotImplementedError

I think to fix this would be adding the get_input_embeddings() to return self.text_model.get_input_embeddings()

def get_input_embeddings(self):
   return self.text_model.get_input_embeddings()

I can give it a shot and make a PR here soon.

Thank you!

Related issues:

• https://github.com/huggingface/peft/issues/1142
• https://github.com/kohya-ss/sd-scripts/issues/323

[vikhyat]

You're right on that being the fix - would love a PR!

[rockerBOO]

Added get_input_embedding() to a clone of the huggingface repo but ran into the following, but not sure if it's related.

*edit: I'm not sure if i'm doing it right either from the clone or using it so could just be how it is "cached". I'm probably not using it properly.

Traceback (most recent call last):
  File "/home/rockerboo/code/caption-train/moondream.py", line 390, in <module>
    main(args)
  File "/home/rockerboo/code/caption-train/moondream.py", line 121, in main
    train(model, tokenizer, args)
  File "/home/rockerboo/code/caption-train/moondream.py", line 303, in train
    loss = compute_loss(batch, accelerator)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/code/caption-train/moondream.py", line 243, in compute_loss
    img_embs = model.vision_encoder.encoder(images)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.cache/huggingface/modules/transformers_modules/moondream2/vision_encoder.py", line 119, in forward
    return self.model["visual"](x)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.cache/huggingface/modules/transformers_modules/moondream2/vision_encoder.py", line 105, in forward
    x = self.patch_embed(x)
        ^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.pyenv/versions/3.11.6/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rockerboo/.cache/huggingface/modules/transformers_modules/moondream2/vision_encoder.py", line 129, in forward
    b, c, hp1, wp2 = x.shape
    ^^^^^^^^^^^^^^
ValueError: not enough values to unpack (expected 4, got 3)

[vikhyat]

I'm getting this too, investigating.

[vikhyat]

Will have a fix for the ValueError: not enough values to unpack (expected 4, got 3) bug pushed shortly.