How to shuffle & optimize Hugging Face datasets for LLM pre-training with StreamingDataset ?

[pietrolesci]

Bug Report

Description

When using StreamingDataset with shuffle=True, accessing individual samples (e.g., dataset[0]) returns the same sample every time, despite shuffle being enabled. This makes it impossible to verify shuffling is working or to access random samples directly from the dataset.

Environment

• LitData version: 0.2.52
• Python version: 3.11.10

Minimal Reproduction Code

from pathlib import Path
import litdata as ld
from litdata.streaming.item_loader import ParquetLoader

INDEX_PATH = Path(".litdata_index")
CACHE_DIR = Path(".litdata_cache")

# Define the Hugging Face dataset URI
hf_dataset_uri = "hf://datasets/pietrolesci/finewebedu-20B/bpe32000minipile"
INDEX_PATH.mkdir(exist_ok=True, parents=True)
ld.index_parquet_dataset(hf_dataset_uri, cache_dir=INDEX_PATH)

dataset = ld.StreamingDataset(
    hf_dataset_uri, 
    item_loader=ParquetLoader(), 
    cache_dir=CACHE_DIR, 
    index_path=INDEX_PATH,
    shuffle=True,
    drop_last=True,
    max_pre_download=5,
)

# This returns the same sample every time
print("Sample 0:", dataset[0])
print("Sample 0 again:", dataset[0])

For context, I ran into this issue while trying to use LitData to transform a Hugging Face dataset into a packed dataset for LLM pretraining. To simplify, I started from a tokenized dataset and wanted to:

• Shuffle the documents
• Pack them into equal-length sequences
• Shuffle the sequences

Any tips about this would be really helpful. Thanks a lot in advance for your attention!

[bhimrazy]

Hi @pietrolesci, thanks for raising this!

In StreamingDataset, shuffling only takes effect during iteration. When you access an item directly by index (e.g., dataset[0]), you’ll always get the same sample at that index, since the shuffle logic is applied at iteration time. For example:

dataset_iter = iter(dataset)
print("Sample 0:", next(dataset_iter))
print("Sample 1:", next(dataset_iter))

You can also inspect the shuffled order of indexes via an internal variable:

print("indexes:", dataset.upcoming_indexes)

If you want different shuffles across runs, you can also pass a different seed when constructing the dataset.

[bhimrazy]

For optimizing the tokenized dataset, you may also want to check out the LLM Pre-training under Features section in the README. Since your dataset is already tokenized, you likely don’t need to apply another tokenizer — you can just return the tokenized ids from the optimize fn.

Please feel free to ask any questions here if anything’s unclear.

[pietrolesci]

Thanks for your swift reply @bhimrazy. Now it's clearer, and shuffling works as expected.

As a follow-up on the optimization: In the case of a HF dataset, what should be passed to optimize, the list of URIs to the parquet files on the Hub or is there a smarter way to do it with the StreamingDataset?

Thanks a lot for your help!

[bhimrazy]

Hi @pietrolesci,

Thanks for your follow-up!

For optimizing a HF dataset , you have a few options to use with LitData's optimize function. Since your dataset is already tokenized and you can access it with StreamingDataset, you can process and pack your data easily. Here are some options to use with optimize:

1. Access by Index from dataset

If you want to get samples by their index, you can pass a list of indexes and fetch each sample inside your optimization function:

import litdata as ld
from functools import partial

dataset = ld.StreamingDataset(...)

def optimize_fn(index, dataset=None):
    data = dataset[index]
    # process and pack data here
    ids = ...
    return ids

ld.optimize(
    fn=partial(optimize_fn, dataset=dataset),
    inputs=list(range(len(dataset))), # these indexes will get passed to optimize fn
    # other arguments ...
)

2. Use dataset as an Iterator

You can also pass an iterator over the dataset to your optimization function.

import litdata as ld
from functools import partial

dataset = ld.StreamingDataset(...)

def optimize_fn(index, dataset=None):
    data = next(dataset)
    # process and pack data here
    ids = ...
    return ids

ld.optimize(
    fn=partial(optimize_fn, dataset=iter(dataset)),
    inputs=list(range(len(dataset))), # this indexes are just to let optimizer know to process how many samples
    # other arguments ...
)

3. Use dataset with StreamingDataLoader

you can also use StreamingDataLoader and pass it to optimize. Your optimize_fn function will get data of batch_size of 1:

import litdata as ld

dataset = ld.StreamingDataset(...)
dataloader = ld.StreamingDataLoader(dataset)  # Default batch_size=1

def optimize_fn(data):
    # data is a batch (default batch size is 1)
    # process and pack batch here
    ids = ...
    return ids

ld.optimize(
    fn=optimize_fn,
    inputs=dataloader,
    # other arguments ...
)

eg:

litData/tests/processing/test_functions.py

Lines 787 to 804 in 8e902df

	dataset = StreamingDataset(parquet_dir, item_loader=ParquetLoader())
	dataloader = StreamingDataLoader(dataset)
	# Verify the dataloader has the expected length
	assert len(dataloader) == num_items, f"Expected dataloader length {num_items}, got {len(dataloader)}"
	# Optimize the dataset using the streaming dataloader as input
	output_dir = os.path.join(tmpdir, "out")
	os.makedirs(output_dir, exist_ok=True)
	optimize(
	fn=optimize_fn,
	inputs=dataloader,
	num_workers=num_workers,
	output_dir=output_dir,
	chunk_bytes="64MB",
	)

If you have any further questions, need clarification, or require help, please feel free to leave a comment here — happy to help 😊

Best,
Bhimraj

[bhimrazy]

I had also created an issue (#558) to add the examples for the same but never got back to it. I think now’s a good time to finally resolve it 😅.

[pietrolesci]

Hi @bhimrazy, thank you so much for your time and swift responses!

I am trying to make it work following the LLM pre-training example in the README but I am not able to make it work:

from functools import partial
from pathlib import Path

import litdata as ld
from litdata.streaming.item_loader import ParquetLoader

INDEX_PATH = Path(".litdata_index")
CACHE_DIR = Path(".litdata_cache")

# Define the Hugging Face dataset URI
hf_dataset_uri = "hf://datasets/pietrolesci/finewebedu-20B/bpe32000minipile"


def optimize_fn(index: int, dataset: ld.StreamingDataset):
    data = dataset[index]
    yield data["input_ids"]


if __name__ == "__main__":
    if not INDEX_PATH.exists():
        INDEX_PATH.mkdir(exist_ok=True, parents=True)
        ld.index_parquet_dataset(hf_dataset_uri, cache_dir=str(INDEX_PATH))

    dataset = ld.StreamingDataset(
        hf_dataset_uri, 
        item_loader=ParquetLoader(), 
        cache_dir=str(CACHE_DIR), 
        index_path=str(INDEX_PATH),
        shuffle=True,
        drop_last=True,
        max_pre_download=2,
    )

    ld.optimize(
        fn=partial(optimize_fn, dataset=dataset),
        inputs=list(range(10)), # test
        chunk_size=(2049 * 2), # test
        # This is important to inform LitData that we are encoding contiguous 1D array (tokens). 
        # LitData skips storing metadata for each sample e.g all the tokens are concatenated to form one large tensor.
        # NOTE: THIS DOES NOT SEEM TO WORK, however it is convenient to avoid storing useless metadata
        item_loader=ld.TokensLoader(), 
    ) # -> I expect the chunks.bin to be a "stream" of tokens that I can load with `StreamingDataset` as sequences and ideally shuffle again (at the sequence level)

I've attempted to debug this myself, but I'm finding it challenging as the documentation for optimize (and the library in general) does not provide much guidance. Would you have any suggestions on what might be going wrong or how I could better approach troubleshooting this? Any guidance would be greatly appreciated—thanks a lot in advance!

---

EDIT: as a minor note, I think it's not super obvious how to pack inside the fn as it works at the individual input element level. So, that's why I was hoping to optimize by simply shuffling documents and storing them, "flattened", into chunks. Then, once I am actually doing training, I can load these chunks (stream of tokens), shape them as (batch_size, seqlen) on the fly, and potentially shuffle them.

---

UPDATE: TokensLoader really wants arrays instead of lists, so:

def optimize_fn(index: int, dataset: ld.StreamingDataset):
    data = dataset[index]
    yield np.array(data["input_ids"], dtype=np.int32)

is now working. Will report back once I test the full pipeline :)

---

UPDATE: This seems to be working and doing what I want. If you have feedback on how to improve/optimise performance, please feel free to let me know :)

import shutil
from collections.abc import Generator
from functools import partial
from pathlib import Path
from typing import Any

import litdata as ld
import numpy as np
from litdata import StreamingDataset, optimize
from litdata.streaming.item_loader import ParquetLoader

# INDEX_PATH = Path(".litdata_index")
CACHE_DIR = Path(".litdata_cache")

# Define the Hugging Face dataset URI
# hf_dataset_uri = "hf://datasets/pietrolesci/finewebedu-20B/bpe32000minipile"
train_uri = "hf://datasets/pietrolesci/minipile/bpe32000minipile/train*.parquet"
val_uri = "hf://datasets/pietrolesci/minipile/bpe32000minipile/validation*.parquet"
EOS_TOKEN_ID = 0  # Example EOS token ID, replace with actual if different

def optimize_fn(index: int, dataset: StreamingDataset) -> Generator[np.ndarray, Any, None]:
    data = next(dataset)  # use next(dataset) instead of dataset[index] for shuffling!
    yield np.array(data["input_ids"] + [EOS_TOKEN_ID], dtype=np.int32)


def process_split(hf_dataset_uri: str) -> None:
    print(f"Processing {hf_dataset_uri}...")
    # NOTE: Indexing is skipped as it does not work with wildcards in the dataset URI
    # if not INDEX_PATH.exists():
    #     INDEX_PATH.mkdir(exist_ok=True, parents=True)
    #     ld.index_parquet_dataset(hf_dataset_uri, cache_dir=str(INDEX_PATH))

    dataset = StreamingDataset(
        hf_dataset_uri,
        item_loader=ParquetLoader(),
        cache_dir=str(CACHE_DIR),
        # index_path=str(INDEX_PATH),
        shuffle=True,
        drop_last=True,
        max_pre_download=2,
    )

    optimize(
        fn=partial(optimize_fn, dataset=iter(dataset)),
        inputs=list(range(len(dataset))),
        chunk_bytes="5GB",
        item_loader=ld.TokensLoader(),
        output_dir=f"./minipile-optimized/{hf_dataset_uri.split('/')[-1].split('*')[0]}",
        num_workers=6,
        compression="zstd",
        # fast_dev_run=True,
    )

    # cleanup
    # shutil.rmtree(INDEX_PATH)
    shutil.rmtree(CACHE_DIR)


if __name__ == "__main__":
    process_split(train_uri)
    process_split(val_uri)

---

Yet another consideration. Since the data are optimized and local, it would be cool to have the possibility to convert StreamingDataset to a regular map-style dataset such that lightning does not complain:

.../lightning/pytorch/utilities/data.py:123: Your `IterableDataset` has `__len__` defined. In combination with multi-process data loading (when num_workers > 1), `__len__` could be inaccurate if each worker is not configured independently to avoid having duplicate data.

[bhimrazy]

That’s great @pietrolesci 🎉, glad to hear you got it working!

On performance: everything looks solid — the only thing I’d point out is the chunk size. For local storage your current 5GB chunks may still work fine, but they’re on the larger side. If you ever plan to stream from cloud/object storage, such large chunks could hurt performance (latency, retries, etc.). We generally recommend keeping chunk sizes ≤500MB, with 64MB being the sweet spot for cloud streaming workloads.

Incase you’d like to try out a multi-node setup, feel free to check out lightning.ai — there’s a LitData plugin available in Studios that makes it easier to scale across nodes.

Some refs: https://lightning.ai/lightning-ai/studios/prepare-the-tinyllama-1t-token-dataset, https://lightning.ai/lightning-ai/studios/pretrain-an-llm-with-pytorch-lightning

Also, about the warning pls feel free to ignore it.

cc: @tchaton in case there are any extra tips worth adding 🙌

[pietrolesci]

Thanks a lot for your feedback and super useful discussion!
I guess the few remaining questions I have are:

• why indexing wildcard URIs does not work (e.g., ld.index_parquet_dataset(train_uri))
• whether it is possible to resume a training run (previously I was handling this with the StatefulDataloader, which is now complaining when I use ld.StreamingDataset. What is the suggested way to do this? I saw this example in the docs showing a manual torch.save(dataloader.state_dict(), "dataloader_state.pt").