Support for Passing task_ids to forward_context for Multi-Task Learning

[FrLdy]

Discussed in #783

^{Originally posted by FrLdy January 20, 2025}

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. 😃