Issues in processing very large datasets
## Describe the bug
I'm trying to add a feature called "subgraph" to CNN/DM dataset (modifications on run_summarization.py of Huggingface Transformers script) --- I'm not quite sure if I'm doing it the right way, though--- but the main problem appears when the training starts where the error ` [OSError: [Errno 12] Cannot allocate memory]`  appears. I suppose this problem roots in RAM issues and how the dataset is loaded during training, but I have no clue of what I can do to fix it.  Observing the dataset's cache directory, I see that it takes ~600GB of memory and that's why I believe special care is needed when loading it into the memory. 


Here are my modifications to `run_summarization.py` code. 


```
# loading pre-computed dictionary where keys are 'id' of article and values are corresponding subgraph
graph_data_train = get_graph_data('train') 
graph_data_validation = get_graph_data('val')
...
...


with training_args.main_process_first(desc="train dataset map pre-processing"):
    train_dataset = train_dataset.map(
        preprocess_function_train,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not data_args.overwrite_cache,
        desc="Running tokenizer on train dataset",
    )

```


And here is the modified preprocessed function:

```
def preprocess_function_train(examples):
        inputs, targets, sub_graphs, ids = [], [], [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] is not None and examples[summary_column][i] is not None:
                # if examples['doc_id'][i] in graph_data.keys():
                inputs.append(examples[text_column][i])
                targets.append(examples[summary_column][i])
                sub_graphs.append(graph_data_train[examples['id'][i]])
                ids.append(examples['id'][i])

        inputs = [prefix + inp for inp in inputs]
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=padding, truncation=True,
                                 sub_graphs=sub_graphs, ids=ids)

            # Setup the tokenizer for targets
        with tokenizer.as_target_tokenizer():
            labels = tokenizer(targets, max_length=max_target_length, padding=padding, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs
```


## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version:  2.1.0
- Platform: Linux Ubuntu
- Python version: 3.6
- PyArrow version: 6.0.1
