Normal baselines

scripts/beaker/debug/llamaish7-normal-launch.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+set -ex
+
+NUM_NODES=32
+
+gantry run \
+  --workspace ai2/OLMo-training \
+  --task-name llamaish7-normal-spike-debug \
+  --description "OLMo medium - 7B - Llamaish Normal Spike Debug" \
+  --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 60m \
+  --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 \
+  --env R2_PROFILE=R2 \
+  --env S3_PROFILE=S3 \
+  --env WEKA_PROFILE=WEKA \
+  --env-secret AWS_CONFIG=PETEW_AWS_CONFIG \
+  --env-secret AWS_CREDENTIALS=PETEW_AWS_CREDENTIALS \
+  --env-secret R2_ENDPOINT_URL=R2_ENDPOINT_URL \
+  --env-secret WEKA_ENDPOINT_URL=WEKA_ENDPOINT_URL \
+  --shared-memory 10GiB \
+  --venv base \
+  --yes \
+  --timeout=-1 \
+  -- /bin/bash -c "scripts/beaker/debug/llamaish7-normal.sh \$BEAKER_LEADER_REPLICA_HOSTNAME ${NUM_NODES} \$BEAKER_REPLICA_RANK"

scripts/beaker/debug/llamaish7-normal.sh

@@ -0,0 +1,54 @@
+#!/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
+
+# Move AWS credentials from env to relevant files
+mkdir -p ~/.aws
+printenv AWS_CONFIG > ~/.aws/config
+printenv AWS_CREDENTIALS > ~/.aws/credentials
+
+export EXPERIMENT=llamaish7-normal-final-spike-rerun-2
+
+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/llamaish7-normal-weka.yaml \
+    --run_name=$EXPERIMENT \
+    --wandb.name=$EXPERIMENT \
+    --wandb.group=$EXPERIMENT \
+    --fsdp.wrapping_strategy=by_block_and_size \
+    --fsdp.sharding_strategy=SHARD_GRAD_OP \
+    --save_folder=runs/ \
+    --activation_checkpointing=fine_grained \
+    --device_train_microbatch_size=2 \
+    --global_train_batch_size=1024 \
+    --save_interval=250 \
+    --eval_interval=250 \
+    --optimizer.metrics_log_interval=1 \
+    --save_overwrite \
+    --save_num_checkpoints_to_keep=3 \
+    --data.num_workers=64 \
+    --fast_forward_batches=1 \
+    --load_path=s3://ai2-llm/checkpoints/OLMo-medium/llamaish7-normal-final/step96750

scripts/beaker/llamaish1-normal-launch.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+set -ex
+
+NUM_NODES=8
+
+gantry run \
+  --workspace ai2/OLMo-training \
+  --task-name llamaish1-normal \
+  --description "OLMo small - 1B - Llamaish Normal Weka" \
+  --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=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 \
+  --env R2_PROFILE=R2 \
+  --env S3_PROFILE=S3 \
+  --env WEKA_PROFILE=WEKA \
+  --env-secret AWS_CONFIG=PETEW_AWS_CONFIG \
+  --env-secret AWS_CREDENTIALS=PETEW_AWS_CREDENTIALS \
+  --env-secret R2_ENDPOINT_URL=R2_ENDPOINT_URL \
+  --env-secret WEKA_ENDPOINT_URL=WEKA_ENDPOINT_URL \
+  --shared-memory 10GiB \
+  --venv base \
+  --yes \
+  --timeout=-1 \
+  -- /bin/bash -c "scripts/beaker/llamaish1-normal.sh \$BEAKER_LEADER_REPLICA_HOSTNAME ${NUM_NODES} \$BEAKER_REPLICA_RANK"

scripts/beaker/llamaish1-normal.sh

@@ -0,0 +1,53 @@
+#!/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
+
+# Move AWS credentials from env to relevant files
+mkdir -p ~/.aws
+printenv AWS_CONFIG > ~/.aws/config
+printenv AWS_CREDENTIALS > ~/.aws/credentials
+
+
+export EXPERIMENT=llamaish1-normal-final
+
+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/llamaish1-normal-weka.yaml \
+    --run_name=$EXPERIMENT \
+    --wandb.name=$EXPERIMENT \
+    --wandb.group=$EXPERIMENT \
+    --fsdp.wrapping_strategy=by_block_and_size \
+    --fsdp.sharding_strategy=SHARD_GRAD_OP \
+    --save_folder=runs/ \
+    --device_train_microbatch_size=4 \
+    --global_train_batch_size=512 \
+    --save_interval=250 \
+    --eval_interval=250 \
+    --optimizer.metrics_log_interval=1 \
+    --save_overwrite \
+    --save_num_checkpoints_to_keep=3 \
+    '--load_path=${path.last_checkpoint:s3://ai2-llm/checkpoints/OLMo-small/llamaish1-normal-final/}'
+    #--load_path=s3://ai2-llm/checkpoints/OLMo-small/llamaish1-normal-shard/step2000

