love2code (3B) recipe

pyproject.toml

@@ -37,7 +37,7 @@ wandb = [
     "wandb",
 ]
 all = [
-    "ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@revert-147-epwalsh/mixture-fix",
+    "ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@learn2code",
     "beaker-py",
     "GitPython>=3.0,<4.0",
     "wandb",

src/cookbook/cli/cli.py

@@ -7,19 +7,14 @@
 import yaml
 from beaker import Beaker
 from beaker.services.job import JobClient
-from olmo_core.utils import generate_uuid, prepare_cli_environment
 from tqdm import tqdm
 from yaspin import yaspin
 
 from cookbook.aliases import ExperimentConfig, LaunchGroup, validate_sources
 from cookbook.cli.eval import convert, evaluate
-from cookbook.utils.config import (
-    build_train_config,
-    config_from_path,
-    mk_experiment_group,
-    mk_launch_configs,
-)
+from cookbook.utils.config import build_train_config, config_from_path, mk_experiment_group, mk_launch_configs
 from cookbook.utils.data import get_token_counts_and_ratios
+from olmo_core.utils import generate_uuid, prepare_cli_environment
 
 logger = logging.getLogger(__name__)
 
@@ -57,7 +52,6 @@ def cli():
 )
 def launch(config: Path, dry_run: bool, no_cache: bool, group_id: Optional[str] = None):
     """Launch an experiment."""
-
     with open(config, "r") as f:
         data = yaml.safe_load(f)
 

src/cookbook/cli/utils.py

@@ -361,8 +361,8 @@ def install_olmo_core(commit_hash: str | None, env: PythonEnv | None = None) ->
 
 def make_destination_dir(input_dir: str, suffix: str, output_dir: str | None = None) -> str:
     if output_dir is None:
-        input_base, input_fn = os.path.split(input_dir)
-        output_dir = os.path.join(input_base, f"{input_fn.rstrip('/')}-{suffix}")
+        input_base, input_fn = os.path.split(input_dir.rstrip("/"))
+        output_dir = os.path.join(input_base, f"{input_fn}-{suffix}")
 
     os.makedirs(output_dir, exist_ok=True)
 

src/cookbook/constants.py

@@ -176,6 +176,12 @@
     "bigcodebench_hard::none",
 ]
 
+ALL_1B_TASKS = [
+    "hellaswag",
+    "piqa",
+] + MMLU_CATEGORIES
+
+
 STARCODER_CODEX_TASKS = [
     "codex_humaneval::starcoder_pass@1",
     "codex_humaneval::starcoder_pass@10",
@@ -200,6 +206,7 @@
     "starcoder": STARCODER_CODEX_TASKS,
     "starcoder::pass@1": STARCODER_PASS_AT_1_TASKS,
     "code-no-bcb": [task for task in ALL_CODEX_TASKS if "bigcodebench" not in task],
+    "1b-evals": ALL_1B_TASKS,
 }
 
 OE_EVAL_GIT_URL = "[email protected]:allenai/oe-eval-internal.git"

src/cookbook/model/aliases.py

