Title: IndexError: Invalid key: 88 is out of bounds for size 0
State: closed

### Describe the bug

This bug generates when I use torch.compile(model) in my code, which seems to raise an error in datasets lib.

### Steps to reproduce the bug

I use the following code to fine-tune Falcon on my private dataset.

```python
import transformers
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer,
    AutoConfig,
    DataCollatorForSeq2Seq,
    Trainer,
    Seq2SeqTrainer,
    HfArgumentParser,
    Seq2SeqTrainingArguments,
    BitsAndBytesConfig,
)

from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
    prepare_model_for_int8_training,
    set_peft_model_state_dict,
)

import torch
import os
import evaluate
import functools
from datasets import load_dataset
import bitsandbytes as bnb
import logging
import json
import copy
from typing import Dict, Optional, Sequence
from dataclasses import dataclass, field


# Lora settings
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT= 0.05
LORA_TARGET_MODULES = ["query_key_value"]


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="Salesforce/codegen2-7B")


@dataclass
class DataArguments:
    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
    train_file: str = field(default=None, metadata={"help": "Path to the evaluation data."})
    eval_file: str = field(default=None, metadata={"help": "Path to the evaluation data."})
    cache_path: str = field(default=None, metadata={"help": "Path to the cache directory."})
    num_proc: int = field(default=4, metadata={"help": "Number of processes to use for data preprocessing."})


@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)."},
    )
    is_lora: bool = field(default=True, metadata={"help": "Whether to use LORA."})


def tokenize(text, tokenizer, max_seq_len=512, add_eos_token=True):
    result = tokenizer(
        text,
        truncation=True,
        max_length=max_seq_len,
        padding=False,
        return_tensors=None,
    )
    if (
        result["input_ids"][-1] != tokenizer.eos_token_id
        and len(result["input_ids"]) < max_seq_len
        and add_eos_token
    ):
        result["input_ids"].append(tokenizer.eos_token_id)
        result["attention_mask"].append(1)

    if add_eos_token and len(result["input_ids"]) >= max_seq_len:
        result["input_ids"][max_seq_len - 1] = tokenizer.eos_token_id
        result["attention_mask"][max_seq_len - 1] = 1

    result["labels"] = result["input_ids"].copy()
    return result


def main():
    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path, 
        cache_dir=data_args.cache_path,
        trust_remote_code=True,
    )

    if training_args.is_lora:
        model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=data_args.cache_path,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            load_in_8bit=True,
            quantization_config=BitsAndBytesConfig(
                load_in_8bit=True,
                llm_int8_threshold=6.0
            ),
        )
        model = prepare_model_for_int8_training(model)

        config = LoraConfig(
            r=LORA_R,
            lora_alpha=LORA_ALPHA,
            target_modules=LORA_TARGET_MODULES,
            lora_dropout=LORA_DROPOUT,
            bias="none",
            task_type="CAUSAL_LM",
        )
        model = get_peft_model(model, config)
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path,
            torch_dtype=torch.float16,
            cache_dir=data_args.cache_path,
            trust_remote_code=True,
        )
    model.config.use_cache = False
    
    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}"
        )
    print_trainable_parameters(model)

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=data_args.cache_path,
        model_max_length=training_args.model_max_length,
        padding_side="left",
        use_fast=True,
        trust_remote_code=True,
    )
    tokenizer.pad_token = tokenizer.eos_token
    
    # Load dataset

    def generate_and_tokenize_prompt(sample):
        input_text = sample["input"]
        target_text = sample["output"] + tokenizer.eos_token
        full_text = input_text + target_text
        tokenized_full_text = tokenize(full_text, tokenizer, max_seq_len=512)
        tokenized_input_text = tokenize(input_text, tokenizer, max_seq_len=512)
        input_len = len(tokenized_input_text["input_ids"]) - 1 # -1 for eos token
        tokenized_full_text["labels"] = [-100] * input_len + tokenized_full_text["labels"][input_len:]
        return tokenized_full_text

    data_files = {}
    if data_args.train_file is not None:
        data_files["train"] = data_args.train_file
    if data_args.eval_file is not None:
        data_files["eval"] = data_args.eval_file
    
    dataset = load_dataset(data_args.data_path, data_files=data_files)
    train_dataset = dataset["train"]
    eval_dataset = dataset["eval"]
    train_dataset = train_dataset.map(generate_and_tokenize_prompt, num_proc=data_args.num_proc)
    eval_dataset = eval_dataset.map(generate_and_tokenize_prompt, num_proc=data_args.num_proc)
    data_collator = DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True)

    # Evaluation metrics
    def compute_metrics(eval_preds, tokenizer):
        metric = evaluate.load('exact_match')
        preds, labels = eval_preds
        # In case the model returns more than the prediction logits
        if isinstance(preds, tuple):
            preds = preds[0]

        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True, clean_up_tokenization_spaces=False)

        # Replace -100s in the labels as we can't decode them
        labels[labels == -100] = tokenizer.pad_token_id
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True, clean_up_tokenization_spaces=False)

        # Some simple post-processing
        decoded_preds = [pred.strip() for pred in decoded_preds]
        decoded_labels = [label.strip() for label in decoded_labels]

        result = metric.compute(predictions=decoded_preds, references=decoded_labels)
        return {'exact_match': result['exact_match']} 
    
    compute_metrics_fn = functools.partial(compute_metrics, tokenizer=tokenizer)

    model = torch.compile(model)

    # Training
    trainer = Trainer(
        model=model, 
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,  
        args=training_args,
        data_collator=data_collator,
        compute_metrics=compute_metrics_fn,
    )
    trainer.train()
    trainer.save_state()
    trainer.save_model(output_dir=training_args.output_dir)
    tokenizer.save_pretrained(save_directory=training_args.output_dir)


if __name__ == "__main__":
    main()
```

