记录一些理解比较模糊的配置的理解,后面有完整的gsm8k_grpo配置示例。
enable_offload :是否开启参数卸载到CPU
启用 torch_memory_saver (TMS) 进行训练时的显存卸载。开启后,模型参数在不参与计算时会被卸载到 CPU 内存,只在前向/反向传播前按需加载回 GPU。这允许在显存不足的情况下训练更大的模型,但会增加 CPU-GPU 数据传输开销。
需要通过
LD_PRELOAD预加载 TMS 的hook库才能生效。当前配置为false,即参数常驻 GPU 显存。
cluster.name_resolve :分布式服务发现(name resolution)机制
在多节点训练中,各个 worker 需要互相找到彼此(如 actor 的 RPC 地址、推理引擎的端口等)。name_resolve 就是一个分布式 KV 存储,worker 启动后将自己的地址写入,其他 worker 通过 key 查询获取。
支持
nfs,etcd3,ray三种后端,目前暂时先不管,默认使用nfs。
rollout.queue_size :采样队列大小
采样任务的异步队列大小。Rollout Controller 内部维护一个 BatchTaskDispatcher,prompt 先进入队列等待推理引擎处理:
qsize = self.config.queue_size or max_concurrent_rollouts * 16设为 null 时默认为 max_concurrent_rollouts × 16。
队列太小会导致提交阻塞(训练侧等待),太大则占用更多内存。一般不需要手动设置。
这个
queue同时作为rollout引擎的input和output队列,input侧受staleness限制后,直接塞满queue,理论上不需要太大,因为rollout引擎能处理的request是有限的
output侧才是重点,如果output队列太小,rollout输出的速度跟不上训练侧的消费数据的速度,rollout数据就会阻塞在队列中。其逻辑如下
submit_task_input()
│
▼
┌─────────────────────┐
│ _pending_inputs │ 无界 deque,永远不阻塞
│ (BatchTaskDispatcher)│
└────────┬────────────┘
│ _commit_loop 线程
│ 受 StalenessManager.get_capacity() 门控
▼
┌─────────────────────┐
│ input_queue │ 有界 Queue(max_queue_size=queue_size) ← 这里可以阻塞
│ (AsyncTaskRunner) │
└────────┬────────────┘
│ asyncio 事件循环取出,创建协程
▼
[ 并发执行推理 ]
│
▼
┌─────────────────────┐
│ output_queue │ 有界 Queue(max_queue_size=queue_size) ← 溢出会崩溃!
│ (AsyncTaskRunner) │
└────────┬────────────┘
│ _fetch_loop 线程消费
▼
┌─────────────────────┐
│ _pending_results │ 无界 dict,训练侧 wait_results() 取走
└─────────────────────┘
rollout.consumer_batch_size:消费者批大小
训练侧从采样队列中消费结果的批大小,即一次凑齐多少条采样结果后才交给训练 step。
在我的理解里,这个应该要和
batch_size相同的
gconfig.n_samples:每个prompt采样次数
实际上就是GRPO的group_size。
gconfig.greedy :是否贪心解码
贪心解码:每步都选概率最高的token,此时temperature/top_p 等采样参数无效,输出在理论上完全确定
在GRPO中必须为
false,不然组内所有sample都完全相同。
actor.mb_spec_max_tokens_per_mb :单次前向/反向的显存峰值
它实际是micro_batch的最大tokens数,AReaL的micro_bacth基于这个值按token划分
micro_batch实际上就是使用梯度累积的结果,直接更新一整个batch的话显存压力太大,于是将batch划分为多个micro_batch,然后累积更新梯度
AReaL是按token数划分而不是按样本数划分,一定程度上更合理,可以想象是在通过这个值调控一次前向/反向的峰值显存,如果OOM了就尝试调小吧。
actor.use_decoupled_loss :是否使用decoupled_loss
AReaL的核心设计之一,由于完全异步的rollout,一个Batch中可能混杂不同版本的数据,decoupled_loss会将trust region中心换到更近的版本。
同时,配合behave_imp_weight_cap 限制重要性权重的上限,避免更新太多
rollout.max_head_offpolicyness :推理引擎最大的版本差异数
实际上就是论文中的max_staleness,如果推理引擎目前与训练引擎的策略版本差异大于这个值,推理引擎就会停下来等训练引擎。
sglang.mem_fraction_static :推理引擎占用GPU显存的比例
如果这个值太大,那留给KV Cache的显存就少了(这个值不包含KV Cache?)
下面是完整的examples/math/gsm8k_grpo.yaml配置:
# ===== 基本信息 =====
experiment_name: gsm8k-grpo # 实验名称
trial_name: trial0 # 试验名称
seed: 1 # 随机种子
enable_offload: false # 是否开启参数卸载到CPU
total_train_epochs: 10 # 总训练轮数
tokenizer_path: ${actor.path} # 分词器路径,复用actor模型路径
# ===== 集群配置 =====
cluster:
n_nodes: 1 # 节点数
n_gpus_per_node: 8 # 每节点GPU数
fileroot: /tmp/areal/experiments # 实验文件根目录
name_resolve:
type: nfs # 服务发现方式:NFS共享文件
nfs_record_root: /tmp/areal/name_resolve
# GPU分配:sglang推理用4卡,训练actor和ref各用4卡(数据并行1,流水线并行1,张量并行1)
allocation_mode: sglang:d4p1t1+d4p1t1
scheduler:
type: null # 不使用集群调度器(本地运行)
# ===== Rollout(采样)配置 =====
rollout:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
max_concurrent_rollouts: 256 # 最大并发采样数
queue_size: null # 采样队列大小,null为不限
consumer_batch_size: ${train_dataset.batch_size} # 消费者批大小
max_head_offpolicyness: 2 # 最大off-policy程度(允许落后训练步数)
enable_rollout_tracing: false # 是否记录采样trace
scheduling_spec: ${actor.scheduling_spec}
fileroot: ${cluster.fileroot}
tokenizer_path: ${tokenizer_path}
dump_to_file: true # 是否将采样结果写入文件
# ===== 生成配置 =====
gconfig:
n_samples: 4 # 每个prompt采样次数
min_new_tokens: 0 # 最少生成token数
max_new_tokens: 1024 # 最多生成token数
greedy: false # 是否贪心解码
temperature: 1.0 # 采样温度
# ===== Actor(策略模型)配置 =====
actor:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
path: Qwen/Qwen2.5-1.5B-Instruct # 模型路径
init_from_scratch: false # 是否从头初始化(否则加载预训练权重)
disable_dropout: true # 关闭dropout(RL训练常用)
gradient_checkpointing: true # 梯度检查点,节省显存
dtype: bfloat16 # 计算精度
mb_spec:
max_tokens_per_mb: 10240 # 每个micro-batch的最大token数
optimizer:
type: adam
lr: 1.70e-5 # 学习率
weight_decay: 0.017 # 权重衰减
beta1: 0.9
beta2: 0.999
eps: 1e-8
lr_scheduler_type: constant # 学习率调度:恒定
gradient_clipping: 1.0 # 梯度裁剪阈值
warmup_steps_proportion: 0.001 # 预热步数比例
eps_clip: 0.4 # PPO/GRPO裁剪范围
temperature: ${gconfig.temperature} # 采样温度
reward_scaling: 10.0 # 奖励缩放系数
reward_bias: -0.5 # 奖励偏置
kl_ctl: 0.0 # KL散度惩罚系数(0表示不使用)
ppo_n_minibatches: 1 # PPO mini-batch数
recompute_logprob: true # 训练时重新计算logprob
use_decoupled_loss: true # 使用解耦损失
behave_imp_weight_cap: 5.0 # 行为重要性权重上限
reward_norm: # 奖励归一化
mean_level: group # 均值归一化级别:组内(同prompt的多个采样)
std_level: group # 标准差归一化级别:组内
group_size: ${gconfig.n_samples}
adv_norm: # 优势函数归一化
mean_level: batch # 均值归一化级别:批次
std_level: batch # 标准差归一化级别:批次
weight_update_mode: xccl # 权重更新方式:集合通信
max_new_tokens: ${gconfig.max_new_tokens}
scheduling_spec: # Worker调度规格
- task_type: worker
port_count: 2
gpu: 1
mem: 32
cmd: python3 -m areal.infra.rpc.rpc_server
env_vars: {}
# ===== Ref(参考模型)配置 =====
ref:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
path: ${actor.path} # 与actor使用相同模型
init_from_scratch: false
disable_dropout: true
dtype: ${actor.dtype}
mb_spec:
max_tokens_per_mb: 10240
optimizer: null # 参考模型不需要优化器(不更新权重)
scheduling_strategy:
type: colocation # 与actor共用GPU
target: actor
scheduling_spec: ${actor.scheduling_spec}
# ===== SGLang 推理引擎配置 =====
sglang:
model_path: ${actor.path}
random_seed: ${seed}
skip_tokenizer_init: true # 跳过分词器初始化(由外部管理)
dtype: ${actor.dtype}
max_running_requests: null # 最大并发请求数,null为自动
context_length: 32768 # 上下文窗口长度
mem_fraction_static: 0.8 # 静态显存占比
# ===== vLLM 推理引擎配置(备选) =====
vllm:
model: ${actor.path}
seed: ${seed}
skip_tokenizer_init: false
dtype: ${actor.dtype}
max_model_len: 32768 # 最大模型长度
gpu_memory_utilization: 0.8 # GPU显存利用率
# ===== 训练数据集 =====
train_dataset:
batch_size: 256 # 批大小
shuffle: true # 是否打乱
pin_memory: true # 固定内存(加速数据传输)
num_workers: 4 # 数据加载线程数
path: openai/gsm8k # 数据集路径(HuggingFace)
type: rl # 数据集类型:强化学习
max_length: 1024 # 最大序列长度
# ===== 验证数据集 =====
valid_dataset:
batch_size: 256
pin_memory: true
num_workers: 4
path: openai/gsm8k
type: rl
# ===== 模型保存 =====
saver:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
fileroot: ${cluster.fileroot}
freq_epochs: 1 # 每1个epoch保存一次
freq_steps: null # 不按步数保存
freq_secs: null # 不按时间保存
# ===== 断点恢复 =====
recover:
mode: disabled # 关闭断点恢复
experiment_name: ${experiment_name}
trial_name: ${trial_name}
fileroot: ${cluster.fileroot}
freq_epochs: 1
freq_steps: null
freq_secs: 3600 # 每3600秒保存恢复点
# ===== 评估 =====
evaluator:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
fileroot: ${cluster.fileroot}
freq_epochs: 1 # 每1个epoch评估一次
freq_steps: null
freq_secs: null
# ===== 统计日志 =====
stats_logger:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
fileroot: ${cluster.fileroot}
wandb:
mode: disabled # 关闭wandb日志
# ===== 性能追踪 =====
perf_tracer:
experiment_name: ${experiment_name}
trial_name: ${trial_name}
fileroot: ${cluster.fileroot}
enabled: false # 关闭性能追踪
session_tracer:
enabled: false