Training loss of BART is going to nan in transformers>=4.21.0
Author: soocheolnoh
Pull request: no
Labels: name: bug

### System Info

transformers==4.20.1 and transformers>=4.21.0
torch==1.12.1

### Who can help?

@patil-suraj

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

Hi, I'm using a huge dataset, so it is hard to show how to reproduce my problem.
I'm using [BART pre-trained model](https://huggingface.co/gogamza/kobart-base-v1) and trying to fine-tune the model as a translation model. But, the training loss is completely differently calculated depending on transformers's versions.

My pseudo code is like:
```python
net = BartForConditionalGeneration.from_pretrained("gogamza/kobart-base-v1").to(rank)
net.train()

with amp.autocast(enabled=True):
    output = net(
        input_ids=input_ids,
        attention_mask=attention_mask,
        decoder_input_ids=decoder_input_ids,,
        decoder_attention_mask=decoder_attention_mask,
    )
    # draw graphs of output.loss
```

I drawed the graph (training loss by iterations) using wandb.

![image](https://user-images.githubusercontent.com/57349134/186881592-594e303c-db8d-4453-b531-a15e0ab2eb0f.png)

`effortless-water-23` (green): `transformers>=4.21.0`
`swept-tree-24` (pink): `transformers==4.20.1`

`swept-tree-24` was slowly coverged to zero, but `effortless-water-23` eventually got nan at 80k+ iterations. (The above graph didn't show that.)
I've searched the difference between `transformers>=4.21.0` and `transformers==4.20.1` especially about BART, and I'm suspicious [this part](https://github.com/huggingface/transformers/commit/d3cb28886ac68beba9a6646b422a4d727b056c0c#diff-5b256a6fcfe7956b744dc1e902f150a79e2692b822a6dbac2c6612f8690846b6).

So as I reverted [the part](https://github.com/huggingface/transformers/blob/main/src/transformers/models/bart/modeling_bart.py#L96) in `transformers>=4.21.0` like:
```python
# mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min))
mask = torch.full((tgt_len, tgt_len), torch.tensor(float("-inf")))
```
the problem was gone. (The result is same as `swept-tree-24` used `transformers==4.20.1`)

Anyway my problem is solved, but I'm wondering what is the real cause of the problem.
Thanks in advance.

### Expected behavior

I explained this part in the reproduction part.

Comments:
This was introduced as part of https://github.com/huggingface/transformers/pull/17306

cc @ydshieh 

Hi @soocheolnoh.

My best guess is that #17306 is not working well with FP16 training (as I see you use `amp`).

- Is it possible for you to share your modified training script, together with the training arguments?

If you have enough computational resource:
  - I know you mentioned you run on a huge dataset. Would it possible for you to check if the issue could be reproduced with a smaller dataset (ideally, with a dataset available on [HF datasets](https://huggingface.co/datasets)?
  - Could you also try with FP32 training and see if the issue appears too?
(You don't need to get NaN - just a training loss graph showing the issue would be enough)

Thank you in advance!







Thanks for the response! @ydshieh 

In my experience, both fp16 and fp32 got nan.

![image](https://user-images.githubusercontent.com/57349134/187511019-42bd417c-889b-4804-88ae-e5b5a0a31419.png)

These experiments were executed using `transformers>=4.21.0` and I don't have any results using `transfomers==4.20.1` under the same condition used the above graph.

Currently I don't have any resource, so I can't try HF datasets right now. But when I get the resource, I will try it and see if the same issue happens.


@soocheolnoh Thank you for the detailed information. I will find some time to investigate this issue. If you are able to share any further findings, it would be very helpful.

Still wondering if it's possible for you to share your modified training script, together with the training arguments? Otherwise I will try to reproduce with the official training script and a HF dataset.

@ydshieh 

When I get GPUs and reproduce the issue again (with my private dataset/HF dataset), I will share my training script with separating essential parts as soon as possible. Currently my scripts aren't an unique file and I'm working on a private repository, so it's hard to share the original scripts.

No problem @soocheolnoh ! Thanks a lot.

Hi, @ydshieh 

With several experiments, I found the reason of the issue. Since I'm a kind of beginner of NLP task, I've used incorrectly a decoder start token(=`decoder_prefix` in args). Originally the model I used was pretrained with \
 decoder start token, but I've used \<pad> decoder start token which led an unexpected result I think.

My training code `train.py` is here:

```python
import argparse
import random
import warnings

import numpy as np
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
import torch.multiprocessing as mp
from datasets import load_dataset
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm
from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast
from transformers.optimization import get_constant_schedule


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--use_wandb", action="store_true")
    parser.add_argument("--wandb_project", default="hg-kobart", type=str)
    parser.add_argument("--num_gpus", default=3, type=int)
    parser.add_argument("--batch_per_gpu", default=24, type=int)
    parser.add_argument("--workers_per_gpu", default=6, type=int)
    parser.add_argument("--epochs", default=1, type=int)
    parser.add_argument("--decoder_prefix", default="<pad>", type=str)
    args = parser.parse_args()
    return args


def seed_everything(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)


class TranslationDataset(Dataset):
    def __init__(self, tokenizer, dataset_path, dataset_split, decoder_prefix, max_len=512, ignore_id=-100):
        super().__init__()
        self.tokenizer = tokenizer
        self.decoder_prefix = decoder_prefix
        self.max_len = max_len
        self.ignore_id = ignore_id
        self.pad_id = tokenizer.pad_token_id

        self.dataset = load_dataset(dataset_path, split=dataset_split)

    def __len__(self):
        return len(self.dataset)

    def _pad(self, input, padding_value=-100):
        return [*input, *(padding_value,) * (self.max_len - len(input))]

    def __getitem__(self, idx):
        data = self.dataset[idx]
        input = self.tokenizer(
            data["premise_original"] + "
", padding="max_length", truncation=True, max_length=self.max_len
        )
        decoder_input = self.tokenizer(
            self.decoder_prefix + data["premise"] + "
",
            padding="max_length",
            truncation=True,
            max_length=self.max_len,
        )
        output = self.tokenizer(data["premise"] + "
")
        output = self._pad(output["input_ids"])

        return {
            "input_ids": np.array(input["input_ids"], dtype=np.int64),
            "decoder_input_ids": np.array(decoder_input["input_ids"], dtype=np.int64),
            "labels": np.array(output, dtype=np.int64),
        }


def train(rank, world_size, args):
    dist.init_process_group(backend="nccl", init_method="tcp://127.0.0.1:33445", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    seed_everything(42 + rank)
    warnings.filterwarnings("ignore")

    tokenizer = PreTrainedTokenizerFast.from_pretrained("gogamza/kobart-base-v1")
    net = BartForConditionalGeneration.from_pretrained("gogamza/kobart-base-v1").to(rank).train()
    net = DDP(net, device_ids=[rank], find_unused_parameters=False)

    param_optimizer = list(net.named_parameters())
    no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {"params": [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], "weight_decay": 0.01},
        {"params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=3e-5)
    scheduler = get_constant_schedule(optimizer)
    scaler = amp.GradScaler(enabled=args.fp16)

    dataset = TranslationDataset(
        tokenizer=tokenizer,
        dataset_path="MoritzLaurer/multilingual-NLI-26lang-2mil7",
        dataset_split="ko_wanli",
        decoder_prefix=args.decoder_prefix,
    )
    sampler = DistributedSampler(dataset, shuffle=True, drop_last=True)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_per_gpu,
        num_workers=args.workers_per_gpu,
        pin_memory=True,
        drop_last=True,
        sampler=sampler,
        persistent_workers=True,
    )

    if args.use_wandb and rank == 0:
        import wandb

        wandb.init(project=args.wandb_project)
        wandb.watch(net)

    for epoch in range(args.epochs):
        if rank == 0:
            print(f"Epoch {epoch+1}/{args.epochs}")
        pbar = tqdm(loader) if rank == 0 else loader
        for data in pbar:
            input_ids = torch.tensor(data["input_ids"]).to(dtype=torch.int64, device=rank)
            attention_mask = torch.ne(input_ids, tokenizer.pad_token_id).to(dtype=torch.float32)
            decoder_input_ids = torch.tensor(data["decoder_input_ids"]).to(dtype=torch.int64, device=rank)
            decoder_attention_mask = torch.ne(decoder_input_ids, tokenizer.pad_token_id).to(dtype=torch.float32)
            labels = torch.tensor(data["labels"]).to(dtype=torch.int64, device=rank)

            optimizer.zero_grad()

            with amp.autocast(enabled=args.fp16):
                output = net(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    decoder_input_ids=decoder_input_ids,
                    decoder_attention_mask=decoder_attention_mask,
                    labels=labels,
                )

            scaler.scale(output.loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

            if rank == 0:
                pbar.set_description(f"Loss: {output.loss.item():.4f}")
            if args.use_wandb and rank == 0:
                wandb.log({"train loss": output.loss.item()})


if __name__ == "__main__":
    args = get_args()
    mp.spawn(train, nprocs=args.num_gpus, args=(args.num_gpus, args))
```

I used 3 * V100 (one node), torch==1.12.1+cu113, and fp16. I added manually special token to data since the original tokenizer didn't post-process the data. The results of `python train.py --fp16 --decoder_prefix 
` and `python train.py --fp16 --decoder_prefix <pad>` with different transformers' versions are belows:

![image](https://user-images.githubusercontent.com/57349134/188361145-cf1889b1-54a3-41ff-ab2d-b5b1a015de83.png)

As the results, If I use correctly `decoder_prefix` as \
, the results are completely same, but if I use \<pad>, the results are different. I'm not sure why the results are different (and if correction of code is needed), but anyway the main reason was an incorrect use of `decoder_prefix`.

Sorry If I bothered you. Thanks!


@soocheolnoh I am very happy that you found the cause and a solution! Also appreciate a lot your effort on the detailed issue description and further investigations!

It's better to use the same decoder start token as the one used in pretraining. Regarding using pad token, it might work in some case, but we should be very careful. I believe in the case of this issue, it might be related to the decoder attention mask. I saw you have
```python
decoder_attention_mask = torch.ne(decoder_input_ids, tokenizer.pad_token_id)
```
So when you used pad token as decoder start token, you prepared a (decoder) attention mask that ignores the decoder start token (which shouldn't be ignored).

Of course, this is just one observation that might be related.

I am going to close the issue. Don't hesitate to reopen if you still have further questions.

Well, just a remark: you mentioned reverting
```python
# mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min))
mask = torch.full((tgt_len, tgt_len), torch.tensor(float("-inf")))
```
also avoided the issue. This might suggests using `-inf` is more robust to this problem, and worth a further investigation from our side.

Thank you! @ydshieh