When I didn't use `torch.cpmpile(model)`, my code worked well. But when I added this line to my code, It produced the following error:

```
Traceback (most recent call last):
  File "falcon_sft.py", line 230, in <module>
    main()
  File "falcon_sft.py", line 223, in main
    trainer.train()
  File "python3.10/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "python3.10/site-packages/transformers/trainer.py", line 1787, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "python3.10/site-packages/accelerate/data_loader.py", line 384, in __iter__
    current_batch = next(dataloader_iter)
  File "python3.10/site-packages/torch/utils/data/dataloader.py", line 633, in __next__
    data = self._next_data()
  File "python3.10/site-packages/torch/utils/data/dataloader.py", line 677, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "python3.10/site-packages/datasets/arrow_dataset.py", line 2807, in __getitems__
    batch = self.__getitem__(keys)
  File "python3.10/site-packages/datasets/arrow_dataset.py", line 2803, in __getitem__
    return self._getitem(key)
  File "python3.10/site-packages/datasets/arrow_dataset.py", line 2787, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
  File "python3.10/site-packages/datasets/formatting/formatting.py", line 583, in query_table
    _check_valid_index_key(key, size)
  File "python3.10/site-packages/datasets/formatting/formatting.py", line 536, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "python3.10/site-packages/datasets/formatting/formatting.py", line 526, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 88 is out of bounds for size 0
```

So I'm confused about why this error was generated, and how to fix it. Is this error produced by datasets or `torch.compile`?

### Expected behavior

I want to use `torch.compile` in my code.

### Environment info

- `datasets` version: 2.14.3
- Platform: Linux-4.18.0-425.19.2.el8_7.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.8
- Huggingface_hub version: 0.16.4
- PyArrow version: 12.0.1
- Pandas version: 2.0.3

Comments:
Hi @TomasAndersonFang,

Have you tried instead to use `torch_compile` in `transformers.TrainingArguments`? https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/trainer#transformers.TrainingArguments.torch_compile

> 

I tried this and got the following error:

```
Traceback (most recent call last):
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    super().run()
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    and self.step()
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1030, in LOAD_ATTR
    result = BuiltinVariable(getattr).call_function(
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 566, in call_function
    result = handler(tx, *args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 931, in call_getattr
    return obj.var_getattr(tx, name).add_options(options)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 124, in var_getattr
    subobj = inspect.getattr_static(base, name)
  File "/apps/Arch/software/Python/3.10.8-GCCcore-12.2.0/lib/python3.10/inspect.py", line 1777, in getattr_static
    raise AttributeError(attr)
AttributeError: config

from user code:
   File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/peft/peft_model.py", line 909, in forward
    if self.base_model.config.model_type == "mpt":

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/llm-copt/fine-tune/falcon/falcon_sft.py", line 228, in <module>
    main()
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/llm-copt/fine-tune/falcon/falcon_sft.py", line 221, in main
    trainer.train()
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/transformers/trainer.py", line 1809, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/transformers/trainer.py", line 2654, in training_step
    loss = self.compute_loss(model, inputs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/transformers/trainer.py", line 2679, in compute_loss
    outputs = model(**inputs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/accelerate/utils/operations.py", line 581, in forward
    return model_forward(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/accelerate/utils/operations.py", line 569, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
    return func(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/cephyr/NOBACKUP/groups/snic2021-23-24/LLM4-CodeOpt/env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    raise InternalTorchDynamoError() from e
torch._dynamo.exc.InternalTorchDynamoError
```

Hi @TomasAndersonFang,

I guess in this case it may be an issue with `transformers` (or `PyTorch`). I would recommend you open an issue on their repo.

@albertvillanova Thanks for your recommendation. I'll do it

@TomasAndersonFang were you able to find a solution to this issue? I would highly appreciate any help. 

Thanks!