#!/usr/bin/env bash

export PYTHONPATH="."
export PATH="/opt/rocm-7.1.1/bin:$PATH"
export ROCM_PATH="/opt/rocm-7.1.1"
export DEV=${DEV:-AMD}
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEVICE_IN_FUNCTION_BUG=1

export DEBUG=${DEBUG:-2}
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export ALL2ALL=${ALL2ALL:-1}
export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
export USE_ATOMICS=${USE_ATOMICS:-1}
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
export MXFP8=${MXFP8:-1}
export ZERO_OPTIM=${ZERO_OPTIM:-1}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
export GBS=$((BS * GRADIENT_ACC_STEPS))

export MODEL="gptoss"
export BASEDIR="/raid/datasets/c4-8b/"
export EVAL_TARGET=3.34 EVAL_FREQ=12288
export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}

export SEED=${SEED:-5760}
export DATA_SEED=${DATA_SEED:-5760}

export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0

export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
if [ -z "$FULL_LAYERS" ]; then
  export LAYERS=${LAYERS:-2}
fi

python3 examples/mlperf/model_train.py
