Add openwebtext dataset for larger scale training without shuffling

torchtrain/datasets/__init__.py

@@ -13,4 +13,5 @@
     "alpaca": build_hf_data_loader,
     "minipile": build_hf_data_loader,
     "c4": build_hf_data_loader,
+    "openwebtext": build_hf_data_loader,
 }

torchtrain/datasets/hf_datasets.py

@@ -17,6 +17,7 @@
     "alpaca": "tatsu-lab/alpaca",
     "minipile": "JeanKaddour/minipile",
     "c4": "allenai/c4",
+    "openwebtext": "Skylion007/openwebtext",
 }
 
 
@@ -32,9 +33,10 @@ class HuggingFaceDataset(IterableDataset):
         rank (int): rank of the current data parallel process
         infinite (bool): whether to loop infinitely over the dataset
 
-    We currently support three datasets:
+    We currently support four datasets:
     alpaca (52K training entries)
-    minipile (1M training entries)
+    minipile (1M training entries, amalgamated from other datasets)
+    openwebtext (1M training entries, same type of data for entire dataset)
     c4 (177M training entries - this dataset is streamed due to the size)
 
     >> Alpaca <<:
@@ -65,6 +67,15 @@ class HuggingFaceDataset(IterableDataset):
     'timestamp': '2019-04-25T12:57:54Z'
     }
 
+    >> OpenWebText <<:
+    OpenWeb crawl, English
+    Example:
+    {
+        'text': "Amazon has launched a new cheaper version of its Echo Dot voice-controlled device today.
+    The launch comes six months after Amazon first introduced two new Echo devices —
+    one of which was the $90 Echo Dot,..."
+    }
+
     Example use (alpaca):
     >>> alpaca_ds = HuggingFaceDataset(dataset_name="alpaca", dataset_path=None, tokenizer=tokenizer)
     >>> for batch in Dataloader(alpaca_ds, batch_size=8):

train_configs/llama_13b.toml

@@ -40,4 +40,4 @@ compile = false
 checkpoint_interval = 3600
 checkpoint_interval_type = "steps"
 checkpoint_folder = ""
-dataset = "minipile"
+dataset = "openwebtext"

train_configs/llama_70b.toml

@@ -40,4 +40,4 @@ compile = false
 checkpoint_interval = 3600
 checkpoint_interval_type = "steps"
 checkpoint_folder = ""
-dataset = "minipile"
+dataset = "openwebtext"