Wrap ForwardContext around full model forward

[calpt]

This PR adapts the ForwardContext to be applied to the full model (including head) forward pass. The original base model forward wrapper is now moved to wrap_base to make sure no second ForwardContext is created for a single forward pass.

This enables passing custom args that are defined in the ForwardContext definition to the top-level model call, as discussed in #783, e.g.:

model = AutoModelForCausalLM.from_pretrained(model_name)
adapters.init(model)
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(["This is a test text"], return_tensors="pt")

# Registers new forward args globally
ForwardContext.context_args.add("task_ids")

# New the new arg name can be used w/o modifying the model's forward method
output = model(**inputs, task_ids=["id_0", "id_1"])

In the example above, the forward context will automatically add the passed context args as attributes, ie. they can be accessed within the foward pass like this:

task_ids = ForwardContext.get_context().task_ids

[FrLdy]

Hello @calpt,
Thanks for the PR !
Just tested it out and it works great !
This definitely makes passing custom args much easier, no need to redefine a new forward for each specific hf model.
I'll rebase my fork once the changes are merged and propose my additions for multi-task support.

[FrLdy]

Hi,
I tried to include a new ForwardContext in my test for multi-task composition:
https://github.com/FrLdy/adapters/blob/6f63c4df35d8deed973f4c791f6b779f8ba4f668/tests/test_misc/test_adapter_composition.py#L177

However, if I don't check whether the new ForwardContext argument (task_ids) is present in context_args, the test runs fine when each TestCase class is executed independently. But when they are run together, only the first one passes, while the subsequent ones fail.

Maybe use a set to store ForwardContext.context_args ?

[calpt]

Maybe use a set to store ForwardContext.context_args ?

makes sense, done that

[FrLdy]

Hello,
I tried to use added context_args with generate method.
Unfortunately, context arguments are rejected by _validate_model_kwargs method from GenerationMixin.

import transformers
import adapters
import torch, torch.nn as nn

model = transformers.T5ForConditionalGeneration(
    transformers.T5Config(
        num_layers=2,
        num_decoder_layers=2
    )
)
adapters.init(model)

model.add_adapter("a")

adapters.ForwardContext.context_args.add("task_ids")
model.generate(input_ids=torch.randint(0, 1000, (3,128)), task_ids=torch.randint(0, 3, (3,)))

which returns

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[20], line 16
     13 model.add_adapter("a", overwrite_ok=True)
     15 adapters.ForwardContext.context_args.add("task_ids")
---> 16 model.generate(input_ids=torch.randint(0, 1000, (3,128)), task_ids=torch.randint(0, 3, (3,)))

File ~/Documents/labo/etr-peft-composition/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/Documents/labo/etr-peft-composition/src/adapters/hf_transformers/src/transformers/generation/utils.py:2009, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2006 assistant_tokenizer = kwargs.pop("assistant_tokenizer", None)  # only used for assisted generation
   2008 generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs)
-> 2009 self._validate_model_kwargs(model_kwargs.copy())
   2010 self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer)
   2012 # 2. Set generation parameters if not already defined

File ~/Documents/labo/etr-peft-composition/src/adapters/hf_transformers/src/transformers/generation/utils.py:1388, in GenerationMixin._validate_model_kwargs(self, model_kwargs)
   1385         unused_model_args.append(key)
   1387 if unused_model_args:
-> 1388     raise ValueError(
   1389         f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"
   1390         " generate arguments will also show up in this list)"
   1391     )

ValueError: The following `model_kwargs` are not used by the model: ['task_ids'] (note: typos in the generate arguments will also show up in this list)

Additional Notes

• No problem using AutoAdapterModel.

[hSterz]

Looks good!