Support for Passing task_ids to forward_context for Multi-Task Learning

[FrLdy]

Hello everyone,

I want to implement the paper MTL-LoRA: Low-Rank Adaptation for Multi-Task Learning.
I read in the documentation that the Parallel composition can be used to enable parallel multi-task training. However, it seems that in this particular case, the BatchSplit composition might be more suitable for routing tasks to the appropriate LoRA module.

I started extending the BatchSplit concept into a dynamic version that uses a task_ids parameter in the forward method, storing it in the context via forward_context.
To follow @calpt's recommendation, I decided to rely on Hugging Face's classes through adapters.init. However, when adding and running tests for test_adapter_composition (which uses transformers.BertForSequenceClassification), I noticed that Hugging Face model classes do not account for **kwargs in their method signatures. This prevents simply passing a new parameter like task_ids for handling it through forward_context.

Additionally, I also tested creating a new class inheriting from ModelBaseAdaptersMixin, inspired by adapters.T5ForConditionalGenerationWithHeadsMixin, to ensure the context is initialized during the first forward call. However, this results in two contexts being created.

Do you think there’s an elegant way to implement this functionality, or would it be necessary to override the forward methods of each model to properly handle the additional parameter?

Any guidance or recommendations would be greatly appreciated. 😃

[calpt]

Hey, this sounds like an interesting project, looking forward to your implementation!

It seems like you already pretty correctly summarized the current limitations. Leveraging the ForwardContext would be the way to go for adding additional forward arguments to existing models. Unfortunately, the ForwardContext is only applied to the forward pass of base model classes (ie AutoModel and friends) currently. Meaning if you wrap a base class instance via adapters.init(), you can easily catch any custom arguments in the forward context. For any "classes with head" (e.g. AutoModelForSequenceClassification), the forward method is not wrapped, so it needs manual overriding at least to pass the args to the base model forward.

tldr: currently, there's no elegant way to add this for all models with prediction heads.

But I think what I could look into is to adapt the current ForwardContext logic in a way that the classes with heads get their forward method wrapped by default while avoiding creation of a second context in the base model forward.

[calpt]

@FrLdy update on this: I've drafted a PR here: #789 to make it easier to pass custom args to a model via the ForwardContext. Feel free to give it a try and let me know if this helps!