Add general purpose LitData streaming data module (#1118)

Author: awaelchli

litgpt/data/__init__.py

@@ -1,6 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-from litgpt.data.base import LitDataModule, SFTDataset, get_sft_collate_fn
+from litgpt.data.base import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.data.alpaca import Alpaca
 from litgpt.data.alpaca_2k import Alpaca2k
 from litgpt.data.alpaca_gpt4 import AlpacaGPT4
@@ -9,6 +9,7 @@
 from litgpt.data.dolly import Dolly
 from litgpt.data.flan import FLAN
 from litgpt.data.lima import LIMA
+from litgpt.data.lit_data import LitData
 from litgpt.data.longform import LongForm
 from litgpt.data.tinyllama import TinyLlama
 from litgpt.data.tinystories import TinyStories
@@ -24,7 +25,8 @@
     "FLAN",
     "JSON",
     "LIMA",
-    "LitDataModule",
+    "LitData",
+    "DataModule",
     "LongForm",
     "OpenWebText",
     "SFTDataset",

litgpt/data/alpaca.py

@@ -9,15 +9,15 @@
 import torch
 from torch.utils.data import random_split, DataLoader
 from lightning_utilities.core.imports import RequirementCache
-from litgpt.data import SFTDataset, get_sft_collate_fn, LitDataModule
+from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
 from litgpt.prompts import PromptStyle
 from litgpt.tokenizer import Tokenizer
 
 _URL = "https://raw.githubusercontent.com/tloen/alpaca-lora/main/alpaca_data_cleaned_archive.json"
 
 
 @dataclass
-class Alpaca(LitDataModule):
+class Alpaca(DataModule):
     """Alpaca data module for supervised finetuning."""
 
     mask_prompt: bool = False

litgpt/data/base.py

@@ -12,7 +12,7 @@
 from litgpt.prompts import PromptStyle
 
 
-class LitDataModule(LightningDataModule):
+class DataModule(LightningDataModule):
     """Base class for all data modules in LitGPT."""
 
     @abstractmethod

litgpt/data/deita.py

@@ -9,12 +9,12 @@
 from torch.utils.data import DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import LitDataModule, SFTDataset, get_sft_collate_fn
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.tokenizer import Tokenizer
 
 
 @dataclass
-class Deita(LitDataModule):
+class Deita(DataModule):
     """Deita data module for supervised finetuning."""
 
     mask_prompt: bool = False

litgpt/data/flan.py

@@ -9,7 +9,7 @@
 from torch.utils.data import DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, LitDataModule
+from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
 from litgpt.data.alpaca import download_if_missing
 from litgpt.tokenizer import Tokenizer
 
@@ -19,7 +19,7 @@
 # TODO: Including all subsets, FLAN is too large to be loaded in memory. Switch the implementation to cache
 #   on disk or use Lightning Data
 @dataclass
-class FLAN(LitDataModule):
+class FLAN(DataModule):
     """FLAN data module for supervised finetuning."""
 
     mask_prompt: bool = False

litgpt/data/json_data.py

@@ -9,12 +9,12 @@
 from torch.utils.data import random_split, DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, LitDataModule
+from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
 from litgpt.tokenizer import Tokenizer
 
 
 @dataclass
-class JSON(LitDataModule):
+class JSON(DataModule):
     """Loads JSON or JSONL data for supervised finetuning."""
 
     json_path: Path

litgpt/data/lima.py

@@ -9,12 +9,12 @@
 from torch.utils.data import random_split, DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import LitDataModule, SFTDataset, get_sft_collate_fn
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.tokenizer import Tokenizer
 
 
 @dataclass
-class LIMA(LitDataModule):
+class LIMA(DataModule):
     """LIMA data module for supervised finetuning."""
 
     mask_prompt: bool = False

litgpt/data/lit_data.py

@@ -0,0 +1,68 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Union, Optional, Tuple
+
+from torch.utils.data import DataLoader
+
+from litgpt import Tokenizer
+from litgpt.data import DataModule
+
+
+@dataclass
+class LitData(DataModule):
+    """Loads data using LitData's StreamingDataset given a path to a folder of preprocessed data (chunks)."""
+
+    data_path: Union[str, Path] = Path("data/")
+    """The path to the data directory containing the preprocessed chunks for the streaming dataset
+    The path can also be a remote path (e.g., s3://). See also ``split_names`` if this path contains subfolders
+    for training- and validation splits."""
+    split_names: Optional[Tuple[str, str]] = None
+    """Optional tuple for names of subfolders for training and validation under ``data_path``. If not provided,
+    all data under data_path will be used for training, and the validation dataloader will be identical to the
+    train dataloader."""
+    seed: int = 42
+    """The random seed for shuffling the dataset."""
+    num_workers: int = 8
+    """How many DataLoader processes to use for loading."""
+
+    batch_size: int = field(init=False, repr=False, default=1)
+    seq_length: int = field(init=False, repr=False, default=2048)
+
+    def __post_init__(self) -> None:
+        if self.split_names is not None and len(self.split_names) != 2:
+            raise ValueError(
+                "If provided `split_names` must be a tuple of two strings, for example: ('train', 'val')."
+            )
+
+    def connect(
+        self,
+        tokenizer: Optional[Tokenizer] = None,
+        batch_size: int = 1,
+        max_seq_length: Optional[int] = None
+    ) -> None:
+        self.batch_size = batch_size
+        self.seq_length = max_seq_length + 1  # Increase by one because we need the next token as well
+
+    def train_dataloader(self) -> DataLoader:
+        input_dir = os.path.join(self.data_path, self.split_names[0]) if self.split_names else str(self.data_path)
+        return self._dataloader(input_dir=input_dir, train=True)
+
+    def val_dataloader(self) -> DataLoader:
+        input_dir = os.path.join(self.data_path, self.split_names[1]) if self.split_names else str(self.data_path)
+        return self._dataloader(input_dir=input_dir, train=False)
+
+    def _dataloader(self, input_dir: str, train: bool):
+        from litdata.streaming import StreamingDataset, TokensLoader
+
+        dataset = StreamingDataset(
+            input_dir=input_dir,
+            item_loader=TokensLoader(block_size=self.seq_length),
+            shuffle=train,
+            drop_last=True,
+        )
+        dataloader = DataLoader(
+            dataset, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_workers, drop_last=True
+        )
+        return dataloader

litgpt/data/longform.py

@@ -9,7 +9,7 @@
 from torch.utils.data import DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, LitDataModule
+from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
 from litgpt.data.alpaca import download_if_missing
 from litgpt.tokenizer import Tokenizer
 
@@ -18,7 +18,7 @@
 
 
 @dataclass
-class LongForm(LitDataModule):
+class LongForm(DataModule):
     """LongForm data module for supervised finetuning."""
 
     mask_prompt: bool = False

litgpt/data/openwebtext.py

@@ -8,11 +8,11 @@
 from torch.utils.data import DataLoader
 
 from litgpt import Tokenizer
-from litgpt.data import LitDataModule
+from litgpt.data import DataModule
 
 
 @dataclass
-class OpenWebText(LitDataModule):
+class OpenWebText(DataModule):
     """The OpenWebText data module for pretraining."""
 
     data_path: Union[str, Path] = Path("data/openwebtext")