pretrain

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 4,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.45.0.dev0",
+  "use_cache": true,
+  "vocab_size": 131072
+}

scripts/backup/base_instruct_datasets.py

@@ -0,0 +1,124 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+core_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    *[
+        {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    *[
+        {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # multilingual instructs
+    #
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    *[
+        {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+]

scripts/backup/base_reason_datasets.py

@@ -0,0 +1,79 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+R1_SYSTEM_PROMPT = '''\
+You are an AI assistant.
+
+Your primary directive is to provide well-reasoned, structured, and extensively detailed responses.
+
+Formatting Requirements:
+- Always structure your replies using: <think>{reasoning}</think>{answer}
+- The <think></think> block should contain at least six reasoning steps when applicable.
+- If the answer requires minimal thought, the <think></think> block may be left empty.
+- The user does not see the <think></think> section. Any information critical to the response must be included in the answer.
+- If you notice that you have engaged in circular reasoning or repetition, immediately terminate {reasoning} with a </think> and proceed to the {answer}
+
+Response Guidelines:
+- Detailed and Structured: Use rich Markdown formatting for clarity and readability.
+- Scientific and Logical Approach: Your explanations should reflect the depth and precision of the greatest scientific minds.
+- Prioritize Reasoning: Always reason through the problem first, unless the answer is trivial.
+- Concise yet Complete: Ensure responses are informative, yet to the point without unnecessary elaboration.
+- Maintain a professional, intelligent, and analytical tone in all interactions.'''
+
+core_reason_datasets = [
+    #
+    # math reason
+    #
+    # 8.43 GB, 450,258
+    *[
+        {'kind': 'instruct', 'path': 'open-r1/OpenR1-Math-220k', 'data_dir': 'data', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general reason
+    #
+    # 3.55 GB, 227,914
+    *[
+        {'kind': 'instruct', 'path': 'open-thoughts/OpenThoughts-114k', 'data_dir': 'data', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['system']}
+        ] + [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in r['conversations']
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 3.98 GB, 814,334
+    # 300k
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-deepseek.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+            *r['messages'],
+            {'role': 'assistant', 'content': '<think>\n' + (r.get('reasoning') or '') + '\n</think>\n' + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 300k
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-flash.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+            *r['messages'],
+            {'role': 'assistant', 'content': '<think>\n' + (r.get('reasoning') or '') + '\n</think>\n' + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 21.1 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K-1.1', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r.get('question') or ''},
+        {'role': 'assistant', 'content': '<think>\n' + (r.get('deepseek_thinking_trajectory') or '') + '\n</think>\n' + (r.get('solution') or '')},
+    ]}
+]

scripts/backup/cpt_base_datasets.py

@@ -0,0 +1,58 @@
+cpt_base_datasets = [
+    #
+    # stem
+    #
+    # 1.44 GB, 63,357
+    {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': 'train', 'format': lambda n: n['abstract']},
+    {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': 'train', 'format': lambda n: n['markdown']},
+
+    #
+    # code
+    #
+    # 1.62 GB, 1,632,309
+    # Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go; SQL, Cypher
+    {'kind': 'base', 'path': 'nampdn-ai/tiny-codes', 'split': 'train', 'format': '{prompt} {response}'},
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # multilingual
+    #
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', 'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train', 'format': lambda n: n['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',
+        ]
+    ],
+
+    #
+    # general knowledge
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'train', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/backup/cpt_base_model.py

@@ -0,0 +1,101 @@
+from unsloth import FastLanguageModel
+import torch
+from transformers import AutoTokenizer
+
+max_seq_length = 4096
+dtype = torch.bfloat16
+load_in_4bit = True
+model_name = '../out/pretrain-base'
+output_dir = '../out/cpt-base'
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=model_name,
+    max_seq_length=max_seq_length,
+    dtype=dtype,
+    load_in_4bit=load_in_4bit,
+)
+
+print('Ignore loaded tokenizer by FastLanguageModel.from_pretrained and using AutoTokenizer.from_pretrained')
+tokenizer = AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True)
+
+print(f'{model=}')
+print(f'{tokenizer=}')
+
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=64, # 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+        "embed_tokens", "lm_head",
+    ], # Add for continual pretraining
+    lora_alpha=16,
+    lora_dropout=0, # Supports any, but = 0 is optimized
+    bias='none',    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing='unsloth', # True or "unsloth" for very long context
+    random_state=23,
+    use_rslora=True,  # We support rank stabilized LoRA
+    loftq_config=None, # And LoftQ
+)
+
+print(f'{model=}')
+
+from datasets import concatenate_datasets
+from cpt_base_datasets import cpt_base_datasets
+from cpt_instruct_datasets import cpt_instruct_datasets
+from unsloth_utils import load_text_dataset, load_chat_dataset
+
+core_datasets = []
+
+for dataset_config in cpt_base_datasets:
+    dataset = load_text_dataset(tokenizer, **dataset_config)
+    print(f'{dataset=}')
+    core_datasets.append(dataset)
+
+# for dataset_config in cpt_instruct_datasets:
+#     dataset = load_chat_dataset(tokenizer, **dataset_config)
+#     print(f'{dataset=}')
+#     core_datasets.append(dataset)
+
+final_dataset = concatenate_datasets(core_datasets)
+print(f'{final_dataset=}')
+
+
+from trl import SFTTrainer
+from transformers import TrainingArguments
+from unsloth import is_bfloat16_supported
+from unsloth import UnslothTrainer, UnslothTrainingArguments
+
+
+trainer = UnslothTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=final_dataset,
+    dataset_text_field='text',
+    max_seq_length=max_seq_length,
+    dataset_num_proc=32,
+
+    args = UnslothTrainingArguments(
+        per_device_train_batch_size=8,
+        gradient_accumulation_steps=8,
+
+        warmup_ratio=0.1,
+        num_train_epochs=1,
+
+        learning_rate=5e-5,
+        embedding_learning_rate=5e-6,
+
+        fp16=not is_bfloat16_supported(),
+        bf16=is_bfloat16_supported(),
+        logging_steps=1,
+        optim='adamw_8bit',
+        weight_decay=0.01,
+        lr_scheduler_type='cosine',
+        seed=23,
+        output_dir=output_dir,
+        report_to='wandb',
+    ),
+)
+
+trainer_stats = trainer.train()

scripts/backup/cpt_instruct_datasets.py

@@ -0,0 +1,119 @@
+import json
+
+
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+cpt_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': 'train', 'field': 'messages'},
+
+    #
+    # multilingual instructs
+    #
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': 'train', 'transform': lambda r: [
+        {'role': 'system', 'content': r['instruction']},
+        {'role': 'user', 'content': r['input']},
+        {'role': 'assistant', 'content': r['output']},
+    ]},
+]

scripts/backup/merge-core-into-base.yaml

@@ -0,0 +1,100 @@
+slices:
+  # 1
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 2
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 3
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 4
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 5
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 6
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 7
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 8
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 9
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 10
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 11
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 12
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 13
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 14
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 15
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 16
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 17
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 18
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 19
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 20
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 21
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 22
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 23
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 24
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+
+merge_method: passthrough
+dtype: bfloat16

scripts/backup/prepare_pretrain_base_datasets.py

@@ -0,0 +1,59 @@
+from functools import partial
+
+from litgpt.tokenizer import Tokenizer
+from litdata import optimize, TokensLoader, StreamingDataset
+from transformers import AutoTokenizer
+
+from utils import tokenize_fn
+from pretrain_base_datasets import pretrain_base_datasets
+from pretrain_instruct_datasets import pretrain_instruct_datasets
+from pretrain_reflection_datasets import pretrain_reflection_datasets
+from pretrain_reasoning_datasets import pretrain_reasoning_datasets
+
+
+#
+# optimize datasets
+#
+for i, (block_size, subchunk_size) in enumerate([(4097, 4000)]):
+    chunk_size = block_size * subchunk_size
+    output_dir = f'../pretrain-base-data-{i}-{block_size}-{subchunk_size}'
+
+    outputs = optimize(
+        fn=partial(
+            tokenize_fn,
+            hf_tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True),
+            tokenizer=Tokenizer('..'),
+        ),
+        inputs=(
+            pretrain_base_datasets +
+            pretrain_instruct_datasets +
+            pretrain_reflection_datasets +
+            pretrain_reasoning_datasets
+        ),
+        output_dir=output_dir,
+        chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+        num_workers=32,
+        reorder_files=False,
+        ## 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.
+        # item_loader=TokensLoader(block_size=block_size),
+    )
+
+#
+# total number of chunks in datasets
+#
+for i, (block_size, subchunk_size) in enumerate([(4097, 4000)]):
+    chunk_size = block_size * subchunk_size
+    input_dir = f'../pretrain-base-data-{i}-{block_size}-{subchunk_size}'
+
+    dataset = StreamingDataset(
+        input_dir=input_dir,
+        item_loader=TokensLoader(block_size=block_size),
+    )
+
+    print(f'{i=}, {block_size=}, {chunk_size=}, {len(dataset)=}, {len(dataset) * block_size=}')
+
+    # total_tokens = sum(len(data) for data in dataset)
+    # print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
+    total_tokens = len(dataset) * block_size
+    print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')

