#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
DATA_PREFIX="${DATA_PREFIX:-}"
DATA_MANIFEST="${DATA_MANIFEST-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_65536_direct_v1/manifest.json}"
TRAIN_ITERS="${TRAIN_ITERS:-10}"
SEQ_LENGTH="${SEQ_LENGTH:-8192}"
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-1024}"
TP="${TP:-1}"
PP="${PP:-1}"
EP="${EP:-1}"
CP="${CP:-1}"
DRY_RUN="${DRY_RUN:-0}"

if [[ -n "$DATA_MANIFEST" ]]; then
  if [[ ! -f "$DATA_MANIFEST" ]]; then
    echo "missing Megatron indexed dataset manifest: $DATA_MANIFEST" >&2
    echo "run scripts/preprocess_megatron_bridge_pretrain_direct.sh first, or set DATA_MANIFEST" >&2
    exit 1
  fi
elif [[ "$DRY_RUN" != "1" && ( -z "$DATA_PREFIX" || ! -f "${DATA_PREFIX}.idx" ) ]]; then
  echo "missing Megatron indexed data prefix: ${DATA_PREFIX:-<empty>}" >&2
  echo "set DATA_MANIFEST or DATA_PREFIX" >&2
  exit 1
fi

DRY_RUN_ARG=""
if [[ "$DRY_RUN" == "1" ]]; then
  DRY_RUN_ARG="--dry-run"
fi

DATA_ARGS=()
if [[ -n "$DATA_MANIFEST" ]]; then
  DATA_ARGS=(--data-manifest "$DATA_MANIFEST")
elif [[ -n "$DATA_PREFIX" ]]; then
  DATA_ARGS=(--data-prefix "$DATA_PREFIX")
fi

DOCKER_MOUNTS=()
if [[ -d /mnt/beegfs ]]; then
  DOCKER_MOUNTS+=(-v /mnt/beegfs:/mnt/beegfs)
fi
if [[ -d /ssd ]]; then
  DOCKER_MOUNTS+=(-v /ssd:/ssd)
fi

docker run --rm --gpus all --ipc=host --network=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  "${DOCKER_MOUNTS[@]}" \
  -w "$REPO_ROOT" \
  "$IMAGE" \
  bash -lc "
    set -euo pipefail
    torchrun --nproc_per_node='$NPROC_PER_NODE' \
      training/megatron_bridge/laoyao_2b_moe_pretrain.py \
      ${DATA_ARGS[*]} \
      --seq-length '$SEQ_LENGTH' \
      --train-iters '$TRAIN_ITERS' \
      --micro-batch-size '$MICRO_BATCH_SIZE' \
      --global-batch-size '$GLOBAL_BATCH_SIZE' \
      --tensor-parallel '$TP' \
      --pipeline-parallel '$PP' \
      --expert-parallel '$EP' \
      --context-parallel '$CP' \
      $DRY_RUN_ARG
  "