Intermittent JSON parse error when streaming the Pile
Author: neelnanda-io | Pull request: no | State: open

## Describe the bug
I have an intermittent error when streaming the Pile, where I get a JSON parse error which causes my program to crash.

This is intermittent - when I rerun the program with the same random seed it does not crash in the same way. The exact point this happens also varied - it happened to me 11B tokens and 4 days into a training run, and now just happened 2 minutes into one, but I can't reliably reproduce it.

I'm using a remote machine with 8 A6000 GPUs via runpod.io

## Expected results
I have a DataLoader which can iterate through the whole Pile

## Actual results

Stack trace:

```
Failed to read file 'zstd://12.jsonl::https://the-eye.eu/public/AI/pile/train/12.jsonl.zst' with error <class 'pyarrow.lib.ArrowInvalid'>: JSON parse error: Invalid value. in row 0
```

I'm currently using HuggingFace accelerate, which also gave me the following stack trace, but I've also experienced this problem intermittently when using DataParallel, so I don't think it's to do with parallelisation 

```
Traceback (most recent call last):                                                                                                          
  File "ddp_script.py", line 1258, in <module>                                                                                              
    main()                                                                                                                                  
  File "ddp_script.py", line 1143, in main                                                                                                  
    for c, batch in tqdm.tqdm(enumerate(data_iter)):                                                                                        
  File "/opt/conda/lib/python3.7/site-packages/tqdm/std.py", line 1195, in __iter__                                                         
    for obj in iterable:                                                                                                                    
  File "/opt/conda/lib/python3.7/site-packages/accelerate/data_loader.py", line 503, in __iter__                                            
    next_batch, next_batch_info, next_skip = self._fetch_batches(main_iterator)                                                             
  File "/opt/conda/lib/python3.7/site-packages/accelerate/data_loader.py", line 454, in _fetch_batches                                      
    broadcast_object_list(batch_info)                                                                                                       
  File "/opt/conda/lib/python3.7/site-packages/accelerate/utils/operations.py", line 333, in broadcast_object_list                          
    torch.distributed.broadcast_object_list(object_list, src=from_process)                                                                  
  File "/opt/conda/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1900, in broadcast_object_list                  
    object_list[i] = _tensor_to_object(obj_view, obj_size)                                                                                  
  File "/opt/conda/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1571, in _tensor_to_object                      
    return _unpickler(io.BytesIO(buf)).load()                                                                                               
_pickle.UnpicklingError: invalid load key, '@'.
```



## Steps to reproduce the bug
```python
from datasets import load_dataset
dataset = load_dataset(
                    cfg["dataset_name"], streaming=True, split="train")
dataset = dataset.remove_columns("meta")
dataset = dataset.map(tokenize_and_concatenate, batched=True)
dataset = dataset.with_format(type="torch")
train_data_loader = DataLoader(
                dataset, batch_size=cfg["batch_size"], num_workers=3)

for batch in train_data_loader:
    continue
```

`tokenize_and_concatenate` is a custom tokenization function I defined on the GPT-NeoX tokenizer to tokenize the text, separated by endoftext tokens, and reshape to have length batch_size, I don't think this is related to tokenization:

```
import numpy as np
import einops
import torch
def tokenize_and_concatenate(examples):
        texts = examples["text"]
        full_text = tokenizer.eos_token.join(texts)
        div = 20
        length = len(full_text) // div
        text_list = [full_text[i * length: (i + 1) * length]
                     for i in range(div)]
        tokens = tokenizer(text_list, return_tensors="np", padding=True)[
            "input_ids"
        ].flatten()
        tokens = tokens[tokens != tokenizer.pad_token_id]
        n = len(tokens)
        curr_batch_size = n // (seq_len - 1)
        tokens = tokens[: (seq_len - 1) * curr_batch_size]
        tokens = einops.rearrange(
            tokens,
            "(batch_size seq) -> batch_size seq",
            batch_size=curr_batch_size,
            seq=seq_len - 1,
        )
        prefix = np.ones((curr_batch_size, 1), dtype=np.int64) * \
            tokenizer.bos_token_id
        return {
            "text": np.concatenate([prefix, tokens], axis=1)
        }
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.4.0
- Platform: Linux-5.4.0-105-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.13
- PyArrow version: 9.0.0
- Pandas version: 1.3.5

ZStandard data:
Version: 0.18.0
Summary: Zstandard bindings for Python
Home-page: https://github.com/indygreg/python-zstandard
Author: Gregory Szorc
Author-email: gregory.szorc@gmail.com
License: BSD
Location: /opt/conda/lib/python3.7/site-packages
Requires: 
Required-by: