File size: 2,097 Bytes
535fb25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | #!/usr/bin/env bash
set -euo pipefail
# 4 GPU radiomics OmniGen LoRA training.
# Override config/port if needed:
# CONFIG_ENTRY=config/grpo.py:general_radiomics_omnigen_4gpu_kl MAIN_PROCESS_PORT=29502 bash scripts/single_node/grpo_4gpu.sh
PROJECT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/../.." && pwd)
export PROJECT_ROOT
WANDB_ENV_FILE="${WANDB_ENV_FILE:-$HOME/.wandb_env}"
if [[ -f "$WANDB_ENV_FILE" ]]; then
# shellcheck disable=SC1090
source "$WANDB_ENV_FILE"
fi
export OMNIGEN_CODE_ROOT=${OMNIGEN_CODE_ROOT:-/home/wenting/gen_joint}
if [ ! -f "$OMNIGEN_CODE_ROOT/OmniGen/__init__.py" ]; then
echo "Missing OmniGen package at $OMNIGEN_CODE_ROOT/OmniGen/__init__.py" >&2
exit 1
fi
export PYTHONPATH="$OMNIGEN_CODE_ROOT:$PROJECT_ROOT:${PYTHONPATH:-}"
export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
export HF_HOME="${HF_HOME:-/home/wenting/.cache/huggingface}"
export HF_HUB_CACHE="${HF_HUB_CACHE:-/home/wenting/.cache/huggingface/hub}"
mkdir -p "$HF_HUB_CACHE"
export WANDB_MODE="${WANDB_MODE:-online}"
if [[ "$WANDB_MODE" == "online" && -z "${WANDB_API_KEY:-}" ]]; then
echo "[ERROR] WANDB_MODE=online but WANDB_API_KEY is not set." >&2
echo "Please export WANDB_API_KEY or put it in $WANDB_ENV_FILE." >&2
exit 1
fi
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
export RADIOMICS_REWARD_FAIL_FAST=${RADIOMICS_REWARD_FAIL_FAST:-1}
export SAVE_FREQ="${SAVE_FREQ:-5}"
if [[ "${SKIP_PREFLIGHT:-0}" != "1" ]]; then
python3 "$PROJECT_ROOT/scripts/debug_preflight_omnigen_env.py"
fi
if [[ "${SKIP_REWARD_SMOKE:-0}" != "1" ]]; then
python3 "$PROJECT_ROOT/scripts/debug_radiomics_reward_smoke.py" \
--config "${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_4gpu_kl}"
fi
python3 -m accelerate.commands.launch \
--config_file scripts/accelerate_configs/multi_gpu.yaml \
--num_processes=${NUM_PROCESSES:-4} \
--main_process_port ${MAIN_PROCESS_PORT:-29501} \
scripts/train_omnigen.py \
--config "${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_4gpu_kl}" \
--config.save_freq="${SAVE_FREQ}"
|