强化学习训练加速
更新时间:2025-03-25
简介
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 百舸启动脚本示例
Bash
1set -x
2# 训练完成后奖励模型保存路径
3SAVE_PATH=/mnt/cluster/models/Llama-2-7b-sft-model-ocra-500k/
4# 基础模型路径
5SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
6# 用于训练奖励模型的数据集路径
7DATASET_PATH=/mnt/cluster/test/dataset/openai/webgpt_comparisons
8# tensorboard日志保存路径
9TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
10
11GPUS_PER_NODE=8
12# Change for multinode config
13MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
14MASTER_PORT=${MASTER_PORT:-29500}
15NNODES=${WORLD_SIZE:-1}
16NODE_RANK=${RANK:-0}
17
18DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
19 --nnodes $NNODES
20 --node_rank $NODE_RANK
21 --master_addr $MASTER_ADDR
22 --master_port $MASTER_PORT"
23
24export NCCL_SOCKET_IFNAME=eth0
25export NCCL_IB_HCA=mlx5
26export NCCL_IB_GID_INDEX=3
27
28read -r -d '' training_commands <<EOF
29/workspace/AIAK-RLHF/examples/train_rm.py \
30 --save_path $SAVE_PATH \
31 --save_steps -1 \
32 --logging_steps 1 \
33 --eval_steps -1 \
34 --train_batch_size 128 \
35 --micro_train_batch_size 1 \
36 --pretrain $SFT_MODEL_PATH \
37 --bf16 \
38 --max_epochs 1 \
39 --max_len 2048 \
40 --zero_stage 3 \
41 --learning_rate 9e-6 \
42 --dataset $DATASET_PATH \
43 --dataset_probs 1 \
44 --flash_attn \
45 --tensorboard $TENSORBOARD_PATH \
46 --gradient_checkpointing
47EOF
48
49GPFS="/workspace/AIAK-RLHF"
50export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
51 && 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启动脚本示例
YAML
1apiVersion: v1
2data:
3 launch.sh: |-
4 #! /bin/bash
5 set -x
6 # 训练完成后奖励模型保存路径
7 SAVE_PATH=/mnt/cluster/models/Llama-2-7b-rm-model-ocra-500k/
8 # 基础模型路径
9 SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
10 # 用于训练奖励模型的数据集路径
11 DATASET_PATH=/mnt/cluster/test/dataset/openai/webgpt_comparisons
12 # tensorboard日志保存路径
13 TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
14
15 GPUS_PER_NODE=8
16 # Change for multinode config
17 MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
18 MASTER_PORT=${MASTER_PORT:-29500}
19 NNODES=${WORLD_SIZE:-1}
20 NODE_RANK=${RANK:-0}
21
22 DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
23 --nnodes $NNODES
24 --node_rank $NODE_RANK
25 --master_addr $MASTER_ADDR
26 --master_port $MASTER_PORT"
27
28 export NCCL_SOCKET_IFNAME=eth0
29 export NCCL_IB_HCA=mlx5
30 export NCCL_IB_GID_INDEX=3
31
32 read -r -d '' training_commands <<EOF
33 /workspace/OpenRLHF/examples/train_rm.py \
34 --save_path $SAVE_PATH \
35 --save_steps -1 \
36 --logging_steps 1 \
37 --eval_steps -1 \
38 --train_batch_size 128 \
39 --micro_train_batch_size 1 \
40 --pretrain $SFT_MODEL_PATH \
41 --bf16 \
42 --max_epochs 1 \
43 --max_len 2048 \
44 --zero_stage 3 \
45 --learning_rate 9e-6 \
46 --dataset $DATASET_PATH \
47 --dataset_probs 1 \
48 --flash_attn \
49 --tensorboard $TENSORBOARD_PATH \
50 --gradient_checkpointing
51 EOF
52
53 GPFS="/workspace/OpenRLHF"
54 export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
55 && python3 -m torch.distributed.launch $DISTRIBUTED_ARGS $training_commands
56
57kind: ConfigMap
58metadata:
59 name: launch-aiak-rm-training
60 namespace: default
61---
62apiVersion: "kubeflow.org/v1"
63kind: PyTorchJob
64metadata:
65 name: aiak-rm-training
66 namespace: default
67spec:
68 pytorchReplicaSpecs:
69 Master:
70 replicas: 1
71 restartPolicy: Never
72 template:
73 metadata:
74 spec:
75 hostNetwork: true
76 dnsPolicy: ClusterFirstWithHostNet
77 schedulerName: volcano
78 containers:
79 - name: pytorch
80 image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
81 imagePullPolicy: Always
82 command:
83 - bash
84 - /workspace/launch.sh
85 env:
86 - name: NCCL_DEBUG
87 value: INFO
88 - name: NCCL_IB_DISABLE
89 value: "0"
90 resources:
91 limits:
92 nvidia.com/gpu: 8
93 rdma/hca: 1
94 securityContext:
95 allowPrivilegeEscalation: true
96 capabilities:
97 add:
98 - IPC_LOCK
99 - SYS_ADMIN
100 - NET_ADMIN
101 privileged: true
102 volumeMounts:
103 - mountPath: /dev/shm
104 name: cache-volume
105 - name: config-volume
106 mountPath: /workspace/launch.sh
107 subPath: launch.sh
108 - name: data
109 mountPath: /mnt/cluster
110 volumes:
111 - name: cache-volume
112 emptyDir:
113 medium: Memory
114 - name: config-volume
115 configMap:
116 name: launch-aiak-rm-training
117 - name: data
118 persistentVolumeClaim:
119 claimName: pvc-pfs
120 Worker:
121 replicas: 1
122 restartPolicy: Never
123 template:
124 metadata:
125 spec:
126 hostNetwork: true
127 dnsPolicy: ClusterFirstWithHostNet
128 schedulerName: volcano
129 containers:
130 - name: pytorch
131 image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
132 imagePullPolicy: Always
133 command:
134 - bash
135 - /workspace/launch.sh
136 env:
137 - name: NCCL_DEBUG
138 value: INFO
139 - name: NCCL_IB_DISABLE
140 value: "0"
141 resources:
142 limits:
143 nvidia.com/gpu: 8
144 rdma/hca: 1
145 securityContext:
146 allowPrivilegeEscalation: true
147 capabilities:
148 add:
149 - IPC_LOCK
150 - SYS_ADMIN
151 - NET_ADMIN
152 privileged: true
153 volumeMounts:
154 - mountPath: /dev/shm
155 name: cache-volume
156 - name: config-volume
157 mountPath: /workspace/launch.sh
158 subPath: launch.sh
159 - name: data
160 mountPath: /mnt/cluster
161 volumes:
162 - name: cache-volume
163 emptyDir:
164 medium: Memory
165 - name: config-volume
166 configMap:
167 name: launch-aiak-rm-training
168 - name: data
169 persistentVolumeClaim:
170 claimName: pvc-pfs
2 PPO训练
2.1 模型准备
用户可以使用自己训练好的SFT基础模型以及奖励模型,也可以用社区开源的模型。以下以社区模型作为示例。
2.1.1 SFT模型
2.1.2 奖励模型
2.2 数据准备
用户可以使用自己准备的数据集作为训练数据,也可以用社区开源的数据集。以下以社区数据集作为示例。
2.2.1 数据集
2.3 启动脚本示例
Bash
1# 根据模型大小设置TP切分,7B模型设置1就好,14B、33B等大小可以尝试改大,取值范围为1~8
2VLLM_TENSOR_PARALLEL_SIZE=1
3SAVE_PATH=/mnt/cluster/llama2-ppo-model
4# 基础模型路径
5SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
6# 奖励模型路径
7REWARD_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-rm-anthropic_hh-lmsys-oasst-webgpt/
8# 用于训练奖励模型的数据集路径
9DATASET_PATH=/mnt/cluster/test/dataset/OpenLLMAI/prompt-collection-v0.1
10# tensorboard日志保存路径
11TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
12
13GPUS_PER_NODE=8
14# Change for multinode config
15MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
16MASTER_PORT=${MASTER_PORT:-29500}
17NNODES=${WORLD_SIZE:-1}
18NODE_RANK=${RANK:-0}
19
20# 计算 VLLM_NUM_ENGINES 的值
21VLLM_NUM_ENGINES=$((GPUS_PER_NODE * NNODES / VLLM_TENSOR_PARALLEL_SIZE))
22
23if [[ $NODE_RANK -eq 0 ]];then
24 ray start --head
25ray job submit --address="http://127.0.0.1:8265" \
26 --runtime-env-json='{"working_dir": "/workspace/AIAK-RLHF"}' \
27 -- python3 examples/train_ppo_ray_optim.py \
28 --num_nodes $NNODES \
29 --num_gpus_per_node $GPUS_PER_NODE \
30 --vllm_num_engines $VLLM_NUM_ENGINES \
31 --vllm_tensor_parallel_size $VLLM_TENSOR_PARALLEL_SIZE \
32 --vllm_gpu_memory_utilization 0.9 \
33 --colocate_all \
34 --ref_reward_offload \
35 --pretrain $SFT_MODEL_PATH \
36 --reward_pretrain $REWARD_MODEL_PATH \
37 --save_path $SAVE_PATH \
38 --micro_train_batch_size 8 \
39 --micro_rollout_batch_size 16 \
40 --rollout_batch_size 1024 \
41 --max_samples 1000000 \
42 --max_epochs 1 \
43 --prompt_max_len 1024 \
44 --generate_max_len 1024 \
45 --zero_stage 3 \
46 --bf16 \
47 --actor_learning_rate 5e-7 \
48 --critic_learning_rate 9e-6 \
49 --init_kl_coef 0.01 \
50 --prompt_data $DATASET_PATH \
51 --input_key context_messages \
52 --apply_chat_template \
53 --normalize_reward \
54 --adam_offload \
55 --flash_attn \
56 --tensorboard $TENSORBOARD_PATH \
57 --gradient_checkpointing
58else
59 max_attempts=300
60 attempt=0
61 while true; do
62 if ray start --address="$MASTER_ADDR:6379" --block; then
63 echo "Connected to master at $MASTER_ADDR."
64 break
65 else
66 attempt=$((attempt+1))
67 echo "Attempt $attempt failed. Retrying in 1 second..."
68 if [[ $attempt -ge $max_attempts ]]; then
69 echo "Failed to connect to master at $MASTER_ADDR after $max_attempts attempts."
70 exit 1
71 fi
72 sleep 1
73 fi
74 done
75fi
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启动脚本
YAML
1apiVersion: v1
2data:
3 launch.sh: |-
4 #! /bin/bash
5 # 根据模型大小设置TP切分,7B模型设置1就好,14B、33B等大小可以尝试改大,取值范围为1~8
6 VLLM_TENSOR_PARALLEL_SIZE=1
7 SAVE_PATH=/mnt/cluster/llama2-ppo-model
8 # 基础模型路径
9 SFT_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-sft-model-ocra-500k/
10 # 奖励模型路径
11 REWARD_MODEL_PATH=/mnt/cluster/test/models/Llama-2-7b-rm-anthropic_hh-lmsys-oasst-webgpt/
12 # 用于训练奖励模型的数据集路径
13 DATASET_PATH=/mnt/cluster/test/dataset/OpenLLMAI/prompt-collection-v0.1
14 # tensorboard日志保存路径
15 TENSORBOARD_PATH=${AIHC_TENSORBOARD_LOG_PATH:-"/mnt/cluster/tensorboard-log"}
16
17 GPUS_PER_NODE=8
18 # Change for multinode config
19 MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
20 MASTER_PORT=${MASTER_PORT:-29500}
21 NNODES=${WORLD_SIZE:-1}
22 NODE_RANK=${RANK:-0}
23
24 # 计算 VLLM_NUM_ENGINES 的值
25 VLLM_NUM_ENGINES=$((GPUS_PER_NODE * NNODES / VLLM_TENSOR_PARALLEL_SIZE))
26
27 if [[ $NODE_RANK -eq 0 ]];then
28 ray start --head
29 ray job submit --address="http://127.0.0.1:8265" \
30 --runtime-env-json='{"working_dir": "/workspace/AIAK-RLHF"}' \
31 -- python3 examples/train_ppo_ray_optim.py \
32 --num_nodes $NNODES \
33 --num_gpus_per_node $GPUS_PER_NODE \
34 --vllm_num_engines $VLLM_NUM_ENGINES \
35 --vllm_tensor_parallel_size $VLLM_TENSOR_PARALLEL_SIZE \
36 --vllm_gpu_memory_utilization 0.9 \
37 --colocate_all \
38 --ref_reward_offload \
39 --pretrain $SFT_MODEL_PATH \
40 --reward_pretrain $REWARD_MODEL_PATH \
41 --save_path $SAVE_PATH \
42 --micro_train_batch_size 8 \
43 --micro_rollout_batch_size 16 \
44 --rollout_batch_size 1024 \
45 --max_samples 1000000 \
46 --max_epochs 1 \
47 --prompt_max_len 1024 \
48 --generate_max_len 1024 \
49 --zero_stage 3 \
50 --bf16 \
51 --actor_learning_rate 5e-7 \
52 --critic_learning_rate 9e-6 \
53 --init_kl_coef 0.01 \
54 --prompt_data $DATASET_PATH \
55 --input_key context_messages \
56 --apply_chat_template \
57 --normalize_reward \
58 --adam_offload \
59 --flash_attn \
60 --tensorboard $TENSORBOARD_PATH \
61 --gradient_checkpointing
62 else
63 max_attempts=300
64 attempt=0
65 while true; do
66 if ray start --address="$MASTER_ADDR:6379" --block; then
67 echo "Connected to master at $MASTER_ADDR."
68 break
69 else
70 attempt=$((attempt+1))
71 echo "Attempt $attempt failed. Retrying in 1 second..."
72 if [[ $attempt -ge $max_attempts ]]; then
73 echo "Failed to connect to master at $MASTER_ADDR after $max_attempts attempts."
74 exit 1
75 fi
76 sleep 1
77 fi
78 done
79 fi
80
81kind: ConfigMap
82metadata:
83 name: launch-aiak-ppo-training
84 namespace: default
85---
86apiVersion: "kubeflow.org/v1"
87kind: PyTorchJob
88metadata:
89 name: aiak-ppo-training
90 namespace: default
91spec:
92 pytorchReplicaSpecs:
93 Master:
94 replicas: 1
95 restartPolicy: Never
96 template:
97 metadata:
98 spec:
99 hostNetwork: true
100 dnsPolicy: ClusterFirstWithHostNet
101 schedulerName: volcano
102 containers:
103 - name: pytorch
104 image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
105 imagePullPolicy: Always
106 command:
107 - bash
108 - /workspace/launch.sh
109 env:
110 - name: NCCL_DEBUG
111 value: INFO
112 - name: NCCL_IB_DISABLE
113 value: "0"
114 resources:
115 limits:
116 nvidia.com/gpu: 8
117 rdma/hca: 1
118 securityContext:
119 allowPrivilegeEscalation: true
120 capabilities:
121 add:
122 - IPC_LOCK
123 - SYS_ADMIN
124 - NET_ADMIN
125 privileged: true
126 volumeMounts:
127 - mountPath: /dev/shm
128 name: cache-volume
129 - name: config-volume
130 mountPath: /workspace/launch.sh
131 subPath: launch.sh
132 - name: data
133 mountPath: /mnt/cluster
134 volumes:
135 - name: cache-volume
136 emptyDir:
137 medium: Memory
138 - name: config-volume
139 configMap:
140 name: launch-aiak-ppo-training
141 - name: data
142 persistentVolumeClaim:
143 claimName: pvc-pfs
144 Worker:
145 replicas: 1
146 restartPolicy: Never
147 template:
148 metadata:
149 spec:
150 hostNetwork: true
151 dnsPolicy: ClusterFirstWithHostNet
152 schedulerName: volcano
153 containers:
154 - name: pytorch
155 image: registry.baidubce.com/cce-ai-native/aiak-rlhf:20240702
156 imagePullPolicy: Always
157 command:
158 - bash
159 - /workspace/launch.sh
160 env:
161 - name: NCCL_DEBUG
162 value: INFO
163 - name: NCCL_IB_DISABLE
164 value: "0"
165 resources:
166 limits:
167 nvidia.com/gpu: 8
168 rdma/hca: 1
169 securityContext:
170 allowPrivilegeEscalation: true
171 capabilities:
172 add:
173 - IPC_LOCK
174 - SYS_ADMIN
175 - NET_ADMIN
176 privileged: true
177 volumeMounts:
178 - mountPath: /dev/shm
179 name: cache-volume
180 - name: config-volume
181 mountPath: /workspace/launch.sh
182 subPath: launch.sh
183 - name: data
184 mountPath: /mnt/cluster
185 volumes:
186 - name: cache-volume
187 emptyDir:
188 medium: Memory
189 - name: config-volume
190 configMap:
191 name: launch-aiak-ppo-training
192 - name: data
193 persistentVolumeClaim:
194 claimName: pvc-pfs
更新日志
v1.0.0.1
- 初始化发布
性能优化:PPO训练流程通过全混部方式进行加速,相比开源提升206%。 全混部方式只支持单个奖励模型。