:hugs: transformers compatibility issues

Hello,

I'm trying to make the `DistributedBloomForCausalLM` work with our library [`inseq`](https://github.com/inseq-team/inseq) to extract feature attributions from BLOOM generations. However, at the moment I am facing some issues that prevent me from using the distributed model:

1. Inseq assumes the possibility of producing a structured output from `model.generate` by passing the `return_dict_in_generate=True` argument, as supported by HuggingFace. In your current implementation, there doesn't seem to be a way to extract such outputs, so when we access the property `sequences` an exception is thrown. To reproduce:

```python
import torch
import inseq
from transformers import BloomTokenizerFast
from petals import DistributedBloomForCausalLM

MODEL_NAME = "bigscience/bloom-petals"
model = DistributedBloomForCausalLM.from_pretrained(MODEL_NAME)
model = model.cuda()
inseq_model = inseq.load_model(model=model, tokenizer="bigscience/bloom-petals", attribution_method="saliency")
out = inseq_model.attribute(
"A cat in French is \"",
generation_args={"max_new_tokens": 3}
)
```

```shell
╭──────────────────────────── Traceback (most recent call last) ────────────────────────────╮
│ :1 in │
│ /usr/local/lib/python3.8/dist-packages/inseq/models/attribution_model.py:184 in attribute │
│ │
│ 181 │ │ │ ) │
│ 182 │ │ if not constrained_decoding: │
│ 183 │ │ │ encoded_input = self.encode(input_texts, return_baseline=True, include_ │
│ ❱ 184 │ │ │ generated_texts = self.generate(encoded_input, return_generation_output │
│ 185 │ │ logger.debug(f"reference_texts={generated_texts}") │
│ 186 │ │ attribution_method = self.get_attribution_method(method, override_default_a │
│ 187 │ │ attributed_fn = self.get_attributed_fn(attributed_fn) │
│ │
│ /usr/local/lib/python3.8/dist-packages/inseq/models/model_decorators.py:13 in │
│ attribution_free_wrapper │
│ │
│ 10 │ │ if self.is_hooked: │
│ 11 │ │ │ was_hooked = True │
│ 12 │ │ │ self.attribution_method.unhook() │
│ ❱ 13 │ │ out = f(self, *args, **kwargs) │
│ 14 │ │ if was_hooked: │
│ 15 │ │ │ self.attribution_method.hook() │
│ 16 │ │ return out │
│ │
│ /usr/local/lib/python3.8/dist-packages/inseq/models/huggingface_model.py:190 in generate │
│ │
│ 187 │ │ │ **kwargs, │
│ 188 │ │ ) │
│ 189 │ │ texts = self.tokenizer.batch_decode( │
│ ❱ 190 │ │ │ generation_out.sequences, │
│ 191 │ │ │ skip_special_tokens=True, │
│ 192 │ │ ) │
│ 193 │ │ if return_generation_output: │
╰───────────────────────────────────────────────────────────────────────────────────────────╯
AttributeError: 'Tensor' object has no attribute 'sequences'
```

2. Using Inseq we can bypass the generation step by attributing a pre-specified generation. In that case, feature attributions will be performed by calling normal forward/backward passes on the model step by step. If I try this by adapting the call to `model.attribute` as:

```python
out = inseq_model.attribute(
"A cat in French is \"",
generated_texts="A cat in French is \"chat\"",
generation_args={"max_new_tokens": 3}
)
```

I get the following error:

```shell
╭──────────────────────────── Traceback (most recent call last) ────────────────────────────╮
│ /usr/local/lib/python3.8/dist-packages/petals/client/remote_model.py:163 in forward │
│ │
│ 160 │ │ attention_mask: Optional[torch.Tensor] = None, │
│ 161 │ │ **kwargs, │
│ 162 │ ): │
│ ❱ 163 │ │ assert attention_mask is None, "DistributedBloomModel does not support atte │
│ 164 │ │ │
│ 165 │ │ for k, v in kwargs.items(): │
│ 166 │ │ │ if not (v is None or v is False): │
╰───────────────────────────────────────────────────────────────────────────────────────────╯
AssertionError: DistributedBloomModel does not support attention masks right now
```

Correct me if I'm wrong, but I believe both `return_dict_in_generate` and `attention_mask` support should be achievable for the `petals` implementation, right? Would you consider supporting such usage? Thanks in advance! :slightly_smiling_face: