Resuming the StreamingDataLoader from a checkpoint makes the iterator hang

[awaelchli]

🐛 Bug

In LitGPT, we currently aren't able to resume correctly from a pretrained checkpoint. The state the dataloader loads causes a stall in the iterator.

To Reproduce

Studio for reproduction.

I made a minimal repro without LitGPT involved:

import torch
from litdata.streaming import CombinedStreamingDataset, StreamingDataLoader, StreamingDataset, TokensLoader


def get_dataloader():
    slim_train_data = StreamingDataset(
        input_dir="s3://tinyllama-template/slimpajama/train",
        item_loader=TokensLoader(block_size=2049),
        shuffle=True,
        drop_last=True,
    )
    train_data = slim_train_data

    train_datasets = [
        slim_train_data,
        StreamingDataset(
            input_dir="s3://tinyllama-template/starcoder",
            item_loader=TokensLoader(block_size=2049),
            shuffle=True,
            drop_last=True,
        ),
    ]

    # Mix SlimPajama data and Starcoder data with these proportions:
    weights = (0.693584, 0.306416)
    train_data = CombinedStreamingDataset(
        datasets=train_datasets, seed=42, weights=weights, iterate_over_all=False
    )
    train_dataloader = StreamingDataLoader(
        train_data, batch_size=4, pin_memory=True, num_workers=8, drop_last=True
    )
    return train_dataloader


def dataloader_state():
    """Dataloader state extracted from the real checkpoint, for reproduction"""
    return {'dataset': {'0': {'num_samples_yielded': 852336, 'num_workers': 8, 'batch_size': 4, 'current_epoch': 1, 'input_dir_path': None, 'input_dir_url': 's3://tinyllama-template/slimpajama/train', 'item_loader': {'block_size': 2049}, 'drop_last': True, 'seed': 42, 'world_size': 1, 'shuffle': True}, '1': {'num_samples_yielded': 376464, 'num_workers': 8, 'batch_size': 4, 'current_epoch': 1, 'input_dir_path': None, 'input_dir_url': 's3://tinyllama-template/starcoder', 'item_loader': {'block_size': 2049}, 'drop_last': True, 'seed': 42, 'world_size': 1, 'shuffle': True}}, 'current_epoch': 0, 'latest_worker_idx': 7, 'num_samples_yielded': {0: [106542, 47058], 1: [106542, 47058], 2: [106542, 47058], 3: [106542, 47058], 4: [106542, 47058], 5: [106542, 47058], 6: [106542, 47058], 7: [106542, 47058]}}


if __name__ == "__main__":
    train_dataloader = get_dataloader()
    # checkpoint = torch.load("dataloader-state.pt")
    train_dataloader.load_state_dict(dataloader_state())
    
    print("start")
    for batch in train_dataloader:  # hangs here
        print(batch)

If executed in a Studio, to avoid authentication errors, the following patch is needed in /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/litdata/streaming/client.py, line 40:

        if has_shared_credentials_file or not _IS_IN_STUDIO or True:
            self._client = boto3.client(
                "s3", config=botocore.config.Config(retries={"max_attempts": 1000, "mode": "adaptive"}, signature_version=botocore.UNSIGNED)
            )

Expected behavior

Continue fetching data from where we left off.

Environment

LitData: 0.2.15

cc @sanyalsunny111

[awaelchli]

In this line of code

https://github.com/Lightning-AI/litdata/blob/14a47f21575abd045115ad3c43a766d08b102f02/src/litdata/streaming/dataloader.py#L444-L448

the range becomes empty because self._loader._latest_worker_idx == self._loader.num_workers == 8. With self._indexes being an empty list, this then causes this if-else logic

https://github.com/Lightning-AI/litdata/blob/14a47f21575abd045115ad3c43a766d08b102f02/src/litdata/streaming/dataloader.py#L462-L476

to go into infinite recursion, because it calls itself when self._indexes evaluates to False.

The logic in _try_put_index is inherently flawed because it is a recursive call into the same method without the state ever changing. Therefore the code path would never change.
It's also not clear to me why it wouldn't be allowed to restore from worker index 0.

[deependujha]

thanks for pinpointing the point of concern. We'll try to fix the issue ASAP.

[deependujha]

the range becomes empty because self._loader._latest_worker_idx == self._loader.num_workers == 8. With self._indexes being an empty list, this then causes this if-else logic

# src/litdata/streaming/dataloader.py#702
self._latest_worker_idx = (obj["latest_worker_idx"] + 1) % (self.num_workers if self.num_workers > 0 else 1)

Now, instead of latest_worker_idx being 8, it will be 0, but the range will still be empty bcoz of else []. But, it doesn't go into infinite recursion. It calls the parent class's _try_put_index.

