Training script "run_mlm.py" doesn't work for certain datasets
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version:
- Platform:3.4.0
- Python version:3.8.3
- PyTorch version (GPU?):3.6.0
- Tensorflow version (GPU?):
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
  @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 Translation: @sshleifer
 Summarization: @sshleifer
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten @TevenLeScao
 blenderbot: @mariamabarham
 Bart: @sshleifer
 Marian: @sshleifer
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao
 RAG: @patrickvonplaten, @lhoestq
 FSTM: @stas00
 examples/seq2seq: @sshleifer
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
@LysandreJik 

## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* [X] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below) MLM 

## To reproduce

Steps to reproduce the behavior:
Basically run the training script in examples/run_mlm.py using wikipedia dataset 
```
python run_mlm.py \ 
--model_name_or_path roberta-base \
--dataset_name wikipedia \
--dataset_config_name 20200501.en \
--do_train \
--output_dir /tmp/test-mlm \
```

## Erorr Message
```
Traceback (most recent call last):
  File "run_mlm2.py", line 388, in <module>
    main()
  File "run_mlm2.py", line 333, in main
    tokenized_datasets = tokenized_datasets.map(
  File "/home/zeyuy/miniconda3/lib/python3.8/site-packages/datasets/dataset_dict.py", line 283, in map
    {
  File "/home/zeyuy/miniconda3/lib/python3.8/site-packages/datasets/dataset_dict.py", line 284, in <dictcomp>
    k: dataset.map(
  File "/home/zeyuy/miniconda3/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1236, in map
    update_data = does_function_return_dict(test_inputs, test_indices)
  File "/home/zeyuy/miniconda3/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1207, in does_function_return_dict    function(*fn_args, indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
  File "run_mlm2.py", line 315, in group_texts
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
  File "run_mlm2.py", line 315, in <dictcomp>
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
TypeError: can only concatenate list (not "str") to list
```

## Expected behavior

The training script should work for all datasets in huggingface datasets. The problem is that the feature other than 'text' (or the feature we train on) interferes when we try to concatenate the tokenized 'text' ('input_ids', 'mask_ids' ...) from each instances.

