Sweeps exp

configs/150M_big_bs/H100.toml

@@ -0,0 +1,16 @@
+name_model = "150M"
+project = "150M_big_bs"
+type_model = "llama2"
+
+[train]
+micro_bs = 64 # change this base on the gpu
+reshard_after_forward = false
+
+[optim]
+batch_size = 2048
+warmup_steps = 500
+total_steps = 4096
+
+
+[optim.optim]
+lr = 8e-4

configs/150M_short/3090.toml

@@ -4,7 +4,7 @@ type_model = "llama2"
 
 [train]
 micro_bs = 16 # change this base on the gpu
-reshard_after_forward = true
+reshard_after_forward = false
 
 [optim]
 batch_size = 512

configs/150M_short/A40.toml

@@ -4,7 +4,7 @@ type_model = "llama2"
 
 [train]
 micro_bs = 32 # change this base on the gpu
-reshard_after_forward = true
+reshard_after_forward = false
 
 
 [optim]

configs/150M_short/H100.toml

@@ -4,7 +4,7 @@ type_model = "llama2"
 
 [train]
 micro_bs = 64 # change this base on the gpu
-reshard_after_forward = true
+reshard_after_forward = false
 
 [optim]
 batch_size = 512

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "setuptools",
     "transformers>=4.44.2",
     "datasets>=3.0.0",
-    "pydantic_config @ git+https://github.com/samsja/pydantic_config.git@74c94ee",
+    "pydantic_config @ git+https://github.com/samsja/pydantic_config.git@b7becc3",
     "torchdata>=0.8.0",
     "fsspec[gcs]>=2024.3.1",
     "ninja",

sweeps/adam.yaml

@@ -0,0 +1,39 @@
+command: 
+ - bash 
+ - scripts/simulate_multi_node_diloco.sh
+ - "1"
+ - "8"
+ - "src/zeroband/train.py"
+ - "@configs/150M_big_bs/H100.toml"
+ - ${args}
+name: adam-sweep
+method: bayes
+early_terminate:
+  type: hyperband
+  max_iter: 27
+  eta: 3
+  s: 2
+metric:
+  goal: minimize
+  name: Perplexity
+parameters:
+  optim.optim.lr:
+    distribution: log_uniform_values
+    min: 1e-4
+    max: 1e-2
+  # optim.optim.betas1:
+  #   distribution: uniform
+  #   min: 0.8
+  #   max: 0.999 
+  optim.optim.betas2:
+    distribution: uniform
+    min: 0.8
+    max: 0.9999
+  optim.optim.weight_decay:
+    distribution: uniform
+    min: 0.0
+    max: 0.4
+  optim.warmup_steps:
+    distribution: q_uniform
+    min: 0
+    max: 2000

sweeps/diloco.yaml

@@ -0,0 +1,43 @@
+command: 
+ - bash 
+ - scripts/simulate_multi_node_diloco.sh
+ - "2"
+ - "4"
+ - "src/zeroband/train.py"
+ - "@configs/150M_big_bs/H100.toml"
+ - "--optim.batch_size"
+ - "1024"
+ - "--diloco.inner_steps"
+ - "25"
+ - ${args}
+name: adam-sweep
+method: bayes
+early_terminate:
+  type: hyperband
+  max_iter: 27
+  eta: 3
+  s: 2
+metric:
+  goal: minimize
+  name: Perplexity
+parameters:
+  optim.optim.lr:
+    distribution: log_uniform_values
+    min: 1e-4
+    max: 1e-2
+  # optim.optim.betas1:
+  #   distribution: uniform
+  #   min: 0.8
+  #   max: 0.999 
+  optim.optim.betas2:
+    distribution: uniform
+    min: 0.8
+    max: 0.9999
+  optim.optim.weight_decay:
+    distribution: uniform
+    min: 0.0
+    max: 0.4
+  optim.warmup_steps:
+    distribution: q_uniform
+    min: 0
+    max: 2000