diff --git a/docs/en/rl/advanced_tutorial/rl_trainer.md b/docs/en/rl/advanced_tutorial/rl_trainer.md index 8257f92e7..6e2072afa 100644 --- a/docs/en/rl/advanced_tutorial/rl_trainer.md +++ b/docs/en/rl/advanced_tutorial/rl_trainer.md @@ -140,7 +140,7 @@ Parameter meanings: | `over_sample_threshold` | The ratio of extra samples that may be generated. A larger value makes the rollout side easier to keep fully loaded, but may produce more samples that are not from the current step. | | `enable_partial_rollout` | Whether rollouts paused before weight synchronization may continue after synchronization. Before using this for tool calling or multi-turn tasks, confirm that the AgentLoop supports continuation. | | `max_staleness` | The number of synchronization cycles by which samples may lag behind the current training progress. A larger value gives more throughput flexibility but weakens the on-policy property. | -| `tail_batch_trigger_size` | When expired samples accumulate to this number, tail batch mode is entered and these samples are retried first. | +| `tail_batch_trigger_size` | Expired-sample retry policy: `-1` disables rerollout, `0` retries immediately without entering tail batch mode, and a positive value enters tail batch mode after that many expired groups accumulate. | `max_staleness` is counted in "weight synchronization cycles". The actual expiration threshold used in code is: @@ -161,12 +161,11 @@ Both oversampling and partial rollout are affected by `max_staleness`: by the earliest model version in the response, so continuation across synchronization cycles also needs room from `max_staleness`. -Tail batch is used to handle samples that have expired during asynchronous production. When the number of `expired` -samples reaches `tail_batch_trigger_size`, `AsyncProduceStrategy` enters tail batch mode: this round no longer -oversamples according to `over_sample_threshold`, only fills the required target, and retries samples from the -expired sample pool first. You can understand it as a non-oversampling synchronous fill-up production. Its goal is -not to improve throughput, but to collect long-tail expired samples again and avoid leaving them in the buffer for -too long. +Expired samples can be rerolled out according to `tail_batch_trigger_size`. `-1` disables rerollout. `0` retries +expired groups as soon as they appear while retaining the normal asynchronous production and oversampling policy. +For a positive value, `AsyncProduceStrategy` waits until the expired pool reaches the configured size, then enters +tail batch mode: this round no longer oversamples according to `over_sample_threshold`, only fills the required +target, and retries samples from the expired sample pool first. Note: it is not recommended to set `max_staleness>0` and `enable_partial_rollout=False` at the same time. With this combination, long-tail oversampled samples may be reset after weight synchronization because partial rollout is not diff --git a/docs/zh_cn/rl/advanced_tutorial/rl_trainer.md b/docs/zh_cn/rl/advanced_tutorial/rl_trainer.md index a32f18b1a..3f7430200 100644 --- a/docs/zh_cn/rl/advanced_tutorial/rl_trainer.md +++ b/docs/zh_cn/rl/advanced_tutorial/rl_trainer.md @@ -131,7 +131,7 @@ produce_strategy_config = AsyncProduceStrategyConfig( | `over_sample_threshold` | 允许额外生成的比例。值越大，rollout 侧越容易保持满载，但也可能产生更多非当前 step 的样本。 | | `enable_partial_rollout` | 权重同步前被暂停的 rollout 是否允许在同步后续跑。工具调用或多轮任务使用前需要确认 AgentLoop 支持续跑。 | | `max_staleness` | 允许样本相对当前训练进度滞后的同步周期数。值越大，吞吐更宽松，on-policy 程度更弱。 | -| `tail_batch_trigger_size` | 过期样本累计到一定数量后，进入 tail batch 模式，优先重试这些样本。 | +| `tail_batch_trigger_size` | 过期样本重试策略：`-1` 关闭 rerollout，`0` 立即重试但不进入 tail batch 模式，正数表示累计到指定 group 数量后进入 tail batch 模式。 | `max_staleness` 按“权重同步周期”计数。代码中实际使用的过期阈值是： @@ -146,7 +146,7 @@ stale_threshold = (max_staleness + 1) * sync_weights_interval - `over_sample_threshold>0` 会为未来 step 提前生成样本。如果这些样本跨过下一次权重同步点，只有 `max_staleness` 允许时才会继续保留为可训练样本。 - `enable_partial_rollout=True` 会让被暂停的 response 在同步后续跑。样本的 staleness 按 response 中最早的模型版本计算，因此跨同步周期续跑时也需要 `max_staleness` 留出空间。 -tail batch 用于处理异步生产中已经过期的样本。当 `expired` 样本数量达到 `tail_batch_trigger_size` 时，`AsyncProduceStrategy` 会进入 tail batch 模式：本轮不再按 `over_sample_threshold` 超发，只补齐必要目标，并优先从过期样本池中取样重试。可以把它理解为一次非超发的同步补齐生产；它的目的不是提高吞吐，而是把长尾过期样本重新收集起来，避免它们长期留在 buffer 中。 +`tail_batch_trigger_size` 控制过期样本的 rerollout。设置为 `-1` 时关闭 rerollout；设置为 `0` 时，过期 group 一出现就立即优先重试，但仍保持普通异步生产和 oversampling 策略；设置为正数时，`AsyncProduceStrategy` 等待 expired pool 累积到指定数量后进入 tail batch 模式，本轮不再按 `over_sample_threshold` 超发，只补齐必要目标，并优先从过期样本池中取样重试。 注意：不建议同时设置 `max_staleness>0` 且 `enable_partial_rollout=False`。这种组合下，长尾超发样本在权重同步后可能因为不支持 partial rollout 被重置（当前在 `RolloutWorker` 中重置样本只保留 prompt 字段）；但由于每次重置过期信息归0，它们不会过期，tail batch 不会及时接管，下一轮同步窗口内仍然可能生成不完并反复重试。当前还没有支持 `tail_batch_max_tries` 机制来按重试次数触发 tail batch。因此 `max_staleness>0` 时，优先开启 `enable_partial_rollout=True`。 diff --git a/examples/v1/config/rl_disagg_multi.py b/examples/v1/config/rl_disagg_multi.py index 394baf0b0..e4aa51038 100644 --- a/examples/v1/config/rl_disagg_multi.py +++ b/examples/v1/config/rl_disagg_multi.py @@ -70,7 +70,7 @@ sync_weights_interval = int(os.environ.get("SYNC_WEIGHTS_INTERVAL", "1")) over_sample_threshold = float(os.environ.get("OVER_SAMPLE_THRESHOLD", "0.0")) partial_rollout = os.environ.get("PARTIAL_ROLLOUT", "0") == "1" -tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "0")) +tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "-1")) max_staleness = int(os.environ.get("MAX_STALENESS", "0")) enable_evaluate = os.environ.get("ENABLE_EVALUATE", "0") == "1" gsm8k_task_weight = float(os.environ.get("GSM8K_TASK_WEIGHT", "3.0")) diff --git a/examples/v1/config/rl_disagg_single.py b/examples/v1/config/rl_disagg_single.py index 6ef2250e3..7878c962e 100644 --- a/examples/v1/config/rl_disagg_single.py +++ b/examples/v1/config/rl_disagg_single.py @@ -80,7 +80,7 @@ sync_weights_interval = int(os.environ.get("SYNC_WEIGHTS_INTERVAL", "1")) over_sample_threshold = float(os.environ.get("OVER_SAMPLE_THRESHOLD", "0.0")) partial_rollout = os.environ.get("PARTIAL_ROLLOUT", "0") == "1" -tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "0")) +tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "-1")) max_staleness = int(os.environ.get("MAX_STALENESS", "0")) prompt_repeat_k = int(os.environ.get("PROMPT_REPEAT_K", "4")) rollout_tp_size = int(os.environ.get("ROLLOUT_TP_SIZE", "1")) diff --git a/tests/rl/test_multi_task_agent_loop_manager.py b/tests/rl/test_multi_task_agent_loop_manager.py index 688932fd7..7432fabcf 100644 --- a/tests/rl/test_multi_task_agent_loop_manager.py +++ b/tests/rl/test_multi_task_agent_loop_manager.py @@ -12,15 +12,19 @@ import asyncio import unittest from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch -from xtuner.v1.data_proto.rl_data import Status +from xtuner.v1.data_proto.rl_data import RolloutState, Status from xtuner.v1.rl.agent_loop_manager.agent_loop_manager import ( AgentLoopManager, AgentLoopManagerConfig, TaskSpecConfig, ) -from xtuner.v1.rl.agent_loop_manager.disagg_agent_loop_manager import DisaggAgentLoopManager +from xtuner.v1.rl.agent_loop_manager.disagg_agent_loop_manager import ( + DisaggAgentLoopManager, + DisaggAgentLoopManagerConfig, + DisaggTaskSpecConfig, +) from xtuner.v1.rl.agent_loop_manager.produce_utils import ( GROUP_GENERATE_TIME_KEY, ProduceBatchStatus, @@ -41,10 +45,12 @@ def __init__( self, cleanup_pause_time_s: float = 0.0, stale_threshold: int = 1, - tail_batch_trigger_size: int = 0, + tail_batch_trigger_size: int = -1, + token_stale_threshold: int | None = None, ): self.cleanup_pause_time_s = cleanup_pause_time_s self.stale_threshold = stale_threshold + self.token_stale_threshold = token_stale_threshold self.tail_batch_trigger_size = tail_batch_trigger_size self.called_batch_sizes: list[int] = [] self.called_train_steps: list[int] = [] @@ -122,6 +128,7 @@ def __init__(self, rollout_states_by_task: dict[str, list[list[Any]]], leftover_ self._leftover_counts = leftover_counts self.refresh_staleness_calls: list[tuple[str, int, int, tuple[Status, ...]]] = [] self.expired_groups_retryable_calls: list[dict[str, bool]] = [] + self.task_token_stale_threshold_calls: list[dict[str, int]] = [] async def get(self, batch_size: int, task_name: str, group_status: Status): assert group_status == Status.COMPLETED @@ -137,11 +144,13 @@ async def refresh_staleness( self, *, task_stale_thresholds: dict[str, int], + task_token_stale_thresholds: dict[str, int] | None = None, expired_groups_retryable_by_task: dict[str, bool] | None = None, current_train_step: int, statuses: list[Status] | None = None, ): self.expired_groups_retryable_calls.append(dict(expired_groups_retryable_by_task or {})) + self.task_token_stale_threshold_calls.append(dict(task_token_stale_thresholds or {})) expired_counts = {} for task_name, stale_threshold in task_stale_thresholds.items(): self.refresh_staleness_calls.append( @@ -191,6 +200,8 @@ def _fake_rollout_controller(): class TestMultiTaskAgentLoopManager(unittest.IsolatedAsyncioTestCase): + """共卡与多 task manager 的 batch 生产、消费和统计行为。""" + def test_manager_config_accepts_single_task_spec(self): # 单 task 配置可以直接传入，兼容最小 AgentLoopManager 配置。 task = TaskSpecConfig.model_construct( @@ -205,11 +216,80 @@ def test_manager_config_accepts_single_task_spec(self): self.assertEqual(manager_config.tasks.task_name, "single_task") + async def test_take_train_batch_applies_token_staleness_mask(self): + # 启用 token staleness 时，公开 produce_batch 路径应返回最终 effective mask。 + state = RolloutState( + rollout_id=1, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=[1, 2], + response_ids=[3, 4], + response_model_steps=[0, 4], + status=Status.COMPLETED, + ) + strategy = _FakeProduceStrategy(token_stale_threshold=4) + replay_buffer = _FakeReplayBuffer( + rollout_states_by_task={"task": [[state]]}, + leftover_counts={}, + ) + manager = AgentLoopManager( + task_runners=[ + _TaskRunner( + task_name="task", + agent_loop=_fake_agent_loop(), + produce_strategy=strategy, + sampler=_FakeSampler(), + weight=1.0, + order=0, + ) + ], + replay_buffer=replay_buffer, + rollout_controller=_fake_rollout_controller(), + ) + + result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + + self.assertEqual(result.rollout_states[0][0].response_mask, [0, 1]) + self.assertEqual(replay_buffer.task_token_stale_threshold_calls, [{"task": 4}]) + + async def test_take_train_batch_skips_agentic_token_staleness_mask(self): + state = RolloutState( + rollout_id=1, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + input_ids=[1, 2], + labels=[-100, 2], + response_mask=[1], + status=Status.COMPLETED, + ) + strategy = _FakeProduceStrategy(token_stale_threshold=4) + manager = AgentLoopManager( + task_runners=[ + _TaskRunner( + task_name="task", + agent_loop=_fake_agent_loop(), + produce_strategy=strategy, + sampler=_FakeSampler(), + weight=1.0, + order=0, + ) + ], + replay_buffer=_FakeReplayBuffer( + rollout_states_by_task={"task": [[state]]}, + leftover_counts={}, + ), + rollout_controller=_fake_rollout_controller(), + ) + + result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + + self.assertEqual(result.rollout_states[0][0].response_mask, [1]) + async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results(self): # 共卡 produce_batch 按 task 权重分配 batch，并按 task 名稳定返回训练数据和 leftover 统计。 strategy_a = _FakeProduceStrategy(tail_batch_trigger_size=2) strategy_b = _FakeProduceStrategy() - strategy_c = _FakeProduceStrategy() + strategy_c = _FakeProduceStrategy(tail_batch_trigger_size=0) replay_buffer = _FakeReplayBuffer( rollout_states_by_task={ "task_a": [["a-0"], ["a-1"]], @@ -268,7 +348,7 @@ async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results self.assertIn("task_c", result.task_results) self.assertEqual( replay_buffer.expired_groups_retryable_calls, - [{"task_b": False, "task_a": True, "task_c": False}], + [{"task_b": False, "task_a": True, "task_c": True}], ) self.assertEqual(strategy_a.called_expired_groups_retryable, [True]) self.assertEqual(strategy_a.cleanup_expired_groups_retryable, [True]) diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index 029a3f030..3c178b476 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -18,7 +18,7 @@ import numpy as np import torch -from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.data_proto.rl_data import RolloutState, Status, reset_rollout_response from xtuner.v1.train.rl_trainer import BaseRLTrainer @@ -111,6 +111,22 @@ def test_text_path_builds_shifted_training_tensors(self): self.assertEqual(info["response_len/mean"], 3.0) self.assertEqual(info["prompt_len/mean"], 3.0) + def test_rerolled_state_without_semantic_mask_uses_all_response_tokens(self): + trainer = self._build_trainer([1.0]) + state = reset_rollout_response(self._state(response_mask=[0, 1, 0])) + state.response = "rerolled response" + state.response_ids = [30, 31] + state.logprobs = [0.1, 0.2] + state.reward = {"score": 1.0} + state.status = Status.COMPLETED + state.finish_reason = "stop" + + data_batches, _ = self._prepare(trainer, [[state]]) + + self.assertIsNone(state.response_mask) + self.assertEqual(data_batches[0]["shifted_labels"].tolist(), [[-100, -100, 30, 31]]) + self.assertEqual(data_batches[0]["advantage"], [1.0, 1.0, 1.0, 1.0, 1.0]) + def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): # 同一个 prompt 下的多个 response 要分别使用自己的 reward 和 advantage。 trainer = self._build_trainer([1.5, -2.0]) diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 17cb83efb..5f6799f7f 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -145,7 +145,8 @@ def _build_context( progress=progress, is_valid_sample_fn=strategy.is_valid_sample_fn, stale_threshold=getattr(strategy, "stale_threshold", None), - expired_groups_retryable=getattr(strategy, "tail_batch_trigger_size", 0) > 0, + token_stale_threshold=getattr(strategy, "token_stale_threshold", None), + expired_groups_retryable=getattr(strategy, "tail_batch_trigger_size", -1) >= 0, ) def _build_disagg_progress( @@ -196,7 +197,8 @@ def _build_disagg_context( progress=progress, is_valid_sample_fn=strategy.is_valid_sample_fn, stale_threshold=getattr(strategy, "stale_threshold", None), - expired_groups_retryable=getattr(strategy, "tail_batch_trigger_size", 0) > 0, + token_stale_threshold=getattr(strategy, "token_stale_threshold", None), + expired_groups_retryable=getattr(strategy, "tail_batch_trigger_size", -1) >= 0, ) async def test_contexts_keep_colocate_and_disagg_control_surface_separate(self): @@ -242,6 +244,53 @@ async def test_contexts_keep_colocate_and_disagg_control_surface_separate(self): update_event.set() self.assertTrue(disagg_ctx.should_abort()) + async def test_disagg_put_uses_consumer_step_for_token_expiry(self): + # 非共卡 put-time check 必须读取 live consumer step，而不是 producer future step。 + task_name = "test_disagg_token_expiry" + strategy = DisaggAsyncProduceStrategyConfig( + max_staleness=3, + max_token_staleness=0, + tail_batch_trigger_size=1, + ).build(sync_weights_interval=1) + progress = self._build_disagg_progress( + task_name, + target=1, + train_step=5, + producer_future_step=9, + target_upto_future_step=9, + ) + ctx = self._build_disagg_context( + strategy, + task_name, + self._build_agent_loop(), + self._build_sampler(), + batch_size=1, + train_step=9, + model_step=3, + progress=progress, + ) + state = RolloutState( + rollout_id=1, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=[1], + tokens=[1, 11], + response="old response", + response_ids=[11], + response_mask=[1], + response_model_steps=[3], + logprobs=[0.1], + finish_reason="stop", + reward={"score": 1.0}, + status=Status.COMPLETED, + ) + + self.assertFalse(await ctx.put_generated_group([state])) + self.assertEqual(await self.replay_buffer.count(task_name, Status.COMPLETED), 0) + self.assertEqual(await self.replay_buffer.count(task_name, Status.EXPIRED), 1) + self.assertEqual(state.status, Status.EXPIRED) + self.assertEqual(state.response_ids, []) + async def test_sampler_with_replay_buffer(self): # 验证 sampler 优先复用 replay buffer 中可重试的 rollout group，耗尽后回退 dataloader。 task_name = "test_task" @@ -628,6 +677,45 @@ async def instrumented_sample(task_name, group_status=None): self.assertEqual(sorted(group[0].group_id for group in completed), [900, 901]) self.assertTrue(all(group[0].seq_staleness == 0 for group in completed)) + async def test_async_produce_strategy_immediate_rerollout_keeps_oversampling(self): + # trigger size 为 0 只立即启用 expired 采样，不进入禁用超发的 tail-batch 模式。 + task_name = "test_immediate_rerollout" + for sample_id in (900, 901): + await self.replay_buffer.put( + [make_rollout_state(sample_id, status=Status.EXPIRED)], + task_name, + expired_groups_retryable=True, + ) + + sampler = self._build_sampler() + original_sample = sampler.sample + sampled_statuses: list[list[Status] | None] = [] + + async def instrumented_sample(task_name, group_status=None): + sampled_statuses.append(group_status) + return await original_sample(task_name=task_name, group_status=group_status) + + sampler.sample = instrumented_sample + strategy = AsyncProduceStrategyConfig( + over_sample_threshold=1.0, + tail_batch_trigger_size=0, + ).build() + ctx = self._build_context( + strategy, + task_name, + self._build_agent_loop(), + sampler, + batch_size=1, + progress=self._build_progress(task_name, target=1), + ) + + await strategy.produce_batch(ctx) + + self.assertEqual( + sampled_statuses, + [[Status.EXPIRED, Status.ABORTED], [Status.EXPIRED, Status.ABORTED]], + ) + async def test_async_produce_strategy_fails_fast_on_invalid_progress(self): # 验证 progress 缺少当前 task key 时 fail fast，避免静默用 0 掩盖调度状态损坏。 task_name = "test_invalid_progress" diff --git a/tests/rl/test_replay_buffer.py b/tests/rl/test_replay_buffer.py index 2dad5c94b..69b5c6a2c 100644 --- a/tests/rl/test_replay_buffer.py +++ b/tests/rl/test_replay_buffer.py @@ -8,15 +8,17 @@ # 会补齐 response_model_steps，并刷新 seq_staleness。 # 4. tail batch disabled 时将 EXPIRED 当作终态，释放重字段且不写入 buffer；enabled 时保留 # prompt/mm_info，只重置 response 和 routed experts 以便重新 rollout。 -# 5. refresh_staleness 的公共契约：可以刷新 completed/aborted 记录，也要尊重显式传入的 +# 5. 写入过期结果时会触发 rerollout：超过 stale_threshold 的 group 会被重置 response 相关字段， +# 并保留 prompt/message 等重新 rollout 所需的输入字段。 +# 6. refresh_staleness 的公共契约：可以刷新 completed/aborted 记录，也要尊重显式传入的 # status 过滤条件。 -# 6. SyncReplayBufferConfig 的采样策略：按 FIFO 顺序返回 group。 -# 7. AsyncReplayBufferConfig 的采样策略：优先返回 seq_staleness 更高的 group； +# 7. SyncReplayBufferConfig 的采样策略：按 FIFO 顺序返回 group。 +# 8. AsyncReplayBufferConfig 的采样策略：优先返回 seq_staleness 更高的 group； # staleness 相同时使用 FIFO 作为 tie-breaker。 -# 8. save/resume 保留采样顺序：sync 恢复后仍是 FIFO，async 恢复后仍按 staleness 排序。 -# 9. save/resume 保留真实 RolloutState 字段：状态、response、tokens、logprobs、reward、 +# 9. save/resume 保留采样顺序：sync 恢复后仍是 FIFO，async 恢复后仍按 staleness 排序。 +# 10. save/resume 保留真实 RolloutState 字段：状态、response、tokens、logprobs、reward、 # error_msg、extra_fields 等字段恢复后应一致。 -# 10. save/resume 保留 Ray ObjectRef：直接 ObjectRef 和 dict(dict(ObjectRef)) 嵌套结构恢复后， +# 11. save/resume 保留 Ray ObjectRef：直接 ObjectRef 和 dict(dict(ObjectRef)) 嵌套结构恢复后， # 解引用得到的内容都应与保存前一致。 import tempfile @@ -45,6 +47,7 @@ def make_rollout_state( response: str | None = None, response_ids: list[int] | None = None, response_model_steps: list[int] | None = None, + response_mask: list[int] | None = None, logprobs: list[float] | None = None, reward: dict | None = None, error_msg: str | None = None, @@ -52,6 +55,8 @@ def make_rollout_state( routed_experts=None, mm_info: dict | None = None, extra_fields: dict | None = None, + input_ids: list[int] | None = None, + labels: list[int] | None = None, ) -> RolloutState: prompt_ids = list(prompt_ids) if prompt_ids is not None else [uid, uid + 1000] response_ids = list(response_ids) if response_ids is not None else [uid + 10] @@ -65,7 +70,7 @@ def make_rollout_state( response=response if response is not None else f"response {uid}", response_ids=response_ids, response_model_steps=list(response_model_steps) if response_model_steps is not None else None, - response_mask=[1 for _ in response_ids], + response_mask=list(response_mask) if response_mask is not None else [1 for _ in response_ids], logprobs=logprobs, routed_experts=routed_experts, finish_reason="stop" if status == Status.COMPLETED else None, @@ -75,6 +80,8 @@ def make_rollout_state( status=status, mm_info=mm_info, extra_fields=dict(extra_fields or {}), + input_ids=input_ids, + labels=labels, ) @@ -257,11 +264,154 @@ async def test_common_put_defaults_to_retryable_expired_group(self): assert reusable.error_msg is None assert reusable.routed_experts is None assert reusable.finish_reason is None - assert reusable.response_mask == [] + assert reusable.response_mask is None assert reusable.mm_info is not None assert reusable.mm_info["pixel_values"] is pixel_values assert reusable.extra_fields == {"train_prompt_ids": [101, 102]} + async def test_common_put_token_expiry_preserves_fresh_group_members(self): + # token expiry 只清理真正全 token 过期的 state，外层 group 仍进入 EXPIRED pool。 + for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: + with self.subTest(replay_buffer_config=config_name): + replay_buffer = replay_buffer_config_cls().build() + expired = make_rollout_state( + 1, + response="expired response", + response_ids=[11, 12], + response_model_steps=[0, 0], + reward={"score": 0.1}, + ) + fresh = make_rollout_state( + 2, + response="fresh response", + response_ids=[21, 22], + response_model_steps=[4, 4], + reward={"score": 0.9}, + ) + + await replay_buffer.put( + [expired, fresh], + "task", + current_train_step=5, + stale_threshold=10, + token_stale_threshold=4, + expired_groups_retryable=True, + ) + + self.assertEqual(await replay_buffer.count("task", Status.COMPLETED), 0) + self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 1) + group = (await replay_buffer.get(1, "task", Status.EXPIRED))[0] + self.assertEqual([item.status for item in group], [Status.EXPIRED, Status.COMPLETED]) + self.assertEqual(group[0].response, "") + self.assertEqual(group[0].response_ids, []) + self.assertIsNone(group[0].reward) + self.assertEqual(group[1].response, "fresh response") + self.assertEqual(group[1].response_ids, [21, 22]) + self.assertEqual(group[1].response_model_steps, [4, 4]) + self.assertEqual(group[1].reward, {"score": 0.9}) + + async def test_common_put_skips_token_expiry_for_agentic_group(self): + for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: + with self.subTest(replay_buffer_config=config_name): + replay_buffer = replay_buffer_config_cls().build() + state = make_rollout_state( + 1, + response="agentic response", + response_model_steps=[0], + input_ids=[1, 2], + labels=[-100, 2], + ) + + await replay_buffer.put( + [state], + "task", + current_train_step=5, + stale_threshold=10, + token_stale_threshold=4, + expired_groups_retryable=True, + ) + + self.assertEqual(await replay_buffer.count("task", Status.COMPLETED), 1) + self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 0) + self.assertEqual(state.status, Status.COMPLETED) + self.assertEqual(state.response, "agentic response") + + async def test_common_put_seq_expiry_preserves_fresh_group_members(self): + # seq expiry 路由整组到 EXPIRED pool，但只标记和清理实际过期的 state。 + for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: + with self.subTest(replay_buffer_config=config_name): + replay_buffer = replay_buffer_config_cls().build() + stale = make_rollout_state( + 1, + response="stale response", + response_model_steps=[0], + ) + fresh = make_rollout_state( + 2, + response="fresh response", + response_model_steps=[4], + ) + + await replay_buffer.put( + [stale, fresh], + "task", + current_train_step=5, + stale_threshold=4, + expired_groups_retryable=True, + ) + + group = (await replay_buffer.get(1, "task", Status.EXPIRED))[0] + self.assertEqual([item.status for item in group], [Status.EXPIRED, Status.COMPLETED]) + self.assertEqual([item.response for item in group], ["", "fresh response"]) + + async def test_common_put_drops_entire_token_expired_group_when_rerollout_is_disabled(self): + # 无 rerollout consumer 时，混合 group 整体终止并释放，缺口由新 prompt 补齐。 + for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: + with self.subTest(replay_buffer_config=config_name): + replay_buffer = replay_buffer_config_cls().build() + expired = make_rollout_state(1, response_model_steps=[0]) + fresh = make_rollout_state(2, response_model_steps=[4]) + + await replay_buffer.put( + [expired, fresh], + "task", + current_train_step=5, + stale_threshold=10, + token_stale_threshold=4, + expired_groups_retryable=False, + ) + + self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 0) + self.assertEqual(len(replay_buffer), 0) + self.assertEqual([item.status for item in (expired, fresh)], [Status.EXPIRED, Status.EXPIRED]) + self.assertIsNone(expired.prompt_ids) + self.assertIsNone(fresh.prompt_ids) + + async def test_common_refresh_token_expiry_moves_mixed_group_to_expired_pool(self): + # batch-start refresh 使用 consumer step 重新判定，并保留仍新鲜 state 的完整训练数据。 + for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: + with self.subTest(replay_buffer_config=config_name): + replay_buffer = replay_buffer_config_cls().build() + expired = make_rollout_state(1, response_model_steps=[0], reward={"score": 0.1}) + fresh = make_rollout_state(2, response_model_steps=[4], reward={"score": 0.9}) + await replay_buffer.put([expired, fresh], "task") + + expired_counts = await replay_buffer.refresh_staleness( + task_stale_thresholds={"task": 10}, + task_token_stale_thresholds={"task": 4}, + expired_groups_retryable_by_task={"task": True}, + current_train_step=5, + ) + + self.assertEqual(expired_counts, {"task": 1}) + self.assertEqual(await replay_buffer.count("task", Status.COMPLETED), 0) + self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 1) + group = (await replay_buffer.get(1, "task", Status.EXPIRED))[0] + self.assertEqual([item.status for item in group], [Status.EXPIRED, Status.COMPLETED]) + self.assertEqual(group[0].response_ids, []) + self.assertEqual(group[1].response_ids, [12]) + self.assertEqual(group[1].reward, {"score": 0.9}) + async def test_common_refresh_staleness_drops_only_terminal_expired_groups(self): # 同一轮 refresh 仍统计两类过期；只删除 terminal EXPIRED，保留 tail batch 可重试项。 for config_name, replay_buffer_config_cls in REPLAY_BUFFER_CONFIGS: diff --git a/tests/rl/test_staleness_policy.py b/tests/rl/test_staleness_policy.py index df97a545f..295c12bdc 100644 --- a/tests/rl/test_staleness_policy.py +++ b/tests/rl/test_staleness_policy.py @@ -1,15 +1,23 @@ +"""Staleness 配置与 mask 行为测试。 + +覆盖整组过期阈值的配置和校验，以及 token 级 staleness mask 的纯函数行为。 +""" + import unittest from pydantic import ValidationError +from xtuner.v1.data_proto.rl_data import RolloutState, calculate_effective_response_mask, reset_rollout_response from xtuner.v1.rl.agent_loop_manager import ( AsyncProduceStrategyConfig, DisaggAsyncProduceStrategyConfig, - calculate_stale_threshold, ) +from xtuner.v1.rl.agent_loop_manager.produce_utils import calculate_stale_threshold class TestStalenessPolicy(unittest.TestCase): + """整组 staleness 阈值、异步策略配置和非法参数校验。""" + def test_max_staleness_zero_uses_sync_interval_as_threshold(self): # max_staleness=0 表示只接受同步间隔内天然存在的最小滞后。 self.assertEqual(calculate_stale_threshold(max_staleness=0, sync_weights_interval=4), 4) @@ -19,18 +27,105 @@ def test_max_staleness_zero_uses_sync_interval_as_threshold(self): self.assertTrue(strategy.is_model_expired(train_step=9, model_step=4)) def test_max_staleness_one_allows_one_extra_sync_interval(self): + # max_staleness=1 表示额外接受一个权重同步周期的滞后。 self.assertEqual(calculate_stale_threshold(max_staleness=1, sync_weights_interval=4), 8) - strategy = DisaggAsyncProduceStrategyConfig(max_staleness=1).build(sync_weights_interval=4) + strategy = DisaggAsyncProduceStrategyConfig( + max_staleness=1, + enable_partial_rollout=True, + ).build(sync_weights_interval=4) self.assertFalse(strategy.is_model_expired(train_step=12, model_step=4)) self.assertTrue(strategy.is_model_expired(train_step=13, model_step=4)) def test_negative_max_staleness_is_invalid(self): + # Pydantic 配置层必须拒绝负的整组 staleness。 with self.assertRaises(ValidationError): AsyncProduceStrategyConfig(max_staleness=-1) with self.assertRaises(ValidationError): DisaggAsyncProduceStrategyConfig(max_staleness=-1) + def test_async_strategies_precompute_token_stale_threshold(self): + # colocated 和 disaggregated 异步策略应使用相同的 token 阈值换算。 + for config_cls in ( + AsyncProduceStrategyConfig, + DisaggAsyncProduceStrategyConfig, + ): + with self.subTest(config_cls=config_cls.__name__): + base_kwargs = {"max_staleness": 1, "enable_partial_rollout": True} + self.assertIsNone(config_cls(**base_kwargs).build(sync_weights_interval=4).token_stale_threshold) + self.assertEqual( + config_cls(max_token_staleness=0, **base_kwargs) + .build(sync_weights_interval=4) + .token_stale_threshold, + 4, + ) + self.assertEqual( + config_cls(max_token_staleness=1, **base_kwargs) + .build(sync_weights_interval=4) + .token_stale_threshold, + 8, + ) + +class TestTokenStalenessMask(unittest.TestCase): + """Token 级 staleness mask 的阈值与 semantic mask 行为。""" + + def test_token_staleness_threshold_can_be_relaxed(self): + # token threshold 放宽一个同步周期后，旧周期 token 应从 masked 变为可训练。 + for token_stale_threshold, expected in ((4, [0, 1]), (8, [1, 1])): + with self.subTest(token_stale_threshold=token_stale_threshold): + state = self._state(response_model_steps=[0, 4]) + + mask = calculate_effective_response_mask( + state, + current_train_step=5, + token_stale_threshold=token_stale_threshold, + ) + + self.assertEqual(mask, expected) + self.assertIsNone(state.response_mask) + + def test_token_staleness_intersects_semantic_response_mask(self): + # 最终 mask 必须同时满足 semantic mask 和 token staleness mask。 + state = self._state(response_model_steps=[0, 4], response_mask=[1, 0]) + + mask = calculate_effective_response_mask( + state, + current_train_step=5, + token_stale_threshold=4, + ) + + self.assertEqual(mask, [0, 0]) + + def test_rerolled_state_without_semantic_mask_uses_token_staleness_only(self): + state = reset_rollout_response(self._state(response_model_steps=[0, 4], response_mask=[0, 1])) + state.response_ids = [3, 4] + state.response_model_steps = [4, 4] + + mask = calculate_effective_response_mask( + state, + current_train_step=5, + token_stale_threshold=4, + ) + + self.assertIsNone(state.response_mask) + self.assertEqual(mask, [1, 1]) + + @staticmethod + def _state( + *, + response_model_steps: list[int] | None, + response_mask: list[int] | None = None, + ) -> RolloutState: + return RolloutState( + rollout_id=1, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=[1, 2], + response_ids=[3, 4], + response_model_steps=response_model_steps, + response_mask=response_mask, + ) + if __name__ == "__main__": unittest.main() diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 7db20f294..8f5652e20 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import Enum -from typing import TYPE_CHECKING, Any, Literal, TypeAlias +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast import numpy as np from pydantic import BaseModel, ConfigDict, Field @@ -246,7 +246,7 @@ def reset_rollout_response(rollout_state: RolloutState) -> RolloutState: rollout_state.logprobs = [] rollout_state.routed_experts = None rollout_state.finish_reason = None - rollout_state.response_mask = [] + rollout_state.response_mask = None rollout_state.response_model_steps = [] rollout_state.reward = None rollout_state.error_msg = None @@ -301,7 +301,7 @@ def update_sample_version(rollout_state: RolloutState, model_step: int) -> Rollo """Append token source model version for newly generated response tokens.""" response_len = len(rollout_state.response_ids or []) - response_model_steps = list(getattr(rollout_state, "response_model_steps", None) or []) + response_model_steps = list(rollout_state.response_model_steps or []) missing_response_steps = max(0, response_len - len(response_model_steps)) if missing_response_steps: response_model_steps.extend([model_step] * missing_response_steps) @@ -313,7 +313,7 @@ def refresh_seq_staleness(group: list[RolloutState], current_train_step: int) -> for rollout_state in group: # response_model_steps 记录每个 response token 的模型版本； # 最早版本决定整条样本的滞后程度。 - response_model_steps = getattr(rollout_state, "response_model_steps", None) or [] + response_model_steps = rollout_state.response_model_steps or [] if response_model_steps: rollout_state.seq_staleness = calculate_seq_staleness(min(response_model_steps), current_train_step) else: @@ -321,25 +321,37 @@ def refresh_seq_staleness(group: list[RolloutState], current_train_step: int) -> return group -def update_expired_status(samples: list[RolloutState], stale_threshold: int) -> list[RolloutState]: - if stale_threshold <= 0: - raise ValueError(f"stale_threshold must be positive, got {stale_threshold}.") - is_group_expired = False - - # 1. 检查组内是否存过期的样本 - for sample in samples: - if sample.status == Status.ABORTED and sample.seq_staleness >= stale_threshold: - logger.debug( - f"Sample {sample.rollout_id} (seq_staleness: {sample.seq_staleness}) exceeded threshold ({stale_threshold}). Triggering group expiration." - ) - is_group_expired = True - break # 一旦发现过期，直接跳出，无需检查剩余样本 - - # 2. 如果存在过期样本，将组内所有样本置为过期 - if is_group_expired: - # NOTE: 当一组数据中有一个样本被标记为过期后，这组数据中就可能出现未超过过期阈值但状态是 aborted 的样本。 - # 这些样本在后续的生成过程中也不应该被继续生成了，所以直接把它们都标记为过期, 才能在preprocess中将之前的response清掉。 - for sample in samples: - sample.status = Status.EXPIRED - - return samples +def calculate_effective_response_mask( + rollout_state: RolloutState, + *, + current_train_step: int, + token_stale_threshold: int, +) -> list[int]: + """Calculate the response mask after applying token staleness. + + Args: + rollout_state (RolloutState): Rollout sample whose response token provenance is evaluated. + current_train_step (int): Trainer step that will consume the sample. + token_stale_threshold (int): Maximum token staleness, measured in trainer steps, allowed for training. + + Returns: + list[int]: The semantic response mask intersected with the token-staleness mask. + """ + response_ids = cast(list[int], rollout_state.response_ids) + response_model_steps = cast(list[int], rollout_state.response_model_steps) + + # semantic mask: 在 agent_loop 中根据是否是 LLM 产生的 token 来 mask 的结果 + semantic_mask = rollout_state.response_mask + if semantic_mask is None: + semantic_mask = [1] * len(response_ids) + + # token_staleness_mask: 根据 token 的新鲜程度来 mask + token_staleness_mask = [ + int(calculate_seq_staleness(response_model_step, current_train_step) < token_stale_threshold) + for response_model_step in response_model_steps + ] + effective_mask = [ + semantic_mask_value * token_staleness_mask_value + for semantic_mask_value, token_staleness_mask_value in zip(semantic_mask, token_staleness_mask) + ] + return effective_mask diff --git a/xtuner/v1/rl/agent_loop_manager/__init__.py b/xtuner/v1/rl/agent_loop_manager/__init__.py index ddd90fd6b..dc64f0a39 100644 --- a/xtuner/v1/rl/agent_loop_manager/__init__.py +++ b/xtuner/v1/rl/agent_loop_manager/__init__.py @@ -17,7 +17,7 @@ DisaggProduceStrategy, DisaggProduceStrategyConfig, ) -from .produce_utils import ProduceBatchResult, ProduceBatchStatus, calculate_stale_threshold +from .produce_utils import ProduceBatchResult, ProduceBatchStatus from .producer import ( AsyncProduceStrategy, AsyncProduceStrategyConfig, @@ -56,7 +56,6 @@ "SyncProduceStrategy", "AsyncProduceStrategy", "DisaggAsyncProduceStrategy", - "calculate_stale_threshold", "SamplerConfig", "Sampler", ] diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index fd79c83ce..07b17e98f 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -226,14 +226,14 @@ async def produce_batch( task_names=self.task_names, target_samples=current_sizes, ) - # 生产前刷新已有 completed / aborted 的 staleness。 + # 生产前刷新可训练/可重试数据的 staleness。 await refresh_for_all_tasks( task_runners=self.task_runners, replay_buffer=self.replay_buffer, logger=self.logger, manager_name=self.name, train_step=train_step, - statuses=[Status.COMPLETED, Status.ABORTED], + statuses=[Status.COMPLETED, Status.ABORTED, Status.EXPIRED], ) produce_start = time.perf_counter() produce_futures = [] @@ -252,6 +252,7 @@ async def produce_batch( progress=local_progress, is_valid_sample_fn=task.is_valid_sample_fn, stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, expired_groups_retryable=task.expired_groups_retryable, ) ) @@ -277,6 +278,7 @@ async def produce_batch( progress=local_progress, is_valid_sample_fn=task.is_valid_sample_fn, stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, expired_groups_retryable=task.expired_groups_retryable, ) ) @@ -287,6 +289,7 @@ async def produce_batch( manager_name=self.name, task_batch_sizes=current_sizes, progress=local_progress, + current_train_step=train_step, pause_time_s=pause_time_s, ) assert result.rollout_states, ( diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index 3518b985c..f2d246989 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -242,6 +242,7 @@ async def _produce_batch_to_buffer( update_event=self._update_event, is_valid_sample_fn=task.is_valid_sample_fn, stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, expired_groups_retryable=task.expired_groups_retryable, ) ) @@ -271,6 +272,7 @@ async def pause_produce(self) -> float: update_event=self._update_event, is_valid_sample_fn=task.is_valid_sample_fn, stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, expired_groups_retryable=task.expired_groups_retryable, ) pause_time_s += await produce_strategy.pause_produce(ctx) @@ -334,7 +336,7 @@ async def get_batch(self, batch_size: int, train_step: int) -> ProduceBatchResul logger=self.logger, manager_name=self.name, train_step=train_step, - statuses=[Status.COMPLETED, Status.ABORTED], + statuses=[Status.COMPLETED, Status.ABORTED, Status.EXPIRED], ) task_batch_sizes = allocate_task_batch_sizes(self.task_runners, batch_size, train_step) current_model_step = train_step - 1 @@ -373,6 +375,7 @@ async def get_batch(self, batch_size: int, train_step: int) -> ProduceBatchResul manager_name=self.name, task_batch_sizes=task_batch_sizes, progress=progress, + current_train_step=train_step, pause_time_s=self._consume_pause_time(), ) if self._status == AgentLoopManagerStatus.EXPIRED_BATCH: @@ -386,7 +389,7 @@ async def get_batch(self, batch_size: int, train_step: int) -> ProduceBatchResul logger=self.logger, manager_name=self.name, train_step=train_step + 1, - statuses=[Status.COMPLETED, Status.ABORTED], + statuses=[Status.COMPLETED, Status.ABORTED, Status.EXPIRED], ) return result await asyncio.sleep(self._STATUS_POLL_INTERVAL_S) diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py index d75c34b7c..43437bbaa 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py @@ -243,12 +243,39 @@ def build( class DisaggAsyncProduceStrategyConfig(DisaggProduceStrategyConfig): - """非共卡异步生产配置。""" + """Configuration for disaggregated asynchronous rollout production. + + Args: + is_valid_sample_fn (IsValidSampleFn): Function used to decide whether a + generated rollout group is trainable. Defaults to + ``default_is_valid_sample_fn``. + should_continue_fn (ShouldContinueFn): Function used to decide whether + production should continue after a group is processed. Defaults to + ``default_should_continue_fn``. + over_sample_threshold (float): Extra completed-sample ratio allowed + before the producer stops. Defaults to 0.0. + enable_partial_rollout (bool): Whether unfinished rollouts can be + continued after a weight sync. Defaults to False. + max_staleness (int): Maximum allowed model-step staleness for replayed + samples. Defaults to 0. + 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. + tail_batch_trigger_size (int): Expired-group rerollout policy. ``-1`` + disables rerollout, ``0`` rerolls out immediately without entering + tail-batch mode, and ``N > 0`` waits until the expired pool contains + at least ``N`` groups before entering tail-batch mode. + """ over_sample_threshold: float = 0.0 enable_partial_rollout: bool = False 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) def build( self, @@ -256,6 +283,11 @@ def build( sync_weights_interval: int = 1, rollout_controller: "Optional[RolloutControllerProxy]" = None, ) -> "DisaggAsyncProduceStrategy": + if self.max_token_staleness is not None and self.max_token_staleness > self.max_staleness: + logger.warning( + "max_token_staleness is greater than max_staleness; token-level masking will not take effect " + "before the group expires." + ) if rollout_controller is not None: import ray @@ -264,6 +296,7 @@ def build( over_sample_threshold=self.over_sample_threshold, enable_partial_rollout=self.enable_partial_rollout, max_staleness=self.max_staleness, + max_token_staleness=self.max_token_staleness, sync_weights_interval=sync_weights_interval, tail_batch_trigger_size=self.tail_batch_trigger_size, is_valid_sample_fn=self.is_valid_sample_fn, @@ -304,6 +337,7 @@ def __init__( enable_partial_rollout: bool, tail_batch_trigger_size: int, max_staleness: int, + max_token_staleness: int | None, sync_weights_interval: int, is_valid_sample_fn: IsValidSampleFn, should_continue_fn: ShouldContinueFn, @@ -318,8 +352,12 @@ def __init__( self.over_sample_threshold = over_sample_threshold self.enable_partial_rollout = enable_partial_rollout self.max_staleness = max_staleness - self.sync_weights_interval = sync_weights_interval self.stale_threshold = calculate_stale_threshold(max_staleness, sync_weights_interval) + self.token_stale_threshold = ( + None + if max_token_staleness is None + else calculate_stale_threshold(max_token_staleness, sync_weights_interval) + ) self.tail_batch_trigger_size = tail_batch_trigger_size self._pending_tasks = _PendingTasks() @@ -363,8 +401,11 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: return ProduceBatchStatus.NORMAL expired_count = await ctx.expired_count() - sample_from_expired = self.tail_batch_trigger_size > 0 and expired_count >= self.tail_batch_trigger_size - if sample_from_expired: + sample_expired = ( + self.tail_batch_trigger_size >= 0 and expired_count > 0 and expired_count >= self.tail_batch_trigger_size + ) + tail_batch_triggered = self.tail_batch_trigger_size > 0 and expired_count >= self.tail_batch_trigger_size + if tail_batch_triggered: logger.info( f"Tail batch trigger condition met: {expired_count} expired samples " f"(threshold: {self.tail_batch_trigger_size}). Enabling tail batch mode." @@ -372,7 +413,7 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: # normal 使用固定超发预算；tail-batch 只补必要缺口。 total_target = ctx.total_target - oversample_budget = 0 if sample_from_expired else math.ceil(self.over_sample_threshold * ctx.task_batch_size) + oversample_budget = 0 if tail_batch_triggered else math.ceil(self.over_sample_threshold * ctx.task_batch_size) scheduled_target = total_target + oversample_budget logger.info( f"Starting produce_batch for task {ctx.task_name} with total_target={total_target}, " @@ -380,7 +421,7 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: ) async def spawn_one() -> asyncio.Task: - rollout_state = await ctx.sample_group(from_expired_pool=sample_from_expired) + rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( ctx.generate_group( rollout_state, diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 25e0541f2..2ea63df39 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import math import time @@ -13,6 +15,7 @@ from xtuner.v1.data_proto.rl_data import ( RolloutState, Status, + calculate_effective_response_mask, discard_rollout_state, get_group_status, ) @@ -44,7 +47,7 @@ def __init__(self, progress_bar: Any | None) -> None: self._tqdm = progress_bar @classmethod - def create(cls, *, strategy_name: str, task_name: str, total: int, initial: int) -> "_ProgressDisplayer": + def create(cls, *, strategy_name: str, task_name: str, total: int, initial: int) -> _ProgressDisplayer: total = max(0, total) initial = min(total, max(0, initial)) if total <= 0 or get_rank() != 0: @@ -92,11 +95,6 @@ def default_should_continue_fn(completed_count: int, batch_size: int, **kwargs) def calculate_stale_threshold(max_staleness: int, sync_weights_interval: int) -> int: - if max_staleness < 0: - raise ValueError(f"max_staleness must be non-negative, got {max_staleness}.") - if sync_weights_interval <= 0: - raise ValueError(f"sync_weights_interval must be positive, got {sync_weights_interval}.") - # max_staleness 按同步周期计数；+1 表示训练天然必须接受的当前同步周期滞后。 return (max_staleness + 1) * sync_weights_interval @@ -122,10 +120,11 @@ class BaseProduceContext: task_name: str train_step: int model_step: int - progress: "ProduceProgress | DisaggProduceProgress" + progress: ProduceProgress | DisaggProduceProgress is_valid_sample_fn: IsValidSampleFn = default_is_valid_sample_fn stale_threshold: int | None = None expired_groups_retryable: bool = True + token_stale_threshold: int | None = None @property def consumer_step(self) -> int: @@ -171,7 +170,7 @@ async def generate_group( return result 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) initial_status = get_group_status(group) discard_status: Status | None = None @@ -209,6 +208,7 @@ async def put_generated_group(self, group: list[RolloutState]) -> bool: model_step=self.model_step, current_train_step=self.consumer_step, stale_threshold=self.stale_threshold, + token_stale_threshold=self.token_stale_threshold, expired_groups_retryable=self.expired_groups_retryable, ) self.progress.add_produced(self.task_name, samples=len(group), tokens=produced_tokens) @@ -265,7 +265,7 @@ class ProduceBatchResult: produced_tokens: int = 0 produce_time_s: float = 0.0 task_batch_sizes: dict[str, int] | None = None - task_results: dict[str, "ProduceBatchResult"] | None = None + task_results: dict[str, ProduceBatchResult] | None = None @dataclass(frozen=True) @@ -287,7 +287,11 @@ def stale_threshold(self) -> int | None: @property def expired_groups_retryable(self) -> bool: - return getattr(self.produce_strategy, "tail_batch_trigger_size", 0) > 0 + return getattr(self.produce_strategy, "tail_batch_trigger_size", -1) >= 0 + + @property + def token_stale_threshold(self) -> int | None: + return getattr(self.produce_strategy, "token_stale_threshold", None) class _TaskSamplerView: @@ -357,7 +361,7 @@ def _fill_leftover_counts(result: ProduceBatchResult, status_counts: dict[Status def _merge_discarded_counts( result: ProduceBatchResult, - progress: "ProduceProgress | DisaggProduceProgress", + progress: ProduceProgress | DisaggProduceProgress, task_name: str, ) -> None: discarded_failed, discarded_filtered = progress.consume_discarded(task_name) @@ -430,14 +434,18 @@ async def refresh_for_all_tasks( statuses: list[Status], ) -> None: task_stale_thresholds: dict[str, int] = {} + task_token_stale_thresholds: dict[str, int] = {} expired_groups_retryable_by_task: dict[str, bool] = {} for task in task_runners: # 没有 stale_threshold 的同步策略按 1 处理。 task_stale_thresholds[task.task_name] = task.stale_threshold if task.stale_threshold is not None else 1 + if task.token_stale_threshold is not None: + task_token_stale_thresholds[task.task_name] = task.token_stale_threshold expired_groups_retryable_by_task[task.task_name] = task.expired_groups_retryable expired_counts = await replay_buffer.refresh_staleness( task_stale_thresholds=task_stale_thresholds, + task_token_stale_thresholds=task_token_stale_thresholds, expired_groups_retryable_by_task=expired_groups_retryable_by_task, current_train_step=train_step, statuses=statuses, @@ -547,7 +555,7 @@ def build_produce_batch_result( task_batch_sizes: dict[str, int], batch_by_task: dict[str, list[list[RolloutState]]], leftover_counts: dict[str, dict[Status, int]], - progress: "ProduceProgress | DisaggProduceProgress", + progress: ProduceProgress | DisaggProduceProgress, pause_time_s: float, ) -> ProduceBatchResult: if len(task_runners) == 1: @@ -599,10 +607,26 @@ async def take_train_batch( logger, manager_name: str, task_batch_sizes: dict[str, int], - progress: "ProduceProgress | DisaggProduceProgress", + progress: ProduceProgress | DisaggProduceProgress, + current_train_step: int, pause_time_s: float = 0.0, ) -> ProduceBatchResult: batch_by_task, consumed_counts = await replay_buffer.take_batch(task_batch_sizes) + + 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, + ) + if hasattr(progress, "mark_consumed"): progress.mark_consumed(consumed_counts) task_names = [task.task_name for task in task_runners] diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 5620be777..c1be6a58f 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -177,7 +177,8 @@ def build( rollout_controller: "Optional[RolloutControllerProxy]" = None, ) -> "SyncProduceStrategy": return SyncProduceStrategy( - is_valid_sample_fn=self.is_valid_sample_fn, should_continue_fn=self.should_continue_fn + is_valid_sample_fn=self.is_valid_sample_fn, + should_continue_fn=self.should_continue_fn, ) @@ -202,8 +203,17 @@ class AsyncProduceStrategyConfig(ProduceStrategyConfig): continued after a weight sync. Defaults to False. max_staleness (int): Maximum allowed model-step staleness for replayed samples. Defaults to 0. - tail_batch_trigger_size (int): Minimum pending tail size that can - trigger a final batch. Defaults to 0. + 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. + tail_batch_trigger_size (int): Expired-group rerollout policy. ``-1`` + disables rerollout, ``0`` rerolls out immediately without entering + tail-batch mode, and ``N > 0`` waits until the expired pool contains + at least ``N`` groups before entering tail-batch mode. **Examples:** @@ -219,7 +229,8 @@ class AsyncProduceStrategyConfig(ProduceStrategyConfig): over_sample_threshold: float = 0.0 enable_partial_rollout: bool = False 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) def build( self, @@ -227,6 +238,11 @@ def build( sync_weights_interval: int = 1, rollout_controller: "Optional[RolloutControllerProxy]" = None, ) -> "AsyncProduceStrategy": + if self.max_token_staleness is not None and self.max_token_staleness > self.max_staleness: + logger.warning( + "max_token_staleness is greater than max_staleness; token-level masking will not take effect " + "before the group expires." + ) if rollout_controller is not None: import ray @@ -235,6 +251,7 @@ def build( over_sample_threshold=self.over_sample_threshold, enable_partial_rollout=self.enable_partial_rollout, max_staleness=self.max_staleness, + max_token_staleness=self.max_token_staleness, sync_weights_interval=sync_weights_interval, tail_batch_trigger_size=self.tail_batch_trigger_size, is_valid_sample_fn=self.is_valid_sample_fn, @@ -315,6 +332,7 @@ def __init__( enable_partial_rollout: bool, tail_batch_trigger_size: int, max_staleness: int, + max_token_staleness: int | None, sync_weights_interval: int, is_valid_sample_fn: IsValidSampleFn, should_continue_fn: ShouldContinueFn, @@ -336,8 +354,12 @@ def __init__( self.over_sample_threshold = over_sample_threshold self.enable_partial_rollout = enable_partial_rollout self.max_staleness = max_staleness - self.sync_weights_interval = sync_weights_interval self.stale_threshold = calculate_stale_threshold(max_staleness, sync_weights_interval) + self.token_stale_threshold = ( + None + if max_token_staleness is None + else calculate_stale_threshold(max_token_staleness, sync_weights_interval) + ) self.tail_batch_trigger_size = tail_batch_trigger_size self._local_pending_tasks: set[asyncio.Task] = set() @@ -362,8 +384,11 @@ async def produce_batch(self, ctx: ProduceContext) -> None: return expired_count = await ctx.expired_count() - sample_from_expired = self.tail_batch_trigger_size > 0 and expired_count >= self.tail_batch_trigger_size - if sample_from_expired: + sample_expired = ( + self.tail_batch_trigger_size >= 0 and expired_count > 0 and expired_count >= self.tail_batch_trigger_size + ) + tail_batch_triggered = self.tail_batch_trigger_size > 0 and expired_count >= self.tail_batch_trigger_size + if tail_batch_triggered: logger.info( f"Tail batch trigger condition met: {expired_count} expired samples " f"(threshold: {self.tail_batch_trigger_size}). Enabling tail batch mode." @@ -371,7 +396,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: # normal 使用固定超发预算；tail-batch 只补必要缺口。 batch_target = ctx.batch_target - oversample_budget = 0 if sample_from_expired else math.ceil(self.over_sample_threshold * ctx.task_batch_size) + oversample_budget = 0 if tail_batch_triggered else math.ceil(self.over_sample_threshold * ctx.task_batch_size) scheduled_target = batch_target + oversample_budget logger.info( f"Starting produce_batch for task {ctx.task_name} with batch_target={batch_target}, " @@ -379,7 +404,7 @@ async def produce_batch(self, ctx: ProduceContext) -> None: ) async def spawn_one() -> asyncio.Task: - rollout_state = await ctx.sample_group(from_expired_pool=sample_from_expired) + rollout_state = await ctx.sample_group(from_expired_pool=sample_expired) return create_task( ctx.generate_group( rollout_state, diff --git a/xtuner/v1/rl/replay_buffer.py b/xtuner/v1/rl/replay_buffer.py index 3ae42a66a..58db4fe8f 100644 --- a/xtuner/v1/rl/replay_buffer.py +++ b/xtuner/v1/rl/replay_buffer.py @@ -13,6 +13,7 @@ from xtuner.v1.data_proto.rl_data import ( RolloutState, Status, + calculate_effective_response_mask, discard_rollout_state, get_group_status, refresh_seq_staleness, @@ -36,19 +37,6 @@ logger = get_logger(__name__) -def maybe_expire_group(group: list[RolloutState], stale_threshold: int) -> None: - if stale_threshold <= 0: - raise ValueError(f"stale_threshold must be positive, got {stale_threshold}.") - - group_status = get_group_status(group) - if group_status not in (Status.COMPLETED, Status.ABORTED): - return - if any(getattr(sample, "seq_staleness", 0) >= stale_threshold for sample in group): - # 生成结果入库前统一做过期翻转，后续存储逻辑只按最终 group status 分类。 - for sample in group: - sample.status = Status.EXPIRED - - @dataclass class StorageItem: # 存储类型 @@ -451,20 +439,65 @@ def __init__(self, policy: ReplayPolicy, storage_backend: StorageBackend): self._storage = storage_backend self._lock = asyncio.Lock() - def _cleanup_expired_group(self, group: list[RolloutState], *, retryable: bool) -> None: - """Release stale state according to whether the group can be - retried.""" + def _apply_staleness_lifecycle( + self, + group: list[RolloutState], + *, + current_train_step: int | None, + stale_threshold: int | None, + token_stale_threshold: int | None, + expired_groups_retryable: bool, + ) -> Status: + """Refresh one group's staleness, expire it when needed, and clean it + up.""" + storage_status = get_group_status(group) + if current_train_step is None or storage_status not in (Status.COMPLETED, Status.ABORTED, Status.EXPIRED): + return storage_status + # NOTE: input_ids/labels 表示 agentic 训练分支，当前暂不支持 agentic token-expired lifecycle。 + is_agentic_group = any(item.input_ids is not None or item.labels is not None for item in group) + # NOTE: An EXPIRED group may still contain COMPLETED states whose responses were preserved. + # Refresh the group again so those states can also expire while waiting for rerollout. + expired_mask = [item.status == Status.EXPIRED for item in group] + + # 1. update seq-level staleness + refresh_seq_staleness(group, current_train_step) + + for index, item in enumerate(group): + if expired_mask[index]: + continue + if stale_threshold is not None and item.seq_staleness >= stale_threshold: + expired_mask[index] = True + continue + if is_agentic_group or token_stale_threshold is None: + continue + if not item.response_ids or (item.response_mask is not None and not any(item.response_mask)): + continue - for item in group: - if retryable: - # Tail batch may reroll this sample. Keep prompt and multimodal - # training inputs, but release the stale response and routed experts. - reset_rollout_response(item) - else: - # No consumer can retry this terminal group. Release all optional - # state before dropping the group's final strong references. + # 2. update token-level staleness + effective_mask = calculate_effective_response_mask( + item, + current_train_step=current_train_step, + token_stale_threshold=token_stale_threshold, + ) + if not any(effective_mask): + expired_mask[index] = True + + # 3. return storage_status when no expired sample in group + if not any(expired_mask): + 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) + else: + for item in group: discard_rollout_state(item) - item.status = Status.EXPIRED + item.status = Status.EXPIRED + + return Status.EXPIRED async def put( self, @@ -474,6 +507,7 @@ async def put( model_step: int | None = None, current_train_step: int | None = None, stale_threshold: int | None = None, + token_stale_threshold: int | None = None, expired_groups_retryable: bool = True, ) -> None: if not items: @@ -481,17 +515,16 @@ async def put( if model_step is not None: for item in items: update_sample_version(item, model_step) - if current_train_step is not None: - refresh_seq_staleness(items, current_train_step) - if stale_threshold is not None: - maybe_expire_group(items, stale_threshold) - - status = get_group_status(items) + status = self._apply_staleness_lifecycle( + items, + current_train_step=current_train_step, + stale_threshold=stale_threshold, + token_stale_threshold=token_stale_threshold, + expired_groups_retryable=expired_groups_retryable, + ) staleness = max(item.seq_staleness for item in items) - if status == Status.EXPIRED: - self._cleanup_expired_group(items, retryable=expired_groups_retryable) - if not expired_groups_retryable: - return + if status == Status.EXPIRED and not expired_groups_retryable: + return storage_item = StorageItem( item=items, uid=0, @@ -519,18 +552,18 @@ async def refresh_staleness( self, *, task_stale_thresholds: dict[str, int], + task_token_stale_thresholds: dict[str, int] | None = None, expired_groups_retryable_by_task: dict[str, bool] | None = None, current_train_step: int, statuses: list[Status] | None = None, ) -> dict[str, int]: - # 刷新可复用样本的 staleness；completed / aborted 都可能来自旧权重，需要按 train_step 淘汰。 - for task_name, stale_threshold in task_stale_thresholds.items(): - if stale_threshold <= 0: - raise ValueError(f"stale_threshold must be positive, got {stale_threshold}.") + # 刷新可复用样本的 staleness；EXPIRED group 中保留的 COMPLETED state + # 在等待 rerollout 期间也可能继续变旧。 if statuses is None: - statuses = [Status.COMPLETED, Status.ABORTED] + statuses = [Status.COMPLETED, Status.ABORTED, Status.EXPIRED] expired_counts: dict[str, int] = {} retryable_by_task = expired_groups_retryable_by_task or {} + token_stale_thresholds = task_token_stale_thresholds or {} async with self._lock: updated_records: list[StorageItem] = [] deleted_uids: list[int] = [] @@ -544,19 +577,20 @@ async def refresh_staleness( records = await self._storage.get(query_dsl) expired_count = 0 for record in records: - refresh_seq_staleness(record.item, current_train_step) + retryable = retryable_by_task.get(task_name, True) + status = self._apply_staleness_lifecycle( + record.item, + current_train_step=current_train_step, + stale_threshold=stale_threshold, + token_stale_threshold=token_stale_thresholds.get(task_name), + expired_groups_retryable=retryable, + ) staleness = max((getattr(item, "seq_staleness", 0) for item in record.item), default=0) - should_expire = any(getattr(item, "seq_staleness", 0) >= stale_threshold for item in record.item) - if should_expire: - retryable = retryable_by_task.get(task_name, True) - self._cleanup_expired_group(record.item, retryable=retryable) - status = Status.EXPIRED + if status == Status.EXPIRED: expired_count += 1 if not retryable: deleted_uids.append(record.uid) continue - else: - status = get_group_status(record.item) updated_records.append(replace(record, status=status, staleness=staleness)) expired_counts[task_name] = expired_count await self._storage.delete(deleted_uids) diff --git a/xtuner/v1/rl/rollout/utils.py b/xtuner/v1/rl/rollout/utils.py index 5c688a0a4..94ccc8023 100644 --- a/xtuner/v1/rl/rollout/utils.py +++ b/xtuner/v1/rl/rollout/utils.py @@ -115,8 +115,8 @@ class PartialRolloutHandler: request. This handler only knows how to continue a single interrupted generation by reusing the previous response as the - next engine input. Agent-loop level multi-turn rollout, including tool messages and response masks, must be handled - by the agent loop itself. + next engine input. Agent-loop level multi-turn rollout, including tool messages, must be handled by the agent loop + itself. """ def __init__(self) -> None: @@ -151,6 +151,7 @@ async def postprocess( prompt_tokens: int, completion_tokens: int, ) -> RolloutState: + """Postprocess a partial rollout using the default semantics.""" rollout_state.finish_reason = finish_reason rollout_state.status = status history_response = rollout_state.response or "" diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index faa6abbc3..9bc1eb420 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -832,7 +832,7 @@ async def generate(self, rollout_state: RolloutState) -> RolloutState: f"No generation needed for request {uid}: max_tokens={payload_max_tokens} or last input_id={last_id} is in eos_token." ) finish_reason = "stop" if is_eos_reached else "length" - # 对于是否开 partial rollout 的情况都直接标记为完成并返回，因为本轮 rollout 未开始，也不需要拼接 + # 本轮 rollout 未开始，不需要执行 partial-rollout 拼接。 rollout_state.finish_reason = finish_reason rollout_state.status = Status.COMPLETED return rollout_state diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 9e17bf0de..a6a85fd76 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1166,7 +1166,7 @@ def _prepare_train_data( response_len_list.append(len(response_ids)) # 根据 response_mask 计算 response_ids 对应的shifted_labels - if not group[i].response_mask: + if group[i].response_mask is None: response_mask = [1] * len(response_ids) response_labels = response_ids else: