general pretrain data generation

scripts/prepare_pretrain_dataset.py

@@ -1,4 +1,5 @@
 import gc
+from typing import Optional
 
 from datasets import load_dataset
 from litdata import optimize, TokensLoader
@@ -6,11 +7,12 @@ from litgpt.tokenizer import Tokenizer
 from functools import partial
 
 
-def batch_iterator(name=None):
+"""
+def batch_iterator_1(path=None):
     # text
-    if name in (None, 'saillab/taco-datasets'):
+    if path in (None, 'saillab/taco-datasets'):
         dataset = (
-            load_dataset(name, data_dir=data_dir, split='train')
+            load_dataset(path, data_dir=data_dir, split='train')
             for data_dir in [
                 'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
                 'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
@@ -32,9 +34,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # text
-    if name in (None, 'xu-song/cc100-samples'):
+    if path in (None, 'xu-song/cc100-samples'):
         dataset = (
-            load_dataset(name, lang, split='train')
+            load_dataset(path, lang, split='train')
             for lang in [
                 'am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br',
                 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es',
@@ -60,8 +62,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # text
-    if name in (None, 'ontocord/fineweb-permissive-multilingual-2m'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'ontocord/fineweb-permissive-multilingual-2m'):
+        dataset = load_dataset(path, split='train')
         
         for row in dataset['text']:
             yield row
@@ -70,8 +72,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # text
-    if name in (None, 'MuskumPillerum/General-Knowledge'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'MuskumPillerum/General-Knowledge'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             if not row['Question'] or not row['Answer']:
@@ -83,9 +85,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # text
-    if name in (None, 'yirenc/general_knowledge_boolean'):
+    if path in (None, 'yirenc/general_knowledge_boolean'):
         for split in ['train', 'validation']:
-            dataset = load_dataset(name, split=split)
+            dataset = load_dataset(path, split=split)
 
             for row in dataset:
                 yield row['question'] + '? ' + str(row['answer']) + '. ' + row['passage']
@@ -94,9 +96,9 @@ def batch_iterator(name=None):
             gc.collect()
 
     # text
-    if name in (None, 'nampdn-ai/tiny-textbooks'):
+    if path in (None, 'nampdn-ai/tiny-textbooks'):
         for split in ['train', 'test']:
-            dataset = load_dataset(name, split=split)
+            dataset = load_dataset(path, split=split)
 
             for row in dataset['textbook']:
                 yield row
@@ -105,8 +107,8 @@ def batch_iterator(name=None):
             gc.collect()
 
     # code
-    if name in (None, 'nampdn-ai/tiny-codes'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'nampdn-ai/tiny-codes'):
+        dataset = load_dataset(path, split='train')
         
         for row in dataset:
             yield (
@@ -119,9 +121,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # code
-    if name in (None, 'bigcode/the-stack-smol-xs'):
+    if path in (None, 'bigcode/the-stack-smol-xs'):
         dataset = (
-            load_dataset(name, lang, split='train', trust_remote_code=True)
+            load_dataset(path, lang, split='train', trust_remote_code=True)
             for lang in [
                 'ada', 'agda', 'alloy', 'antlr', 'applescript', 'assembly',
                 'augeas', 'awk', 'batchfile', 'bison', 'bluespec', 'c',
@@ -150,8 +152,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # text + code
-    if name in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             yield (
@@ -164,8 +166,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # code
-    if name in (None, 'jtatman/python-code-dataset-500k'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'jtatman/python-code-dataset-500k'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             yield (
@@ -178,8 +180,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # code
-    if name in (None, 'iamtarun/python_code_instructions_18k_alpaca'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'iamtarun/python_code_instructions_18k_alpaca'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             yield (
@@ -194,9 +196,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # code
-    if name in (None, 'HuggingFaceH4/CodeAlpaca_20K'):
+    if path in (None, 'HuggingFaceH4/CodeAlpaca_20K'):
         for split in ['train', 'test']:
-            dataset = load_dataset(name, split=split)
+            dataset = load_dataset(path, split=split)
 
             for row in dataset:
                 yield (
@@ -209,8 +211,8 @@ def batch_iterator(name=None):
             gc.collect()
 
     # math
-    if name in (None, 'gair-prox/open-web-math-pro'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'gair-prox/open-web-math-pro'):
+        dataset = load_dataset(path, split='train')
         
         for row in dataset['text']:
             yield row
@@ -219,9 +221,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # math
-    if name in (None, 'rvv-karma/Math-QA'):
+    if path in (None, 'rvv-karma/Math-QA'):
         for split in ['train', 'val', 'test']:
-            dataset = load_dataset(name, split=split)
+            dataset = load_dataset(path, split=split)
             
             for row in dataset:
                 yield (
@@ -234,8 +236,8 @@ def batch_iterator(name=None):
             gc.collect()
 
     # math
-    if name in (None, 'ajibawa-2023/Maths-College'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'ajibawa-2023/Maths-College'):
+        dataset = load_dataset(path, split='train')
         
         for row in dataset:
             yield (
@@ -248,8 +250,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # math
-    if name in (None, 'microsoft/orca-math-word-problems-200k'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'microsoft/orca-math-word-problems-200k'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             yield (
@@ -262,9 +264,9 @@ def batch_iterator(name=None):
         gc.collect()
 
     # math
-    if name in (None, 'fblgit/simple-math'):
+    if path in (None, 'fblgit/simple-math'):
         for split in ['train', 'test']:
-            dataset = load_dataset(name, revision='refs/convert/parquet', split=split)
+            dataset = load_dataset(path, revision='refs/convert/parquet', split=split)
 
             for row in dataset:
                 yield (
@@ -277,8 +279,8 @@ def batch_iterator(name=None):
             gc.collect()
 
     # reasoning
-    if name in (None, 'SkunkworksAI/reasoning-0.01'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'SkunkworksAI/reasoning-0.01'):
+        dataset = load_dataset(path, split='train')
 
         for row in dataset:
             yield (
@@ -293,8 +295,8 @@ def batch_iterator(name=None):
         gc.collect()
 
     # emoji
-    if name in (None, 'badrex/llm-emoji-dataset'):
-        dataset = load_dataset(name, split='train')
+    if path in (None, 'badrex/llm-emoji-dataset'):
+        dataset = load_dataset(path, split='train')
         
         for row in dataset:
             yield (
@@ -311,14 +313,36 @@ def batch_iterator(name=None):
         
         del dataset
         gc.collect()
-
-
-def tokenize_fn(dataset_name, tokenizer=None):
-    for text in batch_iterator(dataset_name):
+"""
+
+def batch_iterator(path: str,
+                   name: Optional[str]=None,
+                   data_dir: Optional[str]=None,
+                   data_files: Optional[str]=None,
+                   revision: Optional[str]=None,
+                   split: str='train',
+                   format: Optional[str]=None):
+    assert format is not None
+    
+    dataset = load_dataset(path=path,
+                           name=name,
+                           data_dir=data_dir,
+                           data_files=data_files,
+                           revision=revision,
+                           split=split,
+                           trust_remote_code=True)
+
+    for row in dataset:
+        text = format.format(**row)
+        yield text
+
+
+def tokenize_fn(datasets_config, tokenizer=None):
+    for text in batch_iterator(**datasets_config):
         text_ids = tokenizer.encode(text, bos=False, eos=True)
         yield text_ids
 
-
+"""
 datasets_names = [
     'saillab/taco-datasets',
     # 'xu-song/cc100-samples',
@@ -340,12 +364,80 @@ datasets_names = [
     # 'SkunkworksAI/reasoning-0.01',
     'badrex/llm-emoji-dataset',
 ]
+"""
+
+datasets_configs = [
+    {'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'},
+    {'path': 'gbharti/wealth-alpaca_lora', 'format': '{instruction} {input} {output}'},
+    *[
+        {'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': 'train[:10%]', 'format': '{instruction} {input} {output}'}
+        for data_dir in [
+            'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
+            'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
+        ]
+    ],
+    *[
+        {'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train[:10%]', 'format': '{text}'}
+        for name in [
+            'am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br',
+            'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es',
+            'et', 'eu', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl',
+            'gn', 'gu', 'ha', 'he', 'hi', 'hi_rom', 'hr', 'ht', 'hu',
+            'hy', 'id', 'ig', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km',
+            'kn', 'ko', 'ku', 'ky', 'la', 'lg', 'li', 'ln', 'lo', 'lt',
+            'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'my_zaw',
+            'ne', 'nl', 'no', 'ns', 'om', 'or', 'pa', 'pl', 'ps', 'pt',
+            'qu', 'rm', 'ro', 'ru', 'sa', 'si', 'sc', 'sd', 'sk', 'sl',
+            'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'ta_rom',
+            'te', 'te_rom', 'th', 'tl', 'tn', 'tr', 'ug', 'uk', 'ur',
+            'ur_rom', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo',
+            'zh-Hans', 'zh-Hant', 'zu',
+        ]
+    ],
+    {'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': 'train[:5%]', 'format': '{text}'},
+    {'path': 'MuskumPillerum/General-Knowledge', 'format': '{Question} {Answer}'},
+    {'path': 'yirenc/general_knowledge_boolean', 'split': 'train+validation', 'format': '{question}? {answer}. {passage}'},
+    {'path': 'nampdn-ai/tiny-textbooks', 'split': 'train+test', 'format': '{textbook}'},
+    {'path': 'nampdn-ai/tiny-codes', 'split': 'train[:5%]', 'format': '{prompt} {response}'},
+    *[
+        {'path': 'bigcode/the-stack-smol-xs', 'name': name, 'format': '{content}'}
+        for name in [
+            'ada', 'agda', 'alloy', 'antlr', 'applescript', 'assembly',
+            'augeas', 'awk', 'batchfile', 'bison', 'bluespec', 'c',
+            'c++', 'c-sharp', 'clojure', 'cmake', 'coffeescript', 'common-lisp',
+            'css', 'cuda', 'dart', 'dockerfile', 'elixir',
+            'elm', 'emacs-lisp','erlang', 'f-sharp', 'fortran', 'glsl', 'go',
+            'groovy', 'haskell','html', 'idris', 'isabelle', 'java', 
+            'java-server-pages', 'javascript', 'julia', 'kotlin', 'lean',
+            'literate-agda', 'literate-coffeescript', 'literate-haskell',
+            'lua', 'makefile', 'maple', 'markdown', 'mathematica', 'matlab',
+            'ocaml', 'pascal', 'perl', 'php', 'powershell', 'prolog',
+            'protocol-buffer', 'python', 'r', 'racket', 'restructuredtext',
+            'rmarkdown', 'ruby', 'rust', 'sas', 'scala', 'scheme', 
+            'shell', 'smalltalk', 'solidity', 'sparql', 'sql', 'stan',
+            'standard-ml', 'stata', 'systemverilog', 'tcl', 'tcsh', 'tex', 
+            'thrift', 'typescript', 'verilog', 'vhdl', 'visual-basic', 'xslt',
+            'yacc', 'zig',
+        ]
+    ],
+    {'path': 'm-a-p/CodeFeedback-Filtered-Instruction', 'split': 'train', 'format': '{query} {answer}'},
+    {'path': 'jtatman/python-code-dataset-500k', 'format': '{instruction} {output}'},
+    {'path': 'iamtarun/python_code_instructions_18k_alpaca', 'format': '{instruction} {input} {output}'},
+    {'path': 'HuggingFaceH4/CodeAlpaca_20K', 'split': 'train+test', 'format': '{prompt} {completion}'},
+    {'path': 'gair-prox/open-web-math-pro', 'split': 'train[:5%]', 'format': '{text}'},
+    {'path': 'rvv-karma/Math-QA', 'split': 'train+val+test', 'format': '{question} {answer}'},
+    {'path': 'ajibawa-2023/Maths-College', 'split': 'train[:10%]', 'format': '{instruction} {output}'},
+    {'path': 'microsoft/orca-math-word-problems-200k', 'format': '{question} {answer}'},
+    {'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'train+test', 'format': '{instruction} = {output}'},
+    {'path': 'SkunkworksAI/reasoning-0.01', 'format': '{instruction} {reasoning} {output}'},
+    {'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'},
+]
 
 outputs = optimize(
     fn=partial(tokenize_fn, tokenizer=Tokenizer('..')),
-    inputs=datasets_names,
+    inputs=datasets_configs,
     output_dir='../pretrain-data/',
     # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
     chunk_size=(2049 * 8012),
-    num_workers=16,
+    num_workers=32,
 )

scripts/pretrain-model.yaml

@@ -1,13 +1,13 @@
 # The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
 # ``model_config``. (type: Optional[str], default: null)
-model_name: "tiny-llama-1.1b"
+model_name: "Llama-3.2-1B"
 
 # A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
 # ``model_config``. (type: Optional[Config], default: null)
 model_config:
   padded_vocab_size: 38400
   vocab_size: 38400
-  block_size: 131072
+  block_size: 8192
   n_layer: 5
   n_head: 32
   head_size: null
@@ -20,7 +20,12 @@ model_config:
   norm_eps: 1e-05
   mlp_class_name: "LLaMAMLP"
   intermediate_size: 3584
-  rope_base: 1000000
+  rope_base: 500000
+  rope_adjustments:
+    factor: 32.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_seq_len: 8192
 
 # Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
 # /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
@@ -71,7 +76,7 @@ train:
 
   # Total number of tokens to train on (type: Optional[int], default: 3000000000000)
   # max_tokens: 3000000000000
-  max_tokens: ??? # ? * 2049 * 3
+  max_tokens: ??? # ? * 2049 * 5
 
   # Limits the number of optimizer steps to run. (type: Optional[int], default: null)
   max_steps: