GitHub issue: Make `BufferShuffledExamplesIterable` resumable
Author: yzhangcs (CONTRIBUTOR)
State: open | Created: 2024-07-22 07:50:02

This PR aims to implement a resumable `BufferShuffledExamplesIterable`.
Instead of saving the entire buffer content, which is very memory-intensive, the newly implemented `BufferShuffledExamplesIterable` saves only the minimal state necessary for recovery, e.g., the random generator states and the state of the first example in the buffer dict.

The idea is that since the buffer size is limited, even if the entire buffer is discarded, we can rebuild it as long as the state of the oldest example is recorded. For buffer size $B$, the expected distance between when an example is pushed and when it is yielded is 
$d = \sum_{k=1}^{\infty} k\frac{1}{B} (1 - \frac{1}{B} )^{k-1} =B$.
Simulation experiments support these claims:
```py
from random import randint

BUFFER_SIZE = 1024

dists = []
buffer = []
for i in range(10000000):
    if i < BUFFER_SIZE:
        buffer.append(i)
    else:
        index = randint(0, BUFFER_SIZE - 1)
        dists.append(i - buffer[index])
        buffer[index] = i

print(f"MIN DIST: {min(dists)}\nMAX DIST: {max(dists)}\nAVG DIST: {sum(dists) / len(dists):.2f}\n")
```
which produces the following output:
```py
MIN DIST: 1
MAX DIST: 15136
AVG DIST: 1023.95

```
The overall time for reconstructing the buffer and recovery should not be too long.
The following code mimics the cases of resuming online tokenization by `datasets` and `StatefulDataLoader` under distributed scenarios, 

```py
import pickle
import time
from itertools import chain
from typing import Any, Dict, List

import torch
from datasets import load_dataset
from torchdata.stateful_dataloader import StatefulDataLoader
from tqdm import tqdm
from transformers import AutoTokenizer, DataCollatorForLanguageModeling

tokenizer = AutoTokenizer.from_pretrained('fla-hub/gla-1.3B-100B')
tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

torch.manual_seed(42)


def tokenize(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
    input_ids = tokenizer(examples['text'])['input_ids']
    input_ids = list(chain(*input_ids))
    total_length = len(input_ids)
    chunk_size = 2048
    total_length = (total_length // chunk_size) * chunk_size
    # the last chunk smaller than chunk_size will be discarded
    return {'input_ids': [input_ids[i: i+chunk_size] for i in range(0, total_length, chunk_size)]}


batch_size = 16
num_workers = 5
context_length = 2048
rank = 1
world_size = 32
prefetch_factor = 2
steps = 2048
path = 'fla-hub/slimpajama-test'
dataset = load_dataset(
    path=path,
    split='train',
    streaming=True,
    trust_remote_code=True
)
dataset = dataset.map(tokenize, batched=True, remove_columns=next(iter(dataset)).keys())
dataset = dataset.shuffle(seed=42)
loader = StatefulDataLoader(dataset=dataset,
                            batch_size=batch_size,
                            collate_fn=data_collator,
                            num_workers=num_workers,
                            persistent_workers=False,
                            prefetch_factor=prefetch_factor)
start = time.time()
for i, batch in tqdm(enumerate(loader)):
    if i == 0:
        print(f'{i}\n{batch["input_ids"]}')
    if i == steps - 1:
        print(f'{i}\n{batch["input_ids"]}')
        state_dict = loader.state_dict()
    if i == steps:
        print(f'{i}\n{batch["input_ids"]}')
        break
print(f"{time.time() - start:.2f}s elapsed")
print(f"{len(pickle.dumps(state_dict)) / 1024**2:.2f}MB states in total")
for worker in state_dict['_snapshot']['_worker_snapshots'].keys():
    print(f"{worker} {len(pickle.dumps(state_dict['_snapshot']['_worker_snapshots'][worker])) / 1024**2:.2f}MB")
print(state_dict['_snapshot']['_worker_snapshots']['worker_0']['dataset_state'])

loader = StatefulDataLoader(dataset=dataset,
                            batch_size=batch_size,
                            collate_fn=data_collator,
                            num_workers=num_workers,
                            persistent_workers=False,
                            prefetch_factor=prefetch_factor)
print("Loading state dict")
loader.load_state_dict(state_dict)
start = time.time()
for batch in loader:
    print(batch['input_ids'])
    break

print(f"{time.time() - start:.2f}s elapsed")

```
and the outputs are
```py
0
tensor([[  909,   395, 19082,  ..., 13088, 16232,   395],
        [  601, 28705, 28770,  ..., 28733,   923,   288],
        [21753, 15071, 13977,  ...,  9369, 28723,   415],
        ...,
        [21763, 28751, 20300,  ..., 28781, 28734,  4775],
        [  354,   396, 10214,  ...,   298,   429, 28770],
        [  333,  6149, 28768,  ...,  2773,   340,   351]])
2047
tensor([[28723,   415,  3889,  ...,   272,  3065,  2609],
        [  403,  3214,  3629,  ...,   403, 21163, 16434],
        [28723,    13, 28749,  ..., 28705, 28750, 28734],
        ...,
        [ 2778,  2251, 28723,  ...,   354,   684,   429],
        [ 5659,   298,  1038,  ...,  5290,   297, 22153],
        [  938, 28723,  1537,  ...,  9123, 28733, 12154]])
2048
tensor([[  769,   278, 12531,  ..., 28721, 19309, 28739],
        [  415, 23347,   622,  ...,  3937,  2426, 28725],
        [28745,  4345, 28723,  ...,   338, 28725,   583],
        ...,
        [ 1670, 28709,  5809,  ..., 28734, 28760,   393],
        [  340,  1277,   624,  ...,   325, 28790,  1329],
        [  523,  1144,  3409,  ...,   359,   359, 17422]])
65.97s elapsed
0.00MB states in total
worker_0 0.00MB
worker_1 0.00MB
worker_2 0.00MB
worker_3 0.00MB
worker_4 0.00MB
{'ex_iterable': {'ex_iterable': {'shard_idx': 0, 'shard_example_idx': 14000}, 'num_examples_since_previous_state': 166, 'previous_state_example_idx': 7394, 'previous_state': {'shard_idx': 0, 'shard_example_idx': 13000}}, 'num_taken': 6560, 'global_example_idx': 7560, 'buffer_state_dict': {'num_taken': 6560, 'global_example_idx': 356, 'index_offset': 0, 'first_state': {'ex_iterable': {'shard_idx': 0, 'shard_example_idx': 1000}, 'num_examples_since_previous_state': 356, 'previous_state_example_idx': 0, 'previous_state': {'shard_idx': 0, 'shard_example_idx': 0}}, 'bit_generator_state': {'state': {'state': 274674114334540486603088602300644985544, 'inc': 332724090758049132448979897138935081983}, 'bit_generator': 'PCG64', 'has_uint32': 0, 'uinteger': 0}}}
Loading state dict
tensor([[  769,   278, 12531,  ..., 28721, 19309, 28739],
        [  415, 23347,   622,  ...,  3937,  2426, 28725],
        [28745,  4345, 28723,  ...,   338, 28725,   583],
        ...,
        [ 1670, 28709,  5809,  ..., 28734, 28760,   393],
        [  340,  1277,   624,  ...,   325, 28790,  1329],
        [  523,  1144,  3409,  ...,   359,   359, 17422]])
24.60s elapsed
```

Not sure if this PR complies with the `datasets` code style. Looking for your help @lhoestq, also very willing to further improve the code if any suggestions are given.