scripts/beaker/llamaish7-normal-launch.sh

@@ -2,22 +2,23 @@
 
 set -ex
 
-NUM_NODES=64
+NUM_NODES=16
 
 gantry run \
   --workspace ai2/OLMo-training \
-  --task-name llamaish7-normal-qk-norm-reorder-zloss \
+  --task-name llamaish7-normal \
   --description "OLMo medium - 7B - Llamaish Normal" \
   --priority urgent \
   --preemptible \
-  --beaker-image shanea/olmo-torch2.3-gantry \
+  --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 15m \
   --env LOG_FILTER_TYPE=local_rank0_only \
@@ -26,6 +27,13 @@ gantry run \
   --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 \
+  --env R2_PROFILE=R2 \
+  --env S3_PROFILE=S3 \
+  --env WEKA_PROFILE=WEKA \
+  --env-secret AWS_CONFIG=PETEW_AWS_CONFIG \
+  --env-secret AWS_CREDENTIALS=PETEW_AWS_CREDENTIALS \
+  --env-secret R2_ENDPOINT_URL=R2_ENDPOINT_URL \
+  --env-secret WEKA_ENDPOINT_URL=WEKA_ENDPOINT_URL \
   --shared-memory 10GiB \
   --venv base \
   --yes \

scripts/beaker/llamaish7-normal.sh

@@ -14,11 +14,20 @@ shift
 # Warm HF cache
 mkdir -p /root/.cache
 pushd /root/.cache
-curl "https://storage.googleapis.com/dirkgr-public/huggingface_cache_v3.tar.gz" | tar --keep-newer-files -xzf -
+curl "https://storage.googleapis.com/hf-cache/huggingface_cache_v4.tar.gz" | tar --keep-newer-files -xzf -
 popd
 export HF_DATASETS_OFFLINE=1
 
-export EXPERIMENT=llamaish7-normal
+# Move AWS credentials from env to relevant files
+mkdir -p ~/.aws
+printenv AWS_CONFIG > ~/.aws/config
+printenv AWS_CREDENTIALS > ~/.aws/credentials
+
+export EXPERIMENT=llamaish7-normal-final
+# export NCCL_IB_HCA=^mlx5_bond
+export NCCL_DEBUG=TRACE
+export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1
+
 
 torchrun \
   --nnodes ${NUM_NODES}:${NUM_NODES} \
@@ -29,26 +38,21 @@ torchrun \
   --node_rank=$BEAKER_REPLICA_RANK \
   --rdzv_conf="read_timeout=420" \
   scripts/train.py \
-  configs/llamaish7-s3.yaml \
+  configs/llamaish7-normal-weka.yaml \
     --run_name=$EXPERIMENT \
     --wandb.name=$EXPERIMENT \
     --wandb.group=$EXPERIMENT \
-    --model.flash_attention=true \
     --fsdp.wrapping_strategy=by_block_and_size \
     --fsdp.sharding_strategy=SHARD_GRAD_OP \
     --save_folder=runs/ \
     --activation_checkpointing=fine_grained \
-    --fused_loss=true \
     --device_train_microbatch_size=2 \
     --global_train_batch_size=1024 \
     --save_interval=250 \
     --eval_interval=250 \
     --optimizer.metrics_log_interval=1 \
     --save_overwrite \
-    --model.init_fn=normal \
-    --model.init_std=0.02 \
-    --model.clip_qkv=null \
     --save_num_checkpoints_to_keep=3 \
-    --scheduler.units=steps \
-    --scheduler.t_warmup=2000
-    # '--load_path=${path.last_checkpoint:s3://ai2-llm/checkpoints/OLMo-medium/llamaish7-normal/}'
+    --data.num_workers=64 \
+    '--load_path=${path.last_checkpoint:s3://ai2-llm/checkpoints/OLMo-medium/llamaish7-normal-final/}'
+    #--load_path=s3://ai2-llm/checkpoints/OLMo-medium/llamaish7-normal/step2000