强化学习训练加速
所有文档
menu

百舸异构计算平台 AIHC

强化学习训练加速

简介

AIAK-RLHF是基于百舸异构计算平台面向大模型强化学习训练场景提供的最佳实践解决方案配套 AI 加速工具,帮助模型开发者高效完成大规模深度学习分布式训练,提升训练效率,在开源代码基础上,对ppo训练流程通过全混部方式进行加速。

  • 支持RM(奖励模型)训练
  • 支持PPO方法训练
  • 端到端性能相比开源提升206%

环境版本

  • 基础依赖
芯片 H800、A800、L20
NVIDIA Driver 450.51 (or later R450)、470.57 (or later R470)、510.47 (or later R510)、525.85 (or later R525)、535.86 (or later R535)、545.23 (or later R545)

百舸资源池资源已默认内置上述驱动,如手动操作更新过驱动需自行确认驱动版本在上述列表

  • 软件包说明
基础镜像 ubuntu22.04
框架 AIAK-RLHF
PyTorch v2.3.0
CUDA v12.1
Python v3.10

在百舸控制台进入 AI加速套件 菜单中可获取最新版本镜像

使用指南

1、奖励模型(RM)训练

1.1 模型准备

用户可以使用自己训练好的SFT基础模型以及奖励模型,也可以用社区开源的模型。以下以社区模型作为示例。

1.1.1 SFT模型

1.2 数据准备

用户可以使用自己准备的数据集作为训练数据,也可以用社区开源的数据集。以下以社区数据集作为示例。

1.2.1 数据集

1.3 百舸启动脚本示例

set -x
# 训练完成后奖励模型保存路径
SAVE_PATH=/mnt/cluster/models/Llama-2-7b-sft-model-ocra-500k/
# 基础模型路径
SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
# 用于训练奖励模型的数据集路径
DATASET_PATH=/mnt/cluster/test/dataset/openai/webgpt_comparisons
# tensorboard日志保存路径
TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
MASTER_PORT=${MASTER_PORT:-29500}
NNODES=${WORLD_SIZE:-1}
NODE_RANK=${RANK:-0}

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
                  --nnodes $NNODES
                  --node_rank $NODE_RANK
                  --master_addr $MASTER_ADDR
                  --master_port $MASTER_PORT"

export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_HCA=mlx5
export NCCL_IB_GID_INDEX=3

read -r -d '' training_commands <<EOF
/workspace/AIAK-RLHF/examples/train_rm.py \
     --save_path $SAVE_PATH \
     --save_steps -1 \
     --logging_steps 1 \
     --eval_steps -1 \
     --train_batch_size 128 \
     --micro_train_batch_size 1 \
     --pretrain $SFT_MODEL_PATH \
     --bf16 \
     --max_epochs 1 \
     --max_len 2048 \
     --zero_stage 3 \
     --learning_rate 9e-6 \
     --dataset $DATASET_PATH \
     --dataset_probs 1 \
     --flash_attn \
     --tensorboard $TENSORBOARD_PATH \
     --gradient_checkpointing
EOF

GPFS="/workspace/AIAK-RLHF"
export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
 && python3 -m torch.distributed.launch $DISTRIBUTED_ARGS $training_commands

1.4 参数说明

参数名 取值类型 是否必选 默认值 枚举值 描述
--pretrain str bigscience/bloomz-1b7 NA 用于训练奖励模型的基础模型路径
--train_batch_size int 128 NA 全局训练批次大小
--micro_train_batch_size int 8 NA 单个GPU训练批次大小
--bf16 bool False NA 使用bf16作为浮点格式,建议使用bf16。
--max_epochs int 1 NA 数据集最大训练轮次
--max_len int 512 NA 单条数据最大长度,超过会被截断
--zero_stage int 2 NA DeepSpeed ZeRO阶段
--flash_attn bool False NA 开启flash attn
--gradient_checkpointing bool False NA 开启重计算
--learning_rate float 1e-5 NA 学习率
--save_path str ./ckpt NA 训练完成后,奖励模型保存路径
--dataset str Dahoas/full-hh-rlh NA 数据集名称,可以传递多个,用逗号分隔
--dataset_probs str 1.0 NA 数据集采样概率,可以传递多个,用逗号分隔。需要跟数据集一一对应,总和为1.0。
--prompt_key str None NA 使用自定义数据集时,数据集中作为prompt的字段
--chosen_key str None NA 使用自定义数据集时,数据集中作为chosen的字段
--rejected_key str None NA 使用自定义数据集时,数据集中作为rejected的字段
--apply_chat_template bool False NA 支持 HuggingFace tokenizer.apply_chat_template 用于数据集处理
--save_steps int -1 NA 训练多少步保存一个临时checkpoint,默认不保存
--logging_steps int 1 NA 日志打印步数
--eval_steps int -1 NA 训练多少步进行一次评估,默认不评估
--tensorboard str None NA 保存tensorboard路径位置,默认不保存

1.5 yaml启动脚本示例

apiVersion: v1
data:
  launch.sh: |-
    #! /bin/bash
    set -x
    # 训练完成后奖励模型保存路径
    SAVE_PATH=/mnt/cluster/models/Llama-2-7b-rm-model-ocra-500k/
    # 基础模型路径
    SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
    # 用于训练奖励模型的数据集路径
    DATASET_PATH=/mnt/cluster/test/dataset/openai/webgpt_comparisons
    # tensorboard日志保存路径
    TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
    
    GPUS_PER_NODE=8
    # Change for multinode config
    MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
    MASTER_PORT=${MASTER_PORT:-29500}
    NNODES=${WORLD_SIZE:-1}
    NODE_RANK=${RANK:-0}
    
    DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
    		  --nnodes $NNODES
    		  --node_rank $NODE_RANK
    		  --master_addr $MASTER_ADDR
    		  --master_port $MASTER_PORT"
    
    export NCCL_SOCKET_IFNAME=eth0
    export NCCL_IB_HCA=mlx5
    export NCCL_IB_GID_INDEX=3
    
    read -r -d '' training_commands <<EOF
    /workspace/OpenRLHF/examples/train_rm.py \
         --save_path $SAVE_PATH \
         --save_steps -1 \
         --logging_steps 1 \
         --eval_steps -1 \
         --train_batch_size 128 \
         --micro_train_batch_size 1 \
         --pretrain $SFT_MODEL_PATH \
         --bf16 \
         --max_epochs 1 \
         --max_len 2048 \
         --zero_stage 3 \
         --learning_rate 9e-6 \
         --dataset $DATASET_PATH \
         --dataset_probs 1 \
         --flash_attn \
         --tensorboard $TENSORBOARD_PATH \
         --gradient_checkpointing
    EOF
    
    GPFS="/workspace/OpenRLHF"
    export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
     && python3 -m torch.distributed.launch $DISTRIBUTED_ARGS $training_commands

kind: ConfigMap
metadata:
  name: launch-aiak-rm-training
  namespace: default
---
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: aiak-rm-training
  namespace: default
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: Never
      template:
        metadata:
        spec:
          hostNetwork: true
          dnsPolicy: ClusterFirstWithHostNet
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
              imagePullPolicy: Always
              command:
                - bash
                - /workspace/launch.sh
              env:
                - name: NCCL_DEBUG
                  value: INFO
                - name: NCCL_IB_DISABLE
                  value: "0"
              resources:
                limits:
                  nvidia.com/gpu: 8
                  rdma/hca: 1
              securityContext:
                allowPrivilegeEscalation: true
                capabilities:
                  add:
                  - IPC_LOCK
                  - SYS_ADMIN
                  - NET_ADMIN
                privileged: true
              volumeMounts:
                - mountPath: /dev/shm
                  name: cache-volume
                - name: config-volume
                  mountPath: /workspace/launch.sh
                  subPath: launch.sh
                - name: data
                  mountPath: /mnt/cluster
          volumes:
            - name: cache-volume
              emptyDir:
                medium: Memory
            - name: config-volume
              configMap:
                name: launch-aiak-rm-training
            - name: data
              persistentVolumeClaim:
                claimName: pvc-pfs
    Worker:
      replicas: 1
      restartPolicy: Never
      template:
        metadata:
        spec:
          hostNetwork: true
          dnsPolicy: ClusterFirstWithHostNet
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
              imagePullPolicy: Always
              command:
                - bash
                - /workspace/launch.sh
              env:
                - name: NCCL_DEBUG
                  value: INFO
                - name: NCCL_IB_DISABLE
                  value: "0"
              resources:
                limits:
                  nvidia.com/gpu: 8
                  rdma/hca: 1
              securityContext:
                allowPrivilegeEscalation: true
                capabilities:
                  add:
                  - IPC_LOCK
                  - SYS_ADMIN
                  - NET_ADMIN
                privileged: true
              volumeMounts:
                - mountPath: /dev/shm
                  name: cache-volume
                - name: config-volume
                  mountPath: /workspace/launch.sh
                  subPath: launch.sh
                - name: data
                  mountPath: /mnt/cluster
          volumes:
            - name: cache-volume
              emptyDir:
                medium: Memory
            - name: config-volume
              configMap:
                name: launch-aiak-rm-training
            - name: data
              persistentVolumeClaim:
                claimName: pvc-pfs

2 PPO训练

2.1 模型准备

用户可以使用自己训练好的SFT基础模型以及奖励模型,也可以用社区开源的模型。以下以社区模型作为示例。

2.1.1 SFT模型
2.1.2 奖励模型

2.2 数据准备

用户可以使用自己准备的数据集作为训练数据,也可以用社区开源的数据集。以下以社区数据集作为示例。

2.2.1 数据集

2.3 启动脚本示例

# 根据模型大小设置TP切分,7B模型设置1就好,14B、33B等大小可以尝试改大,取值范围为1~8
VLLM_TENSOR_PARALLEL_SIZE=1
SAVE_PATH=/mnt/cluster/llama2-ppo-model
# 基础模型路径
SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
# 奖励模型路径
REWARD_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-rm-anthropic_hh-lmsys-oasst-webgpt/
# 用于训练奖励模型的数据集路径
DATASET_PATH=/mnt/cluster/test/dataset/OpenLLMAI/prompt-collection-v0.1
# tensorboard日志保存路径
TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
MASTER_PORT=${MASTER_PORT:-29500}
NNODES=${WORLD_SIZE:-1}
NODE_RANK=${RANK:-0}

# 计算 VLLM_NUM_ENGINES 的值
VLLM_NUM_ENGINES=$((GPUS_PER_NODE * NNODES / VLLM_TENSOR_PARALLEL_SIZE))

if [[ $NODE_RANK -eq 0 ]];then
  ray start --head
ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json='{"working_dir": "/workspace/AIAK-RLHF"}' \
    -- python3 examples/train_ppo_ray_optim.py \
    --num_nodes $NNODES \
    --num_gpus_per_node $GPUS_PER_NODE \
    --vllm_num_engines $VLLM_NUM_ENGINES \
    --vllm_tensor_parallel_size $VLLM_TENSOR_PARALLEL_SIZE \
    --vllm_gpu_memory_utilization 0.9 \
    --colocate_all \
    --ref_reward_offload \
    --pretrain $SFT_MODEL_PATH \
    --reward_pretrain $REWARD_MODEL_PATH \
    --save_path $SAVE_PATH \
    --micro_train_batch_size 8 \
    --micro_rollout_batch_size 16 \
    --rollout_batch_size 1024 \
    --max_samples 1000000 \
    --max_epochs 1 \
    --prompt_max_len 1024 \
    --generate_max_len 1024 \
    --zero_stage 3 \
    --bf16 \
    --actor_learning_rate 5e-7 \
    --critic_learning_rate 9e-6 \
    --init_kl_coef 0.01 \
    --prompt_data $DATASET_PATH \
    --input_key context_messages \
    --apply_chat_template \
    --normalize_reward \
    --adam_offload \
    --flash_attn \
    --tensorboard $TENSORBOARD_PATH \
    --gradient_checkpointing 
else
  max_attempts=300
  attempt=0
  while true; do
    if ray start --address="$MASTER_ADDR:6379" --block; then
      echo "Connected to master at $MASTER_ADDR."
      break
    else
      attempt=$((attempt+1))
      echo "Attempt $attempt failed. Retrying in 1 second..."
      if [[ $attempt -ge $max_attempts ]]; then
        echo "Failed to connect to master at $MASTER_ADDR after $max_attempts attempts."
        exit 1
      fi
      sleep 1
    fi
  done
fi

2.4 参数说明

