基于基础模型和TRL的奖励模型 训练过程加速与显存不足的解决方式

终端 使用

accelerate config

 命令对加速进行设置，主要设置显卡数和数据类型。

设置完成后使用以下命令在终端中运行 奖励模型文件：

accelerate launch \
    /root/autodl-tmp/reward_model_demo.py \
    --model_name_or_path /root/autodl-tmp/Qwen2.5-7B-Instruct \
    --dataset_name /root/autodl-tmp/Data/Data_1 \
    --output_dir /root/autodl-tmp/Qwen2.5-7B-Reward \
    --per_device_train_batch_size 1 \
    --num_train_epochs 1 \
    --gradient_checkpointing True \
    --learning_rate 1.0e-5 \
    --logging_steps 5 \
    --eval_strategy steps \
    --eval_steps 50 \
    --load_best_model_at_end True\
    --max_length 8192 \
    --fp16 \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16

 

 --num_train_epochs 1 表示使用数据训练几遍；

--learning_rate 1.0e-5；学习率，表示微分中的delta值；

--logging_steps 5，日志步长，表示训练几步后打印日志

 --eval_strategy steps \
    --eval_steps 50 \ 表示评估间隔，多少步评估一次，用于查看准确率，损失函数等

--load_best_model_at_end True\ 表示最好的模型就停止

--max_length 8192 表示数据的最大长度，需要与数据集的元素长度匹配，不然数据集就会丢失

如果从运行配置上还无法解决，就需要冻结神经网络的层数来处理，在main函数内，模型运行之前，代码如下：

 #冻结模型前30层
for name, param in model.named_parameters():
        if 'layer' in name and int(name.split('.')[2]) < 30:
           param.requires_grad = False