[Feat] Add off-policy masking for partial rollouts

[YanhuiDua]

背景

在 partial rollout 场景下，同一条 response 中的 token 可能由不同版本的 policy 生成。

现有 sequence staleness 使用 response 中最早的模型版本表示整条样本的 staleness，无法区分：

• 已经过期的旧 policy token；
• 仍然可以参与训练的新 policy token。

本 PR 引入 token staleness，使系统能够：

• 按 token 排除过期的 response token；
• 当一条样本不再包含任何有效训练 token 时，将其标记为过期；
• 只重置实际过期样本的 response；
• 保留同一 rollout group 中其他未过期样本的生成结果。

主要改动

1. Token 级 staleness mask：增加 max_token_staleness 配置，计算方式与 seq staleness 相同，在 take batch 阶段统一更新 response mask
2. Token-expired 生命周期: ReplayBuffer 通过统一的生命周期逻辑处理 sequence 和 token staleness，执行时机为 replay_buffer.put 和 refresh_staleness，与更新 sequence staleness 相同
3. 增强 tail_batch_trigger_size 语义：

• tail_batch_trigger_size = -1: 关闭过期 group 的 rerollout
• tail_batch_trigger_size = 0: 立即优先 rerollout 过期 group，但保持普通异步生产和 oversampling 策略
• tail_batch_trigger_size > 0: 等待 EXPIRED pool 累积到指定 group 数量后进入 tail-batch 模式

说明：这个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

Loading

如何判断一个样本是否过期

 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]

Loading

配置示例

AsyncProduceStrategyConfig(
    max_staleness=2,
    max_token_staleness=0,
    tail_batch_trigger_size=0,
)

该配置表示：

• sequence staleness 允许额外滞后两个权重同步周期；
• token staleness 只接受当前权重同步周期内生成的 token；
• 当样本的有效 token 全部过期后，立即允许 rerollout；
• immediate rerollout 不会使生产策略切换到 tail-batch 模式。

[jayhenry]

@claude review

[YanhuiDua]

@claude review

[claude]

Claude: ## Summary

Adds off-policy masking for partial rollouts: response tokens from prior partial-rollout rounds are masked out of the loss via response_mask, with per-token model-step tracking to refresh masks at consumption time in the replay buffer.

Issues

Critical

• xtuner/v1/rl/rollout/worker.py: Missing else branch in the early-exit path when enable_partial_rollout is False. Execution falls through to the generation loop and makes unnecessary backend requests.

Warning

• xtuner/v1/data_proto/rl_data.py:329: cast on a potentially-None response_mask won't prevent a runtime crash in zip(..., strict=True). Samples without a response_mask (e.g., pre-existing on-policy samples in the replay buffer) will crash when masking is refreshed.
• xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py:143-145: Configuration validation uses assert which is stripped under -O. Should raise ValueError. Same issue in disagg_agent_loop_manager.py.
• xtuner/v1/rl/replay_buffer.py:507: task_mask_offpolicy is now a required parameter — this is a breaking interface change for external callers.
• xtuner/v1/rl/rollout/utils.py:296-302: Off-policy masking path assumes routed_experts from the backend always covers the complete sequence. If the backend returns only new-token routing (as in some LMDeploy configs), this replacement would be incorrect.

Nit

• Style: PR uses from __future__ import annotations in produce_utils.py (commit 3) which changes forward-ref quoting style — fine, but the commit message doesn't mention this cleanup.

ProduceBatchResult impact

Not directly affected — no fields on ProduceBatchResult are added or changed. However, the response_mask field on RolloutState (which feeds into rollout_states inside the result) is now mutated by refresh_response_mask at consumption time — downstream consumers should expect masked-out tokens.

RoutedExperts impact

The off-policy masking path replaces routed_experts instead of concatenating. Old history refs are freed via free_object_refs — ownership transfer is correct for this path. The non-masking path retains the existing concatenation semantics unchanged.

Ray concurrency impact

Not affected — no changes to Ray actor methods, decorators, or concurrency groups.

Verdict

REQUEST_CHANGES — the missing else branch in worker.py is a regression for all users with enable_partial_rollout=False (the default).

[YanhuiDua]

@claude review

[claude]

Claude: ## Summary

PATCH 4/4 addresses all critical and most warning-level issues from the previous review (missing else branch, cast on None, assert → ValueError, breaking interface). Good work on the fix iteration.

One remaining issue:

Warning

• disagg_agent_loop_manager.py:84-87: The assert → ValueError conversion from PATCH 4 was applied to agent_loop_manager.py but missed in disagg_agent_loop_manager.py. Same rationale: assert is stripped under -O.

ProduceBatchResult impact

Not affected — no ProduceBatchResult fields are added or changed. The response_mask on RolloutState (inside rollout_states) is now mutated by refresh_response_mask at consumption time.

RoutedExperts impact

Off-policy path replaces (rather than concatenates) routed_experts, freeing the old ref via free_object_refs. Ownership is correct for this path.

Ray concurrency impact

Not affected.

Verdict

COMMENT — the remaining assert is a minor fix (one-line change in disagg_agent_loop_manager.py). The critical issues from the first review are resolved.

[jayhenry]

对 partial rollout 样本做 off-policy mask 时, off-policy 的定义应该也要跟当前"过期"的定义一致?
或者新增一个字段 max_trainable_response_staleness: int | None = None 替换 off policy partial rollout

• None：所有未过期 token 都可训练，对应原来的 False。
• 0：只训练当前权重同步周期生成的 token。
• N：允许额外 N 个同步周期的 token 参与 loss。
• max_staleness：仍只负责整组过期、reset 和 rerollout，保持修改前语义。

[jayhenry]

@claude review

[YanhuiDua]

@claude review

[jayhenry]

@claude review

[YanhuiDua]

@claude review

[jayhenry]

@claude review

[github-actions]

Claude:

Summary

本 PR 引入 token 级 staleness：新增 max_token_staleness 配置，在 ReplayBuffer 侧用统一的 _apply_staleness_lifecycle 处理 sequence/token 两级过期（只重置真正过期 state 的 response，保留同组新鲜成员），在 take_train_batch 侧按 token 新鲜度重算 response_mask，并把 tail_batch_trigger_size 扩展为 -1 / 0 / N 三态语义（默认值由 0 改为 -1）。

ProduceBatchResult impact: 三处 trainer-visible 变化 —(1) produced_tokens 语义由「整条 response 长度」改为「本轮新增 token 数」；(2) token-expiry 会把原本 COMPLETED 的 group 翻成 EXPIRED，影响 leftover_completed / leftover_expired / produced_samples；(3) tail_batch_trigger_size=0 使 EXPIRED group 变为可重试（此前直接丢弃）。

RoutedExperts impact: token-expiry 只对过期 state 调 reset_rollout_response（内部 free_object_refs），同组保留的 fresh state 会在 EXPIRED pool 中继续持有 routed_experts ObjectRef 直到 rerollout 完成或整组丢弃 —持有窗口变长，但 owner 唯一、非 retryable 分支走 discard_rollout_state，未发现泄漏。

Ray concurrency impact: not affected（无 ray.method、concurrency group 或 actor 构造改动）。

Main Flowchart after this PR

flowchart 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

Loading

核心原理实现与单测

核心实现为三段：calculate_effective_response_mask（semantic mask ∩ token staleness mask）、ReplayBuffer._apply_staleness_lifecycle（put/refresh 两个时机统一做 seq+token 过期判定与清理）、take_train_batch（消费期写回 effective mask）。

真实代码路径已被覆盖的部分：

• mask 纯函数行为（阈值放宽、与 semantic mask 求交、rerollout 后无 semantic mask）：tests/rl/test_staleness_policy.py::TestTokenStalenessMask。
• 生命周期经公开 put / refresh_staleness / get + 真实 Naive/Pandas storage：tests/rl/test_replay_buffer.py 新增 5 例，含 token/seq 过期保留新鲜成员、agentic group 跳过、非 retryable 整组丢弃。
• 消费期 mask 经公开 AgentLoopManager.produce_batch：tests/rl/test_multi_task_agent_loop_manager.py 2 例（普通与 agentic 分支）。
• tail_batch_trigger_size=0 的采样语义与 disagg put-time consumer step：tests/rl/test_producer.py 2 例。

覆盖缺口见「单测建议」。

抽象与信息隐藏评估

• Warning xtuner/v1/rl/agent_loop_manager/producer.py#L210（disagg_producer.py#L265 同）：max_token_staleness 的 docstring 声明「不会使 group 过期或重跑，只收缩 response_mask」，与实现的 token-expired 生命周期相反，配合新默认 tail_batch_trigger_size=-1 会静默丢弃整组。
• Warning xtuner/v1/rl/agent_loop_manager/produce_utils.py#L621：同一条 token-staleness 规则（含 agentic 判定）在 produce_utils 与 replay_buffer 两个 Module 各实现一遍，且 guard 不一致，规则的 Locality 被打散。

单测建议

• Warning xtuner/v1/rl/replay_buffer.py#L490：PR 的头号行为「EXPIRED group 保留 fresh 成员后重跑」只有 buffer 字段级断言，缺少覆盖重跑路径的用例（该行为依赖 worker.generate 的 status == COMPLETED 早退才不会覆写保留的 response）。
• Warning xtuner/v1/rl/agent_loop_manager/produce_utils.py#L173：produced_tokens 改为只统计本轮新增 token，这一 ProduceBatchResult 字段的语义变化没有任何断言覆盖。

其他 Issues

• Warning xtuner/v1/rl/agent_loop_manager/producer.py#L233（disagg_producer.py#L278 同）：tail_batch_trigger_size 旧默认值 0 的语义被翻转（原「丢弃过期组」→ 现「立即 rerollout」），显式写 0 的存量配置行为静默改变且无提示。

Verdict

REQUEST_CHANGES

[claude]

Claude: [测试] produced_tokens 语义改为「仅本轮新增 token」（原为整条 response 长度），这是 ProduceBatchResult 的 trainer-visible 字段，但没有任何断言覆盖。建议补一例：同一 group 连续两轮 partial rollout，断言第二轮只计入新增 token。

[jayhenry]

保留这个验证？

[YanhuiDua]

感觉也没有必要，因为sync_weight_interval 会在RL trainer中检查是否大于0，token和seq的staleness作为AsyncproducerConfig中检查是否大于0

[YifanHe-ailab]

@claude review

[github-actions]

Claude:

Summary

本 PR 引入 token 级 staleness：新增 max_token_staleness 配置，把 seq/token 两级过期统一到 ReplayBuffer._apply_staleness_lifecycle（只重置真正过期的 state，保留同组新鲜成员），在 take_train_batch 侧按 token 新鲜度重算 response_mask，并把 tail_batch_trigger_size 扩展为 -1 / 0 / N 三态语义（默认由 0 改为 -1）。核心正确性逻辑扎实且测试覆盖较完整，遗留一个数据路径性能问题和一个覆盖缺口。

ProduceBatchResult impact: 三处 trainer-visible 变化 —(1) token expiry 会把原本 COMPLETED 的 group 翻成 EXPIRED，影响 leftover_completed / leftover_expired / produced_samples；(2) produced_tokens 语义由「整条 response 长度」改为「本轮新增 token 数」；(3) tail_batch_trigger_size=0 使 EXPIRED group 变为可重试（此前直接丢弃）。

RoutedExperts impact: token expiry 只对过期 state 调 reset_rollout_response（内部 free_object_refs），同组保留的 fresh state 会在 EXPIRED pool 中继续持有 routed_experts object ref 直到 rerollout 完成或整组丢弃 —持有窗口变长，但 owner 唯一、非 retryable 分支走 discard_rollout_state，未发现泄漏。

Ray concurrency impact: not affected（无 ray.method、concurrency group 或 actor 构造改动）。

核心原理实现与单测

核心实现为三段，均有真实代码路径覆盖：

• mask 计算：rl_data.py#L324-L386 的 semantic mask ∩ token staleness mask，以及 agentic / 空 response 的 None 短路 → tests/rl/test_staleness_policy.py::TestTokenStalenessMask（阈值放宽、与 semantic mask 求交、rerollout 后无 semantic mask）。
• 统一生命周期：replay_buffer.py#L442-L494，在 put 与 refresh_staleness 两个时机做 seq+token 过期判定与清理 → tests/rl/test_replay_buffer.py 新增 5 例，全部经公开 put / refresh_staleness / get + 真实 Naive/Pandas storage，覆盖 token/seq 过期保留新鲜成员、agentic group 跳过、非 retryable 整组丢弃。
• 消费期写回 mask：produce_utils.py#L613-L628 → tests/rl/test_multi_task_agent_loop_manager.py 2 例，经公开 AgentLoopManager.produce_batch 覆盖普通与 agentic 两条分支。

此外 tail_batch_trigger_size=0 的采样语义、disagg put-time consumer step、response_mask=None 的 rerollout 训练路径分别由 tests/rl/test_producer.py（2 例）与 tests/rl/test_prepare_train_data.py 覆盖。覆盖缺口见「单测建议」。

单测建议

• Warning xtuner/v1/rl/replay_buffer.py#L456-L458：NOTE 声明的「已 EXPIRED group 中保留的 COMPLETED state 在等待 rerollout 期间也会继续过期」是本次扩大 refresh 状态集的直接原因，但新增 5 例都从非 EXPIRED group 起步，该转换没有任何用例覆盖。

其他 Issues

• Warning xtuner/v1/rl/replay_buffer.py#L463-L477：refresh_staleness 在持锁且无 await 的循环内为每条记录物化完整的 per-token mask 列表，而结果只用于 not any(...) 判空，配合本 PR 扩大的扫描范围（新增 EXPIRED 状态、disagg 每步两次刷新）显著加重了 buffer 锁与事件循环的负担，且未附带 benchmark。

Verdict

COMMENT

[claude]

Claude: [性能] 这里为每条记录物化了完整的 per-token mask 列表（rl_data.py 内每个 state 要分配 2~3 个 len(response_ids) 长度的 list），但结果只用于 L476 的 not any(...) 判空。

该循环位于 refresh_staleness 的 async with self._lock 内且中间没有 await，会阻塞并发的 put/get/count 和宿主 event loop。改动前同一循环只做 refresh_seq_staleness（一次 min()）。同时扫描范围也变大了：refresh 状态集新增了 Status.EXPIRED，disagg 路径每个 train step 还会刷新两次。

建议判空时短路，不要构造 mask（materialize 只保留给 take_train_batch）：

def _has_trainable_token(item, *, current_train_step, token_stale_threshold) -> bool:
    ...
    return any(
        sem and calculate_seq_staleness(step, current_train_step) < token_stale_threshold
        for sem, step in zip(semantic_mask, item.response_model_steps)
    )

按 .claude/CLAUDE.md，数据路径改动需要 before/after benchmark 或说明为何不需要。

ProduceBatchResult impact: 不改变任何字段取值，仅影响 produce_time_s 与整体吞吐。

[claude]

Claude: [测试] 这条 NOTE 描述的转换没有用例覆盖：新增 5 例都从非 EXPIRED group 起步。建议补一例把混合 EXPIRED/COMPLETED 的 group 在更大的 current_train_step 上再过一次 refresh_staleness，断言保留的 fresh state 最终也翻为 EXPIRED 并被 reset。

ProduceBatchResult impact: 该转换决定 leftover_expired 与 leftover_completed 的归属，回归时不会被现有断言发现。