Finetuning with personachat example is not working

I'm trying to run https://github.com/bigscience-workshop/petals/blob/main/examples/prompt-tuning-personachat.ipynb and it is failing on default settings with these exceptions:

```
Feb 08 10:24:01.406 [WARN] [/home/dev/projekty/petals-proj/finetuning/petals/client/sequential_autograd.py.sequential_forward:101] Caught exception when running forward via RemoteSpanInfo(start=16, end=43, peer_id=) (retry in 0 sec): RuntimeError("shape '[2, 272, 14336]' is invalid for input of size 1048576")
```

or

```
Feb 08 10:23:51.349 [WARN] [/home/dev/projekty/petals-proj/finetuning/petals/client/sequential_autograd.py.sequential_forward:101] Caught exception when running forward via RemoteSpanInfo(start=16, end=43, peer_id=) (retry in 0 sec): P2PHandlerError('Failed to call handler `TransformerConnectionHandler.rpc_forward_stream` at 12D3KooWRftAHGeKyYmq35tn5Daqiu4D9767xUfZJX7E6LH2yKs9: ')
```

I don't see such errors when I lower hyperparameters like this. I successfully trained up to ~50% progress without any error, then I stopped it myself:
```
NUM_PREFIX_TOKENS = 4 #16
MODEL_MAX_LENGTH = 50 #256
BATCH_SIZE = 2 # 8
```

Beside this changes, I only changed ```request_timeout=1800``` according to previous discussions on Discord. I'm using current ```git master``` of Petals. Model is Bloomz and at the time of reporting, the health.petals.ml reports "healthy" state of bloomz pool.