interleave_datasets & DataCollatorForLanguageModeling having a conflict ?
Opened by PonteIneptique (open)

### Describe the bug

Hi everyone :)

I have two local & custom datasets (1 "sentence" per line) which I split along the 95/5 lines for pre-training a Bert model. I use a modified version of `run_mlm.py` in order to be able to make use of `interleave_dataset`:

- `tokenize()` runs fine
- `group_text()` runs fine

Everytime, on step 19, I get 

```pytb
  File "env/lib/python3.9/site-packages/transformers/data/data_collator.py", line 779, in torch_mask_tokens
    inputs[indices_random] = random_words[indices_random]
RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Long for the source.
```

I tried:
- training without interleave on dataset 1, it runs
- training without interleave on dataset 2, it runs
- training without `.to_iterable_dataset()`, it hangs then crash
- training without group_text() and padding to max_length seemed to fix the issue, but who knows if this was just because it was an issue that would come much later in terms of steps.

I might have coded something wrong, but I don't get what 

### Steps to reproduce the bug

I have this function:

```py
def build_dataset(path: str, percent: str):
    dataset = load_dataset(
        "text",
        data_files={"train": [path]},
        split=f"train[{percent}]"
    )
    dataset = dataset.map(
        lambda examples: tokenize(examples["text"]),
        batched=True,
        num_proc=num_proc,
    )

    dataset = dataset.map(
        group_texts,
        batched=True,
        num_proc=num_proc,
        desc=f"Grouping texts in chunks of {tokenizer.max_seq_length}",
        remove_columns=["text"]
    )

    print(len(dataset))
    return dataset.to_iterable_dataset()
```

I hardcoded group_text:
```py
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, and if the total_length < max_seq_length  we exclude this batch and return an empty dict.
        # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.
        total_length = (total_length // 512) * 512
        # Split by chunks of max_len.
        result = {
            k: [t[i: i + 512] for i in range(0, total_length, 512)]
            for k, t in concatenated_examples.items()
        }
        # result = {k: [el for el in elements if el] for k, elements in result.items()}
        return result
```

And then I build datasets using the following code:

```py
train1 = build_dataset("d1.txt", ":95%")
train2 = build_dataset("d2.txt", ":95%")
dev1 = build_dataset("d1.txt", "95%:")
dev2 = build_dataset("d2.txt", "95%:")
```

and finally I run
```py
train_dataset = interleave_datasets(
    [train1, train2],
    probabilities=[0.8, 0.2],
    seed=42
)
eval_dataset = interleave_datasets(
    [dev1, dev2],
    probabilities=[0.8, 0.2],
    seed=42
)
```

Then I run the training part which remains mostly untouched:

> CUDA_VISIBLE_DEVICES=1 python custom_dataset.py --model_type bert --per_device_train_batch_size 32 --do_train --output_dir /var/mlm/training-bert/model --max_seq_length 512 --save_steps 10000 --save_total_limit 3 --auto_find_batch_size --logging_dir ./logs-bert --learning_rate 0.0001 --do_train --num_train_epochs 25 --warmup_steps 10000 --max_step 45000 --fp16

### Expected behavior

The model should then train normally, but fails every time at the same step (19).

printing the variables at `inputs[indices_random] = random_words[indices_random]` shows a magnificient empty tensor (, 32) [if I remember well]

### Environment info

transformers[torch] 4.30.2
Ubuntu
A100 0 CUDA 12
Driver Version: 525.116.04

Comments: