blob: 9a5fa8523680cd061752e94dd6e330a616b22096 (
plain)
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
|
#!/bin/bash
# One-shot GEE训练脚本
# 使用方法: bash scripts/train_one_shot_gee.sh
echo "开始One-shot GEE训练..."
# 设置环境变量
export CUDA_VISIBLE_DEVICES=0
export NCCL_TIMEOUT=2700
export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=2700
# 训练参数
MODEL_NAME="Qwen2.5-Math-7B"
MODEL_PATH="/volume/pt-train/models/Qwen2.5-Math-7B" # 请根据实际路径修改
TRAIN_DATA="dataset/1shot_rlvr/pi1_r1280.parquet"
RUN_NAME="one_shot_gee_lambda3"
WANDB_PROJECT="one-shot-gee"
# 检查模型路径
if [ ! -d "$MODEL_PATH" ]; then
echo "错误: 模型路径不存在: $MODEL_PATH"
echo "请修改脚本中的MODEL_PATH变量"
exit 1
fi
# 检查训练数据
if [ ! -f "$TRAIN_DATA" ]; then
echo "错误: 训练数据文件不存在: $TRAIN_DATA"
echo "请检查数据文件路径"
exit 1
fi
echo "模型路径: $MODEL_PATH"
echo "训练数据: $TRAIN_DATA"
echo "运行名称: $RUN_NAME"
# 开始训练
accelerate launch train_gee.py \
--model_name $MODEL_NAME \
--model_path $MODEL_PATH \
--train_data $TRAIN_DATA \
--effective_batch 64 \
--micro_batch_size 2 \
--temperature 0.5 \
--learning_rate 2e-5 \
--max_steps 50 \
--lambda_weight 3.0 \
--auto_anneal \
--balance_dataset \
--log_steps 1 \
--save_steps 1 \
--run_name $RUN_NAME \
--wandb_project $WANDB_PROJECT
echo "训练完成！"
|