But, for some reason it waits indefinitely at 79 count.

---

Then, I tried restoring from index: 0, by modifying the code slightly:

self._indexes = (
            list(range(self._loader._latest_worker_idx, self._loader.num_workers))
            # if self._loader._latest_worker_idx > 0
            # else []
        )

but, everything stays same:

This again goes until 79 and then waits. I think it's either stuck in a for loop or a process waiting for completion of getting tokens. Yet to verify.

[awaelchli]

I found the problem, but it is not in the main process like I showed above. The problem is simpler to see/debug if we reduce the number of workers:

import torch
from litdata.streaming import (
    CombinedStreamingDataset,
    StreamingDataLoader,
    StreamingDataset,
    TokensLoader,
)


def get_dataloader():
    slim_train_data = StreamingDataset(
        input_dir="s3://tinyllama-template/slimpajama/train",
        item_loader=TokensLoader(block_size=2049),
        shuffle=True,
        drop_last=True,
    )
    train_data = slim_train_data

    train_datasets = [
        slim_train_data,
        StreamingDataset(
            input_dir="s3://tinyllama-template/starcoder",
            item_loader=TokensLoader(block_size=2049),
            shuffle=True,
            drop_last=True,
        ),
    ]

    # Mix SlimPajama data and Starcoder data with these proportions:
    weights = (0.693584, 0.306416)
    train_data = CombinedStreamingDataset(
        datasets=train_datasets, seed=42, weights=weights, iterate_over_all=False
    )
    train_dataloader = StreamingDataLoader(
        train_data, batch_size=4, pin_memory=True, num_workers=1, drop_last=True
    )
    return train_dataloader


def dataloader_state():
    """Dataloader state extracted from the real checkpoint, for reproduction"""
    return {
        "dataset": {
            "0": {
                "num_samples_yielded": 852336,
                "num_workers": 1,
                "batch_size": 4,
                "current_epoch": 1,
                "input_dir_path": None,
                "input_dir_url": "s3://tinyllama-template/slimpajama/train",
                "item_loader": {"block_size": 2049},
                "drop_last": True,
                "seed": 42,
                "world_size": 1,
                "shuffle": True,
            },
            "1": {
                "num_samples_yielded": 376464,
                "num_workers": 1,
                "batch_size": 4,
                "current_epoch": 1,
                "input_dir_path": None,
                "input_dir_url": "s3://tinyllama-template/starcoder",
                "item_loader": {"block_size": 2049},
                "drop_last": True,
                "seed": 42,
                "world_size": 1,
                "shuffle": True,
            },
        },
        "current_epoch": 0,
        "latest_worker_idx": 0,
        "num_samples_yielded": {
            0: [106542, 47058],
            # 1: [106542, 47058],
            # 2: [106542, 47058],
            # 3: [106542, 47058],
            # 4: [106542, 47058],
            # 5: [106542, 47058],
            # 6: [106542, 47058],
            # 7: [106542, 47058],
        },
    }


if __name__ == "__main__":
    train_dataloader = get_dataloader()
    # checkpoint = torch.load("dataloader-state.pt")
    train_dataloader.load_state_dict(dataloader_state())

    print("start")
    for batch in train_dataloader:  # hangs here
        print(batch)

What I found so far is that the reader is waiting for a chunk to be available here, but it never is:
https://github.com/Lightning-AI/litdata/blob/c4c9117134e20c6b3f9ca7b071932475ab12da80/src/litdata/streaming/item_loader.py#L132-L134
Chunks are being downloaded, but it seems their names are not the same as what the item loader is waiting for, so it will wait forever

For example, the downloader thread will download a file

'/cache/chunks/c8829d9f5b65f28afca2a0e7e207a8a7/chunk-95-280.bin'

corresponding to a chunk index 6748, but the reader will wait in the while loop for a different file

/cache/chunks/c8829d9f5b65f28afca2a0e7e207a8a7/chunk-120-15.bin

with chunk index 8034

This deadlock suggests that the inconsistent way the filename/chunk index is computed is the problem. But where that originates, I haven't figured out.

[sanyalsunny111]

A quick question and a suggestion? Can we try to make a threshold time say for 50 ms for a particular chunk and move on to the next chunk? I know this can reduce the number of tokens used for training.

[awaelchli]

@sanyalsunny111 It's a systematic problem. We can't just skip chunks and hope it gets better :)

I'm trying to understand the code. One thing I don't understand is what the chunk.chunk_indexes are there for. We put 17000 of them onto the queue here, but why: https://github.com/Lightning-AI/litdata/blob/c4c9117134e20c6b3f9ca7b071932475ab12da80/src/litdata/streaming/reader.py#L262-L263
I doubt that we want to download 17k chunks.