prepare dataset

scripts/model.yaml

@@ -0,0 +1,143 @@
+# 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"
+
+# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
+# ``model_config``. (type: Optional[Config], default: null)
+model_config:
+  padded_vocab_size: 32768
+  vocab_size: 32768
+  block_size: 32768
+  n_layer: 10
+  n_head: 12
+  head_size: null
+  n_embd: 312
+  n_query_groups: 4
+  rotary_percentage: 1.0
+  parallel_residual: false
+  bias: false
+  norm_class_name: "RMSNorm"
+  norm_eps: 1e-05
+  mlp_class_name: "LLaMAMLP"
+  intermediate_size: 1092
+  rope_base: 500000
+
+# 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/contrain/"
+
+# 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: "tangledgroup/tangled-llama-33m-32k-base-v0.1"
+
+# 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: false
+resume: "auto"
+
+# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
+data:
+  class_path: LitData
+
+  init_args:
+    data_path: "../data/"
+    num_workers: 16
+    seq_length: 32768
+
+# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
+train:
+  # Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
+  save_interval: 1000
+
+  # 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: 16
+  # micro_batch_size: 14
+
+  # Number of iterations with learning rate warmup active (type: int, default: 2000)
+  lr_warmup_steps: 2000
+
+  # 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: 3000000000000
+  max_tokens: 9782206713 # 1591379 * 2049 * 3
+
+  # 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: 32768
+
+  # Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
+  tie_embeddings:
+
+  #   (type: Optional[float], default: 1.0)
+  max_norm: 1.0
+
+  #   (type: float, default: 4e-05)
+  min_lr: 4.0e-05
+
+# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
+eval:
+  # Number of optimizer steps between evaluation calls (type: int, default: 1000)
+  interval: 100
+
+  # 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: grokadamw.GrokAdamW
+  # class_path: bitsandbytes.optim.AdamW8bit
+  # class_path: bitsandbytes.optim.PagedAdamW8bit
+  
+  init_args:
+    #   (type: float, default: 0.001)
+    lr: 1.0e-3
+    
+    #   (type: float, default: 0.01)
+    weight_decay: 0.1
+    
+    #   (type: tuple, default: (0.9,0.999))
+    betas:
+      - 0.9
+      - 0.95
+
+# 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: "../"
+
+# 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: 42

scripts/prepare_contrain_dataset.py

@@ -0,0 +1,142 @@
+import gc
+
+from datasets import load_dataset
+from litdata import optimize, TokensLoader
+from litgpt.tokenizer import Tokenizer
+from functools import partial
+
+
+def batch_iterator(name=None):
+    if name in (None, 'Replete-AI/Everything_Instruct_Multilingual'):
+        dataset = load_dataset('Replete-AI/Everything_Instruct_Multilingual', split='train')
+
+        for row in dataset:
+            text = []
+
+            if row['instruction']:
+                text.append(
+                    '<|im_start|>system\n'
+                    f"{row['instruction']}<|im_end|>"
+                )
+
+            if row['input']:
+                text.append(
+                    '<|im_start|>user\n'
+                    f"{row['input']}<|im_end|>"
+                )
+
+            if row['output']:
+                text.append(
+                    '<|im_start|>assistant\n'
+                    f"{row['output']}<|im_end|>"
+                )
+
+            text = '\n'.join(text) + '\n'
+            yield text
+            break
+
+        del dataset
+        gc.collect()
+
+    if name in (None, 'HuggingFaceH4/ultrachat_200k'):
+        dataset = load_dataset('HuggingFaceH4/ultrachat_200k', split='train_sft')
+
+        for row in dataset:
+            text = [
+                f"<|im_start|>{n['role']}\n{n['content']}<|im_end|>"
+                for n in row['messages']
+            ]
+
+            text = '\n'.join(text) + '\n'
+            yield text
+            break
+
+        del dataset
+        gc.collect()
+
+    if name in (None, 'HuggingFaceH4/no_robots'):
+        dataset = load_dataset('HuggingFaceH4/no_robots', split='train')
+
+        for row in dataset:
+            text = [
+                f"<|im_start|>{n['role']}\n{n['content']}<|im_end|>"
+                for n in row['messages']
+            ]
+
+            text = '\n'.join(text) + '\n'
+            yield text
+            break
+
+        del dataset
+        gc.collect()
+
+    if name in (None, 'datatab/ultrafeedback_binarized_serbian'):
+        dataset = load_dataset('datatab/ultrafeedback_binarized_serbian', split='train_sft')
+
+        for row in dataset:
+            text = [
+                f"<|im_start|>{n['role']}\n{n['content']}<|im_end|>"
+                for n in row['chosen']
+            ]
+
+            text = '\n'.join(text) + '\n'
+            yield text
+            break
+
+        del dataset
+        gc.collect()
+
+    if name in (None, 'datatab/alpaca-cleaned-serbian-full'):
+        dataset = load_dataset('datatab/alpaca-cleaned-serbian-full', split='train')
+
+        for row in dataset:
+            text = []
+
+            if row['instruction']:
+                text.append(
+                    '<|im_start|>system\n'
+                    f"{row['instruction']}<|im_end|>"
+                )
+
+            if row['input']:
+                text.append(
+                    '<|im_start|>user\n'
+                    f"{row['input']}<|im_end|>"
+                )
+
+            if row['output']:
+                text.append(
+                    '<|im_start|>assistant\n'
+                    f"{row['output']}<|im_end|>"
+                )
+
+            text = '\n'.join(text) + '\n'
+            yield text
+            break
+
+        del dataset
+        gc.collect()
+
+
+def tokenize_fn(dataset_name, tokenizer=None):
+    for text in batch_iterator(dataset_name):
+        text_ids = tokenizer.encode(text, bos=False, eos=True)
+        yield text_ids
+
+
+datasets_names = [
+    'Replete-AI/Everything_Instruct_Multilingual',
+    'HuggingFaceH4/ultrachat_200k',
+    'HuggingFaceH4/no_robots',
+    'datatab/ultrafeedback_binarized_serbian',
+    'datatab/alpaca-cleaned-serbian-full',
+]
+
+outputs = optimize(
+    fn=partial(tokenize_fn, tokenizer=Tokenizer('..')),
+    inputs=datasets_names,
+    output_dir='../data/',
+    # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+    chunk_size=((32768 + 1) * 500),
+    num_workers=16,
+)

scripts/requirements.in

@@ -0,0 +1,11 @@
+# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+tqdm
+datasets
+jinja2
+transformers
+bitsandbytes
+wandb
+# litgpt[all]
+litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
+litdata
+grokadamw