Add OLMo2-26B config

Author: epwalsh

src/olmo_core/nn/transformer/config.py

@@ -409,6 +409,20 @@ def olmo2_13B(cls, vocab_size: int, **kwargs) -> "TransformerConfig":
             **kwargs,
         )
 
+    @classmethod
+    def olmo2_26B(cls, vocab_size: int, **kwargs) -> "TransformerConfig":
+        """
+        A 26B OLMo model config.
+        """
+        return cls.llama2_26B(
+            vocab_size,
+            block_name=kwargs.pop("block_name", TransformerBlockType.reordered_norm),
+            qk_norm=kwargs.pop("qk_norm", True),
+            rope_theta=kwargs.pop("rope_theta", 500_000),
+            layer_norm_eps=1e-6,
+            **kwargs,
+        )
+
     @classmethod
     def ngpt_271M(cls, vocab_size: int, **kwargs) -> "TransformerConfig":
         """

src/scripts/train/OLMo2-26B.py

@@ -0,0 +1,97 @@
+"""
+Train a 26B OLMo model. Run this script without any arguments to see usage info.
+"""
+
+import logging
+
+from olmo_core.config import DType
+from olmo_core.distributed.parallel import DataParallelType
+from olmo_core.internal.experiment import CommonComponents, main
+from olmo_core.nn.transformer import (
+    TransformerActivationCheckpointingConfig,
+    TransformerActivationCheckpointingMode,
+    TransformerConfig,
+    TransformerDataParallelConfig,
+)
+from olmo_core.optim import AdamWConfig, OptimGroupOverride
+from olmo_core.train import TrainerConfig
+from olmo_core.train.callbacks import CheckpointerCallback, CometCallback, WandBCallback
+
+log = logging.getLogger(__name__)
+
+
+def build_model_config(common: CommonComponents) -> TransformerConfig:
+    return TransformerConfig.olmo2_26B(
+        vocab_size=common.tokenizer.padded_vocab_size(),
+        compile=True,
+        dp_config=TransformerDataParallelConfig(
+            name=DataParallelType.fsdp, param_dtype=DType.bfloat16, reduce_dtype=DType.float32
+        ),
+        ac_config=TransformerActivationCheckpointingConfig(
+            mode=TransformerActivationCheckpointingMode.full
+        ),
+    )
+
+
+def build_optim_config(common: CommonComponents) -> AdamWConfig:
+    del common
+    return AdamWConfig(
+        lr=3e-4,
+        weight_decay=0.1,
+        betas=(0.9, 0.95),
+        group_overrides=[
+            OptimGroupOverride(params=["embeddings.weight"], opts=dict(weight_decay=0.0))
+        ],
+        fused=True,
+    )
+
+
+def build_trainer_config(common: CommonComponents) -> TrainerConfig:
+    return (
+        TrainerConfig(
+            save_folder=common.save_folder,
+            rank_microbatch_size=1 * 4096,
+            save_overwrite=True,
+            metrics_collect_interval=10,
+            cancel_check_interval=1,
+            z_loss_multiplier=1e-5,
+            compile_loss=True,
+        )
+        .with_callback(
+            "checkpointer",
+            CheckpointerCallback(
+                save_interval=10_000,
+                ephemeral_save_interval=250,
+                save_async=True,
+            ),
+        )
+        .with_callback(
+            "comet",
+            CometCallback(
+                name=common.run_name,
+                workspace="ai2",
+                project="OLMo-core-26B",
+                enabled=True,
+                cancel_check_interval=10,
+            ),
+        )
+        .with_callback(
+            "wandb",
+            WandBCallback(
+                name=common.run_name,
+                entity="ai2-llm",
+                project="OLMo-core-26B",
+                enabled=False,
+                cancel_check_interval=10,
+            ),
+        )
+    )
+
+
+if __name__ == "__main__":
+    main(
+        global_batch_size=2048 * 4096,
+        model_config_builder=build_model_config,
+        optim_config_builder=build_optim_config,
+        trainer_config_builder=build_trainer_config,
+    )