[JAX/FLAX]: Language modeling example throws RESOURCE_EXHAUSTED error for large datasets
Hi,

I haven't found any similar issue that is related to my `RESOURCE_EXHAUSTED` error, so I'm creating a new issue here.

Problem is: I can use a dataset that is <= 50GB with a certain batch size for BERT or T5 pre-training. It is working fine.

However, if I want to use a larger dataset > 50GB with the **same** batch size, I always got the following error message:
![20220129_110345](https://user-images.githubusercontent.com/20651387/151709875-c17b8461-a361-4f24-bcc1-d7f15845553b.jpg)

Which is happening here:

https://github.com/huggingface/transformers/blob/0f69b924fbda6a442d721b10ece38ccfc6b67275/examples/flax/language-modeling/run_t5_mlm_flax.py#L442-L450

Under some circumstances, I can train a ~50GB dataset for one epoch, but before entering the second epoch, training also crashes with that `RESOURCE_EXHAUSTED` error.

My current workaround for e.g. a 150GB corpus is: split the corpus into 50GB chunks, and pre-train the model on those chunks.

What would be an effective method to pre-train models on larger datasets with the JAX/FLAX language modeling implementation :thinking: 

Many thanks :heart: 
I've seen the streaming PR approach, but it was closed due to strange pre-training behaviour: https://github.com/huggingface/transformers/pull/12607<|||||>The `generate_batch_splits` function is used here later:

https://github.com/huggingface/transformers/blob/0f69b924fbda6a442d721b10ece38ccfc6b67275/examples/flax/language-modeling/run_t5_mlm_flax.py#L812-L815

As far as I understand it, the list of training sample ids e.g. `[2,3,4,5,6,10,1,0,50,16,12,51]` will be splitted into sublists according to the batch size - let's say 4 - so `train_batch_idx` looks like: `[ [2,3,4,5], [6,10,1,0], [50,16,12,51] ]`.

My question is now: do we really need to to this, or can we simply iterate over the train idx array in batch-size steps, e.g. with:

```python
l = [2,3,4,5,6,10,1,0,50,16,12,51]
batch_size = 4

for i in range(0, len(l), batch_size):
    print(l[i : i + batch_size])
```

in the batch loop:

https://github.com/huggingface/transformers/blob/0f69b924fbda6a442d721b10ece38ccfc6b67275/examples/flax/language-modeling/run_t5_mlm_flax.py#L815

:thinking: <|||||>Will take a look into this, also cc @patrickvonplaten
If possible could you share the command to reproduce this and also post the full stack trace? Thanks<|||||>Wow this looks like a difficult issue! Thanks for posting it here @stefan-it. Maybe we should also include some Flax, JAX guys here in case they've seen this before<|||||>Quick update on it: I was able to use the full 149GB dataset by avoiding the `np.split` logic that splits the list of sample ids into batch-size length sublists:

```diff
+def get_samples_idx(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
+    num_samples = len(samples_idx)
+    samples_to_remove = num_samples % batch_size
+
+    if samples_to_remove != 0:
+        samples_idx = samples_idx[:-samples_to_remove]
+    return samples_idx
+
+
 def write_train_metric(summary_writer, train_metrics, train_time, step):
     summary_writer.scalar("train_time", train_time, step)
 
@@ -809,10 +821,11 @@ def main():
         # Generate an epoch by shuffling sampling indices from the train dataset
         num_train_samples = len(tokenized_datasets["train"])
         train_samples_idx = jax.random.permutation(input_rng, jnp.arange(num_train_samples))
-        train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
+        samples_idx = get_samples_idx(train_samples_idx, train_batch_size)
 
         # Gather the indexes for creating the batch and do a training step
-        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
+        for step, batch_idx in enumerate(tqdm(range(0, len(samples_idx), train_batch_size), desc="Training...", position=1)):
+            batch_idx = train_samples_idx[step : step + train_batch_size]
             samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
             model_inputs = data_collator(samples)
```

This requires minimal changes in the existing code and training is running now with 1,082,358 samples per epoch :)<|||||>Interesting - think we could use this to update the official script no? Seems like you found a much better solution<|||||>Update: I've used this workaround and the OOM error after one epoch also disappears!

So I can prepare a PR for the different example scripts (`run_{clm,mlm,t5_mlm}_flax.py`) :)<|||||>I just experienced the same issue. thanks @stefan-it <|||||>It's also working with the `run_mlm_flax.py` script - I'm currently road-testing that fix to make sure, that it is working.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.