Mixture of datasets

configs/_cluster/entropy.yaml

@@ -19,7 +19,7 @@ infrastructure:
     - 'export HYDRA_FULL_ERROR=1'
 
     # export pixi variables
-    - 'export PIXI_HOME=/storage_nvme_4/nano/pixi_new'
+    - 'export PIXI_HOME=/storage_nvme_4/nano/pixi_md'
     - 'export PATH="$PIXI_HOME/bin:$PATH"'
     - 'export XDG_DATA_HOME="$PIXI_HOME/data"'
     - 'export XDG_CACHE_HOME="$PIXI_HOME/cache"'

configs/_dataset/c4.yaml

@@ -4,27 +4,15 @@ defaults:
 
 trainer:
   train_dataloader:
-    dataset:
-      _target_: src.core.datasets.C4Dataset
-      sequence_length: ${common.sequence_length}
-      tokenize_fn: ???
-      path: ${cluster_switch.train_path_c4}
-      split: train
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
-    num_workers: 8
+    datasets:
+      - path: ${cluster_switch.train_path_c4}
+        weight: 1.0
+    dataset_split: train
+    num_workers: 2
 
   eval_dataloader:
-    dataset:
-      _target_: src.core.datasets.C4Dataset
-      sequence_length: ${common.sequence_length}
-      tokenize_fn: ???
-      path: ${cluster_switch.eval_path_c4}
-      split: validation
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
-    num_workers: 8
+    datasets:
+      - path: ${cluster_switch.eval_path_c4}
+        weight: 1.0
+    dataset_split: validation
+    num_workers: 2

configs/_dataset/default.yaml

@@ -1,12 +1,26 @@
 trainer:
   train_dataloader:
-    _target_: src.core.datasets.get_dataloader
-    dataset: ???
-    total_batch_size: ${common.batch_size}
+    _target_: src.core.datasets.get_mixture_of_datasets_dataloader
+    datasets: ???
+    dataset_split: ???
     num_workers: ???
+    seed: 123
+    sequence_length: ${common.sequence_length}
+    shuffle: true
+    total_batch_size: ${common.batch_size}
+    use_new_sampling_method: true
+    world_size_independent: false
+    tokenize_fn: ???
 
   eval_dataloader:
-    _target_: src.core.datasets.get_dataloader
-    dataset: ???
+    _target_: src.core.datasets.get_mixture_of_datasets_dataloader
+    datasets: ???
+    dataset_split: ???
+    num_workers: ???
+    seed: 123
+    sequence_length: ${common.sequence_length}
+    shuffle: true
     total_batch_size: ${common.batch_size}
-    num_workers: ???
+    use_new_sampling_method: true
+    world_size_independent: false
+    tokenize_fn: ???

configs/_dataset/fineweb.yaml

@@ -2,30 +2,17 @@ defaults:
   - default
   - _self_
 
-
 trainer:
   train_dataloader:
-    dataset:
-      _target_: src.core.datasets.FineWebEduDataset
-      sequence_length: ${common.sequence_length}
-      tokenize_fn: ???
-      path: ${cluster_switch.train_path_fineweb}
-      split: train
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
+    datasets:
+      - path: ${cluster_switch.train_path_fineweb}
+        weight: 1.0
+    dataset_split: train
     num_workers: 2
 
   eval_dataloader:
-    dataset:
-      _target_: src.core.datasets.FineWebEduDataset
-      sequence_length: ${common.sequence_length}
-      tokenize_fn: ???
-      path: ${cluster_switch.eval_path_fineweb}
-      split: train
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
+    datasets:
+      - path: ${cluster_switch.eval_path_fineweb}
+        weight: 1.0
+    dataset_split: train
     num_workers: 2

configs/_dataset/local_dummy.yaml

@@ -4,27 +4,21 @@ defaults:
 
 trainer:
   train_dataloader:
-    dataset:
-      _target_: src.core.datasets.C4Dataset
-      sequence_length: 16
-      tokenize_fn: ???
-      path: data
-      split: train
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
+    datasets:
+      - path: data
+        weight: 1.0
+    dataset_split: train
     num_workers: 0
+    tokenize_fn: 
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
   eval_dataloader:
-    dataset:
-      _target_: src.core.datasets.C4Dataset
-      sequence_length: 16
-      tokenize_fn: ???
-      path: data_eval
-      split: validation
-      seed: 123
-      use_new_sampling_method: true
-      shuffle: true
-      world_size_independent: false
-    num_workers: 0
+    datasets:
+      - path: data_eval
+        weight: 1.0
+    dataset_split: validation
+    num_workers: 0
+    tokenize_fn: 
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2

configs/_dataset/smollm_corpus.yaml

@@ -0,0 +1,38 @@
+defaults:
+  - default
+  - _self_
+
+trainer:
+  train_dataloader:
+    _target_: src.core.datasets.get_mixture_of_datasets_dataloader
+    datasets:
+      - path: /storage_nvme_1/llm-random/datasets/fineweb-edu-dedup/train
+        weight: 0.7
+      - path: /storage_nvme_1/llm-random/datasets/cosmopedia-v2/train
+        weight: 0.15
+      - path: /storage_nvme_2/llm-random/datasets/python-edu
+        weight: 0.08
+      - path: /storage_nvme_2/llm-random/datasets/open-web-math/train
+        weight: 0.07
+    dataset_split: train
+    num_workers: 2
+    tokenize_fn:
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: HuggingFaceTB/SmolLM-1.7B
+
+  eval_dataloader:
+    _target_: src.core.datasets.get_mixture_of_datasets_dataloader
+    datasets:
+      - path: /storage_nvme_1/llm-random/datasets/fineweb-edu-dedup/train
+        weight: 0.7
+      - path: /storage_nvme_1/llm-random/datasets/cosmopedia-v2/train
+        weight: 0.15
+      - path: /storage_nvme_2/llm-random/datasets/python-edu
+        weight: 0.08
+      - path: /storage_nvme_2/llm-random/datasets/open-web-math/train
+        weight: 0.07
+    dataset_split: validation
+    num_workers: 1
+    tokenize_fn:
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: HuggingFaceTB/SmolLM-1.7B