参数名 取值类型 是否必选 默认值 枚举值 描述 备注
--pretrain str None NA 基础模型路径
--reward_pretrain str None NA 奖励模型路径
--micro_train_batch_size int 8 NA 单个GPU训练批次大小
--micro_rollout_batch_size int 8 NA 单次生成批次大小
--rollout_batch_size int 512 NA 全局生成批次大小,需要能被任务中GPU数量以及单次生成批次大小整除
--colocate_all bool True NA 所有模型混合部署在所有GPU
--bf16 bool False NA 使用bf16作为浮点格式,建议使用bf16。
--max_samples int 100000 NA 最大训练样本数,需要根据数据集大小调整
--max_epochs int 1 NA 数据集最大训练轮次
--generate_max_len int 1024 NA 训练过程中,actor模型最大生成长度
--prompt_max_len int 1024 NA 单条样本最大长度,超过会被截断
--zero_stage int 2 NA DeepSpeed ZeRO阶段,7B及以上大小模型建议开启zero3
--actor_learning_rate float 1e-6 NA actor模型学习率
--critic_learning_rate float 9e-6 NA critic模型学习率
--init_kl_coef float 0.02 NA KL散度惩罚系数
--normalize_reward bool False NA 对奖励模型输出的奖励值进行标准化处理
--save_path str ./ckpt NA 训练完成后,奖励模型保存路径
--prompt_data str Dahoas/full-hh-rlh NA 数据集名称,可以传递多个,用逗号分隔
--prompt_data_probs str 1.0 NA 数据集采样概率,可以传递多个,用逗号分隔。需要跟数据集一一对应,总和为1.0。
--input_key str None NA 使用自定义数据集时,数据集中作为input的字段
--input_template str Human: {}\nAssistant: NA 使用自定义数据集时,input的模板
--apply_chat_template bool False NA 支持 HuggingFace tokenizer.apply_chat_template 用于数据集处理
--vllm_num_engines int None NA vllm引擎数量,用于加速PPO中文本生成
--vllm_tensor_parallel_size int None NA 多GPU推理中vLLM引擎的张量并行大小
--vllm_gpu_memory_utilization float 0.9 NA 用于vllm cache的显存比例,默认即可,框架内部自适应调整
--ref_reward_offload bool False NA reference模型、reward模型卸载到内存中
--adam_offload bool False NA adam优化器卸载到内存中
--flash_attn bool False NA 开启flash attn
--gradient_checkpointing bool False NA 开启重计算
--save_steps int -1 NA 训练多少步保存一个临时checkpoint,默认不保存
--logging_steps int 1 NA 日志打印步数
--tensorboard str None NA 保存tensorboard路径位置,默认不保存

2.5 yaml启动脚本

apiVersion: v1
data:
  launch.sh: |-
    #! /bin/bash
    # 根据模型大小设置TP切分,7B模型设置1就好,14B、33B等大小可以尝试改大,取值范围为1~8
    VLLM_TENSOR_PARALLEL_SIZE=1
    SAVE_PATH=/mnt/cluster/llama2-ppo-model
    # 基础模型路径
    SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
    # 奖励模型路径
    REWARD_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-rm-anthropic_hh-lmsys-oasst-webgpt/
    # 用于训练奖励模型的数据集路径
    DATASET_PATH=/mnt/cluster/test/dataset/OpenLLMAI/prompt-collection-v0.1
    # tensorboard日志保存路径
    TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
    
    GPUS_PER_NODE=8
    # Change for multinode config
    MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
    MASTER_PORT=${MASTER_PORT:-29500}
    NNODES=${WORLD_SIZE:-1}
    NODE_RANK=${RANK:-0}
    
    # 计算 VLLM_NUM_ENGINES 的值
    VLLM_NUM_ENGINES=$((GPUS_PER_NODE * NNODES / VLLM_TENSOR_PARALLEL_SIZE))
    
    if [[ $NODE_RANK -eq 0 ]];then
      ray start --head
    ray job submit --address="http://127.0.0.1:8265" \
        --runtime-env-json='{"working_dir": "/workspace/AIAK-RLHF"}' \
        -- python3 examples/train_ppo_ray_optim.py \
        --num_nodes $NNODES \
        --num_gpus_per_node $GPUS_PER_NODE \
        --vllm_num_engines $VLLM_NUM_ENGINES \
        --vllm_tensor_parallel_size $VLLM_TENSOR_PARALLEL_SIZE \
        --vllm_gpu_memory_utilization 0.9 \
        --colocate_all \
        --ref_reward_offload \
        --pretrain $SFT_MODEL_PATH \
        --reward_pretrain $REWARD_MODEL_PATH \
        --save_path $SAVE_PATH \
        --micro_train_batch_size 8 \
        --micro_rollout_batch_size 16 \
        --rollout_batch_size 1024 \
        --max_samples 1000000 \
        --max_epochs 1 \
        --prompt_max_len 1024 \
        --generate_max_len 1024 \
        --zero_stage 3 \
        --bf16 \
        --actor_learning_rate 5e-7 \
        --critic_learning_rate 9e-6 \
        --init_kl_coef 0.01 \
        --prompt_data $DATASET_PATH \
        --input_key context_messages \
        --apply_chat_template \
        --normalize_reward \
        --adam_offload \
        --flash_attn \
        --tensorboard $TENSORBOARD_PATH \
        --gradient_checkpointing 
    else
      max_attempts=300
      attempt=0
      while true; do
        if ray start --address="$MASTER_ADDR:6379" --block; then
          echo "Connected to master at $MASTER_ADDR."
          break
        else
          attempt=$((attempt+1))
          echo "Attempt $attempt failed. Retrying in 1 second..."
          if [[ $attempt -ge $max_attempts ]]; then
            echo "Failed to connect to master at $MASTER_ADDR after $max_attempts attempts."
            exit 1
          fi
          sleep 1
        fi
      done
    fi

kind: ConfigMap
metadata:
  name: launch-aiak-ppo-training
  namespace: default
---
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: aiak-ppo-training
  namespace: default
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: Never
      template:
        metadata:
        spec:
          hostNetwork: true
          dnsPolicy: ClusterFirstWithHostNet
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
              imagePullPolicy: Always
              command:
                - bash
                - /workspace/launch.sh
              env:
                - name: NCCL_DEBUG
                  value: INFO
                - name: NCCL_IB_DISABLE
                  value: "0"
              resources:
                limits:
                  nvidia.com/gpu: 8
                  rdma/hca: 1
              securityContext:
                allowPrivilegeEscalation: true
                capabilities:
                  add:
                  - IPC_LOCK
                  - SYS_ADMIN
                  - NET_ADMIN
                privileged: true
              volumeMounts:
                - mountPath: /dev/shm
                  name: cache-volume
                - name: config-volume
                  mountPath: /workspace/launch.sh
                  subPath: launch.sh
                - name: data
                  mountPath: /mnt/cluster
          volumes:
            - name: cache-volume
              emptyDir:
                medium: Memory
            - name: config-volume
              configMap:
                name: launch-aiak-ppo-training
            - name: data
              persistentVolumeClaim:
                claimName: pvc-pfs
    Worker:
      replicas: 1
      restartPolicy: Never
      template:
        metadata:
        spec:
          hostNetwork: true
          dnsPolicy: ClusterFirstWithHostNet
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
              imagePullPolicy: Always
              command:
                - bash
                - /workspace/launch.sh
              env:
                - name: NCCL_DEBUG
                  value: INFO
                - name: NCCL_IB_DISABLE
                  value: "0"
              resources:
                limits:
                  nvidia.com/gpu: 8
                  rdma/hca: 1
              securityContext:
                allowPrivilegeEscalation: true
                capabilities:
                  add:
                  - IPC_LOCK
                  - SYS_ADMIN
                  - NET_ADMIN
                privileged: true
              volumeMounts:
                - mountPath: /dev/shm
                  name: cache-volume
                - name: config-volume
                  mountPath: /workspace/launch.sh
                  subPath: launch.sh
                - name: data
                  mountPath: /mnt/cluster
          volumes:
            - name: cache-volume
              emptyDir:
                medium: Memory
            - name: config-volume
              configMap:
                name: launch-aiak-ppo-training
            - name: data
              persistentVolumeClaim:
                claimName: pvc-pfs

更新日志

v1.0.0.1

  • 初始化发布

性能优化:PPO训练流程通过全混部方式进行加速,相比开源提升206%。 全混部方式只支持单个奖励模型。

上一篇
训练参数说明
下一篇
更新日志