Skip to content

Run job

提交作业使用AI集群

原生的Slurm作业系统提供高性能的计算资源,主要用于单机多卡和多机多卡的训练任务。

运行训练和推理作业

  • AI算力平台提供了训练作业和推理作业的样例脚本,具体路径如下:
# 训练作业样例脚本
/aifs/public/data/sample_job_script/gpu_train.sh

# 推理作业样例脚本
/aifs/public/data/sample_job_script/gpu_inference.sh

脚本具体内容

/aifs/public/data/sample_job_script/gpu_train.sh

  • 第一部分是作业参数,设置请求的资源,具体可查看计算资源
  • 第二部分是工作负载,包含对环境变量的加载和实际任务的启动,根据实际情况修改
#!/bin/bash

#============== part 1 : job parameters ====================

#SBATCH --partition=gpu
#SBATCH --account=ihepai
#SBATCH --qos=gpunormal         
#SBATCH --mem-per-cpu=4GB
#SBATCH --ntasks-per-node=1
#SBATCH --nodes=1
#SBATCH --gres=gpu:1
#SBATCH --job-name=train_test
#SBATCH --comment="training"

#=================== part 2 : workload =======================
echo "slurm job ${SLURM_JOB_ID} start.."
date

echo "------------------"

# 2.1 Load env
NEW_MODULE_DIR=/cvmfs/slurm.ihep.ac.cn/alma9/modulefiles

# if MODULEPATH doesn't contain NEW_MODULE_DIR
if [[ ! "$MODULEPATH" == *"$NEW_MODULE_DIR"* ]]; then
  module use $NEW_MODULE_DIR
fi

module purge
module load llama_factory/0.9.0-py312

# 2.2 Start training (example)
OUTDIR=$HOME/saves/qwen/lora/sft
mkdir -p $OUTDIR

cd $LLAMA_HOME
llamafactory-cli train \
  --model_name_or_path /aifs/public/data/models/qwen_2.5_7B \
  --stage sft \
  --do_train true \
  --finetuning_type lora \
  --lora_target all \
  --dataset identity,alpaca_en_demo \
  --template qwen \
  --cutoff_len 1024 \
  --max_samples 1000 \
  --overwrite_cache true \
  --preprocessing_num_workers 16 \
  --output_dir $OUTDIR \
  --logging_steps 10 \
  --save_steps 500 \
  --plot_loss true \
  --overwrite_output_dir true \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --learning_rate 1.0e-4 \
  --num_train_epochs 3.0 \
  --lr_scheduler_type cosine \
  --warmup_ratio 0.1 \
  --bf16 true \
  --ddp_timeout 180000000 \
  --val_size 0.1 \
  --per_device_eval_batch_size 1 \
  --eval_strategy steps \
  --eval_steps 500

# 2.3 Finalize
echo "------------------"
echo "slurm job ${SLURM_JOB_ID} done."
date

推理作业样例脚本

/aifs/public/data/sample_job_script/gpu_inference.sh

#!/bin/bash

#============= part 1 : job parameters =========

#SBATCH --partition=gpu
#SBATCH --account=ihepai
#SBATCH --qos=gpunormal
#SBATCH --mem-per-cpu=4GB
#SBATCH --ntasks-per-node=1
#SBATCH --nodes=1
#SBATCH --gres=gpu:1
#SBATCH --job-name=inf_test
#SBATCH --comment="inference"

#============= part 2 : job workload ============
echo "slurm job ${SLURM_JOB_ID} start ..."
date

echo "----------------------------------"
# load env
NEW_MODULE_DIR=/cvmfs/slurm.ihep.ac.cn/alma9/modulefiles

# MODULEPATH doesn't contain NEW_MODULE_DIR
if [[ ! "$MODULEPATH" == *"$NEW_MODULE_DIR"* ]]; then
  module use $NEW_MODULE_DIR
fi

module purge
module load llama_factory/0.9.0-py312

# run inference
cd $LLAMA_HOME
llamafactory-cli api\
        --model_name_or_path /aifs/public/data/models/llama3_8b_instruct \
        --template llama3

echo "----------------------------------"
echo "slurm job ${SLURM_JOB_ID} done."
date