configs/_misc/default.yaml

@@ -6,7 +6,7 @@ infrastructure:
     name: default
     type: wandb
     wandb_entity: ideas_cv
-    project_name: llm-random-test
+    project_name: pc_smollm
     tags:
       - new_wandb_job
 

configs/_model/smollm/base_pc_model.yaml

@@ -23,10 +23,8 @@ model:
 
   embedding:
     _target_: src.projected_compression.model.ProjectedEmbedding
-    embedding:
-      _target_: src.projected_compression.model.Embedding
-      num_embeddings: ${common.vocab_size}
-      embedding_dim: ${common.base_dmodel}
+    num_embeddings: ${common.vocab_size}
+    embedding_dim: ${common.base_dmodel}
     result_out_features: ${common.dmodel}
 
   encoder:
@@ -108,7 +106,7 @@ model:
       result_in_features: ${common.dmodel}
       result_out_features: null
       base_in_features: ${common.base_dmodel}
-      base_out_features: ${model.embedding.num_embeddings}
+      base_out_features: ${common.vocab_size}
     norm_fn:
         _target_: src.core.model.RMSNorm
         _partial_: true

configs/_trainer/llama.yaml

@@ -24,12 +24,11 @@ trainer:
         - ${model.head._target_}
 
   train_dataloader:
-    dataset:
-      tokenize_fn:
-        _target_: src.core.datasets.llama_tokenize_fn
-
+    tokenize_fn:
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: meta-llama/Llama-3.1-8B
 
   eval_dataloader:
-    dataset:
-      tokenize_fn:
-        _target_: src.core.datasets.llama_tokenize_fn
+    tokenize_fn:
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: meta-llama/Llama-3.1-8B

configs/_trainer/llama_distillation.yaml

@@ -151,13 +151,3 @@ trainer:
         - ${model.embedding._target_}
         - ${model.encoder.block_fn._target_}
         - ${model.head._target_}
-
-  train_dataloader:
-    dataset:
-      tokenize_fn:
-        _target_: src.core.datasets.llama_tokenize_fn
-
-  eval_dataloader:
-    dataset:
-      tokenize_fn:
-        _target_: src.core.datasets.llama_tokenize_fn

configs/_trainer/smollm_1700_distillation.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - llama_distillation
+
+common_distillation:
+  dmodel: 2048
+  dff: 8192
+  datt: 2048
+  n_blocks: 24
+  q_heads: 32
+  kv_heads: 32
+  vocab_size: 49152
+
+distillation:
+  load:
+    path: "HuggingFaceTB/SmolLM-1.7B"
+

configs/dataset_mixture_test.yaml

@@ -0,0 +1,34 @@
+defaults:
+  - _cluster@_here_: local
+  - _model@_here_: tiny
+  - _trainer@_here_: llama
+  - _dataset@_here_: smollm_corpus
+  - _checkpoints@_here_: none
+  - _misc@_here_: default
+  - _eval@_here_: default
+
+common:
+  sequence_length: 16
+  batch_size: 2
+
+trainer:
+  gradient_accumulation_steps: 1
+  n_steps: 100
+  learning_rate: 1e-3
+
+  checkpoint:
+    save:
+      type: huggingface
+      path: checkpoint
+
+infrastructure:
+  metric_logger:
+    name: tiny_Local
+    tags:
+      - nano
+      - local
+      - tiny
+
+evaluator:
+  limit: 1
+  device: cpu

configs/pc_project/llama_1B_fine_tune.yaml

@@ -28,7 +28,7 @@ trainer:
   learning_rate: 15
 
   train_dataloader:
-    dataset:
+    datasets:
       seed: 1000
 
   checkpoint:

configs/pc_project/llmr_300.yaml

@@ -38,13 +38,14 @@ trainer:
   distributed: null
 
   train_dataloader:
-    dataset:
-      tokenize_fn: 
-        _target_: src.core.datasets.gpt2_tokenize_fn
+    tokenize_fn: 
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
   eval_dataloader:
     tokenize_fn: 
-      _target_: src.core.datasets.gpt2_tokenize_fn
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
 
 infrastructure:

configs/pc_project/llmr_300_comp.yaml

@@ -41,13 +41,14 @@ trainer:
     path: ??? # CHANGE
 
   train_dataloader:
-    dataset:
-      tokenize_fn: 
-        _target_: src.core.datasets.gpt2_tokenize_fn
+    tokenize_fn: 
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
   eval_dataloader:
     tokenize_fn: 
-      _target_: src.core.datasets.gpt2_tokenize_fn
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
 
 infrastructure:

configs/pc_project/llmr_300_importances.yaml

@@ -25,13 +25,14 @@ trainer:
     path: ??? # CHANGE
 
   train_dataloader:
-    dataset:
-      tokenize_fn: 
-        _target_: src.core.datasets.gpt2_tokenize_fn
+    tokenize_fn: 
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
   eval_dataloader:
     tokenize_fn: 
-      _target_: src.core.datasets.gpt2_tokenize_fn
+      _target_: src.core.datasets.get_tokenize_fn
+      model_name: gpt2
 
 
 infrastructure: