moondream
moondream copied to clipboard
vikhyat
PEFT LoRA with gradient checkpointing
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
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)
Will have a fix for the ValueError: not enough values to unpack (expected 4, got 3) bug pushed shortly.