117 lines
4.6 KiB
Bash
Executable File
117 lines
4.6 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
cd "${ROOT_DIR}"
|
|
|
|
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
|
|
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
|
|
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
|
|
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
|
|
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
|
|
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
|
|
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
|
|
|
|
if [[ -f .venv/bin/activate ]]; then
|
|
source .venv/bin/activate
|
|
elif [[ "${DRY_RUN:-0}" != "1" ]]; then
|
|
echo "Missing .venv. Run ./scripts/setup_env.sh first." >&2
|
|
exit 2
|
|
fi
|
|
|
|
BEEGFS_ROOT="${BEEGFS_ROOT:-/mnt/beegfs/workspace/ti_coding_agent_probe}"
|
|
export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-${BEEGFS_ROOT}/modelscope_cache}"
|
|
MEGATRON_OUTPUT_ROOT="${MEGATRON_OUTPUT_ROOT:-${BEEGFS_ROOT}/megatron_outputs}"
|
|
MEGATRON_LOG_ROOT="${MEGATRON_LOG_ROOT:-${BEEGFS_ROOT}/megatron_logs}"
|
|
mkdir -p "${MODELSCOPE_CACHE}" "${MEGATRON_OUTPUT_ROOT}" "${MEGATRON_LOG_ROOT}" logs
|
|
|
|
TRAIN_DATASET="${TRAIN_DATASET:-${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}}"
|
|
VAL_DATASET="${VAL_DATASET:-${VAL_JSONL:-data/processed/training_probe/validation.jsonl}}"
|
|
MEGATRON_MODEL="${MEGATRON_MODEL:-${QWEN36_27B_MODEL_ID:-Qwen/Qwen3.6-27B}}"
|
|
MEGATRON_RUN_NAME="${MEGATRON_RUN_NAME:-qwen36_27b_megatron_full}"
|
|
MEGATRON_OUTPUT_DIR="${MEGATRON_OUTPUT_DIR:-${MEGATRON_OUTPUT_ROOT}/${MEGATRON_RUN_NAME}}"
|
|
|
|
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
|
|
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
|
|
TENSOR_MODEL_PARALLEL_SIZE="${TENSOR_MODEL_PARALLEL_SIZE:-4}"
|
|
PIPELINE_MODEL_PARALLEL_SIZE="${PIPELINE_MODEL_PARALLEL_SIZE:-1}"
|
|
CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE:-1}"
|
|
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
|
|
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
|
|
MAX_LENGTH="${MAX_LENGTH:-262144}"
|
|
LR="${LR:-1e-5}"
|
|
MIN_LR="${MIN_LR:-1e-6}"
|
|
LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-0.1}"
|
|
NUM_EPOCHS="${NUM_EPOCHS:-1}"
|
|
SAVE_STEPS="${SAVE_STEPS:-1000}"
|
|
EVAL_STEPS="${EVAL_STEPS:-1000}"
|
|
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
|
|
DATASET_NUM_PROC="${DATASET_NUM_PROC:-4}"
|
|
RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}"
|
|
RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}"
|
|
RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}"
|
|
CROSS_ENTROPY_LOSS_FUSION="${CROSS_ENTROPY_LOSS_FUSION:-true}"
|
|
GRADIENT_ACCUMULATION_FUSION="${GRADIENT_ACCUMULATION_FUSION:-false}"
|
|
SAVE_SAFETENSORS="${SAVE_SAFETENSORS:-true}"
|
|
SEQUENCE_PARALLEL="${SEQUENCE_PARALLEL:-true}"
|
|
|
|
require_file() {
|
|
if [[ ! -f "$1" ]]; then
|
|
echo "Missing required file: $1" >&2
|
|
exit 2
|
|
fi
|
|
}
|
|
|
|
run_megatron_sft() {
|
|
require_file "${TRAIN_DATASET}"
|
|
require_file "${VAL_DATASET}"
|
|
mkdir -p "${MEGATRON_OUTPUT_DIR}"
|
|
|
|
local log_file="${MEGATRON_LOG_ROOT}/${MEGATRON_RUN_NAME}.log"
|
|
local cmd=(
|
|
megatron sft
|
|
--model "${MEGATRON_MODEL}"
|
|
--save_safetensors "${SAVE_SAFETENSORS}"
|
|
--dataset "${TRAIN_DATASET}"
|
|
--val_dataset "${VAL_DATASET}"
|
|
--tensor_model_parallel_size "${TENSOR_MODEL_PARALLEL_SIZE}"
|
|
--pipeline_model_parallel_size "${PIPELINE_MODEL_PARALLEL_SIZE}"
|
|
--context_parallel_size "${CONTEXT_PARALLEL_SIZE}"
|
|
--sequence_parallel "${SEQUENCE_PARALLEL}"
|
|
--micro_batch_size "${MICRO_BATCH_SIZE}"
|
|
--global_batch_size "${GLOBAL_BATCH_SIZE}"
|
|
--recompute_granularity "${RECOMPUTE_GRANULARITY}"
|
|
--recompute_method "${RECOMPUTE_METHOD}"
|
|
--recompute_num_layers "${RECOMPUTE_NUM_LAYERS}"
|
|
--finetune true
|
|
--cross_entropy_loss_fusion "${CROSS_ENTROPY_LOSS_FUSION}"
|
|
--gradient_accumulation_fusion "${GRADIENT_ACCUMULATION_FUSION}"
|
|
--lr "${LR}"
|
|
--lr_warmup_fraction "${LR_WARMUP_FRACTION}"
|
|
--min_lr "${MIN_LR}"
|
|
--num_train_epochs "${NUM_EPOCHS}"
|
|
--output_dir "${MEGATRON_OUTPUT_DIR}"
|
|
--save_steps "${SAVE_STEPS}"
|
|
--eval_steps "${EVAL_STEPS}"
|
|
--max_length "${MAX_LENGTH}"
|
|
--dataloader_num_workers "${DATALOADER_NUM_WORKERS}"
|
|
--dataset_num_proc "${DATASET_NUM_PROC}"
|
|
--no_save_optim true
|
|
--no_save_rng true
|
|
--model_author swift
|
|
--model_name ti-coding-agent
|
|
)
|
|
|
|
printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q MODELSCOPE_CACHE=%q ' \
|
|
"${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" "${MODELSCOPE_CACHE}" | tee "${log_file}.cmd"
|
|
printf '%q ' "${cmd[@]}" | tee -a "${log_file}.cmd"
|
|
echo | tee -a "${log_file}.cmd"
|
|
if [[ "${DRY_RUN:-0}" == "1" ]]; then
|
|
return 0
|
|
fi
|
|
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
|
|
NPROC_PER_NODE="${NPROC_PER_NODE}" \
|
|
MODELSCOPE_CACHE="${MODELSCOPE_CACHE}" \
|
|
"${cmd[@]}" 2>&1 | tee "${log_file}"
|
|
}
|