Config for Amberish experiments at 1B

olmo/config.py

@@ -349,6 +349,8 @@ class ModelConfig(BaseConfig):
     to ``False``.
     """
 
+    layer_norm_eps: float = 1e-05
+
     attention_layer_norm_with_affine: bool = True
     """
     Toggle affine transform for the QK norms.

olmo/model.py

@@ -136,11 +136,10 @@ def __init__(
         *,
         size: Optional[int] = None,
         elementwise_affine: Optional[bool] = True,
-        eps: float = 1e-05,
     ):
         super().__init__()
         self.config = config
-        self.eps = eps
+        self.eps = config.layer_norm_eps
         self.normalized_shape = (size or config.d_model,)
         if elementwise_affine or (elementwise_affine is None and self.config.layer_norm_with_affine):
             self.weight = nn.Parameter(torch.ones(self.normalized_shape, device=config.init_device))
@@ -199,9 +198,8 @@ def __init__(
         size: Optional[int] = None,
         low_precision: bool = False,
         elementwise_affine: Optional[bool] = None,
-        eps: float = 1e-05,
     ):
-        super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps)
+        super().__init__(config, size=size, elementwise_affine=elementwise_affine)
         self.low_precision = low_precision
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -230,9 +228,8 @@ def __init__(
         config: ModelConfig,
         size: Optional[int] = None,
         elementwise_affine: Optional[bool] = None,
-        eps: float = 1e-5,
     ):
-        super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps)
+        super().__init__(config, size=size, elementwise_affine=elementwise_affine)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         with torch.autocast(enabled=False, device_type=x.device.type):

pyproject.toml

@@ -14,7 +14,7 @@ requires-python = ">=3.8"
 license = { file = "LICENSE" }
 dependencies = [
     "numpy",
-    "torch>=2.1,<2.3",
+    "torch>=2.1,<=2.3",
     "ai2-olmo-core==0.1.0",
     "omegaconf",
     "rich",

scripts/beaker/llamaish7-normal-launch.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+
+set -ex
+
+NUM_NODES=64
+
+gantry run \
+  --workspace ai2/OLMo-training \
+  --task-name llamaish7-normal \
+  --description "OLMo medium - 7B - Llamaish Normal" \
+  --priority urgent \
+  --preemptible \
+  --beaker-image petew/olmo-torch23-gantry \
+  --cluster ai2/jupiter-cirrascale-2 \
+  --gpus 8 \
+  --replicas "${NUM_NODES}" \
+  --leader-selection \
+  --host-networking \
+  --budget ai2/oe-training \
+  --no-nfs \
+  --propagate-failure \
+  --synchronized-start-timeout 15m \
+  --env LOG_FILTER_TYPE=local_rank0_only \
+  --env OMP_NUM_THREADS=8 \
+  --env OLMO_TASK=model \
+  --env-secret WANDB_API_KEY=AKSHITAB_WANDB_API_KEY \
+  --env-secret AWS_ACCESS_KEY_ID=AKSHITAB_AWS_ACCESS_KEY_ID \
+  --env-secret AWS_SECRET_ACCESS_KEY=AKSHITAB_AWS_SECRET_ACCESS_KEY \
+  --shared-memory 10GiB \
+  --venv base \
+  --yes \
+  --timeout=-1 \
+  -- /bin/bash -c "scripts/beaker/llamaish7-normal.sh \$BEAKER_LEADER_REPLICA_HOSTNAME ${NUM_NODES} \$BEAKER_REPLICA_RANK"

scripts/beaker/llm-360-amber1-launch.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+
+set -ex
+
+NUM_NODES=4
+
+gantry run \
+  --workspace ai2/OLMo-training \
+  --task-name amberish1-base \
+  --description "OLMo small - 1B - Amberish with Amber data" \
+  --priority urgent \
+  --preemptible \
+  --beaker-image petew/olmo-torch23-gantry \
+  --cluster ai2/jupiter-cirrascale-2 \
+  --gpus 8 \
+  --replicas "${NUM_NODES}" \
+  --leader-selection \
+  --host-networking \
+  --budget ai2/oe-training \
+  --no-nfs \
+  --weka oe-training-default:/weka/oe-training-default \
+  --propagate-failure \
+  --synchronized-start-timeout 20m \
+  --env LOG_FILTER_TYPE=local_rank0_only \
+  --env OMP_NUM_THREADS=8 \
+  --env OLMO_TASK=model \
+  --env-secret WANDB_API_KEY=DUSTINS_WANDB_API_KEY \
+  --env-secret AWS_ACCESS_KEY_ID=DUSTINS_AWS_ACCESS_KEY_ID \
+  --env-secret AWS_SECRET_ACCESS_KEY=DUSTINS_AWS_SECRET_ACCESS_KEY \
+  --shared-memory 10GiB \
+  --venv base \
+  --yes \
+  --timeout=-1 \
+  -- /bin/bash -c "scripts/beaker/llm-360-amber1.sh \$BEAKER_LEADER_REPLICA_HOSTNAME ${NUM_NODES} \$BEAKER_REPLICA_RANK"

scripts/beaker/llm-360-amber1.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+set -exuo pipefail
+IFS=$'\n\t'
+
+BEAKER_LEADER_REPLICA_HOSTNAME=$1
+shift
+
+NUM_NODES=$1
+shift
+
+BEAKER_REPLICA_RANK=$1
+shift
+
+# Warm HF cache
+mkdir -p /root/.cache
+pushd /root/.cache
+curl "https://storage.googleapis.com/hf-cache/huggingface_cache_v4.tar.gz" | tar --keep-newer-files -xzf -
+popd
+export HF_DATASETS_OFFLINE=1
+
+
+torchrun \
+  --nnodes ${NUM_NODES}:${NUM_NODES} \
+  --nproc-per-node 8 \
+  --rdzv_id=12347 \
+  --rdzv_backend=static \
+  --rdzv_endpoint=$BEAKER_LEADER_REPLICA_HOSTNAME:29400 \
+  --node_rank=$BEAKER_REPLICA_RANK \
+  --rdzv_conf="read_timeout=420" \
+  scripts/train.py \
+  configs/llm-360-amber1.yaml \
+    --gen1_gc_interval=null \
+    --save_folder=runs/ \
+    --save_interval=1000 \
+    --eval_interval=1000 \
+    --optimizer.metrics_log_interval=1 \
+    --save_overwrite \
+    --save_num_checkpoints_to_keep=3 \
+    '--load_path=s3://ai2-llm/checkpoints/OLMo-small/${run_name}/step69750/'