A quick fix would be on line [295](https://github.com/huggingface/transformers/blob/77a257fc210a56f1fd0d75166ecd654cf58111f3/examples/language-modeling/run_mlm.py#L295), change  

`            remove_columns=[text_column_name],`
to
`            remove_columns=column_names,`

Should I open a PR or someone want to do a quick fix?

<!-- A clear and concise description of what you would expect to happen. -->

> The training script should work for all datasets in huggingface datasets.

To clarify our intent with the examples we provide, this is untrue. We can't support every use case of every user. The script can easily be modified and has plenty of comments to help the user tune them to their needs.

In this particular instance however, I think your proposed fix doesn't break anything and allows us to support more datasets with little change, so I think we can add it. Don't hesitate to suggest a PR and tag me on it.<|||||>Thanks for the reply~ Since this is my first time to do make a pr, I followed on the [guide](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) and run `make test` as suggested. However, most of the test failed.
`947 failed, 188 passed, 124 skipped, 25 warnings in 37.09s`
Those test failed before I even makes any changes to the code. Sorry if this is a newbie question but should we expect all the test to pass?<|||||>Yes, the guide is what is also used by the CI to run all tests, so they should all pass. Make sure you do have the installation from source with all the dev dependencies installed. Apart from that, there is little I can do to help without seeing any error message.<|||||>I finally found out the reason I fail most of the test. It's because I have another training on my gpu which interferes with testing for some reason. 

Now I have all test passed except for 8 of them. Here are those 8 test: 

================================ short test summary info ================================
FAILED tests/test_modeling_ctrl.py::CTRLModelTest::test_multi_gpu_data_parallel_forward
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_multi_gpu_data_parallel_forward
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_multi_gpu_data_parallel_forward
FAILED tests/test_trainer.py::TrainerIntegrationTest::test_dynamic_shapes - IndexError...
FAILED tests/test_trainer.py::TrainerIntegrationTest::test_evaluate - AssertionError: ...
FAILED tests/test_trainer.py::TrainerIntegrationTest::test_load_best_model_at_end - As...
FAILED tests/test_trainer.py::TrainerIntegrationTest::test_number_of_steps_in_training
FAILED tests/test_trainer_distributed.py::TestTrainerDistributed::test_trainer - Runti...
======== 8 failed, 4035 passed, 688 skipped, 608 warnings in 2855.87s (0:47:35) =========

Again, those test are run after I follows the guide by installing from the source with the dev dependencies **but before I made any change to the source code**. So I think just installing from the source with dev dependencies and run "make test" will be enough to reproduce the error.

Here's the traceback error message for 4 of those erorr:

**FAILED tests/test_modeling_ctrl.py::CTRLModelTest::test_multi_gpu_data_parallel_forward**
```

        ...[-1.1481e-01, -8.0696e-02,  7.7946e-02,  ..., -6.4783e-02,
             1.2656e-01,  7.6060e-02]]]]], device='cuda:2'))
dim = 0, destination = 0

    def gather(tensors, dim=0, destination=None, *, out=None):
        r"""Gathers tensors from multiple GPU devices.
    
        Arguments:
            tensors (Iterable[Tensor]): an iterable of tensors to gather.
              Tensor sizes in all dimensions other than :attr:`dim` have to match.
            dim (int, optional): a dimension along which the tensors will be
              concatenated. Default: ``0``.
            destination (torch.device, str, or int, optional): the output device.
              Can be CPU or CUDA. Default: the current CUDA device.
            out (Tensor, optional, keyword-only): the tensor to store gather result.
              Its sizes must match those of :attr:`tensors`, except for :attr:`dim`,
              where the size must equal ``sum(tensor.size(dim) for tensor in tensors)``.
              Can be on CPU or CUDA.
    
        .. note::
            :attr:`destination` must not be specified when :attr:`out` is specified.
    
        Returns:
            - If :attr:`destination` is specified,
                a tensor located on :attr:`destination` device, that is a result of
                concatenating :attr:`tensors` along :attr:`dim`.
            - If :attr:`out` is specified,
                the :attr:`out` tensor, now containing results of concatenating
                :attr:`tensors` along :attr:`dim`.
        """
        if out is None:
            if destination == -1:
                warnings.warn(
                    'Using -1 to represent CPU tensor is deprecated. Please use a '
                    'device object or string instead, e.g., "cpu".')
            destination = _get_device_index(destination, allow_cpu=True, optional=True)
>           return torch._C._gather(tensors, dim, destination)
E           RuntimeError: Input tensor at index 2 has invalid shape [2, 4, 4, 7, 8], but expected [2, 5, 4, 7, 8]

../../test_env_6/lib/python3.8/site-packages/torch/nn/parallel/comm.py:230: RuntimeError

```



**FAILED tests/test_trainer.py::TrainerIntegrationTest::test_dynamic_shapes - IndexError: list index out of range**
```

tests/test_trainer.py:366: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/transformers/trainer.py:1308: in evaluate
    output = self.prediction_loop(
src/transformers/trainer.py:1416: in prediction_loop
    for step, inputs in enumerate(dataloader):
../../test_env_6/lib/python3.8/site-packages/torch/utils/data/dataloader.py:435: in __next__
    data = self._next_data()
../../test_env_6/lib/python3.8/site-packages/torch/utils/data/dataloader.py:475: in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
../../test_env_6/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:44: in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
../../test_env_6/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:44: in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <tests.test_trainer.DynamicShapesDataset object at 0x7f547e1b0250>, i = 48

    def __getitem__(self, i):
>       return {"input_x": self.xs[i], "labels": self.ys[i]}
E       IndexError: list index out of range
```


**FAILED tests/test_trainer.py::TrainerIntegrationTest::test_evaluate - AssertionError: ...**
```

_______________________________________________________________________ TrainerIntegrationTest.test_evaluate ________________________________________________________________________

self = <tests.test_trainer.TrainerIntegrationTest testMethod=test_evaluate>

    def test_evaluate(self):
        trainer = get_regression_trainer(a=1.5, b=2.5, compute_metrics=AlmostAccuracy())
        results = trainer.evaluate()
    
        x, y = trainer.eval_dataset.x, trainer.eval_dataset.ys[0]
        pred = 1.5 * x + 2.5
        expected_loss = ((pred - y) ** 2).mean()
>       self.assertAlmostEqual(results["eval_loss"], expected_loss)
E       AssertionError: 0.37259840965270996 != 0.3790205 within 7 places (0.006422102451324463 difference)

tests/test_trainer.py:312: AssertionError
```


**FAILED tests/test_trainer.py::TrainerIntegrationTest::test_load_best_model_at_end - As...**

```
>           self.check_saved_checkpoints(tmpdir, 64 // self.batch_size, total)

tests/test_trainer.py:565: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/test_trainer.py:211: in check_saved_checkpoints
    self.assertTrue(os.path.isdir(checkpoint))
E   AssertionError: False is not true
```<|||||>Again, this is exactly what the CI does, so those failures are linked to your particular environment. Since you didn't tell us what it is we can't reproduce and fix potential issues.
One way to make sure you don't use your GPU if it's busy is to run `CUDA_VISIBLE_DEVICES='' make tests`.<|||||>Thanks for the reply. Sorry I forgot to list my environment. But as you suggest, `CUDA_VISIBLE_DEVICES=''`, all the test pass when I run them without GPU using. You are right that the test failures are linked to my environment. It's irrelevant to this issue so I think I'll make another post about it.<|||||>Hi newbie questions are better suited for the forum at https://discuss.huggingface.co

We try to keep the issues for bug reports and features/model requests.<|||||>> Hi newbie questions are better suited for the forum at https://discuss.huggingface.co
> 
> We try to keep the issues for bug reports and features/model requests.

Yeah, you are right. I will make another post in the forum instead. 