【AReaL】GRPO配置解析

【AReaL】GRPO配置解析

记录一些理解比较模糊的配置的理解,后面有完整的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引擎的inputoutput队列,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

LICENSED UNDER CC BY-NC-SA 4.0