

TITLE
Longformer version of RoBERTa error

COMMENTS
3

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info

- `transformers` version: 4.2.0
- Platform: Linux-4.15.0-123-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help

 albert, bert, GPT2, XLM: @LysandreJik
 Longformer/Reformer: @patrickvonplaten


## Information

Model I am using script to initialize Longformer starting from [HerBERT](https://huggingface.co/allegro/herbert-klej-cased-v1)

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Install dependencies: `python3 -m pip install -r requirements.txt`.
2. Install apex according to [official documentation](https://github.com/NVIDIA/apex).
3. Run command `CUDA_VISIBLE_DEVICES=0 python3 convert_model_to_longformer.py --finetune_dataset conllu`.

We are using dataset in `.jsonl` format, each line contains 1 CoNLLu entry. It is converted using custom `LineByLineTextDataset` class to LineByLine format from current version of `transformers`. I've added this class to be able to use it in older version (v3.0.2).

Using suggested by author on [allenai/longformer](https://github.com/allenai/longformer) I've used `transformers` in version `3.0.2` and it works fine. But I would like to use recent models to convert them to Long* version and I can't make conversion script work.


## Result

As a result of running command above with `transformers` in version `4.2.0` I've got:

```bash
Traceback (most recent call last):
  File "convert_model_to_longformer.py", line 277, in <module>
    pretrain_and_evaluate(
  File "convert_model_to_longformer.py", line 165, in pretrain_and_evaluate
    eval_loss = trainer.evaluate()
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/trainer.py", line 1442, in evaluate
    output = self.prediction_loop(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/trainer.py", line 1566, in prediction_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/trainer.py", line 1670, in prediction_step
    outputs = model(**inputs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 1032, in forward
    outputs = self.roberta(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 798, in forward
    encoder_outputs = self.encoder(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 498, in forward
    layer_outputs = layer_module(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 393, in forward
    self_attention_outputs = self.attention(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 321, in forward
    self_outputs = self.self(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "convert_model_to_longformer.py", line 63, in forward
    return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions) # v4.2.0
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/longformer/modeling_longformer.py", line 600, in forward
    diagonal_mask = self._sliding_chunks_query_key_matmul(
  File "/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/longformer/modeling_longformer.py", line 789, in _sliding_chunks_query_key_matmul
    batch_size, seq_len, num_heads, head_dim = query.size()
ValueError: too many values to unpack (expected 4)
```

I've changed function `/server/server_1/user/miniconda3/envs/longformer_summary/lib/python3.8/site-packages/transformers/models/longformer/modeling_longformer.py` up to line 789:

```bash
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        is_index_masked=None,
        is_index_global_attn=None,
        is_global_attn=None,
        output_attentions=False,
    ):
        """
        :class:`LongformerSelfAttention` expects `len(hidden_states)` to be multiple of `attention_window`. Padding to
        `attention_window` happens in :meth:`LongformerModel.forward` to avoid redoing the padding on each layer.

        The `attention_mask` is changed in :meth:`LongformerModel.forward` from 0, 1, 2 to:

            * -10000: no attention
            * 0: local attention
            * +10000: global attention
        """
        hidden_states = hidden_states.transpose(0, 1)

        # project hidden states
        query_vectors = self.query(hidden_states)
        key_vectors = self.key(hidden_states)
        value_vectors = self.value(hidden_states)
        print(f"query_vectors: {query_vectors.shape}")
        print(f"key_vectors: {key_vectors.shape}")
        print(f"value_vectors: {value_vectors.shape}")
        print(f"attention_mask: {attention_mask.shape}")


        seq_len, batch_size, embed_dim = hidden_states.size()
        assert (
            embed_dim == self.embed_dim
        ), f"hidden_states should have embed_dim = {self.embed_dim}, but has {embed_dim}"

        # normalize query
        query_vectors /= math.sqrt(self.head_dim)

        query_vectors = query_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1)
        key_vectors = key_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1)

        attn_scores = self._sliding_chunks_query_key_matmul(
            query_vectors, key_vectors, self.one_sided_attn_window_size
        )

        # values to pad for attention probs
        remove_from_windowed_attention_mask = (attention_mask != 0)[:, :, None, None]

        # cast to fp32/fp16 then replace 1's with -inf
        float_mask = remove_from_windowed_attention_mask.type_as(query_vectors).masked_fill(
            remove_from_windowed_attention_mask, -10000.0
        )

        print(f"attn_scores: {attn_scores.shape}")
        print(f"remove_from_windowed_attention_mask: {remove_from_windowed_attention_mask.shape}")
        print(f"float_mask: {float_mask.shape}")

        # diagonal mask with zeros everywhere and -inf inplace of padding
        diagonal_mask = self._sliding_chunks_query_key_matmul(
            float_mask.new_ones(size=float_mask.size()), float_mask, self.one_sided_attn_window_size
        )
```

And as a result I've got:

```bash
attention_mask: torch.Size([2, 1, 1, 1024])
query_vectors: torch.Size([1024, 2, 768])
key_vectors: torch.Size([1024, 2, 768])
value_vectors: torch.Size([1024, 2, 768])
attn_scores: torch.Size([2, 1024, 12, 513])
remove_from_windowed_attention_mask: torch.Size([2, 1, 1, 1, 1, 1024])
float_mask: torch.Size([2, 1, 1, 1, 1, 1024])
```
And after changing version to `3.0.2` and adding print statements I've got:
```bash
attention_mask: torch.Size([2, 1024])
query_vectors: torch.Size([1024, 2, 768])
key_vectors: torch.Size([1024, 2, 768])
value_vectors: torch.Size([1024, 2, 768])
attn_scores: torch.Size([2, 1024, 12, 513])
remove_from_windowed_attention_mask: torch.Size([2, 1024, 1, 1])
float_mask: torch.Size([2, 1024, 1, 1])
```

So maybe it's problem with `_sliding_chunks_query_key_matmul` function?

## Files:

convert_model_to_longformer.py, based on [allenai/longformer/scripts/convert_model_to_long.ipynb](https://github.com/allenai/longformer/blob/master/scripts/convert_model_to_long.ipynb):
```python3
import logging
import os
import math
import copy
import torch
import argparse
from dataclasses import dataclass, field
from transformers import RobertaForMaskedLM, XLMTokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, XLMTokenizer, PreTrainedTokenizer
from transformers import TrainingArguments, HfArgumentParser, XLMTokenizer, RobertaModel, XLMTokenizer
from transformers import LongformerSelfAttention # v4.2.0
# from transformers.modeling_longformer import LongformerSelfAttention # v3.0.2
from conllu import load_conllu_dataset, save_conllu_dataset_in_linebyline_format
from torch.utils.data.dataset import Dataset

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)



class LineByLineTextDataset(Dataset):
    """
    This will be superseded by a framework-agnostic approach
    soon.
    """

    def __init__(self, tokenizer: PreTrainedTokenizer, file_path: str, block_size: int):
        assert os.path.isfile(file_path)
        # Here, we do not cache the features, operating under the assumption
        # that we will soon use fast multithreaded tokenizers from the
        # `tokenizers` repo everywhere =)
        logger.info("Creating features from dataset file at %s", file_path)

        with open(file_path, encoding="utf-8") as f:
            lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        batch_encoding = tokenizer(
            lines,
            add_special_tokens=True,
            truncation=True,
            padding="max_length",
            max_length=block_size,
            pad_to_multiple_of=512)
        self.examples = batch_encoding["input_ids"]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i) -> torch.Tensor:
        return torch.tensor(self.examples[i], dtype=torch.long)


class RobertaLongSelfAttention(LongformerSelfAttention):
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
    ):
        return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)


class RobertaLongForMaskedLM(RobertaForMaskedLM):
    def __init__(self, config):
        super().__init__(config)
        for i, layer in enumerate(self.roberta.encoder.layer):
            # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
            layer.attention.self = RobertaLongSelfAttention(config, layer_id=i)


class RobertaLongModel(RobertaModel):
    def __init__(self, config):
        super().__init__(config)
        for i, layer in enumerate(self.encoder.layer):
            # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
            layer.attention.self = RobertaLongSelfAttention(config, layer_id=i)


def create_long_model(initialization_model, initialization_tokenizer, save_model_to, attention_window, max_pos):
    model = RobertaForMaskedLM.from_pretrained(initialization_model)
    tokenizer = XLMTokenizer.from_pretrained(initialization_tokenizer, model_max_length=max_pos)
    config = model.config

    # extend position embeddings
    tokenizer.model_max_length = max_pos
    tokenizer.init_kwargs['model_max_length'] = max_pos
    current_max_pos, embed_size = model.roberta.embeddings.position_embeddings.weight.shape
    max_pos += 2  # NOTE: RoBERTa has positions 0,1 reserved, so embedding size is max position + 2
    config.max_position_embeddings = max_pos
    assert max_pos > current_max_pos

    # allocate a larger position embedding matrix
    new_pos_embed = model.roberta.embeddings.position_embeddings.weight.new_empty(max_pos, embed_size)

    # copy position embeddings over and over to initialize the new position embeddings
    k = 2
    step = current_max_pos - 2
    while k < max_pos - 1:
        new_pos_embed[k:(k + step)] = model.roberta.embeddings.position_embeddings.weight[2:]
        k += step
    model.roberta.embeddings.position_embeddings.weight.data = new_pos_embed
    model.roberta.embeddings.position_ids.data = torch.tensor([i for i in range(max_pos)]).reshape(1, max_pos) # v4.2.0
    # model.roberta.embeddings.position_ids = torch.tensor([i for i in range(max_pos)]).reshape(1, max_pos) # v3.0.2

    # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
    config.attention_window = [attention_window] * config.num_hidden_layers
    for i, layer in enumerate(model.roberta.encoder.layer):
        longformer_self_attn = LongformerSelfAttention(config, layer_id=i)
        longformer_self_attn.query = copy.deepcopy(layer.attention.self.query)
        longformer_self_attn.key = copy.deepcopy(layer.attention.self.key)
        longformer_self_attn.value = copy.deepcopy(layer.attention.self.value)

        longformer_self_attn.query_global = copy.deepcopy(layer.attention.self.query)
        longformer_self_attn.key_global = copy.deepcopy(layer.attention.self.key)
        longformer_self_attn.value_global = copy.deepcopy(layer.attention.self.value)

        layer.attention.self = longformer_self_attn

    logger.info(f'saving model to {save_model_to}')
    model.save_pretrained(save_model_to)
    tokenizer.save_pretrained(save_model_to)
    return model, tokenizer

def copy_proj_layers(model):
    for i, layer in enumerate(model.roberta.encoder.layer):
        layer.attention.self.query_global = copy.deepcopy(layer.attention.self.query)
        layer.attention.self.key_global = copy.deepcopy(layer.attention.self.key)
        layer.attention.self.value_global = copy.deepcopy(layer.attention.self.value)
    return model

def pretrain_and_evaluate(args, model, tokenizer, eval_only, model_path, max_size):
    val_dataset = LineByLineTextDataset(
        tokenizer=tokenizer,
        file_path=args.val_datapath,
        block_size=max_size,
    )
    if eval_only:
        train_dataset = val_dataset
    else:
        logger.info(f'Loading and tokenizing training data is usually slow: {args.train_datapath}')
        train_dataset = LineByLineTextDataset(
            tokenizer=tokenizer,
            file_path=args.train_datapath,
            block_size=max_size,
        )

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=True,
        mlm_probability=0.15,
    )
    trainer = Trainer(
        model=model,
        args=args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        # prediction_loss_only=True,
    )

    eval_loss = trainer.evaluate()
    eval_loss = eval_loss['eval_loss']
    logger.info(f'Initial eval bpc: {eval_loss/math.log(2)}')

    exit(0)

    if not eval_only:
        trainer = Trainer(
            model=model,
            args=args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            prediction_loss_only=False,
        )

        trainer.train(model_path=model_path)
        trainer.save_model()

        eval_loss = trainer.evaluate()
        eval_loss = eval_loss['eval_loss']
        logger.info(f'Eval bpc after pretraining: {eval_loss/math.log(2)}')

@dataclass
class ModelArgs:
    attention_window: int = field(default=512, metadata={"help": "Size of attention window"})
    max_pos: int = field(default=1024, metadata={"help": "Maximum position"})


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--finetune_dataset", required=True, choices=["conllu"], help="Name of dataset to finetune")
    return parser.parse_args()

if __name__ == "__main__":
    parser = HfArgumentParser((TrainingArguments, ModelArgs,))
    args = parse_args()

    training_args, model_args = parser.parse_args_into_dataclasses(look_for_args_file=False, args=[
        '--output_dir', 'tmp_4.2.0',
        '--warmup_steps', '500',
        '--learning_rate', '0.00003',
        '--weight_decay', '0.01',
        '--adam_epsilon', '1e-6',
        '--max_steps', '3000',
        '--logging_steps', '500',
        '--save_steps', '500',
        '--max_grad_norm', '5.0',
        '--per_device_eval_batch_size', '2',
        '--per_device_train_batch_size', '2',
        '--gradient_accumulation_steps', '4',
        # '--evaluate_during_training',
        '--do_train',
        '--do_eval',
        '--fp16',
        '--fp16_opt_level', 'O2',
    ])

    if args.finetune_dataset == "conllu":
        saved_dataset = '/server/server_1/user/longformer_summary/conllu/'
        if not os.path.exists(saved_dataset):
            os.makedirs(saved_dataset)
            dataset = load_conllu_dataset('/server/server_1/user/conllu_dataset/')
            save_conllu_dataset_in_linebyline_format(dataset, saved_dataset)

    training_args.val_datapath = os.path.join(saved_dataset, 'validation.txt')
    training_args.train_datapath = os.path.join(saved_dataset, 'train.txt')

    initialization_model = 'allegro/herbert-klej-cased-v1'
    initialization_tokenizer = 'allegro/herbert-klej-cased-tokenizer-v1'

    roberta_base = RobertaForMaskedLM.from_pretrained(initialization_model)
    roberta_base_tokenizer = XLMTokenizer.from_pretrained(initialization_tokenizer, model_max_length=512)

    model_path = f'{training_args.output_dir}/{initialization_model}-{model_args.max_pos}'
    if not os.path.exists(model_path):
        os.makedirs(model_path)

    logger.info(f'Converting roberta-base into {initialization_model}-{model_args.max_pos}')
    model, tokenizer = create_long_model(
        initialization_model=initialization_model,
        initialization_tokenizer=initialization_tokenizer,
        save_model_to=model_path,
        attention_window=model_args.attention_window,
        max_pos=model_args.max_pos,
    )

    logger.info(f'Loading the model from {model_path}')
    tokenizer = XLMTokenizer.from_pretrained(model_path)
    model = RobertaLongForMaskedLM.from_pretrained(model_path)

    logger.info(f'Pretraining {initialization_model}-{model_args.max_pos} ... ')

    pretrain_and_evaluate(
        training_args,
        model,
        tokenizer,
        eval_only=False,
        model_path=training_args.output_dir,
        max_size=model_args.max_pos,
    )

    logger.info(f'Copying local projection layers into global projection layers... ')
    model = copy_proj_layers(model)
    logger.info(f'Saving model to {model_path}')
    model.save_pretrained(model_path)

    logger.info(f'Loading the model from {model_path}')
    tokenizer = XLMTokenizer.from_pretrained(model_path)
    model = RobertaLongModel.from_pretrained(model_path)

```

conllu.py
```python3
import re
import glob
import torch
from torch.utils.data import Dataset
import time
import os
import json
from xml.etree.ElementTree import ParseError
import xml.etree.ElementTree as ET
from typing import List, Dict
from sklearn.model_selection import train_test_split


def load_conllu_jsonl(
    path: str,
) -> List[Dict[str, str]]:
    dataset: List[Dict[str, str]] = list()
    with open(path, 'r') as f:
        for jsonl in f.readlines():
            json_file = json.loads(jsonl)
            conllu = json_file['conllu'].split('\n')
            doc_text: str = ""
            utterance: Dict[str, str] = dict()
            for line in conllu:
                try:
                    if line[0].isdigit():
                        if utterance:
                            masked_text = utterance["text"]
                            doc_text = f"{doc_text} {masked_text}.".strip()
                            utterance = dict()
                    elif line[0] == '#':
                        text = line[1:].strip()
                        key = text.split('=')[0].strip()
                        value = text.split('=')[1].strip()
                        utterance[key] = value
                except IndexError:
                    pass

            dataset.append({"text": doc_text})

    return dataset


def load_conllu_dataset(
    path: str,
    train_test_val_ratio: float = 0.1,
) -> Dict[str, List[Dict[str, str]]]:
    dataset: Dict[str, List[Dict[str, str]]] = dict()
    data_dict: Dict[str, List[str]] = dict()

    filepath_list = glob.glob(os.path.join(path, '*.jsonl'))

    train = filepath_list[:int(len(filepath_list)*0.8)]
    test = filepath_list[int(len(filepath_list)*0.8):int(len(filepath_list)*0.9)]
    val = filepath_list[int(len(filepath_list)*0.9):]

    data_dict["test"] = test
    data_dict["train"] = train
    data_dict["validation"] = val

    for key, value in data_dict.items():
        dataset_list: List[Dict[str, str]] = list()
        for filepath in value:
            data = load_conllu_jsonl(path=filepath)
            if data:
                dataset_list.extend(data)

        dataset[key] = dataset_list

    return dataset

def save_conllu_dataset_in_linebyline_format(
    dataset: Dict[str, List[Dict[str, str]]],
    save_dir: str,
) -> None:
    for key, value in dataset.items():
        with open(os.path.join(save_dir, f'{key}.txt'), 'w') as f:
            for line in value:
                # print(line["full"])
                f.write(f'{line["text"]}\n')

```

requirements.txt:
```bash
apex @ file:///server/server_1/user/apex
certifi==2020.12.5
chardet==4.0.0
click==7.1.2
datasets==1.2.0
dill==0.3.3
filelock==3.0.12
idna==2.10
joblib==1.0.0
multiprocess==0.70.11.1
numpy==1.19.4
packaging==20.8
pandas==1.2.0
pyarrow==2.0.0
pyparsing==2.4.7
python-dateutil==2.8.1
pytz==2020.5
regex==2020.11.13
requests==2.25.1
sacremoses==0.0.43
sentencepiece==0.1.94
six==1.15.0
tokenizers==0.8.1rc1
torch==1.7.1
tqdm==4.49.0
transformers==3.0.2
typing-extensions==3.7.4.3
urllib3==1.26.2
xxhash==2.0.0
```

## Expected behavior

Model should be converted, saved and loaded. After that it should be properly fine-tuned and saved on disk.
