Pccl integration

configs/10B/H100_intellect1.toml

@@ -24,19 +24,34 @@ betas1 = 0.9
 betas2 = 0.95
 weight_decay = 0.1
 
+[train.outer_lr_scheduler]
+lr = 0.7
+end_lr = 0.7
+num_decay_steps = 0
+num_warmup_steps = 0
+num_stable_steps = 0
+
+[train.outer_optimizer]
+type = "sgd"
+momentum = 0.9
+nesterov = true
+
 [data]
 seq_length = 8192
-dataset_name_or_paths = "/data/datasets/fineweb-edu,/data/datasets/fineweb,/data/datasets/StackV1-popular,/data/datasets/dclm-baseline-1.0-parquet,/data/datasets/open-web-math"
+dataset_name_or_paths = "/data/datasets/fineweb-edu.bin,/data/datasets/fineweb.bin,/data/datasets/StackV1-popular.bin,/data/datasets/dclm-baseline-1.0-parquet.bin,/data/datasets/open-web-math.bin"
+token_bit_size = 17
 dataset_ratio = "55:10:20:10:5"
 num_workers = 4
 reverse_data_files = true
 split_by_data_rank = false # the 10b training assume that data was already split by datarank. Keeping this for backward compatibility
 
-
 [diloco]
 inner_steps = 100
-compression = "uint8"
+delayed_update = true
 
 [ckpt]
 interval = 100
 path = "/data/10B"
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/10B/H100_simple.toml

@@ -25,3 +25,7 @@ weight_decay = 0.1
 [data]
 seq_length = 8192
 num_workers = 4
+fake = true
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/13B/H100.toml

@@ -13,3 +13,6 @@ batch_size = 512
 [data]
 seq_length = 1024
 dataset_name_or_paths = "datasets/fineweb-edu"
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/150M/A100_debug.toml

@@ -1,12 +1,11 @@
-project = "debug_150m_zero_band"
-
 model_name = "150M"
-model_type = "llama2"
+model_type = "llama3"
 
 wandb = false
+log_all_ranks = true
 
 [hardware]
-micro_batch_size = 64
+micro_batch_size = 32
 torch_compile = true
 
 [train]
@@ -16,6 +15,25 @@ batch_size = 512
 num_warmup_steps = 10
 num_decay_steps = 1000
 
+[train.outer_lr_scheduler]
+lr = 1.0
+end_lr = 1.0
+num_decay_steps = 0
+num_warmup_steps = 0
+num_stable_steps = 0
+
+[train.outer_optimizer]
+type = "sgd"
+
 [data]
-fake = true
+dataset_name_or_paths = 'tests/test_data/parquet/parquet_ds_folder_1,tests/test_data/parquet/parquet_ds_folder_2'
+#dataset_name_or_paths = '/home/mike/IntelliJProjects/dataproctest/working_dir/train_0.bin,/home/mike/IntelliJProjects/dataproctest/working_dir/train_1.bin'
+dataset_ratio = "50:50"
+token_bit_size = 17
+
+[diloco]
+inner_steps = 16
+delayed_update = true
 
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/150M/H100.toml

@@ -13,3 +13,6 @@ batch_size = 512
 [data]
 seq_length = 1024
 dataset_name_or_paths = "datasets/fineweb-edu"
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/150M/H100_best.toml

@@ -24,3 +24,6 @@ weight_decay = 0.24530252977858977
 [data]
 seq_length = 1024
 dataset_name_or_paths = "datasets/fineweb-edu"
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/1B/H100.toml

@@ -13,3 +13,6 @@ batch_size = 512
 [data]
 seq_length = 1024
 dataset_name_or_paths = "datasets/fineweb-edu"
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/70M/H100.toml

@@ -13,3 +13,10 @@ batch_size = 512
 [data]
 seq_length = 1024
 dataset_name_or_paths = "datasets/fineweb-edu"
+
+[diloco]
+inner_steps = 16
+delayed_update = true
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/7B/H100.toml

@@ -1,15 +1,42 @@
-project = "debug_7B_zero_band"
-
 model_name = "7B"
-model_type = "llama2"
+model_type = "llama3"
+
+wandb = true
+log_all_ranks = true
 
 [hardware]