scripts/backup/pretrain_base_datasets.py

@@ -0,0 +1,107 @@
+pretrain_base_datasets = [
+    #
+    # multilingual
+    #
+    # 3.17 GB, 2,226,907
+    *[
+        {'kind': 'base', 'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    # 1.64 GB, 1,001,000
+    *[
+        {'kind': 'base', 'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    # 3.8 GB, 19,454,996
+    *[
+        {'kind': 'base', 'path': 'sentence-transformers/parallel-sentences-wikimatrix', 'data_dir': 'all', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['non_english']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general knowledge
+    #
+    # 65.1 MB, 7,819
+    {'kind': 'base', 'path': 'Sketched33/Cities_Wikipedia_Information', 'format': lambda n: n['wikipedia_content']},
+    # 135 MB, 1,795
+    {'kind': 'base', 'path': 'open-phi/textbooks', 'format': lambda n: n['markdown']},
+    # 631 MB, 111,048
+    {'kind': 'base', 'path': 'open-phi/programming_books_llama', 'format': lambda n: n['markdown']},
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # math
+    #
+    # 12.6 GB, 14M rows
+    *[
+        {'kind': 'base', 'path': 'nvidia/OpenMathInstruct-2', 'split': f'train[{i}%:{i + 10}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # stem
+    #
+    # 1.44 GB, 63,357
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['abstract']}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['markdown']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # code
+    #
+    # 7.81 GB, ~2,804,025
+    *[
+        {'kind': 'base', 'path': 'rombodawg/code_bagel_hermes-2.5', 'split': f'train[{i}%:{i + 10}%]', 'format': '{input} {output}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # multilingual
+    #
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', 'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train', 'format': lambda n: n['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',
+        ]
+    ],
+
+    #
+    # general knowledge
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    *[
+        {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/backup/pretrain_instruct_datasets.py

@@ -0,0 +1,198 @@
+import json
+
+
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 138 MB, 205,568
+    {'kind': 'instruct', 'path': 'CohereForAI/aya_dataset', 'transform': lambda r: [
+        {'role': 'user', 'content': r['inputs']},
+        {'role': 'assistant', 'content': r['targets']},
+    ]},
+
+    # ~3 GB, 4,976,850
+    *[
+        {'kind': 'instruct', 'path': 'saillab/taco-datasets', 'data_dir': name, 'split': 'train', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for name in [
+            'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
+            'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
+        ]
+    ],
+
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    *[
+        {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 4.58 GB, 1,752,473
+    # arcee-ai/The-Tome
+    # - arcee-ai/infini-instruct-top-500k (BAAI/Infinity-Instruct)
+    # - TIGER-Lab/WebInstructSub (top-500k) - IGNORE
+    # - jondurbin/airoboros-3.2
+    # - gardner/glaive-function-calling-v2-sharegpt
+    # - arcee-ai/reasoning-sharegpt (SkunkworksAI/reasoning-0.01)
+    # - arcee-ai/self-instruct-sharegpt (bigcode/self-oss-instruct-sc2-exec-filter-50k)
+    # - cognitivecomputations/ultrainteract_trajectories_sharegpt
+    # - cognitivecomputations/SystemChat-2.0
+    # - arcee-ai/qwen2-72b-magpie-en
+    *[
+        {'kind': 'instruct', 'path': 'arcee-ai/The-Tome', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    *[
+        {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    *[
+        {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # tool/function calling
+    #
+    # 65.7 MB, 11,578
+    {'kind': 'instruct', 'path': 'NousResearch/hermes-function-calling-v1', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    #
+    # agent
+    #
+    # 1.51 GB, 485,874
+    *[
+        {'kind': 'instruct', 'path': 'arcee-ai/agent-data', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 2.21 GB, 1,046,410
+    *[
+        {'kind': 'instruct', 'path': 'microsoft/orca-agentinstruct-1M-v1', 'split': split, 'field': 'messages', 'transform': lambda msgs: json.loads(msgs)}
+        for split in [
+            'creative_content', 'text_modification', 'struct2text_flow', 'rc', 'rag',
+            'text_extraction', 'mcq', 'follow_up', 'analytical_reasoning', 'fermi', 'fs_cot_flow',
+            'code_', 'brain_teaser', 'text_classification', 'open_domain_qa',
+        ]
+    ],
+
+    #
+    # general instructs
+    #
+    # 1.52 GB, 214k (3.98 GB, 814,334)
+    {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-nonreasoning.jsonl', 'split': 'train', 'field': 'messages'},
+    # 4.15 GB, 2,197,730
+    {'kind': 'instruct', 'path': 'HuggingFaceTB/smoltalk', 'name': 'all', 'field': 'messages'},
+]

scripts/backup/pretrain_reasoning_datasets.py

@@ -0,0 +1,98 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_reasoning_datasets = [
+    #
+    # basic reasoning
+    #
+    # 10.8 MB, 15,770
+    {'kind': 'instruct', 'path': 'AtlasUnified/Atlas-Reasoning', 'data_files': 'reasoning.csv', 'transform': lambda r: [
+        {'role': 'user', 'content': r['Prompt']},
+        {'role': 'assistant', 'content': r['Step-by-step reasoning'] + '\n' + r['Solution']},
+    ]},
+    # 1.23 GB, 859,594
+    *[
+        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-CoT', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+    # 148 MB, 72,540
+    *[
+        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-TIR', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # math reasoning
+    #
+    # 1.79 MB, 3,963
+    {'kind': 'instruct', 'path': 'AlgorithmicResearchGroup/math_reasoning_autoformalization_track', 'transform': lambda r: [
+        {'role': 'user', 'content': r['informal_statement']},
+        {'role': 'assistant', 'content': r['informal_proof'] + '\n' + r['formal_proof']},
+    ]},
+    # 307 MB, 19,944
+    {'kind': 'instruct', 'path': 'KingNish/reasoning-base-20k', 'transform': lambda r: [
+        {'role': 'user', 'content': r['user']},
+        {'role': 'assistant', 'content': r['reasoning'] + '\n' + r['assistant']},
+    ]},
+    # 9.45 MB, 10,000
+    {'kind': 'instruct', 'path': 'Aarushhh/math-reasoning-10k', 'transform': lambda r: [
+        {'role': 'user', 'content': r['problem']},
+        {'role': 'assistant', 'content': r['plan'] + '\n' + r['solution']},
+    ]},
+
+    #
+    # cot reasoning
+    #
+    # 11.7 GB, 1,850,809
+    *[
+        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v0', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'user', 'content': r['problem']},
+            {'role': 'assistant', 'content': r['reannotated_assistant_content']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v1', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: r['reannotated_messages']}
+        for i in range(0, 100, 10)
+    ],
+    # 3.85 GB, 300k (3.98 GB, 814,334)
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-deepseek.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            *r['messages'],
+            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
+            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 3.49 GB, 300k (3.98 GB, 814,334)
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-flash.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            *r['messages'],
+            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
+            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.08 GB, 113,957
+    {'kind': 'instruct', 'path': 'open-thoughts/OpenThoughts-114k', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+    # 384 MB, 77,685
+    {'kind': 'instruct', 'path': 'O1-OPEN/OpenO1-SFT', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r['instruction']},
+        {'role': 'assistant', 'content': r['output']},
+    ]},
+    # 6.88 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r['question']},
+        {'role': 'assistant', 'content': '<think>\n' + '\n'.join(r['thinking_trajectories']) + '\n</think>\n' + r['solution']},
+    ]},
+]

scripts/backup/pretrain_reflection_datasets.py

@@ -0,0 +1,39 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_reflection_datasets = [
+    #
+    # reflection
+    #
+    # 4.17 MB, 1,000
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-gpt-4o-judge', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 12.4 MB, 3,000
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-openai-o-mini-judge', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 70.8 MB, 36,549
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-final-dedup', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 30.6 MB, 25,391
+    {'kind': 'instruct', 'path': 'flozi00/reflection-qwen2.5-72b-260924', 'transform': lambda r: [
+        r['system'][0],
+        {'role': 'user', 'content': r['input']},
+        {'role': 'assistant', 'content': r['reflection'] + '\n' + r['output']},
+    ]},
+]

scripts/backup/unsloth_utils.py

@@ -0,0 +1,125 @@
+from typing import Optional, Iterator, Callable, Any
+
+import torch
+from datasets import load_dataset, concatenate_datasets
+from transformers import AutoTokenizer
+
+
+def load_text_dataset(tokenizer: AutoTokenizer,
+                      kind: str,
+                      path: str,
+                      name: Optional[str]=None,
+                      data_dir: Optional[str]=None,
+                      data_files: Optional[str]=None,
+                      keep_in_memory: bool=False,
+                      revision: Optional[str]=None,
+                      split: str='train',
+                      num_proc: Optional[int]=None,
+                      format: Optional[Callable|str]=None) -> Any:
+    assert isinstance(format, str) or callable(format), f'{path=} {format=}'
+    assert kind == 'base'
+
+    dataset = load_dataset(path=path,
+                           name=name,
+                           data_dir=data_dir,
+                           data_files=data_files,
+                           keep_in_memory=keep_in_memory,
+                           revision=revision,
+                           split=split,
+                           trust_remote_code=True,
+                           num_proc=num_proc)
+
+    EOS_TOKEN = tokenizer.eos_token
+
+    def format_dataset(batch):
+        nonlocal EOS_TOKEN
+        nonlocal format
+        texts: list = []
+        rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
+
+        if callable(format):
+            for row in rows:
+                # print(f'{row=}')
+                text = format(row)
+
+                if not text:
+                    text = '[NONE]'
+
+                text += EOS_TOKEN
+                texts.append(text)
+        else:
+            for row in rows:
+                # print(f'{row=}')
+                text = format.format(**row)
+
+                if not text:
+                    text = '[NONE]'
+
+                text += EOS_TOKEN
+                texts.append(text)
+
+        return {'text': texts}
+
+    dataset = dataset.map(format_dataset, batched=True)
+    return dataset
+
+
+def load_chat_dataset(tokenizer: AutoTokenizer,
+                      kind: str,
+                      path: str,
+                      name: Optional[str]=None,
+                      data_dir: Optional[str]=None,
+                      data_files: Optional[str]=None,
+                      keep_in_memory: bool=False,
+                      revision: Optional[str]=None,
+                      split: str='train',
+                      num_proc: Optional[int]=None,
+                      field: Optional[str]=None,
+                      transform: Optional[Callable]=None) -> Any:
+    assert kind == 'instruct'
+
+    dataset = load_dataset(path=path,
+                           name=name,
+                           data_dir=data_dir,
+                           data_files=data_files,
+                           keep_in_memory=keep_in_memory,
+                           revision=revision,
+                           split=split,
+                           trust_remote_code=True,
+                           num_proc=num_proc)
+
+    EOS_TOKEN = tokenizer.eos_token
+
+    def format_dataset(batch):
+        nonlocal EOS_TOKEN
+        nonlocal tokenizer
+        nonlocal field
+        nonlocal transform
+        texts: list = []
+        rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
+
+        if callable(transform):
+            for row in rows:
+                if field:
+                    messages = transform(row[field])
+                else:
+                    messages = transform(row)
+
+                text = tokenizer.apply_chat_template(messages, tokenize=False)
+                text += EOS_TOKEN
+                texts.append(text)
+        else:
+            for row in rows:
+                if field:
+                    messages = row[field]
+                else:
+                    raise ValueError(field)
+
+                text = tokenizer.apply_chat_template(messages, tokenize=False)
+                text += EOS_TOKEN
+                texts.append(text)
+
+        return {'text': texts}
+
+    dataset = dataset.map(format_dataset, batched=True)
+    return dataset

scripts/core_base_datasets.py

@@ -0,0 +1,89 @@
+core_base_datasets = [
+    #
+    # multilingual
+    #
+    # 3.17 GB, 2,226,907
+    *[
+        {'kind': 'base', 'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    # 1.64 GB, 1,001,000
+    *[
+        {'kind': 'base', 'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', 'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train', 'format': lambda n: n['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',
+        ]
+    ],
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # stem
+    #
+    # 12.2 MB, 500,000
+    {'kind': 'base', 'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'train', 'format': '{instruction} = {output}'},
+    {'kind': 'base', 'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'test', 'format': '{instruction} = {output}'},
+    # 125 MB, 1,000,000
+    {'kind': 'base', 'path': 'Gusarich/math-expressions-1m', 'revision': 'refs/convert/parquet', 'split': 'train', 'format': '{expression} = {result}'},
+
+    # 1.44 GB, 63,357
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['abstract']}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['markdown']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # code
+    #
+    # 36.8 MB, 79,013
+    # Rosetta Code currently has 1,203 tasks, 389 draft tasks, and is aware of 883 languages
+    {'kind': 'base', 'path': 'christopher/rosetta-code', 'format': lambda n: n['code']},
+    # 1.62 GB, 1,632,309
+    # Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go; SQL, Cypher
+    *[
+        {'kind': 'base', 'path': 'nampdn-ai/tiny-codes', 'split': f'train[{i}%:{i + 10}%]', 'format': '{prompt} {response}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general knowledge
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    *[
+        {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/core_instruct_datasets.py

@@ -0,0 +1,43 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+R1_SYSTEM_PROMPT = '''\
+You are an AI assistant.
+
+Your primary directive is to provide well-reasoned, structured, and extensively detailed responses.
+
+Formatting Requirements:
+- Always structure your replies using: <think>{reasoning}</think>{answer}
+- The <think></think> block should contain at least six reasoning steps when applicable.
+- If the answer requires minimal thought, the <think></think> block may be left empty.
+- The user does not see the <think></think> section. Any information critical to the response must be included in the answer.
+- If you notice that you have engaged in circular reasoning or repetition, immediately terminate {reasoning} with a </think> and proceed to the {answer}
+
+Response Guidelines:
+- Detailed and Structured: Use rich Markdown formatting for clarity and readability.
+- Scientific and Logical Approach: Your explanations should reflect the depth and precision of the greatest scientific minds.
+- Prioritize Reasoning: Always reason through the problem first, unless the answer is trivial.
+- Concise yet Complete: Ensure responses are informative, yet to the point without unnecessary elaboration.
+- Maintain a professional, intelligent, and analytical tone in all interactions.'''
+
+core_instruct_datasets = [
+    # 65.7 MB, 11,578
+    # 1.89k
+    {'kind': 'instruct', 'path': 'NousResearch/hermes-function-calling-v1', 'data_files': 'func-calling-singleturn.json', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    # 21.1 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K-1.1', 'split': 'train', 'transform': lambda r: [
+        {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+        {'role': 'user', 'content': r.get('question') or ''},
+        {'role': 'assistant', 'content': '<think>\n' + (r.get('deepseek_thinking_trajectory') or '') + '\n</think>\n' + (r.get('solution') or '')},
+    ]}
+]

scripts/prepare_core_datasets.py

@@ -0,0 +1,58 @@
+from functools import partial
+
+from transformers import AutoTokenizer
+from litgpt.tokenizer import Tokenizer
+from litdata import optimize, TokensLoader, StreamingDataset
+
+from utils import tokenize_fn
+from core_base_datasets import core_base_datasets
+from core_instruct_datasets import core_instruct_datasets
+
+
+tokenizer_path = '../tokenizer'
+
+seqs = [
+    (0, 1073741824, 4097, 4000),
+]
+
+#
+# optimize datasets
+#
+for i, (min_len, max_len, block_size, subchunk_size) in enumerate(seqs):
+    chunk_size = block_size * subchunk_size
+    output_dir = f'../core-data-{i}-{min_len}-{max_len}-{block_size}-{subchunk_size}'
+
+    outputs = optimize(
+        fn=partial(
+            tokenize_fn,
+            min_len=min_len,
+            max_len=max_len,
+            hf_tokenizer=AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True, use_fast=True),
+            tokenizer=Tokenizer(tokenizer_path),
+        ),
+        inputs=core_base_datasets + core_instruct_datasets,
+        output_dir=output_dir,
+        chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+        num_workers=32,
+        reorder_files=False,
+        ## 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.
+        # item_loader=TokensLoader(block_size=block_size),
+    )
+
+#
+# total number of chunks in datasets
+#
+for i, (min_len, max_len, block_size, subchunk_size) in enumerate(seqs):
+    chunk_size = block_size * subchunk_size
+    input_dir = f'../core-data-{i}-{min_len}-{max_len}-{block_size}-{subchunk_size}'
+
+    dataset = StreamingDataset(
+        input_dir=input_dir,
+        item_loader=TokensLoader(block_size=block_size),
+    )
+
+    print(f'{i=}, {min_len=}, {max_len=}, {block_size=}, {chunk_size=}, {len(dataset)=}, {len(dataset) * block_size=}')
+    total_tokens = len(dataset) * block_size
+    print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
+    print()

scripts/pretrain_core_model.yaml

@@ -0,0 +1,167 @@
+# 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: 'tangled-alpha-0.8-core'
+
+# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
+# ``model_config``. (type: Optional[Config], default: null)
+model_config:
+  name: 'tangled-alpha-0.8-core'
+  block_size: 131072
+  vocab_size: 131072
+  padded_vocab_size: 131072
+  n_layer: 32
+  n_head: 4
+  n_embd: 512
+  n_query_groups: 4
+  rotary_percentage: 1.0
+  parallel_residual: False
+  bias: False
+  norm_class_name: "RMSNorm"
+  mlp_class_name: "LLaMAMLP"
+  intermediate_size: 2048 # n_embd * 4
+  norm_eps: 1e-5
+  rope_base: 10000
+  head_size: 128 # n_embd / n_head
+
+# 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)
+out_dir: "../out/pretrain-core/"
+
+# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
+# precision: bf16-mixed
+precision: bf16-true
+
+# Optional path to a checkpoint directory to initialize the model from.
+# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
+initial_checkpoint_dir:
+
+# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
+# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
+# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
+# (type: Union[bool, Literal["auto"], Path], default: False)
+resume: "auto"
+
+# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
+data:
+  class_path: LitData
+
+  init_args:
+    data_path: "../core-data-0-0-1073741824-4097-4000/"
+    num_workers: 32
+
+# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
+train:
+  # Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
+  save_interval: 50
+
+  # Number of iterations between logging calls (type: int, default: 1)
+  log_interval: 1
+
+  # Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
+  global_batch_size: 512
+
+  # Number of samples per data-parallel rank (type: int, default: 4)
+  micro_batch_size: 2
+
+  # Number of iterations with learning rate warmup active (type: int, default: 2000)
+  lr_warmup_steps: 500
+
+  # Number of epochs to train on (type: Optional[int], default: null)
+  epochs:
+
+  # Total number of tokens to train on (type: Optional[int], default: 3000000000000)
+  max_tokens: 5274490091
+
+  # Limits the number of optimizer steps to run. (type: Optional[int], default: null)
+  max_steps:
+
+  # Limits the length of samples. Off by default (type: Optional[int], default: null)
+  max_seq_length: 4096
+
+  # Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
+  tie_embeddings: true
+
+  #   (type: Optional[float], default: 1.0)
+  max_norm: 1.0
+
+  #   (type: float, default: 4e-05)
+  min_lr: 3e-5
+
+# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
+eval:
+  # Number of optimizer steps between evaluation calls (type: int, default: 1000)
+  interval: 50
+
+  # Number of tokens to generate (type: Optional[int], default: null)
+  max_new_tokens:
+
+  # Number of iterations (type: int, default: 100)
+  max_iters: 100
+
+  # Whether to evaluate on the validation set at the beginning of the training
+  initial_validation: false
+
+  # Whether to evaluate on the validation set at the end the training
+  final_validation: true
+
+# Optimizer-related arguments
+
+# optimizer:
+#   class_path: torch.optim.AdamW
+#   # class_path: torchao.prototype.low_bit_optim.AdamW8bit
+#   # class_path: torchao.prototype.low_bit_optim.AdamW4bit
+#   # class_path: bitsandbytes.optim.AdamW8bit
+#   # class_path: bitsandbytes.optim.PagedAdamW8bit
+#   init_args:
+#     # (type: float, default: 0.001)
+#     lr: 3e-4
+#     # (type: float, default: 0.01)
+#     weight_decay: 0.01
+#     # (type: tuple, default: (0.9,0.999))
+#     betas:
+#       - 0.9
+#       - 0.999
+
+# optimizer:
+#   class_path: sophia_opt.SophiaG
+#   init_args:
+#     lr: 3e-4
+#     betas:
+#       - 0.9
+#       - 0.95
+#     rho: 0.05
+#     weight_decay: 0.1
+
+# optimizer:
+#   class_path: grokadamw.GrokAdamW
+#   init_args:
+#     # (type: float, default: 0.001)
+#     lr: 3e-4
+#     # (type: float, default: 0.01)
+#     weight_decay: 1e-2
+#     # (type: tuple, default: (0.9,0.999))
+#     betas:
+#       - 0.9
+#       - 0.999
+
+optimizer:
+  class_path: muon.Muon
+  init_args:
+    lr: 3e-4
+    weight_decay: 0.01
+
+# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
+devices: auto
+
+# How many nodes to use. (type: int, default: 1)
+num_nodes: 1
+
+# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
+# module require this. (type: Optional[Path], default: null)
+tokenizer_dir: "../tokenizer"
+
+# The name of the logger to send metrics to. (type: Literal['wandb', 'tensorboard', 'csv'], default: tensorboard)
+logger_name: "wandb"
+
+# The random seed to use for reproducibility. (type: int, default: 42)
+seed: 23

scripts/requirements.in

@@ -0,0 +1,20 @@
+# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+torch>=2.5.0,<2.6.0
+numpy<2.0
+
+tqdm
+Pillow
+datasets
+jinja2
+transformers
+wandb
+litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
+mergekit @ git+https://github.com/arcee-ai/mergekit.git
+# torchao
+# bitsandbytes
+# grokadamw
+# sophia-opt
+# dolphinflow
+muon @ git+https://github.com/cognitivecomputations/Muon
+# unsloth
+lm_eval[ifeval,math]

scripts/train_tokenizer.py

@@ -0,0 +1,115 @@
+import os
+import shutil
+
+from transformers import PreTrainedTokenizerFast
+from tokenizers import Tokenizer, normalizers, pre_tokenizers, processors, decoders
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+
+from utils import batch_dataset_iterator
+from core_base_datasets import core_base_datasets
+from core_instruct_datasets import core_instruct_datasets
+
+
+tokenizer_path = '../tokenizer'
+
+if os.path.exists(tokenizer_path):
+    shutil.rmtree(tokenizer_path)
+
+os.makedirs(tokenizer_path, exist_ok=True)
+
+#
+# special_tokens
+#
+bos_token = '<|endoftext|>'
+eos_token = '<|im_end|>'
+pad_token = '<|pad|>'
+unk_token = '<|unk|>'
+
+special_tokens = [
+    bos_token,
+    eos_token,
+    pad_token,
+    unk_token,
+    '<|im_start|>',
+    '<|im_sep|>',
+    'system',
+    'user',
+    'assistant',
+    '<tools>',
+    '</tools>',
+    '<tool>',
+    '</tool>',
+    '<tool_call>',
+    '</tool_call>',
+    '<tool_response>',
+    '</tool_response>',
+    '<think>',
+    '</think>',
+]
+
+for i in range(64 - len(special_tokens)):
+    special_tokens.append(f'<|reserved_{i}|>')
+
+#
+# BPE Tokenizer
+#
+bpe = BPE(unk_token=None, byte_fallback=True)
+tokenizer = Tokenizer(bpe)
+
+# normalizer
+tokenizer.normalizer = None
+
+# pre-tokenizer
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, trim_offsets=True, use_regex=True)
+
+# post-processor
+tokenizer.post_processor = processors.ByteLevel(add_prefix_space=True, trim_offsets=False, use_regex=True)
+
+# decoder
+tokenizer.decoder = decoders.ByteLevel(add_prefix_space=True, trim_offsets=True, use_regex=True)
+
+#
+# BPE Trainer
+#
+trainer = BpeTrainer(
+    vocab_size=131072, # 128 * 1024
+    min_frequency=3,
+    special_tokens=special_tokens,
+    max_token_length=16,
+)
+
+tokenizer_datasets = core_base_datasets + core_instruct_datasets
+
+tokenizer.train_from_iterator(
+    (batch_dataset_iterator(n) for n in tokenizer_datasets),
+    trainer,
+)
+
+tokenizer.save(os.path.join(tokenizer_path, 'tokenizer.json'))
+tokenizer.model.save(tokenizer_path)
+
+#
+# PreTrainedTokenizerFast
+#
+CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+        "{{'<|im_start|>' + message['role'] + '<|im_sep|>' + message['content'] + '<|im_end|>'}}"
+    "{% endfor %}"
+
+    "{% if add_generation_prompt %}"
+        "{{ '<|im_start|>assistant<|im_sep|>' }}"
+    "{% endif %}"
+)
+
+fast_tokenizer = PreTrainedTokenizerFast(
+    tokenizer_object=tokenizer,
+    chat_template=CHAT_TEMPLATE,
+    bos_token=bos_token,
+    eos_token=eos_token,
+    pad_token=pad_token,
+    unk_token=unk_token,
+    clean_up_tokenization_spaces=False,
+)
+
+fast_tokenizer.save_pretrained(tokenizer_path)

scripts/utils.py

@@ -0,0 +1,151 @@
+import gc
+from typing import Optional, Iterator, Callable
+
+import torch
+from datasets import load_dataset
+from litgpt.tokenizer import Tokenizer
+from transformers import AutoTokenizer
+
+
+def batch_text_iterator(kind: str,
+                        path: str,
+                        name: Optional[str]=None,
+                        data_dir: Optional[str]=None,
+                        data_files: Optional[str]=None,
+                        keep_in_memory: bool=False,
+                        revision: Optional[str]=None,
+                        split: str='train',
+                        num_proc: Optional[int]=None,
+                        format: Optional[Callable|str]=None) -> Iterator[str]:
+    assert isinstance(format, str) or callable(format), f'{path=} {format=}'
+    assert kind == 'base'
+
+    dataset = load_dataset(path=path,
+                           name=name,
+                           data_dir=data_dir,
+                           data_files=data_files,
+                           keep_in_memory=keep_in_memory,
+                           revision=revision,
+                           split=split,
+                           trust_remote_code=True,
+                           num_proc=num_proc)
+
+    if callable(format):
+        for row in dataset:
+            text = format(row)
+
+            if not text:
+                continue
+
+            yield text
+    else:
+        for row in dataset:
+            text = format.format(**row)
+
+            if not text:
+                continue
+
+            yield text
+
+    del dataset
+    gc.collect()
+
+
+def batch_chat_iterator(kind: str,
+                        path: str,
+                        name: Optional[str]=None,
+                        data_dir: Optional[str]=None,
+                        data_files: Optional[str]=None,
+                        keep_in_memory: bool=False,
+                        revision: Optional[str]=None,
+                        split: str='train',
+                        num_proc: Optional[int]=None,
+                        field: Optional[str]=None,
+                        transform: Optional[Callable]=None) -> Iterator[list[dict[str, str]]]:
+    assert kind == 'instruct'
+
+    dataset = load_dataset(path=path,
+                           name=name,
+                           data_dir=data_dir,
+                           data_files=data_files,
+                           keep_in_memory=keep_in_memory,
+                           revision=revision,
+                           split=split,
+                           trust_remote_code=True,
+                           num_proc=num_proc)
+
+    if callable(transform):
+        for row in dataset:
+            if field:
+                messages = transform(row[field])
+            else:
+                messages = transform(row)
+
+            if not messages:
+                continue
+
+            yield messages
+    else:
+        for row in dataset:
+            if field:
+                messages = row[field]
+            else:
+                raise ValueError(field)
+
+            if not messages:
+                continue
+
+            yield messages
+
+    del dataset
+    gc.collect()
+
+
+# NOTE: used only by tokenizer trainer
+def batch_dataset_iterator(dataset_config: dict) -> Iterator[str]:
+    if dataset_config['kind'] == 'base':
+        for text in batch_text_iterator(**dataset_config):
+           yield text
+    elif dataset_config['kind'] == 'instruct':
+        for messages in batch_chat_iterator(**dataset_config):
+            text = '\n'.join(n['content'] for n in messages)
+            yield text
+
+
+def tokenize_text_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    for text in batch_text_iterator(**dataset_config):
+        text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=True)
+        yield text_ids
+
+
+def tokenize_chat_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    for messages in batch_chat_iterator(**dataset_config):
+        text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
+        text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
+        yield text_ids
+
+
+def tokenize_fn(dataset_config: dict, min_len: int, max_len: int, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    if dataset_config['kind'] == 'base':
+        for text in batch_text_iterator(**dataset_config):
+            try:
+                text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=True)
+            except Exception as e:
+                print(f'Skip base raw: {e=} {type(text)=} {text=}')
+                continue
+
+            if min_len <= len(text_ids) <= max_len:
+                yield text_ids
+    elif dataset_config['kind'] == 'instruct':
+        for messages in batch_chat_iterator(**dataset_config):
+            try:
+                text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
+                text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
+            except Exception as e:
+                print(f'Skip instruct row: {e=} {type(messages)=} {messages=}')
+                continue
+
+            if min_len <= len(text_ids) <= max_len:
+                yield text_ids
+    else:
+        raise ValueError(dataset_config['kind'])

tokenizer/special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|im_end|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>"
+}

tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a45777956cfe9ce422244a4cdfafd3aacf7eb4490bde3034fec3dda9c1a98ab
+size 9845205

tokenizer/tokenizer_config.json

@@ -0,0 +1,525 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|unk|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<|im_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "system",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "user",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "assistant",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "</tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<tool>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "</tool>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "17": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "18": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "19": {
+      "content": "<|reserved_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "20": {
+      "content": "<|reserved_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "21": {
+      "content": "<|reserved_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "22": {
+      "content": "<|reserved_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "23": {
+      "content": "<|reserved_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "24": {
+      "content": "<|reserved_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "25": {
+      "content": "<|reserved_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "26": {
+      "content": "<|reserved_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "27": {
+      "content": "<|reserved_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "28": {
+      "content": "<|reserved_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "29": {
+      "content": "<|reserved_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30": {
+      "content": "<|reserved_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "31": {
+      "content": "<|reserved_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32": {
+      "content": "<|reserved_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "33": {
+      "content": "<|reserved_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "34": {
+      "content": "<|reserved_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "35": {
+      "content": "<|reserved_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "36": {
+      "content": "<|reserved_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "37": {
+      "content": "<|reserved_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "38": {
+      "content": "<|reserved_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "39": {
+      "content": "<|reserved_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "40": {
+      "content": "<|reserved_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "41": {
+      "content": "<|reserved_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "42": {
+      "content": "<|reserved_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "43": {
+      "content": "<|reserved_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "44": {
+      "content": "<|reserved_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "45": {
+      "content": "<|reserved_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "46": {
+      "content": "<|reserved_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "47": {
+      "content": "<|reserved_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "48": {
+      "content": "<|reserved_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49": {
+      "content": "<|reserved_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50": {
+      "content": "<|reserved_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "51": {
+      "content": "<|reserved_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "52": {
+      "content": "<|reserved_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "53": {
+      "content": "<|reserved_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "54": {
+      "content": "<|reserved_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "55": {
+      "content": "<|reserved_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "56": {
+      "content": "<|reserved_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "57": {
+      "content": "<|reserved_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "58": {
+      "content": "<|reserved_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "59": {
+      "content": "<|reserved_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "60": {
+      "content": "<|reserved_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "61": {
+      "content": "<|reserved_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "62": {
+      "content": "<|reserved_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "63": {
+      "content": "<|reserved_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '<|im_sep|>' + message['content'] + '<|im_end|>'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant<|im_sep|>' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<|unk|>"
+}