[Feat] Add off-policy masking for partial rollouts - #2003
Conversation
|
@claude review |
1 similar comment
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Claude: ## Summary Adds off-policy masking for partial rollouts: response tokens from prior partial-rollout rounds are masked out of the loss via IssuesCritical
Warning
Nit
ProduceBatchResult impactNot directly affected — no fields on RoutedExperts impactThe off-policy masking path replaces routed_experts instead of concatenating. Old history refs are freed via Ray concurrency impactNot affected — no changes to Ray actor methods, decorators, or concurrency groups. VerdictREQUEST_CHANGES — the missing |
|
@claude review |
|
Claude: ## Summary PATCH 4/4 addresses all critical and most warning-level issues from the previous review (missing One remaining issue: Warning
ProduceBatchResult impactNot affected — no RoutedExperts impactOff-policy path replaces (rather than concatenates) routed_experts, freeing the old ref via Ray concurrency impactNot affected. VerdictCOMMENT — the remaining |
39de861 to
0561379
Compare
|
对 partial rollout 样本做 off-policy mask 时, off-policy 的定义应该也要跟当前"过期"的定义一致?
|
|
@claude review |
04f0f12 to
1f40fff
Compare
|
@claude review |
1f40fff to
8d65892
Compare
|
@claude review |
8d65892 to
43c9e5e
Compare
|
@claude review |
8694013 to
b26bc3f
Compare
|
@claude review |
|
Claude: Summary本 PR 引入 token 级 staleness：新增 ProduceBatchResult impact: 三处 trainer-visible 变化 —(1) RoutedExperts impact: token-expiry 只对过期 state 调 Ray concurrency impact: not affected（无 Main Flowchart after this PRflowchart TD
A[put / refresh_staleness] --> B[_apply_staleness_lifecycle]
B --> C{seq_staleness >= stale_threshold}
C -- 是 --> F[标记该 state EXPIRED]
C -- 否 --> D{agentic group 或 token_stale_threshold 为空}
D -- 是 --> E[保持原状态]
D -- 否 --> G[calculate_effective_response_mask]
G --> H{存在有效 token}
H -- 是 --> E
H -- 否 --> F
F --> I{expired_groups_retryable}
I -- 是 --> J[只重置过期 state 的 response<br/>整组进 EXPIRED pool]
I -- 否 --> K[整组丢弃<br/>含仍新鲜成员]
E --> L[take_batch]
J --> M[从 EXPIRED pool rerollout]
M --> A
L --> N[take_train_batch 重算 response_mask]
N --> O[_prepare_train_data / 训练]
style K fill:#ffcccc,stroke:#cc0000
style N fill:#ffe0b2,stroke:#e65100
style G fill:#e3f2fd,stroke:#1565c0
style J fill:#e3f2fd,stroke:#1565c0
核心原理实现与单测核心实现为三段： 真实代码路径已被覆盖的部分：
覆盖缺口见「单测建议」。 抽象与信息隐藏评估
单测建议
其他 Issues
VerdictREQUEST_CHANGES |
| max_token_staleness (int | None): Maximum extra weight-sync periods a | ||
| response token may lag behind before it is masked out of the loss. | ||
| ``None`` disables token-level masking, ``0`` accepts only tokens | ||
| produced within the current sync period, and ``N`` allows ``N`` | ||
| extra periods. Unlike ``max_staleness``, this does not expire or | ||
| re-roll a group; it only shrinks ``response_mask``. Defaults to | ||
| None. |
There was a problem hiding this comment.
Claude: [设计] docstring 与实现相反：这里声明 max_token_staleness “does not expire or re-roll a group; it only shrinks response_mask”，但 ReplayBuffer._apply_staleness_lifecycle 在 effective mask 全 0 时会把该 state 标为 EXPIRED、整组翻入 EXPIRED pool（replay_buffer.py#L471-L500）。
更关键的是它和新默认值的组合：tail_batch_trigger_size 默认 -1 ⇒ expired_groups_retryable=False ⇒ put 里 if status == Status.EXPIRED and not expired_groups_retryable: return，整组（含仍新鲜的成员，discard_rollout_state 释放全部字段）被直接丢弃。也就是说，用户只按 docstring 打开 max_token_staleness=0 而不动 tail_batch_trigger_size，得到的不是“收缩 mask”，而是静默丢数据。
ProduceBatchResult impact: 原本会计入 COMPLETED 的 group 变成 EXPIRED 或被丢弃，leftover_completed/leftover_expired/produced_samples 都会随之改变。
建议两件事一起做：
- 改写 docstring，明确 token staleness 也会触发 group 级 expire/rerollout（
disagg_producer.py:265同样文案）； - 在
build()里对max_token_staleness is not None and tail_batch_trigger_size < 0给出 warning 或直接校验失败，避免这个组合静默生效。
| for task in task_runners: | ||
| if task.token_stale_threshold is None: | ||
| continue | ||
| for group in batch_by_task.get(task.task_name, []): | ||
| # NOTE: input_ids/labels 表示 agentic 训练分支，当前暂不支持 agentic token-staleness masking。 | ||
| if any(item.input_ids is not None or item.labels is not None for item in group): | ||
| continue | ||
| for rollout_state in group: | ||
| rollout_state.response_mask = calculate_effective_response_mask( | ||
| rollout_state, | ||
| current_train_step=current_train_step, | ||
| token_stale_threshold=task.token_stale_threshold, | ||
| ) |
There was a problem hiding this comment.
Claude: [设计] 同一条 token-staleness 规则被拆到两个 Module 各写一遍，且两处 guard 不一致。
replay_buffer.py 侧（L457-L479）先算 is_agentic_group = any(item.input_ids is not None or item.labels is not None ...)，再对每个 state 加了 if not item.response_ids or (item.response_mask is not None and not any(item.response_mask)): continue 才调用 calculate_effective_response_mask；这里的消费期只重复了 agentic 判定，缺少 response_ids 与全 0 semantic mask 的 guard。
结果是两个 Seam 对「哪些 state 适用 token mask」的判定可以不一致：被 buffer 明确跳过评估的 state，在 take_train_batch 里仍会被送进 calculate_effective_response_mask（该函数用 cast 假设 response_ids/response_model_steps 非 None，已有单独意见）。同时，未来任何一处修规则都必须记得改另一处。
建议把这条规则收敛成 rl_data.py 里紧邻 calculate_effective_response_mask 的一个 Deep 函数，例如：
def apply_effective_response_mask(group, *, current_train_step, token_stale_threshold) -> list[bool]:
"""返回每个 state 是否已无可训练 token；内部统一处理 agentic 跳过与空 response guard。"""_apply_staleness_lifecycle 用它的返回值决定 expire，take_train_batch 用它写回 mask。这样规则、guard 和 agentic 例外集中在一处（Locality），两个调用者只需知道一个小 Interface，也可以直接对这个公开函数补单测。
| max_staleness: int = Field(default=0, ge=0) | ||
| tail_batch_trigger_size: int = 0 | ||
| max_token_staleness: int | None = Field(default=None, ge=0) | ||
| tail_batch_trigger_size: int = Field(default=-1, ge=-1) |
There was a problem hiding this comment.
Claude: [兼容性] 默认值从 0 改为 -1 的同时，0 本身的语义也被翻转了：expired_groups_retryable 由 tail_batch_trigger_size > 0 改成 >= 0，所以显式写 0 的存量配置（如 TAIL_BATCH_TRIGGER_SIZE）从「丢弃过期组」变成「立即 rerollout」。ProduceBatchResult impact: leftover_expired 与 produced_samples 随之改变。建议启动时对 0 打一条语义说明日志。
|
|
||
| async def put_generated_group(self, group: list[RolloutState]) -> bool: | ||
| produced_tokens = sum(len(item.response_ids) for item in group if item.response_ids is not None) | ||
| produced_tokens = sum(len(item.response_ids or []) - len(item.response_model_steps or []) for item in group) |
There was a problem hiding this comment.
Claude: [测试] produced_tokens 语义改为「仅本轮新增 token」（原为整条 response 长度），这是 ProduceBatchResult 的 trainer-visible 字段，但没有任何断言覆盖。建议补一例：同一 group 连续两轮 partial rollout，断言第二轮只计入新增 token。
| return storage_status | ||
|
|
||
| # 4. cleanup sample or cleanup response for expired sample | ||
| if expired_groups_retryable: | ||
| for item, expired in zip(group, expired_mask): | ||
| if expired: | ||
| item.status = Status.EXPIRED | ||
| reset_rollout_response(item) |
There was a problem hiding this comment.
Claude: [测试] 本 PR 的头号行为「只重置过期 state、保留同组新鲜成员」目前只有 buffer 字段级断言，缺少重跑路径覆盖：保留成员不被覆写依赖 worker.generate 的 status == COMPLETED 早退，建议补一例覆盖 EXPIRED group 重跑后新鲜 response 仍在、整组回到 COMPLETED。
2 participants
背景
在 partial rollout 场景下，同一条 response 中的 token 可能由不同版本的 policy 生成。
现有 sequence staleness 使用 response 中最早的模型版本表示整条样本的 staleness，无法区分：
本 PR 引入 token staleness，使系统能够：
主要改动
max_token_staleness配置，计算方式与 seq staleness 相同，在 take batch 阶段统一更新 response maskreplay_buffer.put和refresh_staleness，与更新 sequence staleness 相同说明：这个PR不改动agentic RL的过期语义
token staleness 处理关键阶段
如何采样
flowchart LR A[刷新 staleness] --> B[统计 EXPIRED groups] B --> C{tail_batch_trigger_size} C -- -1 --> D[采样 ABORTED 或新数据] C -- 0 且存在 EXPIRED --> E[优先采样 EXPIRED group] C -- 大于0且达到阈值 --> F[进入 tail batch] C -- 大于0但未达到阈值 --> D E --> G[保持正常异步生产和 oversampling] F --> H[关闭本轮 oversampling] D --> I[执行 rollout] G --> I H --> I如何判断一个样本是否过期
flowchart LR A[刷新 seq staleness] --> B{超过 seq threshold} B -- 是 --> C[state 标记为 EXPIRED] B -- 否 --> D{普通 rollout 且配置 token threshold} D -- 否 --> E[state 保持有效] D -- 是 --> F[计算 effective response mask] F --> G{是否存在有效 token} G -- 否 --> C G -- 是 --> E C --> H[StorageItem 标记为 EXPIRED] H --> I{是否允许 rerollout} I -- 是 --> J[只清空实际过期 state 的 response] I -- 否 --> K[丢弃整个 group]配置示例
该配置表示：