MVP of the dynamic lora allocation framework

[MJHamar]

• inherits LoraLayer and Module with minimal modifications
• the only difference is that we extract the L2 norms of the forward passes and keep a cumulative sum
• how these are handled is TBD
• DynaLoraModel only needs to override the _create_new_module function for now to be compatible (I'm pretty sure)
• I will still need to include functions to handle the activation/deactivation of modules on demand

Have a look, let's have a discussion!

[MJHamar]

self.scheduler.step() is not called because by default, PeftModel.get_base_model returns the backbone model instance (with the LoRAs attached, but without the LoraModel wrapping it). As a result, DynaLoraModel.forward() is never called.

We can fix this by overriding get_base_model in PeftModelWrapper, but it results in all kinds of different errors, because the target model's forward signature is now lost. I'm in the process of fixing it, but may not be able to do it today.

[MJHamar]

Working with the issue that the scheduler is never called.

• The cause of this issue is that the PeftModel, which is PEFT's wrapper around any model implementing peft.tuners.tuners_utils.BaseTuner, provides an api function get_base_model, which by default returns the target model (e.g. roBERTa) instead of the BaseTuner, and as a result the BaseTuner's call method is never invoked
• Overwritten the get_base_model method in utils.wrappers.PeftModelWrapper to force it to return the DynaLoraModel.
• This gave the error ([[#full error]]):

  IndexError: Invalid key: 8431 is out of bounds for size 0

• And is caused within the transformers.Trainer instance, when the dataloader is initialised.
  • Essentially, it drops "unused" columns from the (tokenised) input data, based on the base model's forward signature. In our case, the 'base model' is the DynaLoraInstance and it has forward signature (*args, **kwargs), which does not match any column names.
• To solve this, we need to set
  • hf_args.remove_unused_columns = False in the train.py script, before the Trainer is initialised and remove these columns manually.
  • To remain target-model-agnostic, I added a function that mimics the Trainer class' _remove_unused_columns method, which drops all columns of the input dataset which do not match any of the target model's forward function signature.
  • In case of roBERTa, the signature is

    (input_ids: Optional = Noneattention_mask: Optional = Nonetoken_type_ids: Optional = Noneposition_ids: Optional = Nonehead_mask: Optional = Noneinputs_embeds: Optional = Nonelabels: Optional = Noneoutput_attentions: Optional = Noneoutput_hidden_states: Optional = Nonereturn_dict: Optional = None)

[MJHamar]

full error

  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/src/train.py", line 191, in <module>
    main()
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/src/train.py", line 185, in main
    trainer.train()
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/transformers/trainer.py", line 1859, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/transformers/trainer.py", line 2165, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/accelerate/data_loader.py", line 452, in __iter__
    current_batch = next(dataloader_iter)
                    ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
           ^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/arrow_dataset.py", line 2865, in __getitems__
    batch = self.__getitem__(keys)
            ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/arrow_dataset.py", line 2861, in __getitem__
    return self._getitem(key)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/arrow_dataset.py", line 2845, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/formatting/formatting.py", line 587, in query_table
    _check_valid_index_key(key, size)
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/formatting/formatting.py", line 537, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "/Users/hamarmiklos/Desktop/Egyetem/uva/FoMo-LoRA/.venv/lib/python3.11/site-packages/datasets/formatting/formatting.py", line 527, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 8431 is out of bounds for size 0