Trainer throws error when using torchrun and small dataset
Author: Joshuaclymer
Pull request: no

### System Info

- `transformers` version: 4.28.0.dev0
- Platform: Linux-5.4.0-149-generic-x86_64-with-glibc2.31
- Python version: 3.10.10
- Huggingface_hub version: 0.15.1
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes using torchrun
- Using distributed or parallel set-up in script?: Yes, using torchrun

### Who can help?

@sgugger

### 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

## Command
torchrun --nproc_per_node=2 --master_port=8080 train/train_old_2.py \
    --model_name_or_path /workspace/models/llama7b \
    --data_path /workspace/instruct-generalize/data/test.json \
    --bf16 True \
    --output_dir /workspace/fine-tune-result \
    --num_train_epochs 3 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 2000 \
    --save_total_limit 1 \
    --learning_rate 2e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 True

## Stacktrace excerpt
Note: these are actually two stack traces that are printed over each other because there are two threads.

  File "/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adamw.py", line 171, in step
    adamw(
RuntimeError: The size of tensor a (131078144) must match the size of tensor b (262156288) at non-singleton dimension 0
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adamw.py", line 321, in adamw
    func(
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adamw.py", line 389, in _single_tensor_adamw
    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
RuntimeError: The size of tensor a (131078144) must match the size of tensor b (262156288) at non-singleton dimension 

## Dataset that throws error:
`   
[{"instruction":"Do the addition.", "input": "Hi. What is zero plus zero?","output": "zero"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus one?","output": "one"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus two?","output": "two"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus three?","output": "three"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus four?","output": "four"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus five?","output": "five"},
{"instruction":"Do the addition.","input": "Hi. What is zero plus six?","output": "six"}]
`
## Dataset that doesn't throw error:
Just copy and paste the code above about 10 times to increase the size of the dataset.
  

## Training code
```
import copy
import logging
from dataclasses import dataclass, field
from typing import Dict, Optional, Sequence

import torch
import transformers
import utils
from torch.utils.data import Dataset
from transformers import Trainer

IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "
"
DEFAULT_BOS_TOKEN = "
"
DEFAULT_UNK_TOKEN = "<unk>"
PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Response:"
    ),
}


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")


@dataclass
class DataArguments:
    data_path: str = field(default=None, metadata={"help": "Path to the training data."})


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
    )


def smart_tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg


def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """Tokenize a list of strings."""
    tokenized_list = [
        tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        )
        for text in strings
    ]
    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
    input_ids_lens = labels_lens = [
        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
    ]
    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
    )


def preprocess(
    sources: Sequence[str],
    targets: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
    """Preprocess the data by tokenizing."""
    examples = [s + t for s, t in zip(sources, targets)]
    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
    input_ids = examples_tokenized["input_ids"]
    labels = copy.deepcopy(input_ids)
    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
        label[:source_len] = IGNORE_INDEX
    return dict(input_ids=input_ids, labels=labels)


class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer):
        super(SupervisedDataset, self).__init__()
        logging.warning("Loading data...")
        list_data_dict = utils.jload(data_path)

        logging.warning("Formatting inputs...")
        sources = [f"{example['input']}{tokenizer.eos_token}" for example in list_data_dict["examples"]]
        targets = [f"{example['target']}{tokenizer.eos_token}" for example in list_data_dict["examples"]]

        logging.warning("Tokenizing inputs... This may take some time...")
        data_dict = preprocess(sources, targets, tokenizer)

        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(input_ids=self.input_ids[i], labels=self.labels[i])


@dataclass
class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
        input_ids = torch.nn.utils.rnn.pad_sequence(
            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
        )
        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )


def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, task_name) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    train_dataset = SupervisedDataset(f"/workspace/instruct-generalize/data/{task_name}.json", tokenizer)
    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)


@dataclass
class DataCollatorBCAllTargets(object):
    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]):
        # Get inputs and labels as strings
        input_strings = []
        label_strings = []
        for item in instances:
            if "target" not in item:
                raise Exception(f"The bc_all_targets training strategy requires every example to have a 'target' key. But the following example does not:\n{item}")
            input_strings.append(item["input"])
            if isinstance(item["target"], list):
                for target in item["target"]:
                    label_strings.append(target)
            else:
                label_strings.append(item["target"])
        
        # Combine them and tokenize to produce tokenized inputs
        combined = [input_strings[i] + label_strings[i] for i in range(len(input_strings))]
        encoded_combined = self.tokenizer.batch_encode_plus(combined, padding="longest", return_tensors="pt")

        # append eos tokens
        ones = torch.ones((encoded_combined["input_ids"].shape[0], 1), dtype=int)
        encoded_combined["input_ids"] = torch.cat((encoded_combined["input_ids"], self.tokenizer.eos_token_id*ones), dim=1)
        encoded_combined["attention_mask"] = torch.cat((encoded_combined["attention_mask"], ones), dim=1)

        encoded_labels = self.tokenizer.batch_encode_plus(label_strings, padding="longest", return_tensors="pt")
        label_lengths = [sum(mask) + 1 for mask in encoded_labels["attention_mask"]]

        labels = copy.deepcopy(encoded_combined)["input_ids"]
        for label, label_len in zip(labels, label_lengths):
            label[:-label_len] = IGNORE_INDEX
        print(encoded_combined["input_ids"].shape)
        print(labels.shape)
        print(encoded_combined["attention_mask"].shape)
        return dict(
            input_ids = encoded_combined["input_ids"],
            labels=labels,
            attention_mask=encoded_combined["attention_mask"]
        )

from classes import Model, Task
def train():
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    training_args.remove_unused_columns = False

    yoyo = Model(model_args.model_name_or_path)
    model = yoyo.model
    tokenizer = yoyo.tokenizer

    special_tokens_dict = dict()
    if tokenizer.pad_token is None:
        special_tokens_dict["pad_token"] = DEFAULT_PAD_TOKEN
    if tokenizer.eos_token is None:
        special_tokens_dict["eos_token"] = DEFAULT_EOS_TOKEN
    if tokenizer.bos_token is None:
        special_tokens_dict["bos_token"] = DEFAULT_BOS_TOKEN
    if tokenizer.unk_token is None:
        special_tokens_dict["unk_token"] = DEFAULT_UNK_TOKEN
    #tokenizer.add_special_tokens(special_tokens_dict)
    smart_tokenizer_and_embedding_resize(
        special_tokens_dict=special_tokens_dict,
        tokenizer=tokenizer,
        model=model,
    )
    data_module = make_supervised_data_module(tokenizer=tokenizer, task_name=data_args.data_path)
    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
    trainer.train()
    trainer.save_state()
    trainer.save_model(output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()
```
### Expected behavior

Dataset size should not affect whether the code throws an error.

Comments:
> Dataset size should not affect whether the code throws an error.

I disagree here. The dataset is too small so you will have an error anyway, but it could be clearer I agree. This might be linked to FSDP so tagging @pacman100 

For anyone who happens to be fine-tuning on tiny datasets: you can fix this issue by reducing `gradient_accumulation_steps`. 

I think num_examples has to be greater than `gradient_accumulation_steps*num_GPUs*per_device_train_batch_size`


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.