-micro_batch_size = 64
-reshard_after_forward = false
+micro_batch_size = 8
+reshard_after_forward = true
+torch_compile = false
+attn_fn="sdpa"
 
 [train]
 batch_size = 512
 
+[train.lr_scheduler]
+lr = 3e-4
+end_lr = 0.0
+num_warmup_steps = 8000
+num_decay_steps = 1.2e6
+
+[train.outer_lr_scheduler]
+lr = 0.7
+end_lr = 0.7
+num_decay_steps = 0
+num_warmup_steps = 0
+num_stable_steps = 0
+
+[train.outer_optimizer]
+type = "sgd"
+
 [data]
 seq_length = 1024
-dataset_name_or_paths = "datasets/fineweb-edu"
+dataset_name_or_paths = 'http://65.108.32.176:8080/api/v1/datasets/fineweb-edu-train/stream'
+token_bit_size = 17
+
+[diloco]
+inner_steps = 64
+delayed_update = true
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/debug/diloco.toml

@@ -1,20 +1,39 @@
-model_name = "debugmodel"
-model_type = "llama2"
+model_name = "150M"
+model_type = "llama3"
 
 wandb = false
+log_all_ranks = true
 
 [hardware]
-micro_batch_size = 8
+micro_batch_size = 32
+torch_compile = true
 
 [train]
-batch_size = 16
+batch_size = 512
 
 [train.lr_scheduler]
 num_warmup_steps = 10
-num_decay_steps = 10
+num_decay_steps = 1000
+
+[train.outer_lr_scheduler]
+lr = 1.0
+end_lr = 1.0
+num_decay_steps = 0
+num_warmup_steps = 0
+num_stable_steps = 0
+
+[train.outer_optimizer]
+type = "sgd"
 
 [data]
-fake = true
+#dataset_name_or_paths = 'tests/test_data/parquet/parquet_ds_folder_1,tests/test_data/parquet/parquet_ds_folder_2'
+dataset_name_or_paths = '/home/mike/IntelliJProjects/dataproctest/working_dir/train'
+#dataset_ratio = "50:50"
+token_bit_size = 17
 
 [diloco]
-inner_steps = 5
+inner_steps = 16
+delayed_update = true
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/debug/normal.toml

@@ -7,7 +7,7 @@ wandb = false
 micro_batch_size = 8
 
 [train]
-batch_size = 16
+batch_size = 32
 
 [train.lr_scheduler]
 num_warmup_steps = 10
@@ -16,3 +16,5 @@ num_decay_steps = 10
 [data]
 fake = true
 
+[pccl]
+ccoip_host = "127.0.0.1:48148"

configs/test.toml

@@ -24,3 +24,6 @@ num_warmup_steps = 1000
 lr = 3e-4
 end_lr = 0.0
 num_decay_steps = 80000
+
+[pccl]
+ccoip_host = "127.0.0.1:48148"

pyproject.toml

@@ -11,14 +11,18 @@ dependencies = [
     "transformers>=4.44.2",
     "datasets>=3.0.0",
     "pydantic_config @ git+https://github.com/samsja/pydantic_config.git@b7becc3",
+    "tomli",
     "torchdata>=0.8.0",
     "fsspec[gcs]>=2024.3.1",
     "ninja",
     "zstandard",
     "pyarrow",
     "psutil",
     "wandb",
-    "imageio[ffmpeg]"
+    "numba",
+    "imageio[ffmpeg]",
+    "pccl @ git+https://github.com/PrimeIntellect-ai/pccl.git@main#subdirectory=python/framework",
+    "datasetstream @ git+https://github.com/PrimeIntellect-ai/datasetstream@main"
 ]
 
 [project.optional-dependencies]
@@ -37,4 +41,4 @@ allow-direct-references = true # allow direct references to git repos in depende
 line-length = 120
 
 [tool.uv]
-dev-dependencies = ["ruff>=0.5.0", "pre-commit>=3.0.0","pytest>=7.0.0", "faker", "matplotlib"]
+dev-dependencies = ["ruff>=0.5.0", "pre-commit>=3.0.0", "pytest>=7.0.0", "faker", "matplotlib", "transformers"]