#!/bin/bash
# LOOM at the 1.7B scale (L=15, hidden 1280, 30 routed + 2 shared experts, top-8).
# Paper setting: global batch 1024 x 1024 tokens, 57k steps (~60B tokens), lr 6e-5, WSD.
#
#   H=9 DATA_PATH=data/pretrain-1024 VAL_PATH=data/val-1024 bash scripts/loom_1p7b.sh
#
# MoE capacity is computed per micro-batch, so keep MICRO_BATCH_SAMPLES=8 to reproduce
# the paper runs exactly.
set -euo pipefail
export ARCH_SIZE="${ARCH_SIZE:-T_moe_L15_E32}"
export ARCH_N_LAYERS="${ARCH_N_LAYERS:-15}"
export LR="${LR:-6e-5}"
export TOTAL_STEPS="${TOTAL_STEPS:-57000}"
export LR_WARMUP_RATIO="${LR_WARMUP_RATIO:-0.01}"
export LR_DECAY_RATIO="${LR_DECAY_RATIO:-0.1}"
export LR_MIN_RATIO="${LR_MIN_RATIO:-0.1}"
export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-1024}"
export MICRO_BATCH_SAMPLES="${MICRO_BATCH_SAMPLES:-8}"
export SNAPSHOT_EVERY_STEPS="${SNAPSHOT_EVERY_STEPS:-2000}"
export SNAPSHOT_KEEP_LAST="${SNAPSHOT_KEEP_LAST:-2}"
export RUN_TAG="${RUN_TAG:-loom_1p7b_H${H:?set H (number of loops)}}"
exec bash "$(dirname "${BASH_SOURCE[0]}")/train.sh"