IndexError Not Solving -> IndexError: Invalid key: ?? is out of bounds for size 0 or ??
State: open

### Describe the bug

in <cell line: 1>:1                                                                              │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1537 in train                    │
│                                                                                                  │
│   1534 │   │   inner_training_loop = find_executable_batch_size(                                 │
│   1535 │   │   │   self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size  │
│   1536 │   │   )                                                                                 │
│ ❱ 1537 │   │   return inner_training_loop(                                                       │
│   1538 │   │   │   args=args,                                                                    │
│   1539 │   │   │   resume_from_checkpoint=resume_from_checkpoint,                                │
│   1540 │   │   │   trial=trial,                                                                  │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1789 in _inner_training_loop     │
│                                                                                                  │
│   1786 │   │   │   │   rng_to_sync = True                                                        │
│   1787 │   │   │                                                                                 │
│   1788 │   │   │   step = -1                                                                     │
│ ❱ 1789 │   │   │   for step, inputs in enumerate(epoch_iterator):                                │
│   1790 │   │   │   │   total_batched_samples += 1                                                │
│   1791 │   │   │   │   if rng_to_sync:                                                           │
│   1792 │   │   │   │   │   self._load_rng_state(resume_from_checkpoint)                          │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/accelerate/data_loader.py:377 in __iter__                │
│                                                                                                  │
│   374 │   │   dataloader_iter = super().__iter__()                                               │
│   375 │   │   # We iterate one batch ahead to check when we are at the end                       │
│   376 │   │   try:                                                                               │
│ ❱ 377 │   │   │   current_batch = next(dataloader_iter)                                          │
│   378 │   │   except StopIteration:                                                              │
│   379 │   │   │   yield                                                                          │
│   380                                                                                            │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:633 in __next__           │
│                                                                                                  │
│    630 │   │   │   if self._sampler_iter is None:                                                │
│    631 │   │   │   │   # TODO(https://github.com/pytorch/pytorch/issues/76750)                   │
│    632 │   │   │   │   self._reset()  # type: ignore[call-arg]                                   │
│ ❱  633 │   │   │   data = self._next_data()                                                      │
│    634 │   │   │   self._num_yielded += 1                                                        │
│    635 │   │   │   if self._dataset_kind == _DatasetKind.Iterable and \                          │
│    636 │   │   │   │   │   self._IterableDataset_len_called is not None and \                    │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:677 in _next_data         │
│                                                                                                  │
│    674 │                                                                                         │
│    675 │   def _next_data(self):                                                                 │
│    676 │   │   index = self._next_index()  # may raise StopIteration                             │
│ ❱  677 │   │   data = self._dataset_fetcher.fetch(index)  # may raise StopIteration              │
│    678 │   │   if self._pin_memory:                                                              │
│    679 │   │   │   data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)            │
│    680 │   │   return data                                                                       │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py:49 in fetch             │
│                                                                                                  │
│   46 │   def fetch(self, possibly_batched_index):                                                │
│   47 │   │   if self.auto_collation:                                                             │
│   48 │   │   │   if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:         │
│ ❱ 49 │   │   │   │   data = self.dataset.__getitems__(possibly_batched_index)                    │
│   50 │   │   │   else:                                                                           │
│   51 │   │   │   │   data = [self.dataset[idx] for idx in possibly_batched_index]                │
│   52 │   │   else:                                                                               │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2782 in __getitems__           │
│                                                                                                  │
│   2779 │                                                                                         │
│   2780 │   def __getitems__(self, keys: List) -> List:                                           │
│   2781 │   │   """Can be used to get a batch using a list of integers indices."""                │
│ ❱ 2782 │   │   batch = self.__getitem__(keys)                                                    │
│   2783 │   │   n_examples = len(batch[next(iter(batch))])                                        │
│   2784 │   │   return [{col: array[i] for col, array in batch.items()} for i in range(n_example  │
│   2785                                                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2778 in __getitem__            │
│                                                                                                  │
│   2775 │                                                                                         │
│   2776 │   def __getitem__(self, key):  # noqa: F811                                             │
│   2777 │   │   """Can be used to index columns (by string names) or rows (by integer index or i  │
│ ❱ 2778 │   │   return self._getitem(key)                                                         │
│   2779 │                                                                                         │
│   2780 │   def __getitems__(self, keys: List) -> List:                                           │
│   2781 │   │   """Can be used to get a batch using a list of integers indices."""                │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2762 in _getitem               │
│                                                                                                  │
│   2759 │   │   format_kwargs = kwargs["format_kwargs"] if "format_kwargs" in kwargs else self._  │
│   2760 │   │   format_kwargs = format_kwargs if format_kwargs is not None else {}                │
│   2761 │   │   formatter = get_formatter(format_type, features=self._info.features, **format_kw  │
│ ❱ 2762 │   │   pa_subtable = query_table(self._data, key, indices=self._indices if self._indice  │
│   2763 │   │   formatted_output = format_table(                                                  │
│   2764 │   │   │   pa_subtable, key, formatter=formatter, format_columns=format_columns, output  │
│   2765 │   │   )                                                                                 │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:578 in query_table     │
│                                                                                                  │
│   575 │   │   _check_valid_column_key(key, table.column_names)                                   │
│   576 │   else:                                                                                  │
│   577 │   │   size = indices.num_rows if indices is not None else table.num_rows                 │
│ ❱ 578 │   │   _check_valid_index_key(key, size)                                                  │
│   579 │   # Query the main table                                                                 │
│   580 │   if indices is None:                                                                    │
│   581 │   │   pa_subtable = _query_table(table, key)                                             │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:531 in                 │
│ _check_valid_index_key                                                                           │
│                                                                                                  │
│   528 │   │   │   _check_valid_index_key(min(key), size=size)                                    │
│   529 │   elif isinstance(key, Iterable):                                                        │
│   530 │   │   if len(key) > 0:                                                                   │
│ ❱ 531 │   │   │   _check_valid_index_key(int(max(key)), size=size)                               │
│   532 │   │   │   _check_valid_index_key(int(min(key)), size=size)                               │
│   533 │   else:                                                                                  │
│   534 │   │   _raise_bad_key_type(key)                                                           │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:521 in                 │
│ _check_valid_index_key                                                                           │
│                                                                                                  │
│   518 def _check_valid_index_key(key: Union[int, slice, range, Iterable], size: int) -> None:    │
│   519 │   if isinstance(key, int):                                                               │
│   520 │   │   if (key < 0 and key + size < 0) or (key >= size):                                  │
│ ❱ 521 │   │   │   raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")       │
│   522 │   │   return                                                                             │
│   523 │   elif isinstance(key, slice):                                                           │
│   524 │   │   pass                                             

### Steps to reproduce the bug

``
import json
import os
from pprint import pprint

import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import Dataset,load_dataset

from peft import (
    LoraConfig,
    PeftConfig,
    PeftModel,
    get_peft_model,
    prepare_model_for_kbit_training
)

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,

)

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def print_trainable_parameters(model):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
    )


MODEL_NAME = "tiiuae/falcon-7b"

bnb_config = BitsAndBytesConfig(
    load_in_4bit = True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map = "auto",
    trust_remote_code = True,
    quantization_config = bnb_config
)


tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)


config = LoraConfig(
    r = 16,
    lora_alpha = 32,
    target_modules = ["query_key_value"],
    lora_dropout = 0.05,
    bias = "none",
    task_type = "CASUAL_LM"
)

model = get_peft_model(model,config)
print_trainable_parameters(model)

def generate_prompt(data_point):
  return f"""
<human>: {data_point["question"]}
<assistant>: {data_point["answer"]}  
""".strip()

def generate_and_tokenize_prompt(data_point):
    full_prompt = generate_prompt(data_point)
    tokenized_full_prompt = tokenizer(full_prompt, padding = True, truncation = True,return_tensors = None)
    return dict({
        "input_ids" : tokenized_full_prompt["input_ids"],
        "attention_mask" : tokenized_full_prompt["attention_mask"]

    })


data = data["train"].shuffle().map(generate_and_tokenize_prompt, batched = False)  

OUTPUT_DIR = "experiments"

trainings_args = transformers.TrainingArguments(
    per_device_train_batch_size = 1,
    gradient_accumulation_steps = 4,
    num_train_epochs = 1,
    learning_rate = 2e-4,
    fp16 = True,
    save_total_limit = 3,
    logging_steps = 1,
    output_dir = OUTPUT_DIR,
    max_steps = 80,
    optim = "paged_adamw_8bit",
    lr_scheduler_type = "cosine",
    warmup_ratio = 0.05,
    #remove_unused_columns=True
)

trainer = transformers.Trainer(
    model = model,
    train_dataset = data,
    args = trainings_args, 
    data_collator = transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),

)

model.config.use_cache = False

trainer.train()


IndexError: Invalid key: 32 is out of bounds for size 0

DataSet Format is like : 
[{"question": "How can I create an account?", "answer": "To create an account, click on the 'Sign Up' button on the top right corner of our website and follow the instructions to complete the registration process."}, .... ]

### Expected behavior

-

### Environment info


!pip install -q pip 
!pip install -q bitsandbytes==0.39.0 
!pip install -q torch==2.0.1

!pip install -q git+https://github.com/huggingface/transformers.git 
!pip install -q git+https://github.com/huggingface/peft.git 
!pip install -q git+https://github.com/huggingface/accelerate.git 

!pip install -q datasets 
!pip install -q loralib==0.1.1 
!pip install -q einops==0.6.1 


import json
import os
from pprint import pprint

import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import Dataset,load_dataset

from peft import (
    LoraConfig,
    PeftConfig,
    PeftModel,
    get_peft_model,
    prepare_model_for_kbit_training
)

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,

)

os.environ["CUDA_VISIBLE_DEVICES"] = "0"