Support petals distributed model classes

[gsarti]

Description

This PR adds preliminary support for the AutoDistributedModelForCausalLM class from the petals library. Concretely, it supports bypassing the usage of attention_mask in model.generate and model.forward, while support for attention information in petals is still WIP (see #158 for reference).

Usage

The following code snippet showcases an example of contrastive attribution using the Input X Gradient method and the Llama 65B model (tested on a 6GB RTX3060 machine):

import inseq
from petals import AutoDistributedModelForCausalLM

model_name = "enoch/llama-65b-hf"
model = AutoDistributedModelForCausalLM.from_pretrained(model_name).cuda()
inseq_model = inseq.load_model(model=model, tokenizer=model_name, attribution_method="input_x_gradient")

>>> [INFO] Make sure you follow the LLaMA's terms of use: https://bit.ly/llama2-license for LLaMA 2, https://bit.ly/llama-license for LLaMA 1
>>> [INFO] Using DHT prefix: llama-65b-hf
>>> [INFO] Loading model with input_x_gradient method...

txt = """Option 1: Take a 50 minute bus, then a half hour train, and finally a 10 minute bike ride.
Option 2: Take a 10 minute bus, then an hour train, and finally a 30 minute bike ride.
Which of the options above is faster to get to work?
The answer is Option """

out = inseq_model.attribute(
    txt,
    txt + "1",
    attributed_fn="contrast_prob_diff",
    contrast_targets=txt + "2",
    step_scores=["contrast_prob_diff", "probability"],
)

>>> Attributing with input_x_gradient...: 100%|██████████| 80/80 [00:37<00:00, 37.55s/it]

out.show()

Notes

Method requiring model internals (e.g. attention) are currently not supported and will raise exceptions if used alongside AutoDistributedModelForCausalLM classes.