Intermittent JSON parse error when streaming the Pile
Opened by neelnanda-io (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:

Comments:
Maybe #2838 can help. In this PR we allow to skip bad chunks of JSON data to not crash the training

Did you have warning messages before the error ?

something like this maybe ?
```
03/24/2022 02:19:46 - WARNING - datasets.utils.streaming_download_manager - Got disconnected from remote data host. Retrying in 5sec [1/20]
03/24/2022 02:20:01 - WARNING - datasets.utils.streaming_download_manager - Got disconnected from remote data host. Retrying in 5sec [2/20]
03/24/2022 02:20:09 - ERROR - datasets.packaged_modules.json.json - Failed to read file 'gzip://file-000000000007.json::https://huggingface.co/datasets/lvwerra/codeparrot-clean-train/resolve/1d740acb9d09cf7a3307553323e2c677a6535407/file-000000000007.json.gz' with error <class 'pyarrow.lib.ArrowInvalid'>: JSON parse error: Invalid value. in row 0
```

Ah, thanks! I did get errors like that. Sad that PR wasn't merged in! 

I'm currently just downloading 200GB of the Pile locally to avoid streaming (I have space and it's faster anyway), but that's really useful! I can probably apply the dumb patch of just commenting out the bits that raise the JSON Parse Error lol, based on your code - if I continue the loop should it be fine?

Yup you can get some inspiration from this PR. It simply ignores the bad chunks (a chunk is ~a few MBs of data).
We'll try to merge this PR soon