@@ -0,0 +1,111 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from olmo_core.config import Config
+from olmo_core.data import NumpyDataLoaderConfig, NumpyDatasetConfig, TokenizerConfig
+from olmo_core.distributed.parallel import DataParallelType
+from olmo_core.nn.transformer import TransformerBlockType, TransformerConfig
+from olmo_core.optim import AdamWConfig
+from olmo_core.train import TrainerConfig
+
+
+@dataclass
+class ModelTrainConfig(Config):
+    model: TransformerConfig
+    optim: AdamWConfig
+    dataset: NumpyDatasetConfig
+    data_loader: NumpyDataLoaderConfig
+    trainer: TrainerConfig
+    init_seed: int = 12536
+
+
+@dataclass
+class ModelConfig:
+    compile: bool
+    d_model: int
+    n_heads: int
+    n_layers: int
+    rope_theta: int
+    flash_attention: bool
+    max_sequence_length: int
+    layer_norm_eps: float = 1e-6
+    save_interval: int = 1000
+    eval_interval: int = 200
+    device_batch_size: int = 8
+    batch_divisor: int = 32
+    eps: float = 1e-8
+    betas: tuple = (0.9, 0.95)
+    weight_decay: float = 0.1
+    max_grad_norm: float = 1.0
+    decay_embeddings: bool = False
+    qk_norm: bool = True
+    dp_type: DataParallelType = DataParallelType.fsdp
+    block_type: TransformerBlockType = TransformerBlockType.reordered_norm
+
+    @classmethod
+    def olmo_30m(cls) -> "ModelConfig":
+        return ModelConfig(
+            compile=True,
+            d_model=256,
+            n_heads=8,
+            n_layers=4,
+            rope_theta=500_000,
+            flash_attention=True,
+            max_sequence_length=4096,
+        )
+
+    @classmethod
+    def olmo_190m(cls) -> "ModelConfig":
+        return ModelConfig(
+            compile=True,
+            d_model=768,
+            n_heads=12,
+            n_layers=12,
+            rope_theta=500_000,
+            flash_attention=True,
+            max_sequence_length=4096,
+        )
+
+    @classmethod
+    def olmo_1b(cls) -> "ModelConfig":
+        """
+        OLMo-1b (1_336_035_328 parameters)
+                (1_131_841_536 nonembed params)
+        """
+        return ModelConfig(
+            compile=True,
+            d_model=2048,
+            n_heads=16,
+            n_layers=18,
+            rope_theta=500_000,
+            flash_attention=True,
+            max_sequence_length=4096,
+        )
+
+    @classmethod
+    def love2code_3b(cls) -> "ModelConfig":
+        """
+        num params should be : 3607267840
+        num non_embed parmams should be: 3481438720
+        """
+        return ModelConfig(
+            compile=True,
+            d_model=2560,
+            n_heads=32,
+            n_layers=32,
+            rope_theta=500_000,
+            flash_attention=True,
+            max_sequence_length=2048,
+        )
+
+
+class SupportedModels(Enum):
+    olmo_190m = ModelConfig.olmo_190m()
+    olmo_30m = ModelConfig.olmo_30m()
+    olmo_1b = ModelConfig.olmo_1b()
+    starcoder2_3b = ModelConfig.starcoder_3b()
+
+
+class SupportedTokenizers(Enum):
+    dolma2 = TokenizerConfig.dolma2()
+    gpt_neox = TokenizerConfig.gpt_neox_olmo_dolma_v1_5()

src/cookbook/model/builder.py

@@ -2,13 +2,18 @@
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
-from olmo_core.data import (
-    DataMix,
-    NumpyDataLoaderConfig,
-    NumpyDatasetConfig,
-    NumpyDatasetType,
-    TokenizerConfig,
+from cookbook.aliases import SourceInstance, WandbConfig
+from cookbook.data.dataset import MixtureBuilder
+from cookbook.model.config import (
+    MODEL_TO_LR_MAP,
+    DefaultOptimizerProperties,
+    ModelTrainConfig,
+    SupportedTokenizers,
+    WrappedTransformerConfig,
 )
+from cookbook.model.evaluators import DownstreamEvaluators
+from cookbook.model.schedulers import WSD
+from olmo_core.data import DataMix, NumpyDataLoaderConfig, NumpyDatasetConfig, NumpyDatasetType, TokenizerConfig
 from olmo_core.data.types import NumpyDatasetDType
 from olmo_core.nn.transformer import TransformerConfig
 from olmo_core.optim import AdamWConfig, CosWithWarmup, OptimGroupOverride, Scheduler
@@ -28,18 +33,6 @@
 )
 from olmo_core.train.common import LoadStrategy
 
-from cookbook.aliases import SourceInstance, WandbConfig
-from cookbook.data.dataset import MixtureBuilder
-from cookbook.model.config import (
-    MODEL_TO_LR_MAP,
-    DefaultOptimizerProperties,
-    ModelTrainConfig,
-    SupportedTokenizers,
-    WrappedTransformerConfig,
-)
-from cookbook.model.evaluators import DownstreamEvaluators
-from cookbook.model.schedulers import WSD
-
 logger = logging.getLogger(__name__)
 
 
@@ -199,8 +192,15 @@ def __init__(
         if any(substring in cluster for substring in ["jupiter", "saturn"]) and weka:
             self.root_dir = f"/weka/oe-training-default/ai2-llm"
             logger.info(f"Using Weka bucket as root dir: {self.root_dir}")
-            self.checkpoint_dir = f"{self.root_dir}/checkpoints/{self.beaker_user.lower()}/{self.run_name}"
-
+        elif "augusta" in cluster:
+            try:
+                assert not weka
+            except AssertionError as e:
+                logger.info("Can't be on Augusta and weka!")
+                raise e
+            self.data_dir = self.root_dir = "gs://ai2-llm"
+
+        self.checkpoint_dir = f"{self.root_dir}/checkpoints/{self.beaker_user.lower()}/{self.run_name}"
         self.dataset_cache = f"{self.root_dir}/{self.beaker_user.lower()}/{self.run_name}/dataset-cache"
 
     def get_tokenizer_config(self, tokenizer) -> TokenizerConfig:
@@ -269,7 +269,7 @@ def build_callbacks(self, model: TransformerConfig) -> Dict[str, Callback]:
             "profiler": ProfilerCallback(enabled=self.profile),
             "checkpointer": CheckpointerCallback(
                 save_interval=self.save_interval,
-                ephemeral_save_interval=100,
+                ephemeral_save_interval=20,
                 save_async=True,
             ),
             "wandb": WandBCallback(
@@ -327,6 +327,10 @@ def build_dataset_config(self) -> NumpyDatasetConfig:
             for source in self.sources:
                 source_paths.extend(source.paths)
 
+        # source_paths = []
+        # for source in self.sources:
+        #     source_paths.extend(source.paths)
+
         dataset_config = NumpyDatasetConfig(
             paths=source_paths,
             source_mixture_config=mixture_config,
@@ -387,14 +391,14 @@ def build(self) -> ModelTrainConfig:
             load_path=load_path,
             load_strategy=load_strategy,
             save_folder=self.checkpoint_dir,
+            max_duration=Duration.tokens(self.max_tokens),
             work_dir=self.dataset_cache,
             rank_microbatch_size=rank_microbatch_size,
             save_overwrite=True,
             metrics_collect_interval=10,
             cancel_check_interval=5,
             compile_loss=True,
             z_loss_multiplier=1e-5,
-            max_duration=Duration.tokens(self.max_tokens),
         )
 
         for callback_name, callback in self.build_callbacks(self.transformer_config).items():

src/cookbook/model/config.py

@@ -4,12 +4,8 @@
 
 from olmo_core.config import Config, DType
 from olmo_core.data import NumpyDataLoaderConfig, NumpyDatasetConfig, TokenizerConfig
-from olmo_core.distributed.parallel import DataParallelType
-from olmo_core.nn.transformer import (
-    TransformerBlockType,
-    TransformerConfig,
-    TransformerDataParallelConfig,
-)
+from olmo_core.distributed.parallel import DataParallelConfig, DataParallelType
+from olmo_core.nn.transformer import TransformerBlockType, TransformerConfig
 from olmo_core.optim import AdamWConfig
 from olmo_core.train import TrainerConfig
 
@@ -58,7 +54,7 @@ def olmo_30m(cls, tokenizer: TokenizerConfig) -> TransformerConfig:
             layer_norm_eps=DefaultTransformerProperties.layer_norm_eps,
             qk_norm=DefaultTransformerProperties.qk_norm,
             block_name=DefaultTransformerProperties.block_type,
-            dp_config=TransformerDataParallelConfig(
+            dp_config=DataParallelConfig(
                 name=DefaultTransformerProperties.dp_type,
                 param_dtype=DType.bfloat16,
                 reduce_dtype=DType.float32,
@@ -70,7 +66,7 @@ def olmo2_core_190M(cls, dp_type: Optional[DataParallelType] = None) -> Transfor
         return getattr(TransformerConfig, "olmo2_190M")(
             vocab_size=TokenizerConfig.dolma2().padded_vocab_size(),
             compile=True,
-            dp_config=TransformerDataParallelConfig(
+            dp_config=DataParallelConfig(
                 name=dp_type if dp_type else DefaultTransformerProperties.dp_type,
                 param_dtype=DType.bfloat16,
                 reduce_dtype=DType.float32,
@@ -85,7 +81,19 @@ def olmo2_core_1B(cls, dp_type: Optional[DataParallelType] = None) -> Transforme
         return getattr(TransformerConfig, "olmo2_1B")(
             vocab_size=TokenizerConfig.dolma2().padded_vocab_size(),
             compile=True,
-            dp_config=TransformerDataParallelConfig(
+            dp_config=DataParallelConfig(
+                name=dp_type if dp_type else DefaultTransformerProperties.dp_type,
+                param_dtype=DType.bfloat16,
+                reduce_dtype=DType.float32,
+            ),
+        )
+
+    @classmethod
+    def starcoder2_3B(cls, dp_type: Optional[DataParallelType] = None) -> TransformerConfig:
+        return getattr(TransformerConfig, "starcoder2_3b")(
+            vocab_size=TokenizerConfig.dolma2().padded_vocab_size(),
+            compile=True,
+            dp_config=DataParallelConfig(
                 name=dp_type if dp_type else DefaultTransformerProperties.dp_type,
                 param_dtype=DType.bfloat16,
                 reduce_dtype=DType.float32,
@@ -100,6 +108,8 @@ def from_model_identifier(cls, model_identifier: str) -> TransformerConfig:
             return cls.olmo2_core_190M()
         elif model_identifier == "olmo2_1B":
             return cls.olmo2_core_1B()
+        elif model_identifier == "starcoder2_3b":
+            return cls.starcoder2_3B()
         else:
             raise ValueError(f"Model identifier {model_identifier} is not supported.")