diff --git a/tests/rl/conftest.py b/tests/rl/conftest.py index ebb7714d52..3920ecc867 100644 --- a/tests/rl/conftest.py +++ b/tests/rl/conftest.py @@ -48,8 +48,7 @@ def _pytest_invocation_paths() -> list[Path]: def _should_install_import_stubs() -> bool: paths = _pytest_invocation_paths() return bool(paths) and all( - path.is_file() and path.parent == _RL_TEST_DIR and path.name in _LIGHTWEIGHT_TEST_FILES - for path in paths + path.is_file() and path.parent == _RL_TEST_DIR and path.name in _LIGHTWEIGHT_TEST_FILES for path in paths ) @@ -147,19 +146,16 @@ class TrainingWorker: controller_mod = _new_module("xtuner.v1.rl.trainer.controller") controller_mod.TrainingController = TrainingController - controller_mod.ColateItem = object - controller_mod.__all__ = ["TrainingController", "ColateItem"] + controller_mod.__all__ = ["TrainingController"] worker_mod = _new_module("xtuner.v1.rl.trainer.worker") worker_mod.TrainingWorker = TrainingWorker worker_mod.WorkerConfig = WorkerConfig - worker_mod.WorkerInputItem = dict worker_mod.WorkerLogItem = dict worker_mod.WorkerTrainLogItem = dict worker_mod.__all__ = [ "TrainingWorker", "WorkerConfig", - "WorkerInputItem", "WorkerLogItem", "WorkerTrainLogItem", ] @@ -167,10 +163,8 @@ class TrainingWorker: trainer_pkg.TrainingController = TrainingController trainer_pkg.WorkerConfig = WorkerConfig trainer_pkg.TrainingWorker = TrainingWorker - trainer_pkg.WorkerInputItem = dict trainer_pkg.WorkerLogItem = dict trainer_pkg.WorkerTrainLogItem = dict - trainer_pkg.ColateItem = object sys.modules.setdefault("xtuner.v1.rl.trainer", trainer_pkg) sys.modules.setdefault("xtuner.v1.rl.trainer.controller", controller_mod) diff --git a/tests/rl/test_multi_task_agent_loop_manager.py b/tests/rl/test_multi_task_agent_loop_manager.py index 549ded8d90..4253a8ac46 100644 --- a/tests/rl/test_multi_task_agent_loop_manager.py +++ b/tests/rl/test_multi_task_agent_loop_manager.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.agent_loop.agent_loop import AgentLoop from xtuner.v1.rl.agent_loop_manager.agent_loop_manager import ( AgentLoopManager, AgentLoopManagerConfig, @@ -33,7 +34,7 @@ _TaskRunner, ) from xtuner.v1.rl.agent_loop_manager.producer import SyncProduceStrategyConfig -from xtuner.v1.train.rl_trainer import _migrate_legacy_is_valid_sample_fn +from xtuner.v1.train.rl_trainer import BaseRLTrainer, _migrate_legacy_is_valid_sample_fn class _FakeSampler: @@ -120,10 +121,18 @@ async def pause_produce(self, ctx) -> float: raise RuntimeError("cleanup failure") -class _FakeRolloutState: - def __init__(self, rollout_id: str, group_generate_time_s: float): - self.rollout_id = rollout_id - self.extra_fields = {GROUP_GENERATE_TIME_KEY: group_generate_time_s} +def _named_rollout_state(name: str, group_generate_time_s: float | None = None) -> RolloutState: + extra_fields = {"test_name": name} + if group_generate_time_s is not None: + extra_fields[GROUP_GENERATE_TIME_KEY] = group_generate_time_s + return RolloutState( + rollout_id=int(name.split("-")[-1]), + message=[], + status=Status.COMPLETED, + input_ids=[1, 2], + labels=[-100, 2], + extra_fields=extra_fields, + ) class _FakeReplayBuffer: @@ -312,7 +321,7 @@ def test_trainer_rejects_conflicting_legacy_and_task_validity_checks(self): _migrate_legacy_is_valid_sample_fn(cfg, MagicMock()) async def test_take_train_batch_applies_token_staleness_mask(self): - # 启用 token staleness 时，公开 produce_batch 路径应把过期 token 的监督清出 labels（只清零不恢复）。 + # take_train_batch 先烙制 staleness mask，训练入口检查后保留该 mask。 state = RolloutState( rollout_id=1, group_id=1, @@ -324,6 +333,7 @@ async def test_take_train_batch_applies_token_staleness_mask(self): response_model_steps=[0, 4], status=Status.COMPLETED, ) + AgentLoop._write_train_meta(None, state) strategy = _FakeProduceStrategy(token_stale_threshold=4) replay_buffer = _FakeReplayBuffer( rollout_states_by_task={"task": [[state]]}, @@ -345,6 +355,9 @@ async def test_take_train_batch_applies_token_staleness_mask(self): ) result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, 4]) + trainer = SimpleNamespace(logger=MagicMock(), agent_loop_manager=manager) + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 5, result) self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, 4]) self.assertEqual(replay_buffer.task_token_stale_threshold_calls, [{"task": 4}]) @@ -360,6 +373,7 @@ async def test_take_train_batch_skips_states_without_response_ids(self): logprobs=[0.0, -0.1], status=Status.COMPLETED, ) + AgentLoop._write_train_meta(None, state) strategy = _FakeProduceStrategy(token_stale_threshold=4) manager = AgentLoopManager( task_runners=[ @@ -380,6 +394,9 @@ async def test_take_train_batch_skips_states_without_response_ids(self): ) result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + self.assertEqual(result.rollout_states[0][0].labels, [-100, 2]) + trainer = SimpleNamespace(logger=MagicMock(), agent_loop_manager=manager) + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 5, result) self.assertEqual(result.rollout_states[0][0].labels, [-100, 2]) @@ -398,6 +415,7 @@ async def test_take_train_batch_bakes_full_sequence_agentic_per_token_staleness( labels=[-100, -100, 30, -100, 40], status=Status.COMPLETED, ) + AgentLoop._write_train_meta(None, state) strategy = _FakeProduceStrategy(token_stale_threshold=4) manager = AgentLoopManager( task_runners=[ @@ -418,6 +436,9 @@ async def test_take_train_batch_bakes_full_sequence_agentic_per_token_staleness( ) result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, -100, 40]) + trainer = SimpleNamespace(logger=MagicMock(), agent_loop_manager=manager) + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 5, result) self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, -100, 40]) @@ -433,6 +454,8 @@ async def test_take_train_batch_sync_path_leaves_labels_untouched(self): labels=[-100, -100, 4, 5, -100], status=Status.COMPLETED, ) + state.response_model_steps = [0, 0, 0] + AgentLoop._write_train_meta(None, state) manager = AgentLoopManager( task_runners=[ _TaskRunner( @@ -452,6 +475,9 @@ async def test_take_train_batch_sync_path_leaves_labels_untouched(self): ) result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4) + self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, 4, 5, -100]) + trainer = SimpleNamespace(logger=MagicMock(), agent_loop_manager=manager) + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 5, result) self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, 4, 5, -100]) @@ -462,8 +488,12 @@ async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results strategy_c = _FakeProduceStrategy(tail_batch_trigger_size=0) replay_buffer = _FakeReplayBuffer( rollout_states_by_task={ - "task_a": [["a-0"], ["a-1"]], - "task_b": [["b-0"], ["b-1"], ["b-2"]], + "task_a": [[_named_rollout_state("a-0")], [_named_rollout_state("a-1")]], + "task_b": [ + [_named_rollout_state("b-0")], + [_named_rollout_state("b-1")], + [_named_rollout_state("b-2")], + ], "task_c": [], }, leftover_counts={ @@ -506,7 +536,10 @@ async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results result = await multi_task_manager.produce_batch(batch_size=7, train_step=3, model_step=2) self.assertEqual(result.task_batch_sizes, {"task_a": 5, "task_b": 2, "task_c": 0}) - self.assertEqual(result.rollout_states, [["a-0"], ["a-1"], ["b-0"], ["b-1"]]) + self.assertEqual( + [[item.extra_fields["test_name"] for item in group] for group in result.rollout_states], + [["a-0"], ["a-1"], ["b-0"], ["b-1"]], + ) self.assertEqual(result.leftover_init, 0) self.assertEqual(result.leftover_completed, 1) self.assertEqual(result.leftover_aborted, 2) @@ -529,8 +562,8 @@ async def test_disagg_get_batch_aggregates_multi_task_results_without_colocate_s # 非共卡 get_batch 使用后台 progress 消费 replay buffer，不依赖共卡 produce_batch 继承面。 replay_buffer = _FakeReplayBuffer( rollout_states_by_task={ - "task_a": [["a-0"], ["a-1"]], - "task_b": [["b-0"]], + "task_a": [[_named_rollout_state("a-0")], [_named_rollout_state("a-1")]], + "task_b": [[_named_rollout_state("b-0")]], }, leftover_counts={ ("task_a", Status.COMPLETED): 2, @@ -562,7 +595,10 @@ async def test_disagg_get_batch_aggregates_multi_task_results_without_colocate_s result = await manager.get_batch(batch_size=3, train_step=2) - self.assertEqual(result.rollout_states, [["a-0"], ["a-1"], ["b-0"]]) + self.assertEqual( + [[item.extra_fields["test_name"] for item in group] for group in result.rollout_states], + [["a-0"], ["a-1"], ["b-0"]], + ) self.assertEqual(result.task_batch_sizes, {"task_a": 2, "task_b": 1}) async def test_produce_batch_uses_cleanup_and_reconstructs_group_timing_stats(self): @@ -571,8 +607,8 @@ async def test_produce_batch_uses_cleanup_and_reconstructs_group_timing_stats(se replay_buffer = _FakeReplayBuffer( rollout_states_by_task={ "task_a": [ - [_FakeRolloutState("a-0", 0.5)], - [_FakeRolloutState("a-1", 1.0)], + [_named_rollout_state("a-0", 0.5)], + [_named_rollout_state("a-1", 1.0)], ], }, leftover_counts={}, @@ -652,9 +688,9 @@ async def test_produce_batch_status_stays_normal_for_colocate_flow(self): ], replay_buffer=_FakeReplayBuffer( rollout_states_by_task={ - "task_a": [["a-0"]], - "task_b": [["b-0"]], - "task_c": [["c-0"]], + "task_a": [[_named_rollout_state("a-0")]], + "task_b": [[_named_rollout_state("b-0")]], + "task_c": [[_named_rollout_state("c-0")]], }, leftover_counts={}, ), @@ -664,7 +700,10 @@ async def test_produce_batch_status_stays_normal_for_colocate_flow(self): result = await manager.produce_batch(batch_size=3, train_step=6, model_step=5) self.assertEqual(result.status, ProduceBatchStatus.NORMAL) - self.assertEqual(result.rollout_states, [["a-0"], ["b-0"], ["c-0"]]) + self.assertEqual( + [[item.extra_fields["test_name"] for item in group] for group in result.rollout_states], + [["a-0"], ["b-0"], ["c-0"]], + ) async def test_produce_batch_waits_all_tasks_before_any_pause(self): # fast task 不能在 slow task 仍在生产时提前 pause rollout worker。 @@ -712,7 +751,10 @@ async def pause_produce(self, ctx) -> float: ), ], replay_buffer=_FakeReplayBuffer( - rollout_states_by_task={"task_fast": [["fast-0"]], "task_slow": [["slow-0"]]}, + rollout_states_by_task={ + "task_fast": [[_named_rollout_state("fast-0")]], + "task_slow": [[_named_rollout_state("slow-0")]], + }, leftover_counts={}, ), rollout_controller=_fake_rollout_controller(), @@ -720,7 +762,10 @@ async def pause_produce(self, ctx) -> float: result = await manager.produce_batch(batch_size=2, train_step=4, model_step=3) - self.assertEqual(result.rollout_states, [["fast-0"], ["slow-0"]]) + self.assertEqual( + [[item.extra_fields["test_name"] for item in group] for group in result.rollout_states], + [["fast-0"], ["slow-0"]], + ) self.assertEqual(events[:2], ["fast_produce_done", "slow_produce_done"]) async def test_produce_batch_preserves_original_terminal_exception(self): diff --git a/tests/rl/test_on_policy_distillation.py b/tests/rl/test_on_policy_distillation.py index f41755c8dd..73eff14d37 100644 --- a/tests/rl/test_on_policy_distillation.py +++ b/tests/rl/test_on_policy_distillation.py @@ -11,7 +11,6 @@ build_teacher_launch_server_commands, ) from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status -from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.rl.distillation import ( DistillationConfig, DistillationTrainerAdapter, @@ -23,7 +22,6 @@ TrainTeacherTimings, ) from xtuner.v1.rl.loss import DistillationLossConfig, DistillationLossKwargs -from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.train.rl_trainer import BaseRLTrainer, BaseRLTrainerConfig @@ -609,34 +607,5 @@ async def test_invalid_topk_response_is_retried_before_success(self) -> None: self.assertEqual(request_payload["top_logprobs_num"], 2) -class TestTopKTrainingController(unittest.TestCase): - def test_packs_targets_along_sequence_dimension(self) -> None: - controller = TrainingController(workers=[]) - first = { - "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[1, 2]]),), device="cpu"), - "shifted_labels": torch.tensor([[-100, 2]]), - "advantage": [0.0, 0.0], - "rollout_logprobs": torch.zeros(1, 2), - "teacher_logprobs": torch.tensor([[[-0.1, -0.2], [-0.3, -0.4]]]), - "target_token_ids": torch.tensor([[[1, 2], [3, 4]]]), - } - second = { - "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[3]]),), device="cpu"), - "shifted_labels": torch.tensor([[3]]), - "advantage": [0.0], - "rollout_logprobs": torch.zeros(1, 1), - "teacher_logprobs": torch.tensor([[[-0.5, -0.6]]]), - "target_token_ids": torch.tensor([[[5, 6]]]), - } - - packed = controller._packing([first, second], pack_max_length=4, language_cfg=None) - - self.assertEqual(len(packed), 1) - self.assertEqual(packed[0]["teacher_logprobs"].shape, (1, 4, 2)) - self.assertEqual(packed[0]["target_token_ids"].shape, (1, 4, 2)) - torch.testing.assert_close(packed[0]["teacher_logprobs"][0, 3], torch.zeros(2)) - torch.testing.assert_close(packed[0]["target_token_ids"][0, 3], torch.zeros(2, dtype=torch.long)) - - if __name__ == "__main__": unittest.main() diff --git a/tests/rl/test_pack.py b/tests/rl/test_pack.py new file mode 100644 index 0000000000..b7d5c39c71 --- /dev/null +++ b/tests/rl/test_pack.py @@ -0,0 +1,268 @@ +"""RLDataPacker index-only packing plan 的 contract 测试。 + +Packer 只产出 ``[dp][optimizer_step][pack][sample_index]`` 索引计划与 padding 统计， +不触碰样本数据。测试覆盖设计约束：每个 index 恰好出现一次、pack 不超长、DP/step 维度 +pack 数一致、optimizer steps 多于 pack 数、data-replica 分组、deterministic seed、 +padding 统计，以及 legacy 策略与重构前 controller greedy pack + interleaved DP 分配 + +worker ``iters_per_step`` 重分组的精确等价性。 +""" + +import math +import unittest +from unittest.mock import patch + +from xtuner.v1.rl.trainer.pack import RLDataPacker + + +def _greedy_reference_plan( + lengths: list[int], + pack_max_length: int, + dp_size: int, + optimizer_steps: int, +) -> list[list[list[list[int]]]]: + """重构前行为的参考实现：controller greedy pack + DP-multiple 空 pack 填充 + + interleaved ``[dp::dp]`` 分配 + worker 顺序 ``iters_per_step`` 重分组。""" + packs: list[list[int]] = [] + current: list[int] = [] + current_len = 0 + for index, length in enumerate(lengths): + if current_len + length <= pack_max_length: + current.append(index) + current_len += length + else: + if current: + packs.append(current) + current = [index] + current_len = length + if current: + packs.append(current) + + pad_num = math.ceil(len(packs) / dp_size) * dp_size - len(packs) + packs.extend([[] for _ in range(pad_num)]) + + packs_per_dp = len(packs) // dp_size + iters_per_step = max(1, math.ceil(packs_per_dp / optimizer_steps)) + actual_steps = math.ceil(packs_per_dp / iters_per_step) + return [ + [packs[dp_rank::dp_size][step * iters_per_step : (step + 1) * iters_per_step] for step in range(actual_steps)] + for dp_rank in range(dp_size) + ] + + +class TestRLDataPacker(unittest.TestCase): + PACK_MAX_LENGTH = 32 + + def _make_packer( + self, + world_size: int = 4, + data_replicate_size: int = 1, + optimizer_steps: int = 2, + pack_strategy: str = "legacy", + pack_seed: int | None = None, + ) -> RLDataPacker: + return RLDataPacker( + pack_max_length=self.PACK_MAX_LENGTH, + world_size=world_size, + data_replicate_size=data_replicate_size, + optimizer_steps=optimizer_steps, + pack_strategy=pack_strategy, # type: ignore[arg-type] + pack_seed=pack_seed, + ) + + def _assert_plan_invariants( + self, + plan: list[list[list[list[int]]]], + lengths: list[int], + optimizer_steps: int, + ) -> None: + all_packs = [pack for rank in plan for step in rank for pack in step] + seen = [index for pack in all_packs for index in pack] + self.assertEqual(sorted(seen), list(range(len(lengths)))) + + for rank in plan: + self.assertGreaterEqual(len(rank), 1) + self.assertLessEqual(len(rank), optimizer_steps) + for packs in rank: + for pack in packs: + self.assertLessEqual(sum(lengths[index] for index in pack), self.PACK_MAX_LENGTH) + + steps_per_rank = {len(rank) for rank in plan} + self.assertEqual(len(steps_per_rank), 1) + # 同一 step 中所有 DP rank 的 pack 数一致（按 step 对齐比较）。 + num_steps = next(iter(steps_per_rank)) + for step in range(num_steps): + counts = {len(rank[step]) for rank in plan} + self.assertEqual(len(counts), 1) + + @patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True) + def test_legacy_matches_previous_schedule(self): + # lengths 刻意制造多 pack、非整除与 DP 填充场景。 + lengths = [10, 8, 6, 30, 12, 4, 7, 20, 2, 9] + packer = self._make_packer(world_size=4, optimizer_steps=2) + + plan, padding = packer.pack(lengths) + + self.assertEqual(plan, _greedy_reference_plan(lengths, self.PACK_MAX_LENGTH, dp_size=4, optimizer_steps=2)) + self._assert_plan_invariants(plan, lengths, optimizer_steps=2) + total_packs = sum(len(packs) for rank in plan for packs in rank) + self.assertEqual(padding, total_packs * self.PACK_MAX_LENGTH - sum(lengths)) + + @patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True) + def test_legacy_regroups_when_packs_fewer_than_optimizer_steps(self): + # packs_per_dp=2 < optimizer_steps=4：旧 worker iters_per_step=ceil(2/4)=1， + # 实际只训练 2 个 step。 + lengths = [32, 32, 32, 32] + packer = self._make_packer(world_size=2, optimizer_steps=4) + + plan, padding = packer.pack(lengths) + + self.assertEqual(plan, _greedy_reference_plan(lengths, self.PACK_MAX_LENGTH, dp_size=2, optimizer_steps=4)) + self.assertEqual(len(plan[0]), 2) + self.assertEqual(padding, 4 * self.PACK_MAX_LENGTH - sum(lengths)) + + @patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True) + def test_single_pack_interleaves_across_dp(self): + # 全部样本拼进一个 pack 时，interleaved 分配把唯一 pack 给 rank0，其余 rank 全空 pack。 + lengths = [8, 4] + packer = self._make_packer(world_size=4, optimizer_steps=1) + + plan, padding = packer.pack(lengths) + + self.assertEqual(plan, _greedy_reference_plan(lengths, self.PACK_MAX_LENGTH, dp_size=4, optimizer_steps=1)) + self.assertEqual(plan[0][0], [[0, 1]]) + self.assertEqual(plan[1][0], [[]]) + self.assertEqual(padding, 4 * self.PACK_MAX_LENGTH - sum(lengths)) + + @patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True) + def test_data_replicate_size_divides_world(self): + # world=8、replicate=2 → dp=4，各 replica group 共享同一 plan。 + lengths = [10, 10, 10, 10, 10, 10] + packer = self._make_packer(world_size=8, data_replicate_size=2, optimizer_steps=1) + + plan, _ = packer.pack(lengths) + + self.assertEqual(len(plan), 4) + self.assertEqual(plan, _greedy_reference_plan(lengths, self.PACK_MAX_LENGTH, dp_size=4, optimizer_steps=1)) + + def test_oversize_sample_raises(self): + packer = self._make_packer() + + with self.assertRaisesRegex(ValueError, "exceeds pack_max_length"): + packer.pack([self.PACK_MAX_LENGTH + 1]) + + def test_empty_lengths(self): + packer = self._make_packer() + + plan, padding = packer.pack([]) + + self.assertEqual(plan, []) + self.assertEqual(padding, 0) + + def test_shuffled_indices_use_matching_lengths(self): + lengths = [2, 8, 8] + for strategy in ("legacy", "greedy"): + with self.subTest(strategy=strategy): + packer = RLDataPacker(10, 2, 1, 2, pack_strategy=strategy) + with ( + patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", False), + patch( + "xtuner.v1.rl.trainer.pack.random.shuffle", + side_effect=lambda indices: indices.__setitem__(slice(None), [1, 2, 0]), + ), + ): + plan, padding = packer.pack(lengths) + packs = [pack for rank in plan for step in rank for pack in step] + self.assertEqual(sorted(index for pack in packs for index in pack), [0, 1, 2]) + self.assertTrue(all(sum(lengths[i] for i in pack) <= 10 for pack in packs)) + self.assertEqual(padding, len(packs) * 10 - sum(lengths)) + self.assertEqual(lengths, [2, 8, 8]) + self.assertEqual(len({len(rank) for rank in plan}), 1) + for step in range(len(plan[0])): + self.assertEqual(len({len(rank[step]) for rank in plan}), 1) + + def test_greedy_fills_gaps_from_sft_candidate_buffer(self): + lengths = [6, 7, 4] + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + legacy, old_padding = RLDataPacker(10, 1, 1, 1, pack_strategy="legacy").pack(lengths) + greedy, new_padding = RLDataPacker(10, 1, 1, 1, pack_strategy="greedy").pack(lengths) + self.assertEqual(legacy, [[[[0], [1], [2]]]]) + self.assertEqual(greedy, [[[[0, 2], [1]]]]) + self.assertEqual((old_padding, new_padding), (13, 3)) + + def test_zero_buffer_keeps_sequential_pack_order(self): + lengths = [6, 7, 4] + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + legacy = RLDataPacker(10, 1, 1, 1, pack_strategy="legacy").pack(lengths) + greedy = RLDataPacker(10, 1, 1, 1, pack_strategy="greedy", pack_extra_buffer_size=0).pack(lengths) + self.assertEqual(legacy, greedy) + + def test_unknown_strategy_raises(self): + with self.assertRaisesRegex(ValueError, "Unknown packing strategy"): + self._make_packer(pack_strategy="unknown") + + def test_legacy_matches_main_across_replica_and_step_layouts(self): + # main controller: packs[dp_rank::dp_size]; worker: ceil(packs_per_dp / steps). + for dp_size in (1, 2, 4): + for replicas in (1, 2): + for steps in (1, 2, 4): + for count in (1, 7, 17, 33): + lengths = [32] * count + with self.subTest(dp_size=dp_size, replicas=replicas, steps=steps, count=count): + packer = self._make_packer( + world_size=dp_size * replicas, + data_replicate_size=replicas, + optimizer_steps=steps, + ) + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + plan, _ = packer.pack(lengths) + self.assertEqual(plan, _greedy_reference_plan(lengths, 32, dp_size, steps)) + + @patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True) + def test_greedy_interleaves_instead_of_contiguous_dp_blocks(self): + lengths = [32] * 12 + packer = self._make_packer(pack_strategy="greedy") + + plan, padding = packer.pack(lengths) + + self.assertEqual(plan, [[[[rank], [rank + 4]], [[rank + 8]]] for rank in range(4)]) + self._assert_plan_invariants(plan, lengths, optimizer_steps=2) + total_packs = sum(len(packs) for rank in plan for packs in rank) + self.assertEqual(padding, total_packs * self.PACK_MAX_LENGTH - sum(lengths)) + + def test_balance_strategy_is_deterministic_and_invariants(self): + lengths = [10, 8, 6, 30, 12, 4, 7, 20, 2, 9, 15, 3] + first = self._make_packer(optimizer_steps=2, pack_strategy="balance", pack_seed=7) + second = self._make_packer(optimizer_steps=2, pack_strategy="balance", pack_seed=7) + + plan, padding = first.pack(lengths) + plan_again, _ = second.pack(lengths) + + self.assertEqual(plan, plan_again) + self._assert_plan_invariants(plan, lengths, optimizer_steps=2) + total_packs = sum(len(packs) for rank in plan for packs in rank) + self.assertEqual(padding, total_packs * self.PACK_MAX_LENGTH - sum(lengths)) + + def test_native_strategy_invariants(self): + lengths = [10, 8, 6, 30, 12, 4, 7] + packer = self._make_packer(world_size=4, optimizer_steps=1, pack_strategy="native") + + plan, padding = packer.pack(lengths) + + self._assert_plan_invariants(plan, lengths, optimizer_steps=1) + total_packs = sum(len(packs) for rank in plan for packs in rank) + self.assertEqual(padding, total_packs * self.PACK_MAX_LENGTH - sum(lengths)) + + def test_native_negative_slots_become_empty_packs(self): + # 7 个样本分 4 个 dp：负数 pad 槽从计划中移除后表现为空 pack。 + lengths = [8, 8, 8, 8, 8, 8, 8] + packer = self._make_packer(world_size=4, optimizer_steps=1, pack_strategy="native") + + plan, padding = packer.pack(lengths) + + seen = [index for rank in plan for step in rank for pack in step for index in pack] + self.assertEqual(sorted(seen), list(range(7))) + self.assertGreaterEqual(padding, 4 * self.PACK_MAX_LENGTH - sum(lengths)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/test_prepare_train_data.py b/tests/rl/test_prepare_train_data.py index 78d0f1946c..adaa9497cb 100644 --- a/tests/rl/test_prepare_train_data.py +++ b/tests/rl/test_prepare_train_data.py @@ -1,30 +1,28 @@ -"""RL 训练数据构造的 contract 测试（loop 侧 canonicalize + controller 侧转换）。 +"""RL 训练数据构造的 contract 测试（loop 侧 canonicalize + controller 侧调度准备）。 -重构后训练数据构造分两段，本文件只测两段的公开纯逻辑，不启动 trainer、Ray worker、模型或 -rollout backend： +重构后训练数据构造分三段，本文件只测 loop 侧与 controller 侧的公开纯逻辑，不启动 +trainer、Ray worker、模型或 rollout backend： - loop 侧 ``canonicalize_train_fields``：基类为 prompt+response 型样本构造统一全序列字段 - （``input_ids``/``labels``/``logprobs`` 等长、未 shift）； - localhost/sandbox 各自覆写为 trace 全序列字段的校验（失败置 FAILED，不上抛）。 -- controller 侧 ``TrainingController._rollout_groups_to_colate_items``：消费已 canonicalize 且 labels 已定稿 - （语义洞直接烙在 labels 中）的状态，完成 shift、张量化、组级 advantage、seq_ctx 构造与 - ``data_info`` 统计。 + （``input_ids``/``labels``/``logprobs`` 等长、未 shift）并写轻量 train meta； + localhost/sandbox 复用基类 metadata 计算，训练字段由 trainer 入口检查。 +- controller 侧 ``TrainingController._prepare_rollout_items``：只消费轻量 meta 字段 + （``num_tokens``/``supervised_tokens``/``train_prompt_length`` 等），完成组级 + reward 分组；pack_plan 固定后再计算 advantage，不构造训练张量。 +- token 级张量构造与 pack 物化在 training worker 侧（见 ``test_training_worker_rank.py``）； + index-only 打包计划见 ``test_pack.py``。 当前测试点： -- 文本样本的 input_ids、shifted_labels、rollout_logprobs、advantage 布局。 -- 同一个 prompt 下多个 response 各自使用对应 reward / advantage。 -- VLM 样本 loop 侧使用 train_prompt_ids 构造；controller 侧保留 multimodal 字段与 3D 位置。 -- VLM M-RoPE：get_train_seq_ctx 用 global amax 续写 response position（对齐 SFT get_rope_index_3）。 -- 无效 rollout group 会被跳过；缺 reward 在 task_adv_weight>0 时 fail fast，=0 时零 advantage 参训。 -- 监督语义：labels 是唯一监督载体（语义洞 = labels 中的 -100）；controller 不做任何掩码处理。 -- 多条样本 pack 成一条序列后，input_ids/shifted_labels/advantages/rollout_logprobs（及 teacher 字段） - 与输入位置逐一对齐（regression: advantage 曾比 input_ids 长 1 导致 pack 后整体错位）。 +- 无效 rollout group 由 trainer 入口丢弃；缺 reward 在 task_adv_weight>0 时 fail fast，=0 时零 advantage 参训。 +- 同一个 prompt 下多个 response 各自使用对应 reward / advantage；advantage 统计只计监督位。 +- controller 依赖 write_train_meta 写入的 meta 字段，缺失时 fail fast。 +- shuffle 只打乱调度顺序，items 与 advantages 的配对关系保持不变。 +- 批级 3D position 布局位由 meta 推导（use_3d_position_ids）。 """ import sys import unittest -from typing import cast -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import numpy as np import torch @@ -32,12 +30,11 @@ from xtuner.v1.data_proto.rl_data import ( RolloutState, Status, - TeacherTargets, + validate_train_group, ) -from xtuner.v1.datasets.mllm_tokenize_fn.qwenvl_rope2d import get_rope_index_3 -from xtuner.v1.rl.distillation import DistillationConfig, DistillationTrainerAdapter, RolloutTeacherConfig -from xtuner.v1.rl.loss import DistillationLossConfig -from xtuner.v1.rl.trainer.controller import TrainingController, get_train_seq_ctx +from xtuner.v1.rl.distillation import DistillationTrainerAdapter +from xtuner.v1.rl.trainer.controller import TrainingController + # localhost/sandbox loop 模块顶层导入 lagent；测试环境不依赖其真实实现。 # rate_limiter 必须整路径 stub：父级 MagicMock 没有 __path__，子模块 from-import 无法透传。 @@ -87,6 +84,7 @@ def _state( extra_fields: dict | None = None, input_ids: list[int] | None = None, labels: list[int] | None = None, + position_ids: np.ndarray | None = None, ) -> RolloutState: return RolloutState( rollout_id=uid, @@ -101,6 +99,7 @@ def _state( extra_fields=extra_fields or {}, input_ids=input_ids, labels=labels, + position_ids=position_ids, ) def test_builds_full_sequence_fields(self): @@ -168,6 +167,15 @@ def test_agentic_state_with_input_ids_is_untouched(self): self.assertEqual((state.input_ids, state.labels, state.logprobs), before) loop.tokenizer.assert_not_called() + def test_existing_full_sequence_fields_are_validated_at_admission(self): + loop = self._make_loop() + state = self._state(input_ids=[30, 31, 40], labels=[40], logprobs=None) + loop.canonicalize_train_fields([state]) + self.assertEqual(state.status, Status.COMPLETED) + with self.assertRaisesRegex(ValueError, "labels length mismatch"): + validate_train_group([state], task_name="test") + loop.tokenizer.assert_not_called() + def test_non_completed_state_is_skipped(self): loop = self._make_loop() state = self._state(status=Status.FAILED) @@ -178,29 +186,65 @@ def test_non_completed_state_is_skipped(self): self.assertIsNone(state.labels) self.assertEqual(state.status, Status.FAILED) - def test_missing_prompt_ids_marks_sample_failed(self): - # 构造失败容错：单个样本置 FAILED，不上抛到 producer。 + def test_writes_train_meta_fields(self): + # canonicalize 通过后写下游调度消费的轻量统计：训练长度、prompt/response 长度、 + # 监督 token 数与 position 布局。 + loop = self._make_loop() + state = self._state(response_ids=[20, 21, 22], logprobs=[0.1, 0.2, 0.3]) + + loop.canonicalize_train_fields([state]) + + self.assertEqual(state.num_tokens, 5) + self.assertEqual(state.extra_fields["train_prompt_length"], 3) + self.assertEqual(state.extra_fields["train_response_length"], 3) + self.assertEqual(state.extra_fields["supervised_tokens"], 3) + self.assertEqual(state.extra_fields["position_layout"], "1d") + + def test_train_meta_uses_train_prompt_ids_and_mrope_layout(self): + # VLM 分支 prompt 长度取 train_prompt_ids；3D position_ids 标记 mrope_3d 布局。 + loop = self._make_loop() + state = self._state( + prompt_ids=[1], + response_ids=[102, 103], + extra_fields={"train_prompt_ids": [100, 101]}, + position_ids=np.array([[[0, 1, 2]], [[0, 1, 2]], [[0, 1, 2]]], dtype=np.int64), + ) + + loop.canonicalize_train_fields([state]) + + self.assertEqual(state.num_tokens, 3) + self.assertEqual(state.extra_fields["train_prompt_length"], 2) + self.assertEqual(state.extra_fields["position_layout"], "mrope_3d") + + def test_train_meta_prompt_length_falls_back_to_unsupervised_count(self): + # prompt_ids 缺失时 prompt 长度退回 shifted_labels 的非监督位计数（与旧 data_info 口径一致）。 loop = self._make_loop() state = self._state() state.prompt_ids = None + state.input_ids = [30, 31, 40, 41] + state.labels = [-100, -100, 40, 41] loop.canonicalize_train_fields([state]) - self.assertEqual(state.status, Status.FAILED) - self.assertIn("canonicalize_train_fields failed", state.error_msg) - loop.logger.error.assert_called_once() + self.assertEqual(state.num_tokens, 3) + self.assertEqual(state.extra_fields["train_prompt_length"], 1) + self.assertEqual(state.extra_fields["supervised_tokens"], 2) - def test_failure_is_isolated_within_group(self): + def test_missing_prompt_ids_propagates_construction_error(self): loop = self._make_loop() - bad = self._state(uid=1) - bad.prompt_ids = None - good = self._state(uid=2) - - loop.canonicalize_train_fields([bad, good]) + state = self._state() + state.prompt_ids = None + with self.assertRaisesRegex(AssertionError, "Prompt ids cannot be None or empty"): + loop.canonicalize_train_fields([state]) + self.assertEqual(state.status, Status.COMPLETED) - self.assertEqual(bad.status, Status.FAILED) - self.assertEqual(good.status, Status.COMPLETED) - self.assertEqual(good.input_ids, [10, 11, 12, 20, 21, 22]) + def test_construction_error_stops_the_group(self): + loop = self._make_loop() + bad, good = self._state(uid=1), self._state(uid=2) + bad.prompt_ids = None + with self.assertRaises(AssertionError): + loop.canonicalize_train_fields([bad, good]) + self.assertIsNone(good.input_ids) class TestAgenticLoopCanonicalizeTrainFields(unittest.TestCase): @@ -253,7 +297,7 @@ def test_valid_full_sequence_state_is_unchanged(self): self.assertIsNone(state.error_msg) loop.logger.error.assert_not_called() - def test_labels_length_mismatch_marks_sample_failed(self): + def test_labels_length_mismatch_rejected_at_admission(self): for loop_cls in self.AGENTIC_LOOP_CLASSES: with self.subTest(loop_cls=loop_cls.__name__): loop = self._make_loop(loop_cls) @@ -261,11 +305,11 @@ def test_labels_length_mismatch_marks_sample_failed(self): loop.canonicalize_train_fields([state]) - self.assertEqual(state.status, Status.FAILED) - self.assertIn("canonicalize_train_fields failed", state.error_msg) - loop.logger.error.assert_called_once() + self.assertEqual(state.status, Status.COMPLETED) + with self.assertRaisesRegex(ValueError, "labels length mismatch"): + validate_train_group([state], task_name="test") - def test_missing_labels_marks_sample_failed(self): + def test_missing_labels_rejected_at_admission(self): for loop_cls in self.AGENTIC_LOOP_CLASSES: with self.subTest(loop_cls=loop_cls.__name__): loop = self._make_loop(loop_cls) @@ -273,10 +317,11 @@ def test_missing_labels_marks_sample_failed(self): loop.canonicalize_train_fields([state]) - self.assertEqual(state.status, Status.FAILED) - self.assertIn("labels length mismatch", state.error_msg) + self.assertEqual(state.status, Status.COMPLETED) + with self.assertRaisesRegex(ValueError, "labels length mismatch"): + validate_train_group([state], task_name="test") - def test_logprobs_length_mismatch_marks_sample_failed(self): + def test_logprobs_length_mismatch_rejected_at_admission(self): for loop_cls in self.AGENTIC_LOOP_CLASSES: with self.subTest(loop_cls=loop_cls.__name__): loop = self._make_loop(loop_cls) @@ -288,8 +333,42 @@ def test_logprobs_length_mismatch_marks_sample_failed(self): loop.canonicalize_train_fields([state]) - self.assertEqual(state.status, Status.FAILED) - self.assertIn("logprobs length mismatch", state.error_msg) + self.assertEqual(state.status, Status.COMPLETED) + with self.assertRaisesRegex(ValueError, "logprobs length mismatch"): + validate_train_group([state], task_name="test") + + def test_writes_train_meta_for_valid_states(self): + for loop_cls in self.AGENTIC_LOOP_CLASSES: + with self.subTest(loop_cls=loop_cls.__name__): + loop = self._make_loop(loop_cls) + state = self._agentic_state( + input_ids=[30, 31, 40, 41, 42], + labels=[-100, -100, 40, 41, 42], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + ) + + loop.canonicalize_train_fields([state]) + + self.assertEqual(state.num_tokens, 4) + self.assertEqual(state.extra_fields["train_prompt_length"], 2) + self.assertEqual(state.extra_fields["train_response_length"], 2) + self.assertEqual(state.extra_fields["supervised_tokens"], 3) + self.assertEqual(state.extra_fields["position_layout"], "1d") + + def test_train_meta_skipped_for_invalid_states(self): + # FAILED 与无 input_ids 的样本不写轻量统计。 + for loop_cls in self.AGENTIC_LOOP_CLASSES: + with self.subTest(loop_cls=loop_cls.__name__): + loop = self._make_loop(loop_cls) + failed = self._agentic_state(status=Status.FAILED) + eval_state = self._agentic_state() + + loop.canonicalize_train_fields([failed, eval_state]) + + self.assertIsNone(failed.num_tokens) + self.assertNotIn("supervised_tokens", failed.extra_fields) + self.assertIsNone(eval_state.num_tokens) + self.assertNotIn("supervised_tokens", eval_state.extra_fields) def test_state_without_input_ids_is_skipped(self): # eval 态样本显式置空训练字段且 COMPLETED，不能被构造或校验。 @@ -317,8 +396,13 @@ def test_non_completed_state_is_skipped(self): loop.logger.error.assert_not_called() -class TestConvertRolloutGroups(unittest.TestCase): - """TrainingController._rollout_groups_to_colate_items 合同：shift、advantage、张量与统计。""" +class TestPrepareRolloutItems(unittest.TestCase): + """TrainingController._prepare_rollout_items 合同：组级 advantage 与 data_info 统计。 + + controller 只消费 write_train_meta 写入的轻量 meta 字段（缺失即 fail fast；组级 + 状态/对齐校验由 replay buffer 查询与 canonicalize_train_fields 保证）；shift/张量化/ + pack 物化在 training worker 侧（见 test_training_worker_rank.py），本类不构造任何训练张量。 + """ def _build_controller(self, advantages: list[float], task_adv_weight: float = 1.0) -> TrainingController: controller = TrainingController.__new__(TrainingController) @@ -328,9 +412,16 @@ def _build_controller(self, advantages: list[float], task_adv_weight: float = 1. controller.logger = MagicMock() return controller - def _convert(self, controller, data_groups, pack_max_length=128): - with patch("xtuner.v1.rl.trainer.controller.XTUNER_DETERMINISTIC", True): - return controller._rollout_groups_to_colate_items(data_groups, pack_max_length) + def _prepare(self, controller, data_groups): + prepared = controller._prepare_rollout_items(data_groups, rollout_idx=0) + plan = {"dp_plans": {}, "plan_log": {}} + advantages = controller._prepare_advantages(prepared, plan) + return ( + prepared["rollout_items"], + advantages, + controller._build_trainer_log_info(prepared, pack_plan=plan, log_infos=[], advantages=advantages), + prepared["batch_attr"]["use_3d_position_ids"], + ) def _state( self, @@ -338,20 +429,16 @@ def _state( uid: int = 1, group_id: int = 1, prompt_ids: list[int] | None = None, - response_ids: list[int] | torch.Tensor | None = None, + response_ids: list[int] | None = None, logprobs: list[float] | None = None, supervised_mask: list[int] | None = None, reward: dict | None = None, status: Status = Status.COMPLETED, response: str = "response", - routed_experts=None, position_ids: np.ndarray | None = None, - mm_info: dict | None = None, - extra_fields: dict | None = None, input_ids: list[int] | None = None, labels: list[int] | None = None, - teacher_tokens: list[int] | list[list[int]] | None = None, - teacher_logprobs: list[float] | list[list[float]] | None = None, + write_meta: bool = True, ) -> RolloutState: resolved_prompt_ids = prompt_ids if prompt_ids is not None else [10, 11, 12] resolved_response_ids = response_ids if response_ids is not None else [20, 21, 22] @@ -366,48 +453,26 @@ def _state( reward=reward if reward is not None else {"score": 1.0}, status=status, finish_reason="stop" if status == Status.COMPLETED else "error", - routed_experts=routed_experts, position_ids=position_ids, - mm_info=mm_info, - extra_fields=extra_fields or {}, input_ids=input_ids, labels=labels, - teacher_targets=( - TeacherTargets( - kind="topk" if teacher_tokens and isinstance(teacher_tokens[0], list) else "sampled", - tokens=teacher_tokens, - logprobs=teacher_logprobs, - ) - if teacher_tokens is not None and teacher_logprobs is not None - else None - ), ) - if input_ids is not None: - # agentic 全序列样本由 trace store 直接提供字段，按原样返回。 - return state - # prompt+response 样本：模拟 loop 侧 canonicalize 后的最终形态（语义洞直接烙在 labels）。 - resp_ids = list(resolved_response_ids) - state.input_ids = list(resolved_prompt_ids) + resp_ids - if supervised_mask is None: - state.labels = [-100] * len(resolved_prompt_ids) + resp_ids - else: - state.labels = [-100] * len(resolved_prompt_ids) + [ - resp_id if flag else -100 for resp_id, flag in zip(resp_ids, supervised_mask) - ] - if state.logprobs is not None: - state.logprobs = [0.0] * len(resolved_prompt_ids) + list(state.logprobs) + if input_ids is None: + # prompt+response 样本：模拟 loop 侧 canonicalize 后的最终形态（语义洞直接烙在 labels）。 + resp_ids = list(resolved_response_ids) + state.input_ids = list(resolved_prompt_ids) + resp_ids + if supervised_mask is None: + state.labels = [-100] * len(resolved_prompt_ids) + resp_ids + else: + state.labels = [-100] * len(resolved_prompt_ids) + [ + resp_id if flag else -100 for resp_id, flag in zip(resp_ids, supervised_mask) + ] + if state.logprobs is not None: + state.logprobs = [0.0] * len(resolved_prompt_ids) + list(state.logprobs) + if write_meta: + AgentLoop._write_train_meta(None, state) return state - @staticmethod - def _enable_rollout_distillation(controller, loss_config: DistillationLossConfig) -> None: - controller.distillation = DistillationTrainerAdapter( - DistillationConfig( - loss_config=loss_config, - teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], - data_source_teacher_map={"agent_math": "teacher"}, - ) - ) - def test_noncontiguous_session_reuses_original_advantage(self): controller = self._build_controller([1.5, -2.0]) states = [self._state(uid=i, reward={"score": score}) for i, score in enumerate([1.0, 3.0, 1.0])] @@ -421,61 +486,49 @@ def test_noncontiguous_session_reuses_original_advantage(self): self.assertEqual(indices, [0, 1, 0]) self.assertEqual(controller._compute_group_advantages(rewards, representatives, indices), [1.5, -2.0, 1.5]) - def test_text_path_builds_shifted_training_tensors(self): - # 文本主路径固定 token 布局：input_ids 去掉 response 最后一个 token，label/logprob 对齐预测位置。 + def test_prepare_defers_estimator_until_after_plan(self): + controller = self._build_controller([1.0]) + prepared = controller._prepare_rollout_items([[self._state()]], rollout_idx=0) + self.assertEqual(controller.advantage_estimator.calls, []) + self.assertNotIn("advantages", prepared) + self.assertEqual(controller._prepare_advantages(prepared, {"dp_plans": {}, "plan_log": {}}), [1.0]) + self.assertEqual(len(controller.advantage_estimator.calls), 1) + + def test_returns_items_with_per_sample_advantages(self): + # controller 只做调度准备：states 原样透传，advantage 逐样本对齐，不构造训练张量。 controller = self._build_controller([1.5]) - routed_experts = np.array([[1, 2], [3, 4]]) state = self._state( - prompt_ids=[10, 11, 12], response_ids=[20, 21, 22], logprobs=[0.1, 0.2, 0.3], supervised_mask=[1, 0, 1], reward={"score": 1.0}, - routed_experts=routed_experts, ) - data_batches, info = self._convert(controller, [[state]]) + items, advantages, info, use_3d = self._prepare(controller, [[state]]) - self.assertEqual(len(data_batches), 1) - batch = data_batches[0] - self.assertEqual(batch["seq_ctx"].input_ids.tolist(), [[10, 11, 12, 20, 21]]) - self.assertEqual(batch["shifted_labels"].tolist(), [[-100, -100, 20, -100, 22]]) - torch.testing.assert_close( - batch["rollout_logprobs"], - torch.tensor([[0.0, 0.0, 0.1, 0.2, 0.3]], dtype=torch.float32), - ) - # advantage 与 shifted_labels 逐位置对齐: prompt 段为 0, mask=0 处为 0。 - self.assertEqual(batch["advantage"], [0.0, 0.0, 1.5, 0.0, 1.5]) - self.assertEqual(len(batch["advantage"]), batch["shifted_labels"].numel()) - self.assertIs(batch["seq_ctx"].rollout_routed_experts, routed_experts) + self.assertEqual(len(items), 1) + self.assertIs(items[0], state) + self.assertEqual(advantages, [1.5]) + self.assertFalse(use_3d) self.assertEqual(info["training_samples"], 1) self.assertEqual(info["training_tokens"], 5) self.assertEqual(info["rewards/mean"], 1.0) - # response_len 从 labels 监督位推导（labels 是唯一监督载体）：语义洞不计入。 + # response_len/prompt_len 由 meta 字段推导（labels 是唯一监督载体：语义洞不计入）。 self.assertEqual(info["response_len/mean"], 2.0) self.assertEqual(info["prompt_len/mean"], 3.0) - def test_controller_consumes_labels_as_final_supervision(self): - # controller 不处理任何掩码语义：labels 是什么就消费什么（语义洞已在生成期烙进 labels）。 - controller = self._build_controller([1.0]) - state = self._state(response_ids=[20, 21, 22], logprobs=[0.1, 0.2, 0.3], supervised_mask=[1, 0, 1]) - - data_batches, _ = self._convert(controller, [[state]]) - - self.assertEqual(data_batches[0]["shifted_labels"].tolist(), [[-100, -100, 20, -100, 22]]) - self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.0, 0.0, 1.0]) - def test_multi_sample_group_uses_each_sample_reward_and_advantage(self): # 同一个 prompt 下的多个 response 要分别使用自己的 reward 和 advantage。 controller = self._build_controller([1.5, -2.0]) first = self._state(uid=1, response_ids=[20, 21], reward={"score": 3.0}) second = self._state(uid=2, response_ids=[30, 31], reward={"score": -1.0}) - data_batches, info = self._convert(controller, [[first, second]]) + items, advantages, info, _ = self._prepare(controller, [[first, second]]) - self.assertEqual(len(data_batches), 2) - self.assertEqual(data_batches[0]["advantage"], [0.0, 0.0, 1.5, 1.5]) - self.assertEqual(data_batches[1]["advantage"], [0.0, 0.0, -2.0, -2.0]) + self.assertEqual(len(items), 2) + self.assertIs(items[0], first) + self.assertIs(items[1], second) + self.assertEqual(advantages, [1.5, -2.0]) self.assertEqual(info["batch_size"], 2) self.assertEqual(info["rewards/min"], -1.0) self.assertEqual(info["rewards/max"], 3.0) @@ -503,250 +556,20 @@ def test_advantage_stats_count_only_loss_active_tokens(self): reward={"score": -1.0}, ) - _, info = self._convert(controller, [[plain, agentic]]) + _, _, info, _ = self._prepare(controller, [[plain, agentic]]) # plain: 2 个 label!=-100 的 response token 计入 2.0; agentic: 同理计入 -1.0。 self.assertEqual(info["advantages/mean"], 0.5) self.assertEqual(info["advantages/min"], -1.0) self.assertEqual(info["advantages/max"], 2.0) - def test_vlm_fields_pass_through_with_3d_position_extension(self): - # VLM 样本保留 multimodal 字段；3D position_ids 按 len_response_ids 补 response 段。 - controller = self._build_controller([0.25]) - pixel_values = np.ones((1, 2, 3), dtype=np.float32) - image_grid_thw = np.array([[1, 2, 3]], dtype=np.int32) - position_ids = np.array([[[0, 1]], [[0, 1]], [[0, 1]]], dtype=np.int64) - state = self._state( - prompt_ids=[100, 101], - response_ids=[102, 103], - position_ids=position_ids, - mm_info={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, - ) - - data_batches, _ = self._convert(controller, [[state]]) - - seq_ctx = data_batches[0]["seq_ctx"] - self.assertEqual(seq_ctx.input_ids.tolist(), [[100, 101, 102]]) - self.assertEqual(tuple(seq_ctx.position_ids.shape), (3, 1, 3)) - self.assertEqual(seq_ctx.position_ids.dtype, torch.long) - self.assertIs(seq_ctx.pixel_values, pixel_values) - self.assertEqual(seq_ctx.image_grid_thw.dtype, torch.long) - self.assertEqual(seq_ctx.image_grid_thw.tolist(), [[1, 2, 3]]) - - def test_get_train_seq_ctx_mrope_continues_from_global_amax(self): - """RL 只拿 prompt 的 3D position，续写 response 后须与 SFT get_rope_index_3 一致。 - - Fixture：prompt 以 image tokens 结尾（grid 1x4x4, merge=2 → 2x2），使 per-axis max - 与 global amax 分叉。 - """ - image_token_id = 151655 - vision_start_token_id = 151652 - prompt_ids = [10, vision_start_token_id, image_token_id, image_token_id, image_token_id, image_token_id] - response_ids = [20, 21, 22] - full_ids = torch.tensor([prompt_ids + response_ids], dtype=torch.long) - image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long) - - sft_position_ids = get_rope_index_3( - full_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - ) - prompt_position_ids = get_rope_index_3( - torch.tensor([prompt_ids], dtype=torch.long), - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - ) - - seq_ctx = get_train_seq_ctx( - cast(torch.LongTensor, full_ids), - prompt_position_ids.numpy(), - ) - rl_position_ids = seq_ctx.position_ids - assert rl_position_ids is not None - self.assertEqual(tuple(rl_position_ids.shape), tuple(sft_position_ids.shape)) - torch.testing.assert_close(rl_position_ids, sft_position_ids) - - def test_mixed_agentic_and_vlm_reasoning_use_3d_position_ids(self): - controller = self._build_controller([0.5]) - reasoning_state = self._state( - uid=1, - group_id=1, - prompt_ids=[10, 11, 12], - response_ids=[20, 21, 22], - position_ids=np.arange(3, dtype=np.int64).reshape(1, 1, -1).repeat(3, axis=0), - ) - agentic_state = self._state( - uid=2, - group_id=2, - input_ids=[30, 31, 40, 41, 42], - labels=[-100, -100, 40, 41, 42], - logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], - ) - - data_batches, _ = self._convert(controller, [[reasoning_state], [agentic_state]]) - - self.assertEqual(tuple(data_batches[0]["seq_ctx"].position_ids.shape), (3, 1, 5)) - agentic_position_ids = data_batches[1]["seq_ctx"].position_ids - self.assertEqual(tuple(agentic_position_ids.shape), (3, 1, 4)) - torch.testing.assert_close( - agentic_position_ids, - torch.arange(4, dtype=torch.long).reshape(1, 1, -1).expand(3, -1, -1), - ) - - def test_agentic_topk_targets_include_token_ids_and_logprobs(self): - loss_config = DistillationLossConfig( - policy_loss_cfg={ - "loss_type": "vanilla", - "cliprange_low": 0.2, - "cliprange_high": 0.2, - }, - loss_mode="reverse", - use_policy_gradient=False, - top_k=2, - ) - controller = self._build_controller([0.0]) - self._enable_rollout_distillation(controller, loss_config) - state = self._state( - input_ids=[10, 11, 20, 21, 22], - labels=[-100, -100, 20, -100, 22], - logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], - teacher_tokens=[[100, 101], [102, 103], [104, 105]], - teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], - extra_fields={"origin_data_source": "agent_math"}, - ) - - data_batches, _ = self._convert(controller, [[state]]) - - self.assertEqual(len(data_batches), 1) - batch = data_batches[0] - self.assertEqual(batch["shifted_labels"].tolist(), [[-100, 20, -100, 22]]) - self.assertEqual( - batch["target_token_ids"].tolist(), - [[[0, 0], [100, 101], [102, 103], [104, 105]]], - ) - torch.testing.assert_close( - batch["teacher_logprobs"], - torch.tensor( - [[[0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], - dtype=torch.float32, - ), - ) - with patch("xtuner.v1.rl.loss.distillation_loss.DEVICE", "cpu"): - loss_ctx = loss_config.build( - { - "shifted_labels": batch["shifted_labels"], - "advantages": torch.tensor([batch["advantage"]], dtype=torch.float32), - "old_logprobs": torch.zeros_like(batch["shifted_labels"], dtype=torch.float32), - "teacher_logprobs": batch["teacher_logprobs"], - "target_token_ids": batch["target_token_ids"], - } - ) - assert loss_ctx is not None - type(loss_ctx).build_batches([loss_ctx]) - loss, _ = loss_ctx.loss_fn( - hidden_states=torch.randn(1, 4, 8), - head_weight=torch.randn(128, 8), - head_bias=None, - loss_kwargs=loss_ctx.loss_kwargs, - ) - self.assertTrue(torch.isfinite(loss)) - - def test_plain_topk_targets_include_masked_response_rows(self): - loss_config = DistillationLossConfig( - policy_loss_cfg={"loss_type": "vanilla"}, - loss_mode="forward_kl_topk", - use_policy_gradient=False, - top_k=2, - ) - controller = self._build_controller([0.0]) - self._enable_rollout_distillation(controller, loss_config) - state = self._state( - prompt_ids=[10, 11, 12], - response_ids=[20, 21, 22], - supervised_mask=[0, 1, 1], - teacher_tokens=[[102, 103], [104, 105]], - teacher_logprobs=[[-0.7, -0.8], [-0.9, -1.0]], - extra_fields={"origin_data_source": "agent_math"}, - ) - - data_batches, _ = self._convert(controller, [[state]]) - - batch = data_batches[0] - self.assertEqual(batch["shifted_labels"].tolist(), [[-100, -100, -100, 21, 22]]) - self.assertEqual( - batch["target_token_ids"].tolist(), - [[[0, 0], [0, 0], [0, 0], [102, 103], [104, 105]]], - ) - torch.testing.assert_close( - batch["teacher_logprobs"], - torch.tensor( - [[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [-0.7, -0.8], [-0.9, -1.0]]], - dtype=torch.float32, - ), - ) - - def test_sampled_token_targets_align_for_plain_and_agentic_rollouts(self): - loss_config = DistillationLossConfig( - policy_loss_cfg={"loss_type": "vanilla"}, - loss_mode="k1", - use_policy_gradient=True, - ) - controller = self._build_controller([0.0, 0.0]) - self._enable_rollout_distillation(controller, loss_config) - plain_state = self._state( - uid=1, - group_id=1, - prompt_ids=[10, 11, 12], - response_ids=[20, 21, 22], - supervised_mask=[0, 1, 1], - teacher_tokens=[21, 22], - teacher_logprobs=[-0.7, -0.9], - extra_fields={"origin_data_source": "agent_math"}, - ) - agentic_state = self._state( - uid=2, - group_id=2, - input_ids=[30, 31, 40, 41, 42], - labels=[-100, -100, 40, -100, 42], - logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], - teacher_tokens=[40, 41, 42], - teacher_logprobs=[-1.1, -1.2, -1.3], - extra_fields={"origin_data_source": "agent_math"}, - ) - - data_batches, _ = self._convert(controller, [[plain_state], [agentic_state]]) - - self.assertEqual(len(data_batches), 2) - torch.testing.assert_close( - data_batches[0]["teacher_logprobs"], - torch.tensor([[0.0, 0.0, 0.0, -0.7, -0.9]], dtype=torch.float32), - ) - torch.testing.assert_close( - data_batches[1]["teacher_logprobs"], - torch.tensor([[0.0, -1.1, -1.2, -1.3]], dtype=torch.float32), - ) - self.assertNotIn("target_token_ids", data_batches[0]) - self.assertNotIn("target_token_ids", data_batches[1]) - - def test_invalid_group_is_skipped(self): - # FAILED/FILTERED/ABORTED group 不能进入训练 batch，也不能贡献训练样本数。 - controller = self._build_controller([1.0]) - valid = self._state(uid=1, response_ids=[20, 21], reward={"score": 2.0}) - failed = self._state(uid=2, status=Status.FAILED, response_ids=[30, 31], reward={"score": 4.0}) - - data_batches, info = self._convert(controller, [[valid], [failed]]) - - self.assertEqual(len(data_batches), 1) - self.assertEqual(info["training_samples"], 1) - controller.logger.error.assert_called_once() - def test_missing_reward_score_fails_fast(self): # reward 必须包含 score，否则 advantage 计算前后语义都不明确。 controller = self._build_controller([1.0]) state = self._state(reward={"other": 1.0}) with self.assertRaisesRegex(ValueError, "missing.*score"): - self._convert(controller, [[state]]) + self._prepare(controller, [[state]]) def test_missing_reward_with_zero_task_adv_weight_trains_with_zero_advantage(self): # task_adv_weight=0 时（纯 OPD）缺 reward 不崩溃：样本以零 advantage 参训，不计入 rewards 统计。 @@ -754,211 +577,38 @@ def test_missing_reward_with_zero_task_adv_weight_trains_with_zero_advantage(sel state = self._state() state.reward = None - data_batches, info = self._convert(controller, [[state]]) + items, advantages, info, _ = self._prepare(controller, [[state]]) - self.assertEqual(len(data_batches), 1) - self.assertEqual(data_batches[0]["advantage"], [0.0] * 5) + self.assertEqual(len(items), 1) + self.assertIs(items[0], state) + self.assertEqual(advantages, [0.0]) self.assertEqual(info["training_samples"], 1) self.assertEqual(info["rewards/mean"], 0.0) self.assertEqual(controller.advantage_estimator.calls, []) - def test_group_with_mismatched_full_sequence_logprobs_is_skipped(self): - # 全序列 logprobs 与 input_ids 不对齐的样本由组校验拦截：整组跳过，不进训练。 - controller = self._build_controller([1.0]) - state = self._state( - input_ids=[30, 31, 40, 41, 42], - labels=[-100, -100, 40, 41, 42], - logprobs=[0.0, -0.1, -0.2, -0.3], - ) - - data_batches, info = self._convert(controller, [[state]]) - - self.assertEqual(data_batches, []) - self.assertEqual(info["training_samples"], 0) - - def test_input_ids_must_not_exceed_pack_max_length(self): - # pack_max_length 过小时要在进入 packing 前失败，避免后续训练侧报错难定位。 + def test_missing_train_meta_fails_fast(self): + # 缺 write_train_meta 写入的 meta 字段时 fail fast：这类状态到 worker 转换阶段也会失败。 controller = self._build_controller([1.0]) - state = self._state(prompt_ids=[10, 11, 12], response_ids=[20, 21, 22]) + state = self._state(write_meta=False) - with self.assertRaises(AssertionError): - self._convert(controller, [[state]], pack_max_length=4) - - -class TestConvertRolloutGroupsPackAlignment(unittest.TestCase): - """多条样本 pack 成一条序列后, pg_loss 消费前的张量逐位置对齐。 - - regression: advantage 曾按 `[adv]*len(prompt_ids) + [mask处理]` 构造, 比 input_ids 长 1, - pack 拼接后 advantages 与 shifted_labels/rollout_logprobs 整体错位或形状不匹配。 - """ + with self.assertRaisesRegex(ValueError, "canonicalization"): + self._prepare(controller, [[state]]) - def _make_controller(self, with_teacher_fields: bool = False) -> TrainingController: - controller = TrainingController.__new__(TrainingController) - controller.logger = MagicMock() - controller.task_adv_weight = 1.0 - distillation = MagicMock() - if with_teacher_fields: - # 与真实批次一致的逐位置 teacher 字段形状 (1, seq_len, k) / (1, seq_len)。 - distillation.rollout_teacher_targets.side_effect = lambda state, shifted_labels: { - "teacher_logprobs": torch.zeros(1, len(shifted_labels), 2), - "target_token_ids": torch.zeros(1, len(shifted_labels), 2), - "teacher_indices": torch.full((1, len(shifted_labels)), -1, dtype=torch.int64), - } - else: - distillation.rollout_teacher_targets.return_value = {} - distillation.reward_scalars.return_value = {} - controller.distillation = distillation - controller.advantage_estimator = MagicMock() - # 确定性 advantage: advantage = reward - 0.5, 每个样本值唯一且 float32 精确。 - controller.advantage_estimator.compute.side_effect = lambda rewards, representatives: rewards - 0.5 - return controller - - def _make_sample( - self, - rollout_id: int, - prompt_ids: list[int], - response_ids: list[int], - supervised_mask: list[int], - logprobs: list[float], - reward: float, - ) -> RolloutState: - state = RolloutState( - rollout_id=rollout_id, - message=[], - prompt_ids=prompt_ids, - response_ids=response_ids, - logprobs=logprobs, - response="ok", - reward={"score": reward}, - status=Status.COMPLETED, + def test_use_3d_position_ids_follows_batch_meta(self): + # 批级 3D 布局位：任一样本为 mrope_3d 时整批用 3D position（混合批文本样本会补 3D 轴）。 + controller = self._build_controller([1.0, 1.0]) + text = self._state(uid=1, response_ids=[20, 21]) + vlm = self._state( + uid=2, + response_ids=[30, 31], + position_ids=np.array([[[0, 1]], [[0, 1]], [[0, 1]]], dtype=np.int64), ) - # 模拟 loop 侧 canonicalize 后的最终形态（语义洞直接烙在 labels）。 - state.input_ids = list(prompt_ids) + list(response_ids) - state.labels = [-100] * len(prompt_ids) + [ - resp_id if flag else -100 for resp_id, flag in zip(response_ids, supervised_mask) - ] - state.logprobs = [0.0] * len(prompt_ids) + list(logprobs) - return state - @staticmethod - def _expected_segment( - sample: RolloutState, adv_val: float - ) -> tuple[list[int], list[int], list[float], list[float]]: - """单个样本 pack 前的正确布局: input_ids 去掉 response 末位(EOS 只作 label)。""" - assert sample.prompt_ids is not None - assert sample.response_ids is not None - assert sample.labels is not None - assert sample.logprobs is not None - prompt_len = len(sample.prompt_ids) - input_ids = list(sample.prompt_ids) + list(sample.response_ids)[:-1] - response_labels = list(sample.labels)[len(sample.labels) - len(sample.response_ids) :] - labels = [-100] * (prompt_len - 1) + response_labels - advantages = [0.0] * (prompt_len - 1) + [0.0 if label == -100 else adv_val for label in response_labels] - # canonical logprobs 为全序列对齐（prompt 段补 0），shift 后即 pack 布局。 - logprobs = list(sample.logprobs)[1:] - return input_ids, labels, advantages, logprobs - - @staticmethod - def _pack_and_extract(data_batches: list[dict], pack_max_length: int) -> dict: - controller = TrainingController.__new__(TrainingController) - packed_batches = controller._packing(data_batches, pack_max_length, None) - assert len(packed_batches) == 1, "all samples are expected to fit into a single pack" - packed = packed_batches[0] - result = { - "input_ids": packed["seq_ctx"].input_ids.squeeze(0).tolist(), - "labels": packed["shifted_labels"].squeeze(0).tolist(), - "advantages": packed["advantages"].squeeze(0).tolist(), - "rollout_logprobs": packed["rollout_logprobs"].squeeze(0).tolist(), - } - for key in ("teacher_logprobs", "target_token_ids", "teacher_indices"): - if packed.get(key) is not None: - result[key] = packed[key].squeeze(0).tolist() - return result - - def _build_data_groups(self) -> list[list[RolloutState]]: - """两组样本; 组内共享 prompt(与真实 RL 组一致), 长度/监督位/advantage 刻意互不相同。""" - s1 = self._make_sample( - 0, - [101, 102, 103, 104], - [1001, 1002, 1003, 1004, 1005], - [1, 1, 0, 1, 1], - [-0.5, -1.0, -1.5, -2.0, -2.5], - 1.0, - ) - s2 = self._make_sample(1, [101, 102, 103, 104], [2001, 2002, 2003], [1, 0, 1], [-0.5, -1.0, -1.5], 2.0) - s3 = self._make_sample( - 2, [301, 302, 303, 304, 305], [3001, 3002, 3003, 3004], [1, 1, 1, 0], [-0.5, -1.0, -1.5, -2.0], 4.0 - ) - s4 = self._make_sample(3, [301, 302, 303, 304, 305], [4001, 4002], [1, 0], [-0.5, -1.0], 8.0) - return [[s1, s2], [s3, s4]] - - def _convert(self, controller, data_groups, pack_max_length: int): - with patch("xtuner.v1.rl.trainer.controller.XTUNER_DETERMINISTIC", True): - return controller._rollout_groups_to_colate_items(data_groups, pack_max_length) - - def _assert_real_token_alignment( - self, samples: list[RolloutState], packed: dict, total_len: int, packed_len: int - ) -> None: - self.assertEqual(len(packed["input_ids"]), packed_len) - self.assertEqual(len(packed["labels"]), packed_len) - self.assertEqual(len(packed["advantages"]), packed_len) - self.assertEqual(len(packed["rollout_logprobs"]), packed_len) - - offset = 0 - for sample in samples: - adv_val = sample.reward["score"] - 0.5 - exp_ids, exp_labels, exp_advs, exp_logprobs = self._expected_segment(sample, adv_val) - seg = slice(offset, offset + len(exp_ids)) - self.assertEqual(packed["input_ids"][seg], exp_ids) - self.assertEqual(packed["labels"][seg], exp_labels) - self.assertEqual(packed["advantages"][seg], exp_advs) - self.assertEqual(packed["rollout_logprobs"][seg], exp_logprobs) - offset += len(exp_ids) - self.assertEqual(offset, total_len) - - def test_packed_fields_align_without_padding(self) -> None: - controller = self._make_controller() - data_groups = self._build_data_groups() - samples = [sample for group in data_groups for sample in group] - total_len = sum(len(s.prompt_ids) + len(s.response_ids) - 1 for s in samples) - - data_batches, _ = self._convert(controller, data_groups, pack_max_length=total_len) - - self.assertEqual(len(data_batches), len(samples)) - for item in data_batches: - self.assertEqual(len(item["advantage"]), item["shifted_labels"].numel()) - - packed = self._pack_and_extract(data_batches, total_len) - self._assert_real_token_alignment(samples, packed, total_len, total_len) - - def test_packed_fields_align_with_padding(self) -> None: - controller = self._make_controller() - data_groups = self._build_data_groups() - samples = [sample for group in data_groups for sample in group] - total_len = sum(len(s.prompt_ids) + len(s.response_ids) - 1 for s in samples) - pack_max_length = (total_len // 16 + 1) * 16 - - data_batches, _ = self._convert(controller, data_groups, pack_max_length=pack_max_length) - packed = self._pack_and_extract(data_batches, pack_max_length) - - self._assert_real_token_alignment(samples, packed, total_len, pack_max_length) - # pad 区: label 为 -100, advantage 为 controller 的 pad 值 -100。 - self.assertTrue(all(label == -100 for label in packed["labels"][total_len:])) - self.assertTrue(all(adv == -100 for adv in packed["advantages"][total_len:])) - - def test_packed_fields_align_with_teacher_fields(self) -> None: - controller = self._make_controller(with_teacher_fields=True) - data_groups = self._build_data_groups() - samples = [sample for group in data_groups for sample in group] - total_len = sum(len(s.prompt_ids) + len(s.response_ids) - 1 for s in samples) - - data_batches, _ = self._convert(controller, data_groups, pack_max_length=total_len) - packed = self._pack_and_extract(data_batches, total_len) - - self._assert_real_token_alignment(samples, packed, total_len, total_len) - self.assertEqual(len(packed["teacher_logprobs"]), total_len) - self.assertEqual(len(packed["target_token_ids"]), total_len) - self.assertEqual(len(packed["teacher_indices"]), total_len) + _, _, _, text_only = self._prepare(controller, [[text]]) + self.assertFalse(text_only) + + _, _, _, mixed = self._prepare(controller, [[text], [vlm]]) + self.assertTrue(mixed) if __name__ == "__main__": diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 6c7c493a23..782e2b65bf 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -385,6 +385,28 @@ def is_valid_sample_fn(samples): self.assertEqual(result.failed_samples, 1) self.assertEqual(result.filtered_samples, 1) + async def test_producer_does_not_require_training_fields(self): + strategy = SyncProduceStrategyConfig().build() + ctx = self._build_context( + strategy, "eval_or_custom_loop", self._build_agent_loop(), self._build_sampler(), batch_size=1 + ) + state = make_rollout_state(1) + self.assertIsNone(state.input_ids) + self.assertIsNone(state.labels) + self.assertTrue(await ctx.put_generated_group([state])) + self.assertEqual(len(self.replay_buffer), 1) + + async def test_inactive_loop_result_is_rejected_before_insertion(self): + strategy = SyncProduceStrategyConfig().build() + agent_loop = self._build_agent_loop() + state = make_rollout_state(1, status=Status.INIT) + agent_loop.generate_group = AsyncMock(return_value=[state]) + ctx = self._build_context(strategy, "bad_status", agent_loop, self._build_sampler(), batch_size=1) + result = await ctx.generate_group([state]) + with self.assertRaisesRegex(ValueError, "invalid generated status"): + await ctx.put_generated_group(result) + self.assertEqual(len(self.replay_buffer), 0) + async def test_put_generated_group_releases_non_retryable_trace_sessions(self): # FAILED / FILTERED 不进入 replay buffer，必须在丢弃 RolloutState 前释放对应 trace session。 cases = ( diff --git a/tests/rl/test_rl_colocate_trainer_integration.py b/tests/rl/test_rl_colocate_trainer_integration.py index b4c9580f6f..1f3bf93666 100644 --- a/tests/rl/test_rl_colocate_trainer_integration.py +++ b/tests/rl/test_rl_colocate_trainer_integration.py @@ -1,24 +1,20 @@ import os -import unittest import shutil import tempfile -import ray +import unittest from pathlib import Path -from xtuner.v1.rl.utils import AcceleratorResourcesConfig, CPUResourcesConfig +import ray + +from transformers import AutoTokenizer from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig -from xtuner.v1.model import get_model_config_from_hf +from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig from xtuner.v1.datasets.rl_tokenize_fn import RLTextTokenizeFnConfig -from xtuner.v1.train.trainer import LoadCheckpointConfig -from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig -from xtuner.v1.rl.trainer import WorkerConfig -from xtuner.v1.rl.loss import GRPOLossConfig -from xtuner.v1.rl.rollout.worker import RolloutConfig -from xtuner.v1.rl.judger import GSM8KJudgerConfig -from xtuner.v1.loss import CELossConfig from xtuner.v1.datasets.sft_tokenize_fn import OpenaiTokenizeFunctionConfig -from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig +from xtuner.v1.loss import CELossConfig +from xtuner.v1.model import get_model_config_from_hf +from xtuner.v1.rl.agent_loop.agent_loop import AgentLoop from xtuner.v1.rl.agent_loop import SingleTurnAgentLoopConfig from xtuner.v1.rl.agent_loop_manager import ( AgentLoopManagerConfig, @@ -27,8 +23,15 @@ TaskSpecConfig, ) from xtuner.v1.rl.evaluator import EvaluatorConfig -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status -from transformers import AutoTokenizer +from xtuner.v1.rl.judger import GSM8KJudgerConfig +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trainer import WorkerConfig +from xtuner.v1.rl.utils import AcceleratorResourcesConfig, CPUResourcesConfig +from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig +from xtuner.v1.train.trainer import LoadCheckpointConfig + QWEN3_PATH = os.environ["QWEN3_PATH"] ALPACA_PATH = os.environ["ALPACA_PATH"] @@ -105,13 +108,12 @@ def build_trainer_config(self, work_dir, checkpoint_interval=1, checkpoint_maxke ) # SFT configs for WorkerConfig - sft_dataset_config = [{ - "dataset": DatasetConfig(name='alpaca', anno_path=data_path), - "tokenize_fn": OpenaiTokenizeFunctionConfig( - chat_template='qwen3', - max_length=2048 - ) - }] + sft_dataset_config = [ + { + "dataset": DatasetConfig(name="alpaca", anno_path=data_path), + "tokenize_fn": OpenaiTokenizeFunctionConfig(chat_template="qwen3", max_length=2048), + } + ] sft_dataloader_cfg = DataloaderConfig( dataset_config_list=sft_dataset_config, pack_max_length=2048, @@ -241,21 +243,21 @@ def test_rl_train_with_sft(self): prompts = ["What is 2+2?", "What is the capital of France?"] responses = [ ["4", "Four", "2+2=4", "The answer is 4"], - ["Paris", "The capital is Paris", "Paris, France", "It's Paris"] + ["Paris", "The capital is Paris", "Paris, France", "It's Paris"], ] group_rewards = [1.0, 0.8, 0.9, 0.7] rollout_groups = [] for group_idx, (prompt, response_list) in enumerate(zip(prompts, responses)): - prompt_ids = tokenizer(prompt, return_tensors='pt')['input_ids'].flatten().tolist() + prompt_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].flatten().tolist() group = [] for i, response in enumerate(response_list): - response_ids = tokenizer(response, return_tensors='pt')['input_ids'].flatten().tolist() + response_ids = tokenizer(response, return_tensors="pt")["input_ids"].flatten().tolist() # Canonical full-sequence convention (AgentLoop.canonicalize_train_fields): # - input_ids keeps the whole prompt+response sequence - # - labels supervise the full response; TrainingController applies the + # - labels supervise the full response; the training worker applies the # one-position shift at conversion time - group.append(RolloutState( + state = RolloutState( rollout_id=group_idx * len(response_list) + i, group_id=group_idx, message=[{"role": "user", "content": prompt}], @@ -266,7 +268,10 @@ def test_rl_train_with_sft(self): status=Status.COMPLETED, input_ids=prompt_ids + response_ids, labels=[-100] * len(prompt_ids) + response_ids, - )) + ) + # Controller scheduling only consumes the lightweight train meta fields. + AgentLoop._write_train_meta(None, state) + group.append(state) rollout_groups.append(group) # RLColocateTrainer initializes by offloading train workers to CPU. @@ -283,8 +288,8 @@ def test_rl_train_with_sft(self): log_infos, _ = train_controller.fit(rollout_groups, pack_max_length=1024, rollout_idx=1) efficient_attn_ratio_list = [] for log_info in log_infos: - efficient_attn_ratio_list.append(log_info['sft_train_metrics']['efficient_attn_ratio']) - self.assertTrue(all([ratio > 0 for ratio in efficient_attn_ratio_list])) + efficient_attn_ratio_list.append(log_info["sft_train_metrics"]["efficient_attn_ratio"]) + self.assertTrue(all(ratio > 0 for ratio in efficient_attn_ratio_list)) # Kill and rebuild del trainer @@ -303,9 +308,7 @@ def test_rl_train_with_sft(self): # Resume and verify load_checkpoint_cfg = LoadCheckpointConfig( - checkpoint_path=checkpoint_path, - load_optimizer_states=False, - load_optimizer_args=False + checkpoint_path=checkpoint_path, load_optimizer_states=False, load_optimizer_args=False ) train_controller.resume(load_checkpoint_cfg) @@ -313,7 +316,7 @@ def test_rl_train_with_sft(self): log_infos, _ = train_controller.fit(rollout_groups, pack_max_length=1024, rollout_idx=1) new_efficient_attn_ratio_list = [] for log_info in log_infos: - new_efficient_attn_ratio_list.append(log_info['sft_train_metrics']['efficient_attn_ratio']) + new_efficient_attn_ratio_list.append(log_info["sft_train_metrics"]["efficient_attn_ratio"]) efficient_attn_ratio_list.sort() new_efficient_attn_ratio_list.sort() diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 73b99ee23d..2bc5661531 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -30,6 +30,8 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch +from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.agent_loop.agent_loop import AgentLoop from xtuner.v1.rl.agent_loop_manager import ( AgentLoopManagerStatus, ProduceBatchResult, @@ -39,6 +41,23 @@ from xtuner.v1.train.rl_trainer import RLDisaggregatedTrainer, _validate_sync_intervals +def make_train_state(session_id=None): + state = RolloutState( + group_id=1, + rollout_id=1, + session_id=session_id, + message=[], + status=Status.COMPLETED, + prompt_ids=[1], + response_ids=[2], + response_model_steps=[0], + input_ids=[1, 2], + labels=[-100, 2], + ) + AgentLoop._write_train_meta(None, state) + return state + + class _FakeManager: def __init__(self, get_batch_results): self._results = list(get_batch_results) @@ -191,7 +210,7 @@ def _minimal_train_info(self, *, training_samples: int, training_tokens: int, be def test_fit_persists_checkpoint_for_completed_model_step(self): # 验证 checkpoint 以 fit 完成的 model_step 为准，并通过 async manager.save 落盘。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=None) + train_sample = make_train_state() manager = _FakeManager([ProduceBatchResult(rollout_states=[[train_sample]])]) manager.save = AsyncMock() trainer = self._make_trainer(manager) @@ -223,7 +242,7 @@ def test_fit_persists_checkpoint_for_completed_model_step(self): def test_fit_retries_same_step_after_empty_expired_skip(self): # 验证空 expired batch 只同步上一版模型，不推进 train_step，并重试同一步。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=None) + train_sample = make_train_state() manager = _FakeManager( [ ProduceBatchResult(rollout_states=[], status=ProduceBatchStatus.EXPIRED_BATCH), @@ -249,7 +268,7 @@ def test_fit_retries_same_step_after_empty_expired_skip(self): def test_fit_trains_non_empty_expired_batch_then_syncs_current_step(self): # 验证非空 expired batch 仍会训练，并用当前完成的 model_step 恢复 producer。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=None) + train_sample = make_train_state() manager = _FakeManager( [ProduceBatchResult(rollout_states=[[train_sample]], status=ProduceBatchStatus.EXPIRED_BATCH)] ) @@ -264,7 +283,7 @@ def test_fit_trains_non_empty_expired_batch_then_syncs_current_step(self): def test_fit_rebinds_weight_update_with_rollout_update_address(self): # 验证非共卡后续同步权重时继续沿用 rollout config 中的 NCCL update 地址。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=None) + train_sample = make_train_state() manager = _FakeManager([ProduceBatchResult(rollout_states=[[train_sample]])]) trainer = self._make_trainer(manager) trainer._rollout_config = SimpleNamespace(weight_update_host="10.0.0.1", weight_update_port=23456) @@ -311,7 +330,7 @@ def blocking_train_one_batch(*args, **kwargs): def test_train_batch_releases_only_consumed_trace_sessions_for_disaggregated_rollout(self): # 后台 producer 在 learner 训练期间仍会创建 trace；训练结束只能释放当前消费 batch 的 session。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=101) + train_sample = make_train_state(session_id=101) trainer = self._make_trainer(_FakeManager([])) trainer._release_trace_sessions = RLDisaggregatedTrainer._release_trace_sessions.__get__( trainer, RLDisaggregatedTrainer @@ -409,7 +428,7 @@ def test_fit_observes_background_producer_failure_before_training_waited_batch(s def test_fit_runs_eval_before_reset_and_stops_producer(self): # 验证 eval 在 producer 恢复前执行，避免生产侧提前抢占 rollout 资源。 # 确定性排序依赖 RolloutState 的 group_id 和 rollout_id，测试用轻量对象模拟即可。 - train_sample = SimpleNamespace(group_id=1, rollout_id=1, session_id=None) + train_sample = make_train_state() eval_sample = SimpleNamespace(group_id=2, rollout_id=2, session_id=None) manager = _FakeManager([ProduceBatchResult(rollout_states=[[train_sample]], status=ProduceBatchStatus.NORMAL)]) trainer = self._make_trainer(manager) diff --git a/tests/rl/test_rollout_to_train_item.py b/tests/rl/test_rollout_to_train_item.py new file mode 100644 index 0000000000..030864f50c --- /dev/null +++ b/tests/rl/test_rollout_to_train_item.py @@ -0,0 +1,275 @@ +"""TrainingWorker rollout-to-train-item 转换的 contract 测试。 + +``TrainingWorker._convert_rollout_items_to_train_items`` 是 RolloutState -> (seq_ctx, loss_inputs) +的唯一入口：shift、scalar advantage 逐监督位广播、rollout logprob shift、position 布局统一、 +seq_ctx 构造、teacher target 对齐与 CPU per-item loss_inputs 构建都在这一步完成。本文件只测该 +纯转换逻辑，不启动 Ray、模型或 rollout backend。 + +当前测试点： +- 文本样本：input_ids 去尾、shifted_labels、rollout_logprobs shift 与形状对齐。 +- scalar advantage 广播：prompt 占位与 labels=-100 的语义洞位为 0，监督位为 group advantage。 +- use_3d_position_ids=True 时文本样本广播出三轴 MRoPE 位置；False 时保持 1D。 +- VLM 样本：prompt 段 3D position 续写 response、mm_info/pixel_values 透传、routed_experts 挂到 seq_ctx。 +- 蒸馏：rollout teacher target 对齐结果进入 loss_inputs（sampled 无 target_token_ids，topk 有）。 +- token advantage 与 shifted labels 等长，语义洞与 prompt 不参与监督。 +""" + +import unittest +from types import SimpleNamespace + +import numpy as np +import torch + +from xtuner.v1.data_proto.rl_data import RolloutState, Status, TeacherTargets +from xtuner.v1.rl.distillation import DistillationConfig, DistillationTrainerAdapter, RolloutTeacherConfig +from xtuner.v1.rl.loss import DistillationLossConfig, GRPOLossConfig +from xtuner.v1.rl.trainer.worker import TrainBatchAttr, TrainingWorker + + +def _batch_attr(use_3d: bool) -> TrainBatchAttr: + return { + "rollout_idx": 0, + "use_3d_position_ids": use_3d, + "pack_loss_keys": [], + "has_routed_experts": False, + } + + +class TestConvertRolloutItemsToTrainItems(unittest.TestCase): + def _make_worker(self, distillation: DistillationTrainerAdapter | None = None) -> TrainingWorker: + # 生产契约：蒸馏开启时 train_worker_cfg.loss_cfg 即 distillation_config.loss_config。 + worker = TrainingWorker.__new__(TrainingWorker) + worker.config = SimpleNamespace(loss_cfg=GRPOLossConfig(policy_loss_cfg={"loss_type": "vanilla"})) + worker._distillation = distillation if distillation is not None else DistillationTrainerAdapter(None) + return worker + + def _state( + self, + *, + uid: int = 1, + prompt_ids: list[int] | None = None, + response_ids: list[int] | None = None, + logprobs: list[float] | None = None, + supervised_mask: list[int] | None = None, + routed_experts=None, + position_ids: np.ndarray | None = None, + mm_info: dict | None = None, + extra_fields: dict | None = None, + input_ids: list[int] | None = None, + labels: list[int] | None = None, + teacher_tokens: list[int] | list[list[int]] | None = None, + teacher_logprobs: list[float] | list[list[float]] | None = None, + ) -> RolloutState: + resolved_prompt_ids = prompt_ids if prompt_ids is not None else [10, 11, 12] + resolved_response_ids = response_ids if response_ids is not None else [20, 21, 22] + state = RolloutState( + rollout_id=uid, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=resolved_prompt_ids, + response="response", + response_ids=resolved_response_ids, + logprobs=logprobs, + status=Status.COMPLETED, + finish_reason="stop", + routed_experts=routed_experts, + position_ids=position_ids, + mm_info=mm_info, + extra_fields=extra_fields or {}, + input_ids=input_ids, + labels=labels, + teacher_targets=( + TeacherTargets( + kind="topk" if teacher_tokens and isinstance(teacher_tokens[0], list) else "sampled", + tokens=teacher_tokens, + logprobs=teacher_logprobs, + ) + if teacher_tokens is not None and teacher_logprobs is not None + else None + ), + ) + if input_ids is not None: + return state + state.input_ids = list(resolved_prompt_ids) + list(resolved_response_ids) + if supervised_mask is None: + state.labels = [-100] * len(resolved_prompt_ids) + list(resolved_response_ids) + else: + state.labels = [-100] * len(resolved_prompt_ids) + [ + resp_id if flag else -100 for resp_id, flag in zip(resolved_response_ids, supervised_mask) + ] + if state.logprobs is not None: + state.logprobs = [0.0] * len(resolved_prompt_ids) + list(state.logprobs) + return state + + def test_text_path_builds_shifted_rl_train_item(self): + worker = self._make_worker() + state = self._state( + response_ids=[20, 21, 22], + logprobs=[0.1, 0.2, 0.3], + supervised_mask=[1, 0, 1], + ) + + items = worker._convert_rollout_items_to_train_items([state], [1.5], _batch_attr(False)) + + self.assertEqual(len(items), 1) + seq_ctx = items[0]["seq_ctx"] + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertEqual(seq_ctx.input_ids.tolist(), [[10, 11, 12, 20, 21]]) + self.assertEqual(loss_kwargs.shifted_labels.tolist(), [[-100, -100, 20, -100, 22]]) + torch.testing.assert_close( + loss_kwargs.advantages, + torch.tensor([[0.0, 0.0, 1.5, 0.0, 1.5]], dtype=torch.float32), + ) + torch.testing.assert_close( + loss_kwargs.rollout_logprobs, + torch.tensor([[0.0, 0.0, 0.1, 0.2, 0.3]], dtype=torch.float32), + ) + self.assertIsNone(getattr(loss_kwargs, "teacher_logprobs", None)) + self.assertIsNone(getattr(loss_kwargs, "target_token_ids", None)) + self.assertIsNone(getattr(loss_kwargs, "teacher_indices", None)) + for tensor in (loss_kwargs.shifted_labels, loss_kwargs.advantages, loss_kwargs.rollout_logprobs): + self.assertEqual(tensor.shape[1], seq_ctx.input_ids.shape[1]) + + def test_advantages_follow_each_sample_scalar(self): + worker = self._make_worker() + first = self._state(uid=1, response_ids=[20, 21], supervised_mask=[1, 1]) + second = self._state(uid=2, response_ids=[30, 31], supervised_mask=[1, 0]) + + items = worker._convert_rollout_items_to_train_items([first, second], [1.5, -2.0], _batch_attr(False)) + + torch.testing.assert_close( + items[0]["loss_inputs"].get("advantages"), + torch.tensor([[0.0, 0.0, 1.5, 1.5]], dtype=torch.float32), + ) + torch.testing.assert_close( + items[1]["loss_inputs"].get("advantages"), + torch.tensor([[0.0, 0.0, -2.0, 0.0]], dtype=torch.float32), + ) + + def test_text_positions_stay_1d_without_3d_flag(self): + worker = self._make_worker() + state = self._state(response_ids=[20, 21]) + + items = worker._convert_rollout_items_to_train_items([state], [0.0], _batch_attr(False)) + + torch.testing.assert_close(items[0]["seq_ctx"].position_ids, torch.arange(4).reshape(1, -1)) + + def test_text_positions_broadcast_to_3d_in_mixed_batches(self): + worker = self._make_worker() + state = self._state(response_ids=[20, 21]) + + items = worker._convert_rollout_items_to_train_items([state], [0.0], _batch_attr(True)) + + position_ids = items[0]["seq_ctx"].position_ids + self.assertEqual(tuple(position_ids.shape), (3, 1, 4)) + torch.testing.assert_close( + position_ids, + torch.arange(4, dtype=torch.long).reshape(1, 1, -1).expand(3, -1, -1), + ) + + def test_vlm_position_extension_and_multimodal_passthrough(self): + worker = self._make_worker() + pixel_values = np.ones((1, 2, 3), dtype=np.float32) + image_grid_thw = np.array([[1, 2, 3]], dtype=np.int32) + routed_experts = np.array([[1, 2], [3, 4]]) + state = self._state( + prompt_ids=[100, 101], + response_ids=[102, 103], + position_ids=np.array([[[0, 1]], [[0, 1]], [[0, 1]]], dtype=np.int64), + mm_info={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + routed_experts=routed_experts, + ) + + items = worker._convert_rollout_items_to_train_items([state], [0.25], _batch_attr(True)) + + seq_ctx = items[0]["seq_ctx"] + self.assertEqual(seq_ctx.input_ids.tolist(), [[100, 101, 102]]) + self.assertEqual(tuple(seq_ctx.position_ids.shape), (3, 1, 3)) + self.assertIs(seq_ctx.pixel_values, pixel_values) + self.assertEqual(seq_ctx.image_grid_thw.tolist(), [[1, 2, 3]]) + self.assertIs(seq_ctx.rollout_routed_experts, routed_experts) + + def _enable_rollout_distillation(self, worker: TrainingWorker, loss_config: DistillationLossConfig) -> None: + worker._distillation = DistillationTrainerAdapter( + DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"agent_math": "teacher"}, + ) + ) + worker.config.loss_cfg = loss_config + + def test_topk_teacher_targets_land_in_loss_inputs(self): + worker = self._make_worker() + self._enable_rollout_distillation( + worker, + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ), + ) + state = self._state( + input_ids=[10, 11, 20, 21, 22], + labels=[-100, -100, 20, -100, 22], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + teacher_tokens=[[100, 101], [102, 103], [104, 105]], + teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], + extra_fields={"origin_data_source": "agent_math"}, + ) + + items = worker._convert_rollout_items_to_train_items([state], [0.0], _batch_attr(False)) + + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertEqual(loss_kwargs.shifted_labels.tolist(), [[-100, 20, -100, 22]]) + self.assertEqual( + loss_kwargs.target_token_ids.tolist(), + [[[0, 0], [100, 101], [102, 103], [104, 105]]], + ) + torch.testing.assert_close( + loss_kwargs.teacher_logprobs, + torch.tensor([[[0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], dtype=torch.float32), + ) + + def test_sampled_teacher_targets_omit_target_token_ids(self): + worker = self._make_worker() + self._enable_rollout_distillation( + worker, + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="k1", + use_policy_gradient=True, + ), + ) + state = self._state( + response_ids=[20, 21], + logprobs=[0.1, 0.2], + teacher_tokens=[20, 21], + teacher_logprobs=[-0.7, -0.9], + extra_fields={"origin_data_source": "agent_math"}, + ) + + items = worker._convert_rollout_items_to_train_items([state], [0.0], _batch_attr(False)) + + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertIsNone(getattr(loss_kwargs, "target_token_ids", None)) + torch.testing.assert_close( + loss_kwargs.teacher_logprobs, + torch.tensor([[0.0, 0.0, -0.7, -0.9]], dtype=torch.float32), + ) + + def test_token_advantages_mask_prompt_and_semantic_holes(self): + worker = self._make_worker() + state = self._state(supervised_mask=[1, 0, 1]) + items = worker._convert_rollout_items_to_train_items([state], [[1, 2, 3, 4, 5]], _batch_attr(False)) + torch.testing.assert_close(items[0]["loss_inputs"]["advantages"], torch.tensor([[0.0, 0.0, 3.0, 0.0, 5.0]])) + + def test_token_advantages_require_shifted_length(self): + worker = self._make_worker() + with self.assertRaisesRegex(ValueError, "Token advantage length mismatch"): + worker._convert_rollout_items_to_train_items([self._state()], [[1.0]], _batch_attr(False)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/test_train_batch_contract.py b/tests/rl/test_train_batch_contract.py new file mode 100644 index 0000000000..c9d09c5e84 --- /dev/null +++ b/tests/rl/test_train_batch_contract.py @@ -0,0 +1,289 @@ +"""Canonical field validation, staleness metadata and reusable packing layouts.""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +from xtuner.v1.data_proto.rl_data import RolloutState, Status, validate_rollout_group, validate_train_group +from xtuner.v1.rl.agent_loop.agent_loop import AgentLoop +from xtuner.v1.rl.agent_loop_manager.produce_utils import ProduceBatchResult, _TaskRunner, take_train_batch +from xtuner.v1.rl.agent_loop_manager.producer import ProduceProgress +from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig +from xtuner.v1.rl.trainer.controller import TrainingController +from xtuner.v1.train.rl_trainer import BaseRLTrainer + + +def canonical_state(uid=1): + state = RolloutState( + rollout_id=uid, + group_id=uid, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=[10, 11], + response_ids=[20, 21, 22, 23], + input_ids=[10, 11, 20, 21, 22, 23], + labels=[-100, -100, 20, -100, 22, 23], + response="answer", + status=Status.COMPLETED, + reward={"score": 1.0}, + ) + state.response_model_steps = [1] * 4 + AgentLoop._write_train_meta(None, state) + return state + + +class TestTrainFields(unittest.TestCase): + def admit(self, group): + return validate_rollout_group(group, task_name="math") + + def test_optional_logprobs(self): + state = canonical_state() + validate_train_group([state], task_name="math") + state.logprobs = [0.0] * 6 + AgentLoop._write_train_meta(None, state) + validate_train_group([state], task_name="math") + + def test_invalid_training_fields_include_context(self): + for field, value in ( + ("input_ids", None), + ("input_ids", [1]), + ("labels", None), + ("labels", [1]), + ("logprobs", [0.0]), + ): + with self.subTest(field=field): + state = canonical_state() + setattr(state, field, value) + self.admit([state]) + with self.assertRaisesRegex(ValueError, f"task='math'.*rollout_id=1.*{field}"): + validate_train_group([state], task_name="math") + + def test_public_validation_does_not_fill_versions(self): + state = canonical_state() + state.response_model_steps = [0] + self.assertIsNone(self.admit([state])) + self.assertEqual(state.response_model_steps, [0]) + + def test_training_validation_does_not_infer_missing_historical_versions(self): + state = canonical_state() + state.response_model_steps = [0] + with self.assertRaisesRegex(ValueError, "response_model_steps length"): + validate_train_group([state], task_name="math") + self.assertEqual(state.response_model_steps, [0]) + + def test_eval_and_partial_results_need_no_train_payload(self): + for status in (Status.COMPLETED, Status.ABORTED, Status.EXPIRED): + with self.subTest(status=status): + state = RolloutState(message=[], status=status, response_ids=[1]) + self.admit([state]) + self.assertIsNone(state.input_ids) + self.assertIsNone(state.response_model_steps) + + def test_producer_rejects_inactive_result_without_mutating_it(self): + for status in (Status.INIT, Status.ARCHIVED): + state = RolloutState(message=[], status=status, response_ids=[1]) + with self.assertRaisesRegex(ValueError, "invalid generated status"): + self.admit([state]) + self.assertEqual(state.status, status) + self.assertIsNone(state.response_model_steps) + + def test_training_validation_rejects_missing_or_stale_metadata(self): + for key, invalid in ( + ("train_prompt_length", -1), + ("train_response_length", 100), + ("supervised_tokens", 100), + ("position_layout", "mrope_3d"), + ("has_rollout_logprobs", True), + ("has_routed_experts", True), + ): + for missing in (False, True): + with self.subTest(key=key, missing=missing): + state = canonical_state() + if missing: + state.extra_fields.pop(key) + else: + state.extra_fields[key] = invalid + self.admit([state]) + with self.assertRaisesRegex(ValueError, key): + validate_train_group([state], task_name="math") + + def test_training_validation_rejects_stale_token_count(self): + state = canonical_state() + state.num_tokens += 1 + with self.assertRaisesRegex(ValueError, "num_tokens"): + validate_train_group([state], task_name="math") + + def test_public_validation_rejects_versions_longer_than_response(self): + state = canonical_state() + state.response_model_steps = [1] * 5 + with self.assertRaisesRegex(ValueError, "response_model_steps length=5.*response_ids=4"): + self.admit([state]) + self.assertEqual(state.response_model_steps, [1] * 5) + + +class TestTrainBatchContract(unittest.IsolatedAsyncioTestCase): + async def take(self, buffer, threshold): + runner = _TaskRunner( + task_name="math", + agent_loop=MagicMock(), + sampler=MagicMock(), + produce_strategy=SimpleNamespace(token_stale_threshold=threshold), + ) + result = await take_train_batch( + task_runners=[runner], + replay_buffer=buffer, + logger=MagicMock(), + manager_name="test", + task_batch_sizes={"math": 1}, + progress=ProduceProgress.build(task_names=["math"], target_samples={"math": 1}), + current_train_step=2, + ) + + trainer = SimpleNamespace(logger=MagicMock()) + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 2, result) + return result + + async def test_staleness_refreshes_only_supervision_count_idempotently(self): + buffer = SyncReplayBufferConfig().build() + state = canonical_state() + # Two fresh supervised tokens, one stale token, and one existing semantic hole. + state.response_model_steps = [0, 0, 2, 2] + old_meta = dict(state.extra_fields) + old_length = state.num_tokens + for _ in range(2): + await buffer.put([state], "math") + result = await self.take(buffer, 1) + state = result.rollout_states[0][0] + self.assertEqual(state.labels, [-100, -100, -100, -100, 22, 23]) + self.assertEqual(state.extra_fields["supervised_tokens"], 2) + self.assertEqual(state.num_tokens, old_length) + self.assertEqual(state.extra_fields, {**old_meta, "supervised_tokens": 2}) + + async def test_training_boundary_discards_entire_invalid_group(self): + bad, sibling, good = canonical_state(1), canonical_state(2), canonical_state(3) + bad.labels = [20] + invalid_group = [bad, sibling] + trainer = SimpleNamespace(logger=MagicMock()) + batch = [invalid_group, [good]] + with patch("xtuner.v1.train.rl_trainer.release_and_discard_rollout_groups", new=AsyncMock()) as release: + dropped = await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, batch, 2) + self.assertEqual(dropped, (1, 2)) + self.assertEqual(batch, [[good]]) + release.assert_awaited_once_with([invalid_group]) + trainer.logger.warning.assert_called_once() + + async def test_all_invalid_groups_are_released_before_empty_batch_error(self): + state = canonical_state() + state.extra_fields["supervised_tokens"] = 100 + trainer = SimpleNamespace(logger=MagicMock()) + batch = [[state]] + with patch("xtuner.v1.train.rl_trainer.release_and_discard_rollout_groups", new=AsyncMock()) as release: + with self.assertRaisesRegex(ValueError, "No valid training data"): + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, batch, 2) + release.assert_awaited_once() + self.assertEqual(batch, []) + + async def test_training_check_preserves_existing_label_masks(self): + state = canonical_state() + state.task_name = "math" + state.response_model_steps = [0, 0, 2, 2] + state.labels = [-100, -100, -100, -100, 22, 23] + AgentLoop._write_train_meta(None, state) + trainer = SimpleNamespace(logger=MagicMock()) + batch = [[state]] + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, batch, 2) + self.assertEqual(state.labels, [-100, -100, -100, -100, 22, 23]) + self.assertEqual(state.extra_fields["supervised_tokens"], 2) + + async def test_invalid_groups_removed_from_task_results(self): + bad, good = canonical_state(1), canonical_state(2) + bad.labels = None + result = ProduceBatchResult( + rollout_states=[[bad], [good]], + task_results={"math": ProduceBatchResult(rollout_states=[[bad], [good]])}, + ) + trainer = SimpleNamespace(logger=MagicMock()) + with patch("xtuner.v1.train.rl_trainer.release_and_discard_rollout_groups", new=AsyncMock()): + await asyncio.to_thread(BaseRLTrainer._prepare_train_batch, trainer, result.rollout_states, 2, result) + self.assertEqual(result.rollout_states, [[good]]) + self.assertEqual(result.task_results["math"].rollout_states, [[good]]) + + async def test_shared_batch_retrieval_does_not_require_or_mutate_train_fields(self): + buffer = SyncReplayBufferConfig().build() + state = RolloutState(message=[], status=Status.COMPLETED, response_ids=[1]) + await buffer.put([state], "math", model_step=0) + runner = _TaskRunner( + task_name="math", + agent_loop=MagicMock(), + sampler=MagicMock(), + produce_strategy=SimpleNamespace(token_stale_threshold=1), + ) + result = await take_train_batch( + task_runners=[runner], + replay_buffer=buffer, + logger=MagicMock(), + manager_name="eval", + task_batch_sizes={"math": 1}, + progress=ProduceProgress.build(task_names=["math"], target_samples={"math": 1}), + current_train_step=2, + ) + self.assertEqual(result.rollout_states, [[state]]) + self.assertIsNone(state.labels) + self.assertEqual(state.extra_fields, {}) + + async def test_replay_fills_partial_versions_before_validation(self): + buffer = SyncReplayBufferConfig().build() + state = canonical_state() + state.response_model_steps = [0] + await buffer.put([state], "math", model_step=2, current_train_step=2) + self.assertEqual(state.response_model_steps, [0, 2, 2, 2]) + result = await self.take(buffer, None) + self.assertEqual(result.rollout_states[0][0].response_model_steps, [0, 2, 2, 2]) + + async def test_bad_versions_are_discarded_at_training_boundary(self): + buffer = SyncReplayBufferConfig().build() + state = canonical_state() + state.response_model_steps = [2] * 5 + await buffer.put([state], "math", model_step=2, current_train_step=2) + self.assertEqual(len(buffer), 1) + with self.assertRaisesRegex(ValueError, "No valid training data"): + await self.take(buffer, None) + self.assertEqual(len(buffer), 0) + + +class TestPackPlanDispatch(unittest.TestCase): + def test_layout_without_advantages_preserves_mapping_when_targets_change(self): + controller = TrainingController(workers=[object()] * 4) + states = [canonical_state(i) for i in range(5)] + for state, length in zip(states, [6, 4, 7, 3, 8]): + state.num_tokens = length + config = SimpleNamespace( + optimizer_steps=2, pack_strategy="legacy", pack_seed=None, pack_extra_buffer_size=1000 + ) + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + plan = controller._build_pack_plan(states, 10, config, 2) + self.assertEqual(plan["dp_plans"][0]["global_indices"], [0, 1, 4]) + self.assertEqual(plan["dp_plans"][0]["plan"], [[[0, 1]], [[2]]]) + self.assertEqual(plan["dp_plans"][1]["plan"], [[[0, 1]], [[]]]) + for advantages in ([1.0, 2.0, 3.0, 4.0, 5.0], [-5.0, -4.0, -3.0, -2.0, -1.0]): + dispatches = controller._build_dispatches(states, advantages, plan) + for rank, dp_plan in plan["dp_plans"].items(): + dispatch = dispatches[rank] + self.assertIs(dispatch["plan"], dp_plan["plan"]) + self.assertEqual(dispatch["advantages"], [advantages[i] for i in dp_plan["global_indices"]]) + self.assertEqual(dispatch["rollout_items"], [states[i] for i in dp_plan["global_indices"]]) + with self.assertRaisesRegex(ValueError, "lengths differ"): + controller._build_dispatches(states, [], plan) + token_advantages = [[float(index)] * state.num_tokens for index, state in enumerate(states)] + for rank, dispatch in controller._build_dispatches(states, token_advantages, plan).items(): + self.assertEqual( + dispatch["advantages"], [token_advantages[index] for index in plan["dp_plans"][rank]["global_indices"]] + ) + with self.assertRaisesRegex(ValueError, "Token advantage length mismatch"): + controller._build_dispatches(states, [[0.0]] * len(states), plan) + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + padding_plan = controller._build_pack_plan(states[:1], 10, config, 1) + dispatches = controller._build_dispatches(states[:1], [1.0], padding_plan) + self.assertEqual(dispatches[3]["rollout_items"], []) + self.assertEqual(dispatches[3]["advantages"], []) + self.assertEqual(dispatches[3]["plan"], [[[]]]) diff --git a/tests/rl/test_training_worker_rank.py b/tests/rl/test_training_worker_rank.py new file mode 100644 index 0000000000..533174516d --- /dev/null +++ b/tests/rl/test_training_worker_rank.py @@ -0,0 +1,775 @@ +"""TrainingWorker 侧训练数据构造的 contract 测试（RolloutState → (seq_ctx, loss_inputs) → pack）。 + +重构后 token 级张量构造从 controller 移到 training worker，本文件只测两段公开纯逻辑， +不启动 Ray、模型或 rollout backend： + +- ``_convert_rollout_items_to_train_items``：shift、advantage 张量化、seq_ctx 构造、 + distillation teacher targets 附加，保留 CPU 原始 loss_inputs + （use_3d_position_ids 时给文本样本补 3D 轴）。 +- ``_pack_train_items`` / ``_pack_one_batch``：按 controller 的 index-only 打包计划把 + (seq_ctx, loss_inputs) 拼成 ``pack_max_length`` 定长的 pack：loss_kwargs 张量沿序列维 cat、 + 模板键（rollout_logprobs / teacher 字段）在空 pack 上以固定 padding 值填充，最后重建 + pack 级 loss_ctx（仍在 CPU，设备迁移与 sp 切分留给 ``_fit``），保证跨 rank collective + 形状一致。 + +当前测试点： +- shift 布局、advantage 与 shifted_labels 逐位置对齐（regression: advantage 曾比 input_ids + 长 1 导致 pack 后整体错位）。 +- 混合批文本样本补 3D MRoPE 轴；VLM multimodal 字段与 3D position 原样保留。 +- topk / sampled teacher targets 与 shifted 布局对齐（prompt 段补零行，语义洞保留行）。 +- pack 模板要求的 loss 键必须在所有真实样本中存在；空 pack 的模板键用 padding 值填充。 +- padding 区取值表：labels=-100、advantages/rollout_logprobs=0、teacher_indices=-1、 + teacher 张量按 top_k 补零；3D position 补 arange、routed experts 补 dummy id。 +""" + +import unittest +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +import numpy as np +import ray +import torch + +from xtuner.v1.data_proto.rl_data import RolloutState, Status, TeacherTargets +from xtuner.v1.data_proto.sequence_context import SequenceContext +from xtuner.v1.datasets.mllm_tokenize_fn.qwenvl_rope2d import get_rope_index_3 +from xtuner.v1.rl.agent_loop.agent_loop import AgentLoop +from xtuner.v1.rl.distillation import DistillationConfig, DistillationTrainerAdapter, RolloutTeacherConfig +from xtuner.v1.rl.loss import DistillationLossConfig, GRPOLossConfig +from xtuner.v1.rl.trainer.pack import RLDataPacker +from xtuner.v1.rl.trainer.worker import TrainBatchAttr, TrainingWorker, get_train_seq_ctx + + +def _rollout_distillation_config(loss_config: DistillationLossConfig) -> DistillationConfig: + return DistillationConfig( + loss_config=loss_config, + teachers=[RolloutTeacherConfig(name="teacher", endpoints=["http://teacher"])], + data_source_teacher_map={"agent_math": "teacher"}, + ) + + +def _make_worker(distillation_config=None) -> TrainingWorker: + # 真实配置下 rollout teacher 依赖 loss_cfg 为 DistillationLossConfig，teacher 模板键才能进入 loss_ctx。 + loss_cfg = ( + distillation_config.loss_config + if distillation_config is not None + else GRPOLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + ) + worker = TrainingWorker.__new__(TrainingWorker) + worker.config = SimpleNamespace( + pack_max_length=16, + loss_cfg=loss_cfg, + distillation_config=distillation_config, + model_cfg=SimpleNamespace(n_routed_experts=8), + ) + worker._distillation = DistillationTrainerAdapter(distillation_config) + return worker + + +def _make_state( + *, + uid: int = 1, + prompt_ids: list[int] | None = None, + response_ids: list[int] | None = None, + logprobs: list[float] | None = None, + supervised_mask: list[int] | None = None, + input_ids: list[int] | None = None, + labels: list[int] | None = None, + position_ids: np.ndarray | None = None, + mm_info: dict | None = None, + routed_experts=None, + teacher_tokens: list[int] | list[list[int]] | None = None, + teacher_logprobs: list[float] | list[list[float]] | None = None, +) -> RolloutState: + resolved_prompt_ids = prompt_ids if prompt_ids is not None else [10, 11, 12] + resolved_response_ids = response_ids if response_ids is not None else [20, 21, 22] + state = RolloutState( + rollout_id=uid, + group_id=1, + message=[{"role": "user", "content": "prompt"}], + prompt_ids=resolved_prompt_ids, + response="response", + response_ids=resolved_response_ids, + logprobs=logprobs, + reward={"score": 1.0}, + status=Status.COMPLETED, + finish_reason="stop", + routed_experts=routed_experts, + position_ids=position_ids, + mm_info=mm_info, + extra_fields={"origin_data_source": "agent_math"} if teacher_tokens is not None else {}, + input_ids=input_ids, + labels=labels, + teacher_targets=( + TeacherTargets( + kind="topk" if isinstance(teacher_tokens[0], list) else "sampled", + tokens=teacher_tokens, + logprobs=teacher_logprobs, + ) + if teacher_tokens is not None and teacher_logprobs is not None + else None + ), + ) + if input_ids is None: + # prompt+response 样本：模拟 loop 侧 canonicalize 后的最终形态（语义洞直接烙在 labels）。 + resp_ids = list(resolved_response_ids) + state.input_ids = list(resolved_prompt_ids) + resp_ids + if supervised_mask is None: + state.labels = [-100] * len(resolved_prompt_ids) + resp_ids + else: + state.labels = [-100] * len(resolved_prompt_ids) + [ + resp_id if flag else -100 for resp_id, flag in zip(resp_ids, supervised_mask) + ] + if state.logprobs is not None: + state.logprobs = [0.0] * len(resolved_prompt_ids) + list(state.logprobs) + AgentLoop._write_train_meta(None, state) + return state + + +def _batch_attr( + use_3d: bool = False, + pack_loss_keys: list[str] | None = None, + has_routed_experts: bool = False, +) -> TrainBatchAttr: + return { + "rollout_idx": 0, + "use_3d_position_ids": use_3d, + "pack_loss_keys": pack_loss_keys or [], + "has_routed_experts": has_routed_experts, + } + + +def _convert(worker: TrainingWorker, states: list[RolloutState], advantages: list[float], use_3d: bool = False): + return worker._convert_rollout_items_to_train_items(states, advantages, _batch_attr(use_3d=use_3d)) + + +class TestConvertRolloutItemsToTrainItems(unittest.TestCase): + """RolloutState → (seq_ctx, loss_ctx)：shift、advantage 张量化、seq_ctx、teacher targets 与 CPU loss_inputs。""" + + def test_text_path_shifts_and_aligns_loss_inputs(self): + # 文本主路径固定 token 布局：input_ids 去掉 response 最后一个 token，label/logprob 对齐预测位置。 + worker = _make_worker() + routed_experts = np.array([[1, 2], [3, 4]]) + state = _make_state( + response_ids=[20, 21, 22], + logprobs=[0.1, 0.2, 0.3], + supervised_mask=[1, 0, 1], + routed_experts=routed_experts, + ) + + items = _convert(worker, [state], [1.5]) + + self.assertEqual(len(items), 1) + seq_ctx = items[0]["seq_ctx"] + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertEqual(seq_ctx.input_ids.tolist(), [[10, 11, 12, 20, 21]]) + self.assertEqual(loss_kwargs.shifted_labels.tolist(), [[-100, -100, 20, -100, 22]]) + torch.testing.assert_close( + loss_kwargs.rollout_logprobs, + torch.tensor([[0.0, 0.0, 0.1, 0.2, 0.3]], dtype=torch.float32), + ) + # advantage 与 shifted_labels 逐位置对齐: prompt 段为 0, mask=0 处为 0。 + torch.testing.assert_close( + loss_kwargs.advantages, + torch.tensor([[0.0, 0.0, 1.5, 0.0, 1.5]], dtype=torch.float32), + ) + self.assertIs(seq_ctx.rollout_routed_experts, routed_experts) + + def test_conversion_defers_loss_context_construction_until_packing(self): + worker = _make_worker() + with patch.object(type(worker.config.loss_cfg), "build", side_effect=AssertionError("premature loss context")): + items = _convert(worker, [_make_state()], [1.0]) + self.assertEqual(set(items[0]), {"seq_ctx", "loss_inputs"}) + packed = worker._pack_train_items(items, [[[0]]], _batch_attr()) + self.assertEqual(packed[0][0][1].loss_kwargs.shifted_labels.shape[1], worker.config.pack_max_length) + + def test_mixed_batch_extends_text_positions_to_3d(self): + # use_3d_position_ids=True 时 1D 文本样本沿三轴重复 arange，与 VLM 样本 pack 在一起。 + worker = _make_worker() + reasoning = _make_state( + uid=1, + position_ids=np.arange(3, dtype=np.int64).reshape(1, 1, -1).repeat(3, axis=0), + ) + agentic = _make_state( + uid=2, + input_ids=[30, 31, 40, 41, 42], + labels=[-100, -100, 40, 41, 42], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + ) + + items = _convert(worker, [reasoning, agentic], [0.5, 0.5], use_3d=True) + + self.assertEqual(tuple(items[0]["seq_ctx"].position_ids.shape), (3, 1, 5)) + agentic_position_ids = items[1]["seq_ctx"].position_ids + self.assertEqual(tuple(agentic_position_ids.shape), (3, 1, 4)) + torch.testing.assert_close( + agentic_position_ids, + torch.arange(4, dtype=torch.long).reshape(1, 1, -1).expand(3, -1, -1), + ) + + def test_vlm_keeps_mm_fields_and_3d_positions(self): + # VLM 样本保留 multimodal 字段；3D position_ids 按 response 段续写。 + worker = _make_worker() + pixel_values = np.ones((1, 2, 3), dtype=np.float32) + image_grid_thw = np.array([[1, 2, 3]], dtype=np.int32) + position_ids = np.array([[[0, 1]], [[0, 1]], [[0, 1]]], dtype=np.int64) + state = _make_state( + prompt_ids=[100, 101], + response_ids=[102, 103], + position_ids=position_ids, + mm_info={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + ) + + items = _convert(worker, [state], [0.25], use_3d=True) + + seq_ctx = items[0]["seq_ctx"] + self.assertEqual(seq_ctx.input_ids.tolist(), [[100, 101, 102]]) + self.assertEqual(tuple(seq_ctx.position_ids.shape), (3, 1, 3)) + self.assertEqual(seq_ctx.position_ids.dtype, torch.long) + self.assertIs(seq_ctx.pixel_values, pixel_values) + self.assertEqual(seq_ctx.image_grid_thw.dtype, torch.long) + self.assertEqual(seq_ctx.image_grid_thw.tolist(), [[1, 2, 3]]) + + def test_get_train_seq_ctx_mrope_continues_from_global_amax(self): + """RL 只拿 prompt 的 3D position，续写 response 后须与 SFT get_rope_index_3 一致。 + + Fixture：prompt 以 image tokens 结尾（grid 1x4x4, merge=2 → 2x2），使 per-axis max + 与 global amax 分叉。 + """ + image_token_id = 151655 + vision_start_token_id = 151652 + prompt_ids = [10, vision_start_token_id, image_token_id, image_token_id, image_token_id, image_token_id] + response_ids = [20, 21, 22] + full_ids = torch.tensor([prompt_ids + response_ids], dtype=torch.long) + image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long) + + sft_position_ids = get_rope_index_3( + full_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + ) + prompt_position_ids = get_rope_index_3( + torch.tensor([prompt_ids], dtype=torch.long), + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + ) + + seq_ctx = get_train_seq_ctx(cast(torch.LongTensor, full_ids), prompt_position_ids.numpy()) + rl_position_ids = seq_ctx.position_ids + assert rl_position_ids is not None + self.assertEqual(tuple(rl_position_ids.shape), tuple(sft_position_ids.shape)) + torch.testing.assert_close(rl_position_ids, sft_position_ids) + + def test_topk_teacher_targets_align_shifted_layout(self): + # topk targets 补齐 prompt 段零行后与 shifted_labels 逐行对齐，并可直接构建 loss。 + loss_config = DistillationLossConfig( + policy_loss_cfg={ + "loss_type": "vanilla", + "cliprange_low": 0.2, + "cliprange_high": 0.2, + }, + loss_mode="reverse", + use_policy_gradient=False, + top_k=2, + ) + worker = _make_worker(_rollout_distillation_config(loss_config)) + state = _make_state( + input_ids=[10, 11, 20, 21, 22], + labels=[-100, -100, 20, -100, 22], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + teacher_tokens=[[100, 101], [102, 103], [104, 105]], + teacher_logprobs=[[-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]], + ) + + items = _convert(worker, [state], [0.0]) + + self.assertEqual(len(items), 1) + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertEqual(loss_kwargs.shifted_labels.tolist(), [[-100, 20, -100, 22]]) + self.assertEqual( + loss_kwargs.target_token_ids.tolist(), + [[[0, 0], [100, 101], [102, 103], [104, 105]]], + ) + torch.testing.assert_close( + loss_kwargs.teacher_logprobs, + torch.tensor([[[0.0, 0.0], [-0.5, -0.6], [-0.7, -0.8], [-0.9, -1.0]]], dtype=torch.float32), + ) + with patch("xtuner.v1.rl.loss.distillation_loss.DEVICE", "cpu"): + loss_ctx = loss_config.build( + { + "shifted_labels": loss_kwargs.shifted_labels, + "advantages": loss_kwargs.advantages, + "old_logprobs": torch.zeros_like(loss_kwargs.shifted_labels, dtype=torch.float32), + "teacher_logprobs": loss_kwargs.teacher_logprobs, + "target_token_ids": loss_kwargs.target_token_ids, + } + ) + assert loss_ctx is not None + type(loss_ctx).build_batches([loss_ctx]) + loss, _ = loss_ctx.loss_fn( + hidden_states=torch.randn(1, 4, 8), + head_weight=torch.randn(128, 8), + head_bias=None, + loss_kwargs=loss_ctx.loss_kwargs, + ) + self.assertTrue(torch.isfinite(loss)) + + def test_topk_targets_include_masked_response_rows(self): + # labels=-100 的 response 行也保留 teacher 行（loss 侧自行忽略），只在前面补零行。 + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + worker = _make_worker(_rollout_distillation_config(loss_config)) + state = _make_state( + prompt_ids=[10, 11, 12], + response_ids=[20, 21, 22], + supervised_mask=[0, 1, 1], + teacher_tokens=[[102, 103], [104, 105]], + teacher_logprobs=[[-0.7, -0.8], [-0.9, -1.0]], + ) + + items = _convert(worker, [state], [0.0]) + + loss_kwargs = SimpleNamespace(**items[0]["loss_inputs"]) + self.assertEqual(loss_kwargs.shifted_labels.tolist(), [[-100, -100, -100, 21, 22]]) + self.assertEqual( + loss_kwargs.target_token_ids.tolist(), + [[[0, 0], [0, 0], [0, 0], [102, 103], [104, 105]]], + ) + torch.testing.assert_close( + loss_kwargs.teacher_logprobs, + torch.tensor([[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [-0.7, -0.8], [-0.9, -1.0]]], dtype=torch.float32), + ) + + def test_sampled_targets_have_no_token_ids(self): + # sampled（k1）目标只有逐位置 teacher logprobs，不产出 target_token_ids 键。 + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="k1", + use_policy_gradient=True, + ) + worker = _make_worker(_rollout_distillation_config(loss_config)) + plain_state = _make_state( + uid=1, + prompt_ids=[10, 11, 12], + response_ids=[20, 21, 22], + supervised_mask=[0, 1, 1], + teacher_tokens=[21, 22], + teacher_logprobs=[-0.7, -0.9], + ) + agentic_state = _make_state( + uid=2, + input_ids=[30, 31, 40, 41, 42], + labels=[-100, -100, 40, -100, 42], + logprobs=[0.0, -0.1, -0.2, -0.3, -0.4], + teacher_tokens=[40, 41, 42], + teacher_logprobs=[-1.1, -1.2, -1.3], + ) + + items = _convert(worker, [plain_state, agentic_state], [0.0, 0.0]) + + self.assertEqual(len(items), 2) + torch.testing.assert_close( + items[0]["loss_inputs"].get("teacher_logprobs"), + torch.tensor([[0.0, 0.0, 0.0, -0.7, -0.9]], dtype=torch.float32), + ) + torch.testing.assert_close( + items[1]["loss_inputs"].get("teacher_logprobs"), + torch.tensor([[0.0, -1.1, -1.2, -1.3]], dtype=torch.float32), + ) + self.assertIsNone(items[0]["loss_inputs"].get("target_token_ids")) + self.assertIsNone(items[1]["loss_inputs"].get("target_token_ids")) + + +class TestPackTrainItems(unittest.TestCase): + """(seq_ctx, loss_inputs) 列表 → 定长 pack：拼接、padding 值表、模板键与 pack 级 loss_ctx 重建。""" + + def _pack(self, worker, items, plan, use_3d=False, pack_loss_keys=None, has_routed_experts=False): + return worker._pack_train_items(items, plan, _batch_attr(use_3d, pack_loss_keys, has_routed_experts)) + + def test_packs_concatenate_along_sequence(self): + # 两个样本按 plan 顺序拼进一个 pack，padding 区补满 pack_max_length。 + worker = _make_worker() + first = _make_state(uid=1, response_ids=[20, 21], logprobs=[0.1, 0.2]) + second = _make_state(uid=2, prompt_ids=[30, 31], response_ids=[40], logprobs=[-0.3]) + items = _convert(worker, [first, second], [1.0, 2.0]) + + step_batches = self._pack(worker, items, [[[0, 1]]]) + + self.assertEqual(len(step_batches), 1) + seq_ctx, loss_ctx = step_batches[0][0] + seq_len = seq_ctx.input_ids.shape[1] + self.assertEqual(seq_len, worker.config.pack_max_length) + # first: prompt 3 + response 2 → 4 个训练位; second: 2 个训练位; padding 10。 + self.assertEqual(seq_ctx.num_padding, 16 - 6) + loss_kwargs = loss_ctx.loss_kwargs + # pack 阶段产物保持 CPU，设备迁移由 _fit 统一执行。 + self.assertEqual(loss_kwargs.shifted_labels.device.type, "cpu") + # 真实段逐位置保序；padding 区 labels=-100、advantages=0。 + self.assertEqual(loss_kwargs.shifted_labels[0, :4].tolist(), [-100, -100, 20, 21]) + self.assertEqual(loss_kwargs.shifted_labels[0, 4:6].tolist(), [-100, 40]) + torch.testing.assert_close( + loss_kwargs.advantages[0, :6], + torch.tensor([0.0, 0.0, 1.0, 1.0, 0.0, 2.0]), + ) + pad = loss_kwargs.shifted_labels[0, 6:] + self.assertTrue(torch.all(pad == -100)) + self.assertTrue(torch.all(loss_kwargs.advantages[0, 6:] == 0)) + + def test_pack_fields_round_trip_to_original_samples(self): + # 各字段共用同一份真实 pack plan；按样本边界拆回后应逐元素一致。 + # 长度和字段值各不相同，避免相同值掩盖样本重排、漏样本或字段错位。 + worker = _make_worker( + _rollout_distillation_config( + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + ) + ) + worker.config.pack_max_length = 10 + lengths = [6, 7, 4, 10, 3, 8, 5] + states = [] + for index, length in enumerate(lengths): + base = (index + 1) * 100 + response = list(range(base + 2, base + length + 1)) + states.append( + _make_state( + uid=index + 1, + prompt_ids=[base, base + 1], + response_ids=response, + supervised_mask=[int(position % 3 != 1) for position in range(len(response))], + logprobs=[-(base + position) / 1000 for position in range(len(response))], + teacher_tokens=[[token, token + 1000] for token in response], + teacher_logprobs=[[-token / 1000, -token / 2000] for token in response], + ) + ) + items = _convert(worker, states, [float(index + 1) for index in range(len(states))]) + originals = [ + { + "input_ids": item["seq_ctx"].input_ids.clone(), + "position_ids": item["seq_ctx"].position_ids.clone(), + **{key: value.clone() for key, value in item["loss_inputs"].items() if value is not None}, + } + for item in items + ] + for strategy in ("legacy", "greedy", "balance", "native"): + with self.subTest(strategy=strategy): + with patch("xtuner.v1.rl.trainer.pack.XTUNER_DETERMINISTIC", True): + plans, _ = RLDataPacker(10, 3, 1, 2, pack_strategy=strategy, pack_seed=7).pack(lengths) + restored = {} + for plan in plans: + batches = self._pack(worker, items, plan, pack_loss_keys=["rollout_logprobs"]) + self.assertEqual(len(batches), len(plan)) + for step_plan, step_batches in zip(plan, batches): + self.assertEqual(len(step_batches), len(step_plan)) + for indices, (seq, loss) in zip(step_plan, step_batches): + fields = { + "input_ids": seq.input_ids, + "position_ids": seq.position_ids, + **{ + key: getattr(loss.loss_kwargs, key) + for key in originals[0] + if key not in ("input_ids", "position_ids") + }, + } + offset = 0 + for index in indices: + self.assertNotIn(index, restored, "sample was packed more than once") + restored[index] = { + key: value[:, offset : offset + lengths[index]].clone() + for key, value in fields.items() + } + offset += lengths[index] + self.assertEqual(seq.num_padding, worker.config.pack_max_length - offset) + for value in fields.values(): + self.assertEqual(value.shape[1], worker.config.pack_max_length) + self.assertEqual(sorted(restored), list(range(len(states)))) + for index, original in enumerate(originals): + for key, expected in original.items(): + with self.subTest(sample=index, field=key): + torch.testing.assert_close(restored[index][key], expected, rtol=0, atol=0) + # Packing must also leave its source fields untouched. + source = ( + getattr(items[index]["seq_ctx"], key) + if key in ("input_ids", "position_ids") + else items[index]["loss_inputs"][key] + ) + torch.testing.assert_close(source, expected, rtol=0, atol=0) + + def test_single_item_without_padding_preserves_fields(self): + # 单样本满包仍经过 cat，保留字段和值并统一 packed 表示。 + worker = _make_worker() + state = _make_state( + uid=1, + prompt_ids=[10, 11, 12, 10, 11, 12, 10, 11, 12, 10, 11, 12, 10, 11, 12, 10], + response_ids=[20], + ) + items = _convert(worker, [state], [1.0]) + self.assertEqual(items[0]["seq_ctx"].input_ids.shape[1], worker.config.pack_max_length) + + step_batches = self._pack(worker, items, [[[0]]]) + + seq_ctx, _ = step_batches[0][0] + torch.testing.assert_close(seq_ctx.input_ids, items[0]["seq_ctx"].input_ids) + self.assertEqual(seq_ctx.num_padding, 0) + + def test_empty_pack_materializes_all_padding(self): + # 调度占位空 pack：整包都是 padding（全零 input_ids + labels=-100 + advantages=0）。 + worker = _make_worker() + + step_batches = self._pack(worker, [], [[[]]]) + + seq_ctx, loss_ctx = step_batches[0][0] + seq_len = seq_ctx.input_ids.shape[1] + self.assertEqual(seq_len, worker.config.pack_max_length) + self.assertEqual(seq_ctx.num_padding, 16) + self.assertTrue(torch.all(seq_ctx.input_ids == 0)) + loss_kwargs = loss_ctx.loss_kwargs + self.assertTrue(torch.all(loss_kwargs.shifted_labels == -100)) + self.assertTrue(torch.all(loss_kwargs.advantages == 0)) + + def test_rollout_logprobs_template_key_on_empty_pack(self): + # pack_loss_keys 声明的批级模板键在空 pack 上也要物化，保证跨 rank collective 形状一致。 + worker = _make_worker() + + step_batches = self._pack(worker, [], [[[]]], pack_loss_keys=["rollout_logprobs"]) + + _, loss_ctx = step_batches[0][0] + rollout_logprobs = loss_ctx.loss_kwargs.rollout_logprobs + assert rollout_logprobs is not None + torch.testing.assert_close( + rollout_logprobs, + torch.zeros(1, worker.config.pack_max_length, dtype=torch.float32), + ) + + def test_teacher_template_keys_padded_on_empty_pack(self): + # 配置 distillation 后空 pack 也要带 teacher 模板键（top_k 形状补零），供全 pad rank 对齐。 + loss_config = DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + worker = _make_worker(_rollout_distillation_config(loss_config)) + + step_batches = self._pack(worker, [], [[[]]]) + + _, loss_ctx = step_batches[0][0] + loss_kwargs = loss_ctx.loss_kwargs + torch.testing.assert_close( + loss_kwargs.teacher_logprobs, + torch.zeros(1, worker.config.pack_max_length, 2, dtype=torch.float32), + ) + torch.testing.assert_close( + loss_kwargs.target_token_ids, + torch.zeros(1, worker.config.pack_max_length, 2, dtype=torch.long), + ) + + def test_mixed_presence_within_pack_raises(self): + # 同一 pack 模板要求的 loss 键必须在所有真实样本中存在，混出即 fail fast。 + worker = _make_worker() + items = [ + { + "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[1, 2]]),), device="cpu"), + "loss_inputs": { + "shifted_labels": torch.tensor([[-100, 2]]), + "advantages": torch.zeros(1, 2), + "rollout_logprobs": torch.zeros(1, 2), + }, + }, + { + "seq_ctx": SequenceContext.from_input_ids((torch.tensor([[3, 4]]),), device="cpu"), + "loss_inputs": {"shifted_labels": torch.tensor([[3, 4]]), "advantages": torch.zeros(1, 2)}, + }, + ] + + with self.assertRaisesRegex(ValueError, "must be present on every sample"): + self._pack(worker, items, [[[0, 1]]], pack_loss_keys=["rollout_logprobs"]) + + def test_oversize_pack_raises(self): + # pack 内真实 token 数超过 pack_max_length 时 fail fast（controller 侧 packer 同样拦截）。 + worker = _make_worker() + long_ids = list(range(100, 120)) + state = _make_state( + uid=1, + input_ids=long_ids, + labels=[-100] * 19 + [42], + logprobs=[0.0] * 20, + ) + items = _convert(worker, [state], [1.0]) + + with self.assertRaises(AssertionError): + self._pack(worker, items, [[[0]]]) + + def test_padding_seq_ctx_uses_3d_positions_when_configured(self): + # use_3d_position_ids=True 时 padding 区 position_ids 为三轴 arange。 + worker = _make_worker() + + step_batches = self._pack(worker, [], [[[]]], use_3d=True) + + position_ids = step_batches[0][0][0].position_ids + self.assertEqual(tuple(position_ids.shape), (3, 1, worker.config.pack_max_length)) + torch.testing.assert_close( + position_ids, + torch.arange(worker.config.pack_max_length, dtype=torch.long).reshape(1, 1, -1).expand(3, -1, -1), + ) + + def test_padding_seq_ctx_gets_dummy_routed_experts(self): + # has_routed_experts=True 时 padding 区补 dummy expert id，形状与真实 rollout_routed_experts 对齐。 + worker = _make_worker() + + step_batches = self._pack(worker, [], [[[]]], has_routed_experts=True) + + routed_experts = step_batches[0][0][0].rollout_routed_experts[0] + self.assertEqual(tuple(routed_experts.shape), (worker.config.pack_max_length, 1, 1)) + self.assertTrue(torch.all(routed_experts >= 0)) + self.assertTrue(torch.all(routed_experts < worker.config.model_cfg.n_routed_experts)) + + def test_required_loss_field_missing_from_every_real_sample_raises(self): + worker = _make_worker() + items = _convert(worker, [_make_state()], [1.0]) + with self.assertRaisesRegex(ValueError, "Required pack field rollout_logprobs"): + self._pack(worker, items, [[[0]]], pack_loss_keys=["rollout_logprobs"]) + + def test_sampled_teacher_padding_preserves_real_targets_and_empty_dp(self): + worker = _make_worker( + _rollout_distillation_config(DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"})) + ) + state = _make_state(teacher_tokens=[20, 21, 22], teacher_logprobs=[-0.1, -0.2, -0.3]) + items = _convert(worker, [state], [0.0]) + real_length = items[0]["seq_ctx"].input_ids.shape[1] + # Two DP ranks: one real sample and one scheduling-only empty pack. + from xtuner.v1.rl.trainer.pack import RLDataPacker + + plans, _ = RLDataPacker(16, 2, 1, 1).pack([real_length]) + for dp_rank, plan in enumerate(plans): + with self.subTest(dp_rank=dp_rank): + seq, loss = self._pack(worker, items if dp_rank == 0 else [], plan)[0][0] + kwargs = loss.loss_kwargs + self.assertEqual(tuple(kwargs.teacher_logprobs.shape), (1, 16)) + self.assertEqual(kwargs.teacher_logprobs.device.type, "cpu") + self.assertIsNone(kwargs.target_token_ids) + self.assertEqual(seq.input_ids.shape[1], kwargs.shifted_labels.shape[1]) + start = real_length if dp_rank == 0 else 0 + if dp_rank == 0: + torch.testing.assert_close( + kwargs.teacher_logprobs[:, :real_length], items[0]["loss_inputs"].get("teacher_logprobs") + ) + self.assertTrue(torch.all(kwargs.teacher_logprobs[:, start:] == 0)) + self.assertTrue(torch.all(kwargs.shifted_labels[:, start:] == -100)) + self.assertTrue(torch.all(kwargs.advantages[:, start:] == 0)) + + def test_short_pack_pads_topk_targets_and_rollout_logprobs(self): + worker = _make_worker( + _rollout_distillation_config( + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + ) + ) + state = _make_state( + logprobs=[-0.3, -0.4, -0.5], + teacher_tokens=[[20, 30], [21, 31], [22, 32]], + teacher_logprobs=[[-0.1, -0.2]] * 3, + ) + items = _convert(worker, [state], [0.0]) + real_length = items[0]["seq_ctx"].input_ids.shape[1] + seq, loss = self._pack(worker, items, [[[0]]], pack_loss_keys=["rollout_logprobs"])[0][0] + for key in ("shifted_labels", "advantages", "rollout_logprobs", "teacher_logprobs", "target_token_ids"): + with self.subTest(key=key): + packed = getattr(loss.loss_kwargs, key) + original = getattr(SimpleNamespace(**items[0]["loss_inputs"]), key) + self.assertEqual(packed.shape[1], seq.input_ids.shape[1]) + self.assertEqual(packed.device.type, "cpu") + torch.testing.assert_close(packed[:, :real_length], original) + pad_value = -100 if key == "shifted_labels" else 0 + self.assertTrue(torch.all(packed[:, real_length:] == pad_value)) + + def test_short_pack_pads_training_teacher_indices(self): + worker = _make_worker() + loss_config = DistillationLossConfig(policy_loss_cfg={"loss_type": "vanilla"}) + worker.config.loss_cfg = loss_config + seq = SequenceContext.from_input_ids((torch.tensor([[1, 2, 3]]),), device="cpu") + item = { + "seq_ctx": seq, + "loss_inputs": { + "shifted_labels": torch.tensor([[-100, 2, 3]]), + "advantages": torch.tensor([[0.0, 1.0, 1.0]]), + "teacher_indices": torch.tensor([[2, 2, 2]]), + }, + } + _, packed = self._pack(worker, [item], [[[0]]], pack_loss_keys=["teacher_indices"])[0][0] + torch.testing.assert_close(packed.loss_kwargs.teacher_indices[:, :3], torch.tensor([[2, 2, 2]])) + self.assertTrue(torch.all(packed.loss_kwargs.teacher_indices[:, 3:] == -1)) + self.assertEqual(tuple(packed.loss_kwargs.teacher_indices.shape), (1, 16)) + + def test_single_full_pack_materializes_single_ref_and_segment_refs(self): + for segmented in (False, True): + with self.subTest(segmented=segmented): + worker = _make_worker() + worker.config.model_cfg = SimpleNamespace( + n_routed_experts=8, num_hidden_layers=1, num_experts_per_tok=1 + ) + worker.config.free_rollout_routed_experts_in_worker = False + ref = ray.ObjectRef(b"a" * 28) + payload = [ref] if segmented else ref + state = _make_state(prompt_ids=list(range(16)), response_ids=[20], routed_experts=payload) + items = _convert(worker, [state], [1.0]) + seq, _ = self._pack(worker, items, [[[0]]], has_routed_experts=True)[0][0] + self.assertEqual(seq.rollout_routed_experts, [payload]) + values = (np.arange(17 if segmented else 16) % 8).astype(np.int32).reshape(-1, 1) + with patch.object(ray, "get", return_value=values): + worker._add_rollout_routed_experts(seq, seq.rollout_routed_experts) + expected = values[1:] if segmented else values + np.testing.assert_array_equal(seq.rollout_routed_experts.numpy().reshape(-1, 1), expected) + + def test_configured_teacher_loss_keys(self): + # 模板键由 distillation 配置推导：rollout topk → logprobs+token_ids；sampled → 仅 logprobs； + # train teachers → teacher_indices；未配置 → 空。 + topk_worker = _make_worker( + _rollout_distillation_config( + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="forward_kl_topk", + use_policy_gradient=False, + top_k=2, + ) + ) + ) + self.assertEqual(topk_worker._configured_teacher_loss_keys(), ["teacher_logprobs", "target_token_ids"]) + + sampled_worker = _make_worker( + _rollout_distillation_config( + DistillationLossConfig( + policy_loss_cfg={"loss_type": "vanilla"}, + loss_mode="k1", + use_policy_gradient=True, + ) + ) + ) + self.assertEqual(sampled_worker._configured_teacher_loss_keys(), ["teacher_logprobs"]) + + train_teacher_config = SimpleNamespace( + rollout_teachers=[], + train_teachers=[object()], + loss_config=SimpleNamespace(uses_sampled_token_targets=True), + ) + train_teacher_worker = _make_worker() + train_teacher_worker.config.distillation_config = train_teacher_config + self.assertEqual(train_teacher_worker._configured_teacher_loss_keys(), ["teacher_indices"]) + + self.assertEqual(_make_worker()._configured_teacher_loss_keys(), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 1173f31e88..0fdb3af5c6 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -285,63 +285,6 @@ def reset_rollout_response(rollout_state: RolloutState) -> RolloutState: return rollout_state -def is_valid_for_training(group_data_items: list[RolloutState], logger) -> bool: - """Checks if a group of rollout states is valid for a training step. - - Args: - group_data_items (list[RolloutState]): A list of RolloutState objects. - logger (logging.Logger): Logger used to report the invalid reason. - - Returns: - bool: True if the group is valid, False otherwise. - - NOTE: Why this check is needed: - - For system fault tolerance, this check is performed at rollout / dataflow - time, but we still do it here to ensure training data integrity. - - 'filtered'/'failed': These items are fundamentally broken or incomplete and - should not be used for training. - - 'aborted': These items represent rollouts that were stopped - prematurely. Using such partial data could lead the model to learn - undesirable behaviors (e.g., stopping generation too early). - - Empty response/response_ids: The model's generated response is the core - of the training data for RL algorithms like PPO. If the response is - missing, there is nothing to compute rewards on or to train the model with. - """ - is_abort = any(item.status == Status.ABORTED for item in group_data_items) - is_filtered = any(item.status == Status.FILTERED for item in group_data_items) - is_failed = any(item.status == Status.FAILED for item in group_data_items) - if is_filtered or is_failed or is_abort: - logger.warning( - f"Invalid dataflow group found during training, rollout state filtered: {is_filtered}, failed: {is_failed}, aborted: {is_abort}." - ) - return False - for item in group_data_items: - if item.input_ids is not None: - input_ids_valid = len(item.input_ids) > 1 - labels_valid = item.labels is not None and len(item.labels) == len(item.input_ids) - logprobs_valid = item.logprobs is None or len(item.logprobs) == len(item.input_ids) - if not input_ids_valid or not labels_valid or not logprobs_valid: - logger.warning( - "Invalid dataflow item found during training: input_ids, labels, and logprobs lengths mismatch." - ) - return False - continue - - response_valid = item.response is not None and len(item.response) > 0 - ids_valid = item.response_ids is not None and len(item.response_ids) > 0 - if not ids_valid: - # NOTE: `response_ids` is the critical field for token-in-token-out mode, so we ensure it's not empty. - logger.warning( - "Invalid dataflow item found during training: no response or response_ids and skip this item." - ) - return False - if not response_valid: - # NOTE: check valid response string for judger inputs - logger.warning("Invalid dataflow item found during training: empty response string and skip this item.") - return False - return True - - def get_group_status(rollout_states: list[RolloutState]) -> Status: """Get the group status based on the individual rollout states. @@ -464,8 +407,7 @@ def calculate_group_effective_response_masks( policy. The function is pure: ``group`` is inspected, not modified. Each returned mask is the effective - response mask aligned with ``response_ids``; callers bake them into ``labels`` when the train - batch is taken, clearing zero-mask supervised positions to ``-100``. The clearing only ever + response mask aligned with ``response_ids``; callers bake them into ``labels`` in ``take_train_batch``, clearing zero-mask supervised positions to ``-100``. The clearing only ever extends, so masks computed at growing trainer steps converge to the same labels. ``None`` means token staleness is disabled or does not apply to that state (no labels, or no ``response_ids`` to align with). @@ -495,3 +437,72 @@ def calculate_group_effective_response_masks( ) ) return masks + + +def validate_rollout_group(group: list[RolloutState], *, task_name: str) -> None: + """Check producer result structure, final status and existing version + bounds. + + This function is read-only: it neither counts tokens nor fills model versions. + """ + allowed = {Status.COMPLETED, Status.FAILED, Status.FILTERED, Status.ABORTED, Status.EXPIRED} + for state in group: + context = f"task={task_name!r}, group_id={state.group_id}, rollout_id={state.rollout_id}" + if state.status not in allowed: + raise ValueError(f"{context}: invalid generated status {state.status}") + response_length = len(state.response_ids or []) + version_length = len(state.response_model_steps or []) + if version_length > response_length: + raise ValueError( + f"{context}: response_model_steps length={version_length}, expected at most response_ids={response_length}" + ) + + +def validate_train_group(group: list[RolloutState], *, task_name: str) -> None: + """Check a completed group's full fields immediately before training. + + Raise ValueError so the training boundary can discard the entire group. This function does not fill source versions + or modify labels or metadata. + """ + if not group or any(not isinstance(state, RolloutState) or state.status != Status.COMPLETED for state in group): + raise ValueError(f"task={task_name!r}: training requires a nonempty COMPLETED rollout group") + for state in group: + context = f"task={task_name!r}, group_id={state.group_id}, rollout_id={state.rollout_id}" + input_ids = state.input_ids + if input_ids is None or len(input_ids) < 2: + length = None if input_ids is None else len(input_ids) + raise ValueError(f"{context}: input_ids must contain at least 2 tokens, got {length}") + for field in ("labels", "logprobs"): + value = getattr(state, field) + if field == "logprobs" and value is None: + continue + length = None if value is None else len(value) + if length != len(input_ids): + raise ValueError(f"{context}: {field} length mismatch: {length}, expected input_ids={len(input_ids)}") + labels = state.labels + assert labels is not None + response_ids = state.response_ids or [] + version_length = len(state.response_model_steps or []) + if version_length != len(response_ids): + raise ValueError( + f"{context}: response_model_steps length={version_length}, expected response_ids={len(response_ids)}" + ) + position_ids = state.position_ids + expected = { + "train_response_length": len(response_ids), + "supervised_tokens": sum(label != -100 for label in labels[1:]), + "position_layout": "mrope_3d" if position_ids is not None and position_ids.ndim == 3 else "1d", + "has_rollout_logprobs": state.logprobs is not None, + "has_routed_experts": state.routed_experts is not None, + } + if state.num_tokens != len(input_ids) - 1: + raise ValueError( + f"{context}: num_tokens={state.num_tokens}, expected={len(input_ids) - 1}; missing or stale train meta" + ) + prompt_length = state.extra_fields.get("train_prompt_length") + if not isinstance(prompt_length, int) or prompt_length < 0: + raise ValueError(f"{context}: missing or invalid train_prompt_length={prompt_length}") + for key, value in expected.items(): + actual = state.extra_fields.get(key) + if actual != value: + raise ValueError(f"{context}: train meta {key}={actual!r}, expected={value!r}") diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index b7e8736a5b..c9f65ce2cd 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -11,7 +11,12 @@ from ray.actor import ActorClass, ActorProxy from ray.util.placement_group import PlacementGroup -from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status, get_group_status +from xtuner.v1.data_proto.rl_data import ( + RolloutState, + SampleParams, + Status, + get_group_status, +) from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.rollout.constants import AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY @@ -284,66 +289,84 @@ async def generate_one(state: RolloutState) -> RolloutState: group = await self._teacher_scorer.on_group_ready(group) return self.canonicalize_train_fields(group) - def canonicalize_train_fields(self, group: list[RolloutState]) -> list[RolloutState]: - """Canonicalize the training token fields of one rollout group. - - Prompt+response samples without ``input_ids`` are assembled in place into the - unified full-sequence convention: ``input_ids``/``labels``/``logprobs`` share one - length, where ``labels[j]`` is the prediction target of the token at position - ``j`` and the train controller applies the one-position shift at conversion time. - Samples that already carry ``input_ids`` satisfy the convention and are left - untouched. Non-completed samples are skipped. Failures mark the single sample as - ``Status.FAILED`` instead of raising into the producer. Loops that assemble their - own full sequences (e.g. the agentic localhost/sandbox loops) override this - method with their own loop-specific implementation instead of relying on this - default. + def _write_train_meta(self, rollout_state: RolloutState) -> None: + """Write the controller-facing lightweight train metadata in place. + + Agent loops call this once the canonical training fields (``input_ids``/``labels``/ + ``logprobs``) are finalized, so downstream scheduling only consumes scalar statistics + instead of token-level fields. Samples that are not ``COMPLETED`` or that lack aligned + ``input_ids``/``labels`` are left untouched; the function is idempotent. Args: - group (list[RolloutState]): One group of rollout states after generation, - judging, and filtering. + rollout_state (RolloutState): Rollout state updated in place: ``num_tokens`` + becomes the shifted training length (``len(input_ids) - 1``) and + ``extra_fields`` gains ``train_prompt_length``, ``train_response_length``, + ``supervised_tokens``, ``position_layout``, ``has_rollout_logprobs`` and + ``has_routed_experts``. + """ + if rollout_state.status != Status.COMPLETED: + return + input_ids = rollout_state.input_ids + labels = rollout_state.labels + if input_ids is None or labels is None or len(input_ids) != len(labels) or len(input_ids) < 2: + return + shifted_labels = labels[1:] + supervised_tokens = sum(1 for label in shifted_labels if label != -100) + # prompt_len 统计原始输入 prompt 长度（VLM 的 prompt 在 train_prompt_ids）；prompt_ids + # 缺失时退回 shifted_labels 的非监督位计数，与 controller 既有 data_info 口径一致。 + prompt_ids = rollout_state.extra_fields.get("train_prompt_ids") or rollout_state.prompt_ids + prompt_length = len(prompt_ids) if prompt_ids else len(shifted_labels) - supervised_tokens + position_ids = rollout_state.position_ids + rollout_state.num_tokens = len(input_ids) - 1 + rollout_state.extra_fields["train_prompt_length"] = prompt_length + rollout_state.extra_fields["train_response_length"] = len(rollout_state.response_ids or []) + rollout_state.extra_fields["supervised_tokens"] = supervised_tokens + rollout_state.extra_fields["position_layout"] = ( + "mrope_3d" if position_ids is not None and position_ids.ndim == 3 else "1d" + ) + # 批级 presence 标志：controller 只读布尔位即可判定 pack 模板（rollout logprobs 的 + # IS 指标 collective、routed experts dummy 填充），不需要触碰 token 级字段。 + rollout_state.extra_fields["has_rollout_logprobs"] = rollout_state.logprobs is not None + rollout_state.extra_fields["has_routed_experts"] = rollout_state.routed_experts is not None + + def canonicalize_train_fields(self, group: list[RolloutState]) -> list[RolloutState]: + """Build full, unshifted training fields and write lightweight + metadata. - Returns: - list[RolloutState]: The same group with canonical training fields. + Already assembled sequences are preserved. A construction error marks that sample failed so the producer can + discard its group without aborting the produce window. """ for state in group: - if state.status != Status.COMPLETED or state.input_ids is not None: + if state.status != Status.COMPLETED: continue try: - is_vlm_model = "train_prompt_ids" in state.extra_fields - if is_vlm_model: - # TODO(hha): VLM, 不好的设计，后续要去掉 - prompt_ids = state.extra_fields["train_prompt_ids"] - else: - prompt_ids = state.prompt_ids - assert prompt_ids is not None and len(prompt_ids) > 0, ( - f"Prompt ids cannot be None or empty in data: {state}" - ) - - response_ids: list[int] - if state.response_ids is not None: - resp_ids_raw = state.response_ids - if isinstance(resp_ids_raw, torch.Tensor): - response_ids = resp_ids_raw.flatten().tolist() - else: - response_ids = cast(list[int], resp_ids_raw) - else: - assert state.response is not None, "response item cannot be None" - response_ids = self.tokenizer(state.response, return_tensors="pt")["input_ids"].flatten().tolist() - # 归一化写回：Tensor 展平为 list、tokenizer 兜底结果也落到 response_ids，后续消费者零分支。 - state.response_ids = response_ids - - if state.logprobs is not None: - assert len(state.logprobs) == len(response_ids), ( - f"{len(state.logprobs)} vs {len(response_ids)}, data: {state}" + if state.input_ids is None: + prompt_ids = state.extra_fields.get("train_prompt_ids", state.prompt_ids) + assert prompt_ids is not None and len(prompt_ids) > 0, ( + f"Prompt ids cannot be None or empty in data: {state}" ) - # 只有 response 部分有 logprobs, 前面补 0 对齐全序列 - state.logprobs = [0.0] * len(prompt_ids) + list(state.logprobs) - - # 返回的 routed_experts 不包括 eos 的值，实际上也不需要，训练输入会去掉最后一位 - state.input_ids = list(prompt_ids) + response_ids - state.labels = [-100] * len(prompt_ids) + list(response_ids) + if state.response_ids is not None: + response_ids = ( + state.response_ids.flatten().tolist() + if isinstance(state.response_ids, torch.Tensor) + else list(state.response_ids) + ) + else: + assert state.response is not None, "response item cannot be None" + response_ids = ( + self.tokenizer(state.response, return_tensors="pt")["input_ids"].flatten().tolist() + ) + state.response_ids = response_ids + if state.logprobs is not None: + assert len(state.logprobs) == len(response_ids), ( + f"{len(state.logprobs)} vs {len(response_ids)}, data: {state}" + ) + state.logprobs = [0.0] * len(prompt_ids) + list(state.logprobs) + state.input_ids = list(prompt_ids) + response_ids + state.labels = [-100] * len(prompt_ids) + response_ids + self._write_train_meta(state) except Exception as exc: - self.logger.error(f"Canonicalize train fields failed for rollout_id={state.rollout_id}: {exc}") + self.logger.error(f"Canonicalize failed for rollout_id={state.rollout_id}: {exc}") state.status = Status.FAILED state.error_msg = f"canonicalize_train_fields failed: {exc}" return group diff --git a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py index bdd7f8c117..bc303be9b0 100644 --- a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py +++ b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py @@ -163,9 +163,6 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou resp_id if mask_value else -100 for resp_id, mask_value in zip(final_response_ids, final_response_mask) ] rollout_state.response = self.tokenizer.decode(rollout_state.response_ids) - assert len(rollout_state.input_ids) == len(rollout_state.labels) == len(rollout_state.logprobs), ( - f"{len(rollout_state.input_ids)} vs {len(rollout_state.labels)} vs {len(rollout_state.logprobs)}" - ) if self.judger is not None and not self.enable_batch_judge: rollout_state = await self.run_judger(rollout_state) return rollout_state diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index 509eca487f..0d2bbeea80 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -156,37 +156,13 @@ async def generate_one(state: RolloutState) -> RolloutState: return self.canonicalize_train_fields(samples) def canonicalize_train_fields(self, group: list[RolloutState]) -> list[RolloutState]: - """Validate the canonical training fields of one localhost agentic - rollout group. - - This loop fills ``input_ids``/``labels``/``logprobs`` from the trace store as one - unshifted full sequence, so there is nothing left to assemble. Canonicalization - only re-checks the unified convention (``labels`` present and of equal length, - ``logprobs`` either ``None`` or of equal length) and marks a violating sample - ``Status.FAILED`` instead of raising into the producer. States without - ``input_ids`` (eval-mode states) carry no training fields and are skipped. - - Args: - group (list[RolloutState]): One group of rollout states after generation, - judging, and filtering. - - Returns: - list[RolloutState]: The same group with validated training fields. + """Preserve trace-exported unshifted fields and write train metadata. + + Eval states without training fields are left untouched. The trainer checks full training fields before + consuming a group. """ for state in group: - if state.status != Status.COMPLETED or state.input_ids is None: - continue - try: - labels = state.labels - if labels is None or len(labels) != len(state.input_ids): - expected = 0 if labels is None else len(labels) - raise ValueError(f"labels length mismatch: {expected} vs {len(state.input_ids)}") - if state.logprobs is not None and len(state.logprobs) != len(state.input_ids): - raise ValueError(f"logprobs length mismatch: {len(state.logprobs)} vs {len(state.input_ids)}") - except Exception as exc: - self.logger.error(f"Canonicalize train fields failed for rollout_id={state.rollout_id}: {exc}") - state.status = Status.FAILED - state.error_msg = f"canonicalize_train_fields failed: {exc}" + self._write_train_meta(state) return group async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 25a4809d60..ed1a9d0ccd 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -252,38 +252,13 @@ async def generate_one(state: RolloutState) -> list[RolloutState]: return self.canonicalize_train_fields(samples) def canonicalize_train_fields(self, group: list[RolloutState]) -> list[RolloutState]: - """Validate the canonical training fields of one sandbox agentic - rollout group. - - Each segment state carries ``input_ids``/``labels``/``logprobs`` exported from - the trace store as one unshifted full sequence, so there is nothing left to - assemble. Canonicalization only re-checks the unified convention (``labels`` - present and of equal length, ``logprobs`` either ``None`` or of equal length) - per segment state and marks a violating sample ``Status.FAILED`` instead of - raising into the producer. States without ``input_ids`` (eval-mode states) carry - no training fields and are skipped. - - Args: - group (list[RolloutState]): One flattened group of segment rollout states - after generation, judging, and filtering. - - Returns: - list[RolloutState]: The same group with validated training fields. + """Preserve trace-exported unshifted fields and write train metadata. + + Eval states without training fields are left untouched. The trainer checks full training fields before + consuming a group. """ for state in group: - if state.status != Status.COMPLETED or state.input_ids is None: - continue - try: - labels = state.labels - if labels is None or len(labels) != len(state.input_ids): - expected = 0 if labels is None else len(labels) - raise ValueError(f"labels length mismatch: {expected} vs {len(state.input_ids)}") - if state.logprobs is not None and len(state.logprobs) != len(state.input_ids): - raise ValueError(f"logprobs length mismatch: {len(state.logprobs)} vs {len(state.input_ids)}") - except Exception as exc: - self.logger.error(f"Canonicalize train fields failed for rollout_id={state.rollout_id}: {exc}") - state.status = Status.FAILED - state.error_msg = f"canonicalize_train_fields failed: {exc}" + self._write_train_meta(state) return group # NOTE: A single sandbox session may yield multiple trainable segments, so this returns a list diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 810f972cff..407ffee44e 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -17,6 +17,7 @@ Status, calculate_group_effective_response_masks, get_group_status, + validate_rollout_group, ) from xtuner.v1.rl.agent_loop import AgentLoopSpec from xtuner.v1.rl.replay_buffer import ReplayBuffer @@ -151,19 +152,23 @@ async def generate_group( enable_partial_rollout=enable_partial_rollout, ) elapsed = time.perf_counter() - start - for item in result: - extra_fields = getattr(item, "extra_fields", None) - if extra_fields is None: - extra_fields = {} - setattr(item, "extra_fields", extra_fields) - extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed + if isinstance(result, list): + for item in result: + if isinstance(item, RolloutState): + item.extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed return result async def put_generated_group(self, group: list[RolloutState]) -> bool: - produced_tokens = sum(len(item.response_ids or []) - len(item.response_model_steps or []) for item in group) + produced_tokens = sum( + max(0, len(item.response_ids or []) - len(item.response_model_steps or [])) for item in group + ) initial_status = get_group_status(group) - if initial_status in (Status.COMPLETED, Status.FILTERED): + allowed_statuses = {Status.COMPLETED, Status.FAILED, Status.FILTERED, Status.ABORTED, Status.EXPIRED} + if all(item.status in allowed_statuses for item in group) and initial_status in ( + Status.COMPLETED, + Status.FILTERED, + ): rewards_sum = 0.0 rewards_count = 0 for item in group: @@ -174,6 +179,15 @@ async def put_generated_group(self, group: list[RolloutState]) -> bool: rewards_count += 1 self.progress.add_raw_rewards(self.task_name, rewards_sum, rewards_count) + try: + validate_rollout_group(group, task_name=self.task_name) + except ValueError as exc: + logger.warning(f"Discard invalid generated group for task={self.task_name!r}: {exc}") + self.progress.add_produced(self.task_name, samples=len(group), tokens=produced_tokens) + self.progress.add_discarded(self.task_name, Status.FAILED, samples=len(group)) + await release_and_discard_rollout_groups([group]) + return False + if initial_status in (Status.FAILED, Status.FILTERED): # 失败样本和业务过滤样本都不进入 replay buffer。 self.progress.add_produced(self.task_name, samples=len(group), tokens=produced_tokens) @@ -587,7 +601,14 @@ async def take_train_batch( current_train_step: int, pause_time_s: float = 0.0, ) -> ProduceBatchResult: + """Take rollout results and apply the existing per-task token-staleness + masks.""" batch_by_task, consumed_counts = await replay_buffer.take_batch(task_batch_sizes) + # Use the existing task field to retain attribution through sorting/debug saves. + for task_name, groups in batch_by_task.items(): + for group in groups: + for state in group: + state.task_name = task_name for task in task_runners: if task.token_stale_threshold is None: @@ -604,6 +625,7 @@ async def take_train_batch( labels = cast(list[int], item.labels) offset = len(labels) - len(mask) labels[offset:] = [label if mask_value else -100 for label, mask_value in zip(labels[offset:], mask)] + item.extra_fields["supervised_tokens"] = sum(label != -100 for label in labels[1:]) if hasattr(progress, "mark_consumed"): progress.mark_consumed(consumed_counts) diff --git a/xtuner/v1/rl/loss/base_loss.py b/xtuner/v1/rl/loss/base_loss.py index 7805b967f1..870764b5f6 100644 --- a/xtuner/v1/rl/loss/base_loss.py +++ b/xtuner/v1/rl/loss/base_loss.py @@ -102,6 +102,7 @@ def build( self, data: dict, sp_mesh: DeviceMesh | None = None, + device: str | torch.device = DEVICE, ) -> "BaseRLLossContext | None": """Build RL loss context from data dict. @@ -114,6 +115,8 @@ def build( - rollout_is_weights (torch.Tensor | None): Importance sampling weights - ref_logprobs (torch.Tensor | None): Reference model log probabilities sp_mesh (DeviceMesh | None): Sequence parallel device mesh + device (str | torch.device): Target device of the built loss kwargs. Defaults to the + accelerator device; pass ``"cpu"`` to build the context before any device transfer. Returns: BaseRLLossContext | None: The built loss context, or None if required fields are missing @@ -138,7 +141,7 @@ def build( rollout_logprobs=rollout_logprobs, is_weights=rollout_is_weights, ref_logprobs=ref_logprobs, - ).to(DEVICE) + ).to(device) if sp_mesh is not None and sp_mesh.size() > 1: loss_kwargs = loss_kwargs.sp_split(sp_mesh) diff --git a/xtuner/v1/rl/loss/distillation_loss.py b/xtuner/v1/rl/loss/distillation_loss.py index 4f8292f687..084d671dbb 100644 --- a/xtuner/v1/rl/loss/distillation_loss.py +++ b/xtuner/v1/rl/loss/distillation_loss.py @@ -117,6 +117,7 @@ def build( self, data: dict, sp_mesh: DeviceMesh | None = None, + device: str | torch.device = DEVICE, ) -> "DistillationLossContext | None": if "shifted_labels" not in data or "advantages" not in data: return None @@ -136,7 +137,7 @@ def build( # local slice to the accelerator. if sp_mesh is not None and sp_mesh.size() > 1: loss_kwargs = loss_kwargs.sp_split(sp_mesh) - loss_kwargs = loss_kwargs.to(DEVICE) + loss_kwargs = loss_kwargs.to(device) return self.loss_ctx_cls(self, loss_kwargs) diff --git a/xtuner/v1/rl/trainer/__init__.py b/xtuner/v1/rl/trainer/__init__.py index 79a02a0313..66319ee9ae 100644 --- a/xtuner/v1/rl/trainer/__init__.py +++ b/xtuner/v1/rl/trainer/__init__.py @@ -5,12 +5,12 @@ compute_rollout_importance_weights, merge_rollout_is_metrics, ) -from .controller import ColateItem, TrainingController -from .worker import TrainingWorker, WorkerConfig, WorkerInputItem, WorkerLogItem, WorkerTrainLogItem +from .controller import TrainingController +from .worker import TrainBatchAttr, TrainingWorker, WorkerConfig, WorkerLogItem, WorkerTrainLogItem __all__ = [ - "ColateItem", + "TrainBatchAttr", "TrainingController", "RolloutImportanceSampling", "compute_rollout_importance_weights", @@ -18,7 +18,6 @@ "compute_mismatch_metrics", "merge_rollout_is_metrics", "WorkerConfig", - "WorkerInputItem", "WorkerTrainLogItem", "WorkerLogItem", "TrainingWorker", diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index 2ab1f80b89..e536d2ea21 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -1,24 +1,20 @@ from __future__ import annotations -import math import os -import random -from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast +import time +from typing import TYPE_CHECKING, Any, Literal, Sequence, TypedDict, cast -import numpy as np import ray import torch -from typing_extensions import NotRequired -from xtuner.v1.data_proto.rl_data import RolloutState, is_valid_for_training -from xtuner.v1.data_proto.sequence_context import SequenceContext -from xtuner.v1.model.compose.base import BaseComposeConfig +from xtuner.v1.data_proto.rl_data import RolloutState from xtuner.v1.rl.distillation import DistillationTrainerAdapter from xtuner.v1.rl.utils import free_object_refs from xtuner.v1.train.trainer import LoadCheckpointConfig -from xtuner.v1.utils import XTUNER_DETERMINISTIC, get_logger +from xtuner.v1.utils import get_logger -from .worker import TrainingWorker, WorkerLogItem +from .pack import DPRankPackIndices, RLDataPacker +from .worker import TrainBatchAttr, TrainingWorker, WorkerConfig, WorkerLogItem if TYPE_CHECKING: @@ -27,85 +23,46 @@ TRAIN_RAY_GET_TIMEOUT = os.getenv("XTUNER_TRAIN_RAY_GET_TIMEOUT", 5 * 3600) # default 5 hours -class ColateItem(TypedDict): - seq_ctx: SequenceContext - shifted_labels: torch.Tensor - advantage: list[float] - rollout_logprobs: torch.Tensor | None - teacher_logprobs: NotRequired[torch.Tensor | None] - target_token_ids: NotRequired[torch.Tensor | None] - teacher_indices: NotRequired[torch.Tensor | None] - - -class PackedBatch(TypedDict): - """Output of ``TrainingController._packing``: samples concatenated into one - sequence.""" - - seq_ctx: SequenceContext - shifted_labels: torch.Tensor - advantages: torch.Tensor - rollout_logprobs: torch.Tensor | None - teacher_logprobs: torch.Tensor | None - target_token_ids: torch.Tensor | None - teacher_indices: torch.Tensor | None - - -class _SampleObservation(TypedDict): - """Per-sample metrics observed while converting one rollout state.""" - - advantage: float - supervised_positions: int - prompt_len: int - tool_turns: int | None - training_tokens: int - - -def get_train_seq_ctx( - input_ids: torch.LongTensor, - position_ids: np.ndarray | None = None, - multimodal_train_info: dict | None = None, -) -> SequenceContext: - """Build a CPU ``SequenceContext`` for one training sample. - - Args: - input_ids (torch.LongTensor): Model input tokens with shape ``(1, seq_len)``. - position_ids (np.ndarray | None): Optional position ids. A 3D array triggers - the VLM MRoPE layout; a prompt-only position segment is extended to cover - the whole input. - multimodal_train_info (dict | None): Optional multimodal payload with - ``pixel_values``, ``image_grid_thw`` and ``num_img_tokens``. - - Returns: - SequenceContext: The CPU sequence context of the sample. - """ - seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cpu") - position_ids = _to_cpu_tensor(position_ids, dtype=torch.long) - if position_ids is not None and len(position_ids.shape) == 3: - # Match get_rope_index_3: response text continues from a single global - # max(T, H, W), not per-axis maxima. Per-axis max diverges when the - # prompt ends on image tokens (T≈0 while H/W are large). The extension - # length is derived from the position deficit so samples whose positions - # already cover the whole input (full-sequence agentic samples) need no - # extra shape information. - num_extend = input_ids.size(-1) - position_ids.size(-1) - if num_extend > 0: - max_value = position_ids.amax() - response_position_ids = ( - (torch.arange(1, num_extend + 1, device=position_ids.device, dtype=position_ids.dtype) + max_value) - .view(1, 1, -1) - .expand(3, 1, -1) - ) - position_ids = torch.cat([position_ids, response_position_ids], dim=-1) - seq_ctx.position_ids = position_ids # type: ignore[assignment] - assert position_ids.size(-1) == input_ids.size(-1) +class _DPRankDispatch(TypedDict): + """Per-DP-rank payload of one training batch, before the Ray upload in + ``fit``.""" + + plan: DPRankPackIndices + advantages: list[float | list[float]] + rollout_items: list[RolloutState] + + +class _DPRankPlan(TypedDict): + """Global sample mapping and local pack indices for one DP rank.""" + + global_indices: list[int] + plan: DPRankPackIndices + + +class _PackPlan(TypedDict): + """Index-only layout and packing metrics, independent of training + targets.""" + + dp_plans: dict[int, _DPRankPlan] + plan_log: dict[str, float] + - if multimodal_train_info: - seq_ctx.pixel_values = multimodal_train_info.get("pixel_values") - seq_ctx.image_grid_thw = _to_cpu_tensor(multimodal_train_info.get("image_grid_thw"), dtype=torch.long) - num_img_tokens = multimodal_train_info.get("num_img_tokens") - if num_img_tokens is not None: - seq_ctx.num_img_tokens = [num_img_tokens] - return seq_ctx +class _GroupRewardInfo(TypedDict): + """Reward clusters and sample mapping for one original rollout group.""" + + rewards: list[float] + representatives: list[RolloutState] + cluster_indices: list[int | None] + + +class _PreparedBatch(TypedDict): + """Outputs of ``_prepare_rollout_items`` consumed by the fit phases.""" + + rollout_items: list[RolloutState] + reward_groups: list[_GroupRewardInfo] + batch_attr: TrainBatchAttr + cluster_rewards: list[float] + distillation_reward_observations: list[tuple[RolloutState, float]] def _stat_tensor(values: list[float]) -> torch.Tensor: @@ -135,23 +92,25 @@ def _summarize_process_group_results(results: list[dict[str, Any]]) -> str: return summary -def _verify_packed_alignment(packed_batch: PackedBatch) -> None: - # Position-wise fields must stay aligned with input_ids; a mismatch would silently shift - # the element-wise policy loss or crash on shape mismatch in the loss function. - assert packed_batch["seq_ctx"].input_ids is not None - seq_len = packed_batch["seq_ctx"].input_ids.shape[1] - shifted_labels = packed_batch["shifted_labels"] - advantages = packed_batch["advantages"] - assert shifted_labels.shape[1] == seq_len, f"{shifted_labels.shape[1]} vs {seq_len}" - assert advantages.shape[1] == seq_len, f"{advantages.shape[1]} vs {seq_len}" - for field_name, optional_field in ( - ("rollout_logprobs", packed_batch["rollout_logprobs"]), - ("teacher_logprobs", packed_batch["teacher_logprobs"]), - ("target_token_ids", packed_batch["target_token_ids"]), - ("teacher_indices", packed_batch["teacher_indices"]), - ): - if optional_field is not None: - assert optional_field.shape[1] == seq_len, f"{field_name}: {optional_field.shape[1]} vs {seq_len}" +def _state_uses_3d_position_ids(state: RolloutState) -> bool: + layout = state.extra_fields.get("position_layout") + if layout is not None: + return layout == "mrope_3d" + return state.position_ids is not None and state.position_ids.ndim == 3 + + +def _state_has_rollout_logprobs(state: RolloutState) -> bool: + flag = state.extra_fields.get("has_rollout_logprobs") + if flag is None: + return state.logprobs is not None + return bool(flag) + + +def _state_has_routed_experts(state: RolloutState) -> bool: + flag = state.extra_fields.get("has_routed_experts") + if flag is None: + return state.routed_experts is not None + return bool(flag) class TrainingController: @@ -167,271 +126,6 @@ def __init__( self.task_adv_weight = self.distillation.task_adv_weight self.logger = get_logger() - # TODO(hha): 这个逻辑不够通用，应该复用 sft 函数，从而支持 expand soft pack - def _get_pack_infos(self, dataset, num_tokens, target, random=None): - inds = list(range(len(dataset))) - if random is not None: - random.shuffle(inds) - - item_buffer = [] - length_buffer = [] - longest = 0 - - pack_infos = [] - for shfl_i in inds: - if num_tokens[shfl_i] + sum(length_buffer) <= target: - item_buffer.append(shfl_i) - length_buffer.append(num_tokens[shfl_i]) - longest = max(longest, num_tokens[shfl_i]) - else: - if len(item_buffer) > 0: - info = { - "indices": item_buffer, - "longest": int(longest), - } - pack_infos.append(info) - - item_buffer = [shfl_i] - length_buffer = [num_tokens[shfl_i]] - longest = num_tokens[shfl_i] - - if len(item_buffer) > 0: - info = { - "indices": item_buffer, - "longest": int(longest), - } - - pack_infos.append(info) - - return pack_infos - - # TODO(hha): 这个逻辑不够通用，和模型绑定了 - def _packing(self, data_batches, pack_max_length, language_cfg): - pack_infos = self._get_pack_infos( - data_batches, - [data["seq_ctx"].input_ids.numel() for data in data_batches], - pack_max_length, - ) - packed_data_batches = [] - - is_qwen3_vl = any(data["seq_ctx"].position_ids.ndim == 3 for data in data_batches) - - has_rollout_routed_experts = False - if data_batches[0]["seq_ctx"].rollout_routed_experts is not None: - assert language_cfg is not None - has_rollout_routed_experts = True - n_routed_experts = language_cfg.n_routed_experts - - for pack_info in pack_infos: - indices = pack_info["indices"] - total_len = sum([data_batches[i]["seq_ctx"].input_ids.shape[1] for i in indices]) - pad_len = pack_max_length - total_len - seq_ctx_list = [data_batches[i]["seq_ctx"] for i in indices] - label_list = [data_batches[i]["shifted_labels"] for i in indices] - advantage_list = [data_batches[i]["advantage"] for i in indices] - - rollout_logprobs_list = None - if "rollout_logprobs" in data_batches[0] and data_batches[0]["rollout_logprobs"] is not None: - rollout_logprobs_list = [data_batches[i]["rollout_logprobs"] for i in indices] - - teacher_logprobs_list = None - if "teacher_logprobs" in data_batches[0] and data_batches[0]["teacher_logprobs"] is not None: - teacher_logprobs_list = [data_batches[i]["teacher_logprobs"] for i in indices] - - target_token_ids_list = None - if "target_token_ids" in data_batches[0] and data_batches[0]["target_token_ids"] is not None: - target_token_ids_list = [data_batches[i]["target_token_ids"] for i in indices] - - teacher_indices_list = None - if "teacher_indices" in data_batches[0] and data_batches[0]["teacher_indices"] is not None: - teacher_indices_list = [data_batches[i]["teacher_indices"] for i in indices] - - if pad_len > 0: - # Reduce the attn calculation time by using multiple short sequence packs - pad_tokens = tuple( - torch.zeros(1, 1024, dtype=data_batches[0]["seq_ctx"].input_ids.dtype, device="cpu") - for _ in range(pad_len // 1024) - ) - if pad_len % 1024 > 0: - pad_tokens = pad_tokens + ( - torch.zeros(1, pad_len % 1024, dtype=data_batches[0]["seq_ctx"].input_ids.dtype, device="cpu"), - ) - pad_seq_ctx = SequenceContext.from_input_ids(pad_tokens, device="cpu") - pad_seq_ctx.num_padding = pad_len - pad_labels = torch.full( - (1, pad_len), - -100, - dtype=data_batches[0]["shifted_labels"].dtype, - device=data_batches[0]["shifted_labels"].device, - ) - pad_advantages = [-100] * pad_len - if is_qwen3_vl: - _position_ids_list = [] - for pad_token in pad_tokens: - _position_ids = torch.arange(pad_token.size(-1)).view(1, 1, -1).expand(3, 1, -1) - _position_ids_list.append(_position_ids) - pad_seq_ctx.position_ids = torch.cat(_position_ids_list, dim=-1) - - if has_rollout_routed_experts: - pad_rand_index = torch.randint(low=0, high=n_routed_experts, size=(pad_len, 1, 1)) - pad_seq_ctx.rollout_routed_experts = pad_rand_index - - seq_ctx_list.append(pad_seq_ctx) - label_list.append(pad_labels) - advantage_list.append(pad_advantages) - if rollout_logprobs_list is not None: - pad_rollout_logprobs = torch.zeros( - 1, - pad_len, - dtype=data_batches[0]["rollout_logprobs"].dtype, - device=data_batches[0]["shifted_labels"].device, - ) - rollout_logprobs_list.append(pad_rollout_logprobs) - if teacher_logprobs_list is not None: - teacher_logprobs_shape = data_batches[0]["teacher_logprobs"].shape[2:] - pad_teacher_logprobs = torch.zeros( - (1, pad_len, *teacher_logprobs_shape), - dtype=data_batches[0]["teacher_logprobs"].dtype, - device=data_batches[0]["shifted_labels"].device, - ) - teacher_logprobs_list.append(pad_teacher_logprobs) - if target_token_ids_list is not None: - target_token_ids_shape = data_batches[0]["target_token_ids"].shape[2:] - pad_target_token_ids = torch.zeros( - (1, pad_len, *target_token_ids_shape), - dtype=data_batches[0]["target_token_ids"].dtype, - device=data_batches[0]["shifted_labels"].device, - ) - target_token_ids_list.append(pad_target_token_ids) - if teacher_indices_list is not None: - pad_teacher_indices = torch.full( - (1, pad_len), - -1, - dtype=data_batches[0]["teacher_indices"].dtype, - device=data_batches[0]["teacher_indices"].device, - ) - teacher_indices_list.append(pad_teacher_indices) - - seq_ctx = SequenceContext.cat(seq_ctx_list) - shifted_labels = torch.cat(label_list, dim=1) # (1, max_len) - advantage_flat = [item for sublist in advantage_list for item in sublist] - advantages = torch.tensor(advantage_flat, dtype=torch.float32).unsqueeze(0) - - rollout_logprobs = None - if rollout_logprobs_list is not None: - rollout_logprobs = torch.cat(rollout_logprobs_list, dim=1) # (1, max_len) - - teacher_logprobs = None - if teacher_logprobs_list is not None: - teacher_logprobs = torch.cat(teacher_logprobs_list, dim=1) # (1, max_len) - - target_token_ids = None - if target_token_ids_list is not None: - target_token_ids = torch.cat(target_token_ids_list, dim=1) - - teacher_indices = None - if teacher_indices_list is not None: - teacher_indices = torch.cat(teacher_indices_list, dim=1) # (1, max_len) - - packed_batch = { - "seq_ctx": seq_ctx, - "shifted_labels": shifted_labels, - "advantages": advantages, - "rollout_logprobs": rollout_logprobs, - "teacher_logprobs": teacher_logprobs, - "target_token_ids": target_token_ids, - "teacher_indices": teacher_indices, - } - _verify_packed_alignment(packed_batch) - packed_data_batches.append(packed_batch) - return packed_data_batches - - def _grouped_by_max_length(self, packed_data_batches): - # sort 过后可能第一个 batch 会有很多 pad tokens，因为最后一个 pack 可能只有少量真实数据。 - # 比如组成了 16 个 pack，第 16 个 pack 可能只有几条真实数据，剩下的都是 pad tokens。 - # 排序后这条 pack 会被放在最前面，导致 rank0 的第一个 step 消耗的有效 token 数往往少于其他 rank，是正常现象。 - return sorted(packed_data_batches, key=lambda x: x["seq_ctx"].max_length_q, reverse=True) - - def _rollout_groups_to_colate_items( - self, - rollout_groups: list[list[RolloutState]], - pack_max_length: int, - raw_rewards_sum: float = 0.0, - raw_rewards_count: int = 0, - ) -> tuple[list[ColateItem], dict[str, float]]: - """Convert rollout groups into packable training items and step - statistics. - - Samples arrive canonicalized from the agent loop (``input_ids``/``labels``/ - ``logprobs`` share one full-sequence length) with token staleness already baked - into ``labels``. The conversion runs in phases: (1) cluster group rewards by - session, (2) estimate group-level advantages, (3) shift and tensorize each - rollout state into a ``ColateItem``, and (4) summarize the ``data_info`` - metrics. - - Args: - rollout_groups (list[list[RolloutState]]): Rollout groups selected for training. - pack_max_length (int): Maximum token length per sample accepted for packing. - raw_rewards_sum (float): Producer-side raw reward sum used by ``raw_rewards/mean``. - raw_rewards_count (int): Producer-side raw reward count used by ``raw_rewards/mean``. - - Returns: - tuple[list[ColateItem], dict[str, float]]: Packable training items and the step - statistics dictionary. - """ - has_3d_position = any( - state.position_ids is not None and state.position_ids.ndim == 3 - for group in rollout_groups - for state in group - ) - - cluster_rewards_list: list[float] = [] - distillation_reward_observations: list[tuple[RolloutState, float]] = [] - observations: list[_SampleObservation] = [] - training_samples = 0 - - data_batches: list[ColateItem] = [] - for group in rollout_groups: - if not is_valid_for_training(group, self.logger): - self.logger.error(f"Skip one data group {group} due to rollout failed or empty response.") - continue - training_samples += len(group) - - # Phase 1: session reward clustering. Collect rewards independently from - # task-advantage computation: pure OPD may omit rewards entirely; when rewards - # are present they remain useful observability signals. - cluster_rewards, session_representatives, cluster_indices = self._cluster_group_rewards(group) - cluster_rewards_list.extend(cluster_rewards) - distillation_reward_observations.extend(zip(session_representatives, cluster_rewards)) - - # Phase 2: group-level advantage estimation. - sample_advantages = self._compute_group_advantages( - cluster_rewards, session_representatives, cluster_indices - ) - - # Phase 3: per-sample shift, tensorization and seq_ctx construction. - for i, state in enumerate(group): - data_batch, observation = self._convert_rollout_state(state, sample_advantages[i], has_3d_position) - assert observation["training_tokens"] <= pack_max_length, ( - f"{observation['training_tokens']} vs {pack_max_length}" - ) - data_batches.append(data_batch) - observations.append(observation) - - if not XTUNER_DETERMINISTIC: - random.shuffle(data_batches) - - # Phase 4: step statistics summarization. - info_dict = self._summarize_data_info( - observations, - cluster_rewards_list, - distillation_reward_observations, - training_samples, - raw_rewards_sum=raw_rewards_sum, - raw_rewards_count=raw_rewards_count, - ) - return data_batches, info_dict - def _cluster_group_rewards( self, group: list[RolloutState] ) -> tuple[list[float], list[RolloutState], list[int | None]]: @@ -484,106 +178,36 @@ def _compute_group_advantages( return [float(cluster_advantages[index].item()) if index is not None else 0.0 for index in cluster_indices] return [0.0] * len(cluster_indices) - def _convert_rollout_state( - self, - state: RolloutState, - advantage_val: float, - has_3d_position: bool, - ) -> tuple[ColateItem, _SampleObservation]: - """Shift, tensorize and build the sequence context of one rollout - state. - - Args: - state (RolloutState): Canonicalized rollout state with full-sequence ``input_ids``/``labels``. - advantage_val (float): Group-level advantage broadcast to every supervised position. - has_3d_position (bool): Whether any sample in the batch carries 3D MRoPE position ids. - - Returns: - tuple[ColateItem, _SampleObservation]: The packable training item and the per-sample - metrics observed during conversion. - """ - input_ids = state.input_ids - labels = state.labels - assert input_ids is not None and labels is not None and len(input_ids) == len(labels), ( - f"Rollout state is not canonicalized for training: {state}" - ) - input_len = len(input_ids) - shifted_labels = labels[1:] - - input_ids_t = cast(torch.LongTensor, torch.tensor(input_ids[:-1], dtype=torch.int64).unsqueeze(0)) - shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) - - rollout_logprobs: torch.Tensor | None = None - if state.logprobs is not None: - assert len(state.logprobs) == input_len, f"{len(state.logprobs)} vs {input_len}, data: {state}" - rollout_logprobs = torch.tensor(state.logprobs[1:], dtype=torch.float32).unsqueeze(0) - assert rollout_logprobs.size() == shifted_labels_t.size(), ( - f"{rollout_logprobs.size()} vs {shifted_labels_t.size()}" + def _prepare_advantages(self, prepared: _PreparedBatch, pack_plan: _PackPlan) -> list[float] | list[list[float]]: + """Prepare targets after packing indices are fixed, without changing + sample order or placing targets inside the plan.""" + advantages: list[float] = [] + for group in prepared["reward_groups"]: + advantages.extend( + self._compute_group_advantages(group["rewards"], group["representatives"], group["cluster_indices"]) ) + return advantages - # Keep the advantage layout aligned with input_ids (response excludes EOS). - # Prompt positions predict prompt tokens whose shifted labels are -100, so their - # advantage is 0; the last entry of response_ids (EOS) is only a label, never an input. - actual_advantages = [0.0 if label == -100 else advantage_val for label in shifted_labels] - - # prompt_len 统计原始输入 prompt 长度（VLM 的 prompt 在 train_prompt_ids）； - # response_len 统计 LLM 生成的 token 数（labels 非 -100 的监督位）。环境/工具 - # 插入的语义洞不计入两者；prompt_ids 缺失时 prompt_len 退回洞计数。 - prompt_ids = state.extra_fields.get("train_prompt_ids") or state.prompt_ids - prompt_len = len(prompt_ids) if prompt_ids else sum(label == -100 for label in shifted_labels) - response_len = sum(label != -100 for label in shifted_labels) - - position_ids = state.position_ids - if has_3d_position and (position_ids is None or position_ids.ndim != 3): - seq_len = input_ids_t.size(-1) - text_position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1) - # Mixed batches use three-axis MRoPE position IDs. Repeat the normal text - # positions on every axis so text samples have the same rank as the VLM - # samples when their sequence contexts are packed together. - position_ids = np.broadcast_to(text_position_ids, (3, 1, seq_len)).copy() - multimodal_train_info = cast(dict | None, state.mm_info) - seq_ctx = get_train_seq_ctx(input_ids_t, position_ids, multimodal_train_info) - seq_ctx.rollout_routed_experts = state.routed_experts # type: ignore[assignment] # n,layer*expert - - teacher_fields = self.distillation.rollout_teacher_targets(state, shifted_labels=shifted_labels) - - data_dict: ColateItem = { - "seq_ctx": seq_ctx, - "shifted_labels": shifted_labels_t, - "advantage": actual_advantages, - "rollout_logprobs": rollout_logprobs, - } - data_dict.update(cast(ColateItem, teacher_fields)) - - # 有可能有重复，但是没有其他更好办法 - turns = state.extra_fields.get("agent_tool_turns") - observation: _SampleObservation = { - "advantage": advantage_val, - "supervised_positions": response_len, - "prompt_len": prompt_len, - "tool_turns": turns if isinstance(turns, int) else None, - "training_tokens": input_len - 1, - } - return data_dict, observation - - def _summarize_data_info( + def _build_trainer_log_info( self, - observations: list[_SampleObservation], - cluster_rewards_list: list[float], - distillation_reward_observations: list[tuple[RolloutState, float]], - training_samples: int, - raw_rewards_sum: float, - raw_rewards_count: int, + prepared: _PreparedBatch, + pack_plan: _PackPlan, + log_infos: list[WorkerLogItem], + advantages: Sequence[float | list[float]], + raw_rewards_sum: float = 0.0, + raw_rewards_count: int = 0, ) -> dict[str, float]: - """Reduce phase-1/phase-3 observations into the step ``data_info`` - metrics. + """Assemble the complete ``data_info`` step log from one fit's phases. + + All log-related logic of ``fit`` lives here: data statistics of the prepared + batch, plan-derived packing metrics and the maxima of the per-worker + conversion/materialize timings. Args: - observations (list[_SampleObservation]): Per-sample metrics collected in phase 3. - cluster_rewards_list (list[float]): Session-clustered rewards collected in phase 1. - distillation_reward_observations (list[tuple[RolloutState, float]]): Session - representative states paired with their cluster reward. - training_samples (int): Number of states in the valid groups selected for training. + prepared (_PreparedBatch): Output of ``_prepare_rollout_items``. + pack_plan (_PackPlan): Output of ``_build_pack_plan``. + log_infos (list[WorkerLogItem]): Per-worker training logs of the batch. + advantages: Final scalar or token advantages in global sample order. raw_rewards_sum (float): Producer-side raw reward sum used by ``raw_rewards/mean``. raw_rewards_count (int): Producer-side raw reward count used by ``raw_rewards/mean``. @@ -595,22 +219,31 @@ def _summarize_data_info( response_len_list: list[float] = [] tool_turns_list: list[int] = [] training_tokens = 0 - for observation in observations: + for state, advantage in zip(prepared["rollout_items"], advantages): # Advantage metrics count one entry per supervised position, matching the - # per-position advantage layout built in phase 3. - advantages_list.extend([observation["advantage"]] * observation["supervised_positions"]) - prompt_len_list.append(observation["prompt_len"]) - response_len_list.append(observation["supervised_positions"]) - if observation["tool_turns"] is not None: - tool_turns_list.append(observation["tool_turns"]) - training_tokens += observation["training_tokens"] + # per-position advantage layout built by the worker conversion. + supervised_positions = int(state.extra_fields["supervised_tokens"]) + if isinstance(advantage, list): + labels = state.labels + assert labels is not None + advantages_list.extend(value for value, label in zip(advantage, labels[1:]) if label != -100) + else: + advantages_list.extend([advantage] * supervised_positions) + prompt_len_list.append(int(state.extra_fields["train_prompt_length"])) + response_len_list.append(supervised_positions) + turns = state.extra_fields.get("agent_tool_turns") + if isinstance(turns, int): + tool_turns_list.append(turns) + assert state.num_tokens is not None + training_tokens += state.num_tokens # rewards/* report the per-session reward distribution; batch_size/training_samples - # count the valid rollout segments included in training. - rewards_t = _stat_tensor(cluster_rewards_list) + # count the rollout segments included in training. + rewards_t = _stat_tensor(prepared["cluster_rewards"]) advantages_t = _stat_tensor(advantages_list) prompt_len_t = _stat_tensor(prompt_len_list) response_len_t = _stat_tensor(response_len_list) + training_samples = len(prepared["rollout_items"]) raw_rewards_mean = raw_rewards_sum / raw_rewards_count if raw_rewards_count > 0 else rewards_t.mean().item() info_dict: dict[str, float] = { @@ -632,14 +265,177 @@ def _summarize_data_info( "prompt_len/min": prompt_len_t.min().item(), "prompt_len/max": prompt_len_t.max().item(), } - info_dict.update(self.distillation.reward_scalars(distillation_reward_observations)) + info_dict.update(self.distillation.reward_scalars(prepared["distillation_reward_observations"])) if tool_turns_list: tool_turns_t = torch.tensor(tool_turns_list, dtype=torch.float32) info_dict["tool_turns/mean"] = tool_turns_t.mean().item() info_dict["tool_turns/min"] = float(tool_turns_t.min().item()) info_dict["tool_turns/max"] = float(tool_turns_t.max().item()) + info_dict.update(pack_plan["plan_log"]) + # WorkerLogItem is a TypedDict; dynamic-key reads below need a plain dict view. + worker_logs = cast(list[dict[str, Any]], log_infos) + for key, target_key in ( + ("packing_conversion_time_s", "packing/conversion_time_s_max"), + ("packing_pack_time_s", "packing/pack_time_s_max"), + ): + worker_values = [log_info[key] for log_info in worker_logs if key in log_info] + if worker_values: + info_dict[target_key] = max(worker_values) return info_dict + def _prepare_rollout_items(self, rollout_groups: list[list[RolloutState]], rollout_idx: int) -> _PreparedBatch: + """Flatten samples and retain reward/group mappings without estimating + advantages or building training tensors. + + Groups arrive pre-filtered: the replay buffer selects by ``Status.COMPLETED`` + and sample alignment is guaranteed by the trainer entry check, so no group-level validation happens here; missing train meta + fields still fail fast. + + Args: + rollout_groups (list[list[RolloutState]]): Rollout groups selected for training. + rollout_idx (int): Global rollout (train step) id written into ``batch_attr``. + + Returns: + _PreparedBatch: Flat rollout items in training order with original + reward/group mappings, reward clusters for statistics, and the batch-level + ``TrainBatchAttr`` derived in the same single pass over the states. + """ + cluster_rewards_list: list[float] = [] + distillation_reward_observations: list[tuple[RolloutState, float]] = [] + rollout_items: list[RolloutState] = [] + reward_groups: list[_GroupRewardInfo] = [] + use_3d_position_ids = False + has_rollout_logprobs = False + has_routed_experts = False + for group in rollout_groups: + # Phase 1: session reward clustering keeps one reward per agentic session so + # rewards/* and advantages are not amplified by per-session segment counts. + cluster_rewards, session_representatives, cluster_indices = self._cluster_group_rewards(group) + cluster_rewards_list.extend(cluster_rewards) + distillation_reward_observations.extend(zip(session_representatives, cluster_rewards)) + + reward_groups.append( + { + "rewards": cluster_rewards, + "representatives": session_representatives, + "cluster_indices": cluster_indices, + } + ) + for state in group: + rollout_items.append(state) + use_3d_position_ids = use_3d_position_ids or _state_uses_3d_position_ids(state) + has_rollout_logprobs = has_rollout_logprobs or _state_has_rollout_logprobs(state) + has_routed_experts = has_routed_experts or _state_has_routed_experts(state) + # 只读 AgentLoop 规范化时写入的轻量 meta 字段（AgentLoop canonicalization），不触碰 + # input_ids/labels 等重载荷；meta 缺失说明样本未过 canonicalize，直接 fail fast。 + if ( + state.extra_fields.get("supervised_tokens") is None + or state.extra_fields.get("train_prompt_length") is None + or state.num_tokens is None + ): + raise ValueError( + "Rollout state is missing the train meta fields written by AgentLoop canonicalization: " + f"rollout_id={state.rollout_id}" + ) + + return { + "rollout_items": rollout_items, + "reward_groups": reward_groups, + "batch_attr": { + "rollout_idx": rollout_idx, + "use_3d_position_ids": use_3d_position_ids, + "pack_loss_keys": ["rollout_logprobs"] if has_rollout_logprobs else [], + "has_routed_experts": has_routed_experts, + }, + "cluster_rewards": cluster_rewards_list, + "distillation_reward_observations": distillation_reward_observations, + } + + def _build_pack_plan( + self, + rollout_items: list[RolloutState], + pack_max_length: int, + worker_cfg: WorkerConfig, + data_replicate_size: int, + ) -> _PackPlan: + """Build an index-only per-DP-rank plan without training targets. + + Args: + rollout_items (list[RolloutState]): Flat rollout states in training order. + pack_max_length (int): Maximum token budget of one pack. + worker_cfg (WorkerConfig): Training worker config providing packing options. + data_replicate_size (int): Number of workers holding one data replica. + + Returns: + _PackPlan: Global sample indices and local pack indices per DP rank, + plus packing metrics. Payload assembly happens separately, once the + algorithm has produced its training targets. + """ + plan_begin = time.perf_counter() + lengths = [] + real_tokens = 0 + for state in rollout_items: + assert state.num_tokens is not None, ( + f"Rollout state is missing num_tokens written by AgentLoop canonicalization: rollout_id={state.rollout_id}" + ) + lengths.append(state.num_tokens) + real_tokens += state.num_tokens + data_packer = RLDataPacker( + pack_max_length=pack_max_length, + world_size=len(self.workers), + data_replicate_size=data_replicate_size, + optimizer_steps=worker_cfg.optimizer_steps, + pack_strategy=worker_cfg.pack_strategy, + pack_seed=worker_cfg.pack_seed, + pack_extra_buffer_size=worker_cfg.pack_extra_buffer_size, + ) + packed_plan, padding_tokens = data_packer.pack(lengths) + + dp_plans: dict[int, _DPRankPlan] = {} + for dp_rank, dp_plan in enumerate(packed_plan): + global_indices = sorted({index for step in dp_plan for pack in step for index in pack}) + global_to_local = {global_index: local_index for local_index, global_index in enumerate(global_indices)} + dp_plans[dp_rank] = { + "plan": [[[global_to_local[index] for index in pack] for pack in step] for step in dp_plan], + "global_indices": global_indices, + } + total_packs = sum(len(step) for dp_plan in packed_plan for step in dp_plan) + total_tokens = real_tokens + padding_tokens + plan_log = { + "packing/real_tokens": float(real_tokens), + "packing/padding_tokens": float(padding_tokens), + "packing/efficiency": real_tokens / total_tokens if total_tokens > 0 else 0.0, + "packing/num_packs": float(total_packs), + "packing/num_optimizer_steps": float(len(packed_plan[0])), + "packing/plan_time_s": time.perf_counter() - plan_begin, + } + return {"dp_plans": dp_plans, "plan_log": plan_log} + + def _build_dispatches( + self, + rollout_items: list[RolloutState], + advantages: Sequence[float | list[float]], + pack_plan: _PackPlan, + ) -> dict[int, _DPRankDispatch]: + """Attach states and training targets to an already fixed sample + mapping.""" + if len(rollout_items) != len(advantages): + raise ValueError(f"rollout_items and advantages lengths differ: {len(rollout_items)} vs {len(advantages)}") + for state, advantage in zip(rollout_items, advantages): + if isinstance(advantage, list) and len(advantage) != state.num_tokens: + raise ValueError( + f"Token advantage length mismatch for rollout_id={state.rollout_id}: " + f"{len(advantage)} vs num_tokens={state.num_tokens}" + ) + return { + dp_rank: { + "plan": dp_plan["plan"], + "rollout_items": [rollout_items[index] for index in dp_plan["global_indices"]], + "advantages": [advantages[index] for index in dp_plan["global_indices"]], + } + for dp_rank, dp_plan in pack_plan["dp_plans"].items() + } + def fit( self, rollout_groups: list[list[RolloutState]], @@ -648,143 +444,58 @@ def fit( raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, ) -> tuple[list[WorkerLogItem], dict[str, float]]: - data_batches, data_info = self._rollout_groups_to_colate_items( - rollout_groups, - pack_max_length, - raw_rewards_sum=raw_rewards_sum, - raw_rewards_count=raw_rewards_count, - ) - has_rollout_routed_experts = False - language_cfg = None - if data_batches[0]["seq_ctx"].rollout_routed_experts is not None: - model_cfg = ray.get(self.workers[0].get_model_cfg.remote()) # type: ignore[attr-defined] - has_rollout_routed_experts = True - language_cfg = model_cfg - if isinstance(model_cfg, BaseComposeConfig): - language_cfg = model_cfg.text_config - - packed_data_batches = self._packing(data_batches, pack_max_length, language_cfg) - # packed_data_batches = self._grouped_by_max_length(packed_data_batches) - - # TODO(hha): 这个逻辑不够通用，和模型绑定了 - is_qwen3_vl = False - if len(packed_data_batches[0]["seq_ctx"].position_ids.shape) == 3: - is_qwen3_vl = True - - # todo: support round up - num_packed_data_batches = len(packed_data_batches) + prepared = self._prepare_rollout_items(rollout_groups, rollout_idx) + rollout_items = prepared["rollout_items"] + if not rollout_items: + raise ValueError(f"Rollout {rollout_idx} has no valid training samples") + + worker_cfg = ray.get(self.workers[0].get_worker_cfg.remote()) # type: ignore[attr-defined] data_replicate_size = ray.get(self.workers[0].get_data_replicate_size.remote()) # type: ignore[attr-defined] - dp_size = len(self.workers) // data_replicate_size - pad_num = math.ceil(num_packed_data_batches / dp_size) * dp_size - num_packed_data_batches - if pad_num > 0: - # Reduce the attn calculation time by using multiple short sequence packs - assert data_batches[0]["seq_ctx"].input_ids is not None - pad_tokens = tuple( - torch.zeros(1, 1024, dtype=data_batches[0]["seq_ctx"].input_ids.dtype, device="cpu") - for _ in range(pack_max_length // 1024) - ) - if pack_max_length % 1024 > 0: - assert data_batches[0]["seq_ctx"].input_ids is not None - pad_tokens = pad_tokens + ( - torch.zeros( - 1, pack_max_length % 1024, dtype=data_batches[0]["seq_ctx"].input_ids.dtype, device="cpu" - ), - ) - pad_seq_ctx = SequenceContext.from_input_ids(pad_tokens, device="cpu") # type: ignore - pad_seq_ctx.num_padding = pack_max_length - if is_qwen3_vl: - _position_ids_list = [] - for pad_token in pad_tokens: - _position_ids = torch.arange(pad_token.size(-1)).view(1, 1, -1).expand(3, 1, -1) - _position_ids_list.append(_position_ids) - pad_seq_ctx.position_ids = torch.cat(_position_ids_list, dim=-1) # type: ignore - - pad_shifted_labels = torch.full( - (1, pack_max_length), - -100, - dtype=packed_data_batches[0]["shifted_labels"].dtype, - device="cpu", - ) - pad_advantages = torch.full( - (1, pack_max_length), - -100, - dtype=packed_data_batches[0]["advantages"].dtype, - device="cpu", - ) - if has_rollout_routed_experts: - pad_rand_index = torch.randint( - low=0, - high=1, - size=(1, 1, 1), # add dummy data, true data will be initialized in train worker.fit - ) - pad_seq_ctx.rollout_routed_experts = pad_rand_index + pack_plan = self._build_pack_plan(rollout_items, pack_max_length, worker_cfg, data_replicate_size) - pad_rollout_logprobs = None - if "rollout_logprobs" in packed_data_batches[0] and packed_data_batches[0]["rollout_logprobs"] is not None: - pad_rollout_logprobs = torch.zeros( - 1, pack_max_length, dtype=packed_data_batches[0]["rollout_logprobs"].dtype, device="cpu" - ) - pad_teacher_logprobs = None - if "teacher_logprobs" in packed_data_batches[0] and packed_data_batches[0]["teacher_logprobs"] is not None: - teacher_logprobs_shape = packed_data_batches[0]["teacher_logprobs"].shape[2:] - pad_teacher_logprobs = torch.zeros( - (1, pack_max_length, *teacher_logprobs_shape), - dtype=packed_data_batches[0]["teacher_logprobs"].dtype, - device="cpu", - ) - pad_target_token_ids = None - if "target_token_ids" in packed_data_batches[0] and packed_data_batches[0]["target_token_ids"] is not None: - target_token_ids_shape = packed_data_batches[0]["target_token_ids"].shape[2:] - pad_target_token_ids = torch.zeros( - (1, pack_max_length, *target_token_ids_shape), - dtype=packed_data_batches[0]["target_token_ids"].dtype, - device="cpu", - ) - pad_teacher_indices = None - if "teacher_indices" in packed_data_batches[0] and packed_data_batches[0]["teacher_indices"] is not None: - pad_teacher_indices = torch.full( - (1, pack_max_length), - -1, - dtype=packed_data_batches[0]["teacher_indices"].dtype, - device="cpu", - ) - pad_data = { - "seq_ctx": pad_seq_ctx, - "shifted_labels": pad_shifted_labels, - "advantages": pad_advantages, - "rollout_logprobs": pad_rollout_logprobs, - "teacher_logprobs": pad_teacher_logprobs, - "target_token_ids": pad_target_token_ids, - "teacher_indices": pad_teacher_indices, + # Pixel-value object refs must stay alive until every worker has consumed the + # dispatch, so the finally clause releases them on both the success and the + # failure path. + handles: list[ray.ObjectRef] = [] + try: + advantages = self._prepare_advantages(prepared, pack_plan) + dp_dispatches = self._build_dispatches(rollout_items, advantages, pack_plan) + # Nested in a list so Ray does not dereference it before TrainingWorker.fit + # explicitly fetches it; one DP rank's replicas share the ref. + item_refs_by_dp = { + dp_rank: [ray.put(dispatch["rollout_items"])] for dp_rank, dispatch in dp_dispatches.items() } - pad_data_samples = [pad_data for _ in range(pad_num)] - packed_data_batches = packed_data_batches + pad_data_samples - - handles = [] - data_batch_refs = {} - for worker_idx, worker in enumerate(self.workers): - dp_idx = worker_idx // data_replicate_size - if dp_idx not in data_batch_refs: - data_batch_refs[dp_idx] = ray.put(packed_data_batches[dp_idx::dp_size]) - handles.append( - worker.fit.remote( # type: ignore[attr-defined] - data_batches=data_batch_refs[dp_idx], - rollout_idx=rollout_idx, + for worker_idx, worker in enumerate(self.workers): + dp_rank = worker_idx // data_replicate_size + dp_dispatch = dp_dispatches[dp_rank] + handles.append( + worker.fit.remote( # type: ignore[attr-defined] + rollout_item_refs=item_refs_by_dp[dp_rank], + advantages=dp_dispatch["advantages"], + pack_plan=dp_dispatch["plan"], + batch_attr=prepared["batch_attr"], + ) ) - ) - try: log_infos = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) finally: - # free pixel values ref free_pixel_value_refs: list[ray.ObjectRef] = [] - for data in packed_data_batches: - if data["seq_ctx"].pixel_values is not None: - free_pixel_value_refs.extend(data["seq_ctx"].pixel_values) - if len(free_pixel_value_refs) > 0: + for state in rollout_items: + pixel_values = state.mm_info.get("pixel_values") if state.mm_info is not None else None + if pixel_values is not None: + free_pixel_value_refs.extend(pixel_values) + if free_pixel_value_refs: free_object_refs(free_pixel_value_refs) - del data_batch_refs - del packed_data_batches + del rollout_items + + data_info = self._build_trainer_log_info( + prepared, + pack_plan, + log_infos, + advantages, + raw_rewards_sum=raw_rewards_sum, + raw_rewards_count=raw_rewards_count, + ) return log_infos, data_info def offload(self, target: Literal["model", "optimizer", "all"] = "all"): @@ -870,10 +581,3 @@ def save(self, dcp_dir: str, no_save_optimizer: bool = False): handles = [worker.save.remote(dcp_dir, no_save_optimizer) for worker in self.workers] # type: ignore ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) return - - -def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None) -> torch.Tensor | None: - if value is None: - return None - assert isinstance(value, np.ndarray), f"Expected np.ndarray, got {type(value)}" - return torch.as_tensor(value, dtype=dtype, device="cpu") diff --git a/xtuner/v1/rl/trainer/pack.py b/xtuner/v1/rl/trainer/pack.py new file mode 100644 index 0000000000..40af331057 --- /dev/null +++ b/xtuner/v1/rl/trainer/pack.py @@ -0,0 +1,308 @@ +import math +import random +from collections.abc import Sequence +from typing import Literal, TypeAlias + +import numpy as np +import torch + +from xtuner.v1.datasets.packing import get_pack_chunk_infos, get_pack_infos_by_soft_split +from xtuner.v1.datasets.sampler import get_length_grouped_indices +from xtuner.v1.utils import XTUNER_DETERMINISTIC + + +PackIndices: TypeAlias = list[int] +OptimizerStepPackIndices: TypeAlias = list[PackIndices] +DPRankPackIndices: TypeAlias = list[OptimizerStepPackIndices] +PackedDataIndices: TypeAlias = list[DPRankPackIndices] + +_BALANCE_DEFAULT_SEED = 42 + + +def get_greedy_pack_infos(data_indices: Sequence[int], num_tokens: Sequence[int], target: int) -> list[PackIndices]: + """Group sample indices into sequential greedy packs without touching + sample data. + + Samples are visited in the given order and appended to the current pack while they + fit; the first non-fitting sample starts a new pack. + + Args: + data_indices (Sequence[int]): Sample indices to schedule, in visit order. + num_tokens (Sequence[int]): Training token length of each sample, aligned with + ``data_indices``. + target (int): Maximum token budget of one pack. + + Returns: + list[PackIndices]: The greedy pack layout; every input index appears exactly + once. + """ + if len(data_indices) != len(num_tokens): + raise ValueError("data_indices and num_tokens must have the same length") + + infos = get_pack_infos_by_soft_split( + inds=list(range(len(data_indices))), + dataset_id=0, + num_tokens=np.asarray(num_tokens), + pack_max_length=target, + ) + return [[int(data_indices[position]) for position in info["indices"]] for info in infos] + + +class RLDataPacker: + """Build an index-only packing plan for RL training data. + + The packer owns scheduling decisions only. It never creates padding tensors, + concatenates ``SequenceContext`` objects, or otherwise materializes a pack. Empty + packs (``[]``) are scheduling-only padding slots: the training worker materializes + them as all-padding packs of ``pack_max_length`` tokens. + + Args: + pack_max_length (int): Maximum token budget of one pack. + world_size (int): Total number of training workers. + data_replicate_size (int): Number of workers holding one data replica. + optimizer_steps (int): Configured optimizer steps per training iteration. + pack_strategy (Literal["legacy", "greedy", "balance", "native"]): Packing + strategy. ``legacy`` reproduces main's sequential greedy pack, + interleaved DP allocation and optimizer-step grouping; ``greedy`` fills + spare capacity from a candidate buffer; ``balance`` length-groups samples + before partitioning; ``native`` splits samples evenly + across DP ranks before packing. Defaults to "legacy". + pack_extra_buffer_size (int): Candidate buffer used by greedy gap filling. + Zero disables gap filling. Defaults to 1000. + pack_seed (int | None): Seed for randomized strategies such as ``balance``. + Defaults to None. + + Ordering is a per-strategy decision: ``legacy`` and ``greedy`` shuffle the visit + order before packing unless ``XTUNER_DETERMINISTIC`` is set, ``balance`` performs + its own seeded length grouping, and ``native`` keeps the input order. + """ + + def __init__( + self, + pack_max_length: int, + world_size: int, + data_replicate_size: int, + optimizer_steps: int, + pack_strategy: Literal["legacy", "greedy", "balance", "native"] = "legacy", + pack_seed: int | None = None, + pack_extra_buffer_size: int = 1000, + ): + self.pack_max_length = pack_max_length + self.world_size = world_size + self.data_replicate_size = data_replicate_size + self.optimizer_steps = optimizer_steps + self.dp_size = self.world_size // self.data_replicate_size + if pack_extra_buffer_size < 0: + raise ValueError("pack_extra_buffer_size must be non-negative") + self.pack_extra_buffer_size = pack_extra_buffer_size + self.pack_seed = pack_seed + self.split_size = 1024 + strategy_map = { + "legacy": self._legacy_pack, + "greedy": self._greedy_pack, + "balance": self._balance_pack, + "native": self._native_pack, + } + if pack_strategy not in strategy_map: + raise ValueError(f"Unknown packing strategy: {pack_strategy}") + self._impl = strategy_map[pack_strategy] + + def pack(self, data_lengths: Sequence[int]) -> tuple[PackedDataIndices, int]: + """Return the packed plan ``[dp][optimizer_step][pack][sample_index]`` + and the padding token count. + + Args: + data_lengths (Sequence[int]): Training token length of every sample. + + Returns: + tuple[PackedDataIndices, int]: The packed plan and the total padding tokens + across all packs (including empty all-padding packs). + + Raises: + ValueError: A single sample exceeds ``pack_max_length``. + """ + if not data_lengths: + return [], 0 + for data_index, data_length in enumerate(data_lengths): + if data_length > self.pack_max_length: + raise ValueError( + f"Single sample {data_index} seq len {data_length} exceeds pack_max_length {self.pack_max_length}" + ) + + data_indices = list(range(len(data_lengths))) + packed_data_indices = self._impl(data_indices, data_lengths) + padding_tokens = self._count_padding_tokens(packed_data_indices, data_lengths) + return packed_data_indices, padding_tokens + + def _legacy_pack(self, data_indices: list[int], data_lengths: Sequence[int]) -> PackedDataIndices: + if not XTUNER_DETERMINISTIC: + random.shuffle(data_indices) + visit_lengths = [data_lengths[index] for index in data_indices] + total_pack_indices = get_greedy_pack_infos(data_indices, visit_lengths, self.pack_max_length) + # Interleaved DP allocation over a DP-multiple pack list, then per-rank + # sequential optimizer-step grouping: this is the exact schedule the previous + # controller + `iters_per_step` worker regrouping produced. + pad_num = math.ceil(len(total_pack_indices) / self.dp_size) * self.dp_size - len(total_pack_indices) + total_pack_indices.extend([[] for _ in range(pad_num)]) + + packs_per_dp = len(total_pack_indices) // self.dp_size + iters_per_step = max(1, math.ceil(packs_per_dp / self.optimizer_steps)) + actual_optimizer_steps = math.ceil(packs_per_dp / iters_per_step) + + packed_data_indices: PackedDataIndices = [] + for dp_rank in range(self.dp_size): + dp_pack_indices = total_pack_indices[dp_rank :: self.dp_size] + packed_data_indices.append( + [ + dp_pack_indices[step * iters_per_step : (step + 1) * iters_per_step] + for step in range(actual_optimizer_steps) + ] + ) + return packed_data_indices + + def _greedy_pack(self, data_indices: list[int], data_lengths: Sequence[int]) -> PackedDataIndices: + if not XTUNER_DETERMINISTIC: + random.shuffle(data_indices) + visit_lengths = [data_lengths[index] for index in data_indices] + lengths = np.asarray(visit_lengths) + infos = get_pack_chunk_infos( + # SFT consumes this list from its tail; reverse a private position list. + inds=list(reversed(range(len(data_indices)))), + dataset_id=0, + target=self.pack_max_length, + pack_extra_buffer_size=self.pack_extra_buffer_size, + num_tokens=lengths, + proxy_attn_flops=lengths, + ) + total_pack_indices = [[data_indices[position] for position in info["indices"]] for info in infos] + # Interleaved DP allocation over a DP-multiple pack list, then per-rank + # sequential optimizer-step grouping: this is the exact schedule the previous + # controller + `iters_per_step` worker regrouping produced. + pad_num = math.ceil(len(total_pack_indices) / self.dp_size) * self.dp_size - len(total_pack_indices) + total_pack_indices.extend([[] for _ in range(pad_num)]) + + packs_per_dp = len(total_pack_indices) // self.dp_size + iters_per_step = max(1, math.ceil(packs_per_dp / self.optimizer_steps)) + actual_optimizer_steps = math.ceil(packs_per_dp / iters_per_step) + + packed_data_indices: PackedDataIndices = [] + for dp_rank in range(self.dp_size): + dp_pack_indices = total_pack_indices[dp_rank :: self.dp_size] + packed_data_indices.append( + [ + dp_pack_indices[step * iters_per_step : (step + 1) * iters_per_step] + for step in range(actual_optimizer_steps) + ] + ) + return packed_data_indices + + def _balance_pack( + self, + data_indices: list[int], + data_lengths: Sequence[int], + ) -> PackedDataIndices: + seed = self.pack_seed if self.pack_seed is not None else _BALANCE_DEFAULT_SEED + torch_generator = torch.Generator().manual_seed(seed) + random_generator = random.Random(seed) + grouped_indices = get_length_grouped_indices( + max_lengths=list(data_lengths), + group_batch_size=len(data_indices), + group_size=self.dp_size, + torch_generator=torch_generator, + random_generator=random_generator, + ) + + partitioned_indices: list[list[PackIndices]] = [ + [[] for _ in range(self.optimizer_steps)] for _ in range(self.dp_size) + ] + for i, data_index in enumerate(grouped_indices): + dp_rank = i % self.dp_size + step_idx = (i // self.dp_size) % self.optimizer_steps + partitioned_indices[dp_rank][step_idx].append(int(data_index)) + + packed_data_indices: PackedDataIndices = [ + [[] for _ in range(self.optimizer_steps)] for _ in range(self.dp_size) + ] + max_packs_per_step = [0] * self.optimizer_steps + for dp_rank in range(self.dp_size): + for step_idx in range(self.optimizer_steps): + packed_step_indices = self._pack_indices( + partitioned_indices[dp_rank][step_idx], + data_lengths, + ) + packed_data_indices[dp_rank][step_idx] = packed_step_indices + max_packs_per_step[step_idx] = max(max_packs_per_step[step_idx], len(packed_step_indices)) + + self._align_pack_count(packed_data_indices, max_packs_per_step) + return packed_data_indices + + def _native_pack( + self, + data_indices: list[int], + data_lengths: Sequence[int], + ) -> PackedDataIndices: + # Use private negative indices as scheduling-only padding slots. They are + # removed from the returned plan and materialized as padding by workers. + scheduled_indices = data_indices.copy() + if len(scheduled_indices) % self.dp_size != 0: + pad_num = self.dp_size - (len(scheduled_indices) % self.dp_size) + scheduled_indices.extend(-(index + 1) for index in range(pad_num)) + + batches_per_dp_group = np.array_split(scheduled_indices, self.dp_size) + actual_optimizer_steps = min(len(batches_per_dp_group[0]), self.optimizer_steps) + packed_data_indices: PackedDataIndices = [ + [[] for _ in range(actual_optimizer_steps)] for _ in range(self.dp_size) + ] + max_packs_per_step = [0] * actual_optimizer_steps + + for dp_rank, dp_worker_indices in enumerate(batches_per_dp_group): + indices_for_optim_steps = np.array_split(dp_worker_indices, actual_optimizer_steps) + for step_idx, step_indices_array in enumerate(indices_for_optim_steps): + step_indices = [int(index) for index in step_indices_array] + packed_step_indices = self._pack_indices(step_indices, data_lengths) + packed_data_indices[dp_rank][step_idx] = packed_step_indices + max_packs_per_step[step_idx] = max(max_packs_per_step[step_idx], len(packed_step_indices)) + + self._align_pack_count(packed_data_indices, max_packs_per_step) + return packed_data_indices + + def _pack_indices(self, data_indices: Sequence[int], data_lengths: Sequence[int]) -> OptimizerStepPackIndices: + scheduled_lengths = [self._get_scheduled_length(data_index, data_lengths) for data_index in data_indices] + if sum(scheduled_lengths) > self.pack_max_length: + packs = get_greedy_pack_infos(data_indices, scheduled_lengths, self.pack_max_length) + else: + packs = [[int(data_index) for data_index in data_indices]] + return [[data_index for data_index in pack if data_index >= 0] for pack in packs] + + def _get_scheduled_length(self, data_index: int, data_lengths: Sequence[int]) -> int: + if data_index >= 0: + return data_lengths[data_index] + return min(self.split_size, self.pack_max_length) + + def _align_pack_count(self, packed_data_indices: PackedDataIndices, max_packs_per_step: Sequence[int]) -> None: + for step_idx, max_packs in enumerate(max_packs_per_step): + for dp_rank in range(self.dp_size): + missing_packs = max_packs - len(packed_data_indices[dp_rank][step_idx]) + packed_data_indices[dp_rank][step_idx].extend([[] for _ in range(missing_packs)]) + + def _count_padding_tokens( + self, + packed_data_indices: PackedDataIndices, + data_lengths: Sequence[int], + ) -> int: + total_packs = 0 + scheduled_tokens = 0 + seen_indices: list[int] = [] + for dp_rank_indices in packed_data_indices: + for step_indices in dp_rank_indices: + total_packs += len(step_indices) + for pack_indices in step_indices: + seen_indices.extend(pack_indices) + pack_tokens = sum(data_lengths[data_index] for data_index in pack_indices) + if pack_tokens > self.pack_max_length: + raise RuntimeError(f"Pack tokens {pack_tokens} exceed pack_max_length {self.pack_max_length}") + scheduled_tokens += pack_tokens + + if sorted(seen_indices) != list(range(len(data_lengths))): + raise RuntimeError("Packing plan must contain every data index exactly once") + return total_packs * self.pack_max_length - scheduled_tokens diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index d0a2345cd7..f753225d00 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -2,7 +2,6 @@ import contextlib import json -import math import os import time from contextlib import contextmanager @@ -10,6 +9,7 @@ from typing import ( TYPE_CHECKING, Iterable, + Literal, Sequence, TypedDict, cast, @@ -26,7 +26,7 @@ import torch import torch.distributed as dist from mmengine.runner import set_random_seed -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from ray.actor import ActorClass, ActorProxy from torch.distributed.device_mesh import init_device_mesh from typing_extensions import NotRequired @@ -34,6 +34,7 @@ from transformers import AutoTokenizer from xtuner.v1.config.fsdp import FSDPConfig from xtuner.v1.config.optim import LRConfig, OptimConfig +from xtuner.v1.data_proto.rl_data import RolloutState from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.datasets.config import DataloaderConfig from xtuner.v1.datasets.dataloader import Dataloader @@ -46,6 +47,7 @@ from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo from xtuner.v1.profiler import profiling_memory, profiling_time from xtuner.v1.rl.distillation import ( + DistillationConfig, DistillationTrainerAdapter, TrainTeacherManager, TrainTeacherManagerConfig, @@ -73,6 +75,7 @@ from xtuner.v1.utils.nccl_process_group import resume_nccl_process_groups, suspend_nccl_process_groups from ..rollout_is import merge_rollout_is_metrics +from .pack import DPRankPackIndices DEVICE = get_device() @@ -80,6 +83,61 @@ _UINT16_CAPACITY = 1 << 16 +def get_train_seq_ctx( + input_ids: torch.LongTensor, + position_ids: np.ndarray | None = None, + multimodal_train_info: dict | None = None, +) -> SequenceContext: + """Build a CPU ``SequenceContext`` for one training sample. + + Args: + input_ids (torch.LongTensor): Model input tokens with shape ``(1, seq_len)``. + position_ids (np.ndarray | None): Optional position ids. A 3D array triggers + the VLM MRoPE layout; a prompt-only position segment is extended to cover + the whole input. + multimodal_train_info (dict | None): Optional multimodal payload with + ``pixel_values``, ``image_grid_thw`` and ``num_img_tokens``. + + Returns: + SequenceContext: The CPU sequence context of the sample. + """ + seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cpu") + position_ids = _to_cpu_tensor(position_ids, dtype=torch.long) + if position_ids is not None and len(position_ids.shape) == 3: + # Match get_rope_index_3: response text continues from a single global + # max(T, H, W), not per-axis maxima. Per-axis max diverges when the + # prompt ends on image tokens (T≈0 while H/W are large). The extension + # length is derived from the position deficit so samples whose positions + # already cover the whole input (full-sequence agentic samples) need no + # extra shape information. + num_extend = input_ids.size(-1) - position_ids.size(-1) + if num_extend > 0: + max_value = position_ids.amax() + response_position_ids = ( + (torch.arange(1, num_extend + 1, device=position_ids.device, dtype=position_ids.dtype) + max_value) + .view(1, 1, -1) + .expand(3, 1, -1) + ) + position_ids = torch.cat([position_ids, response_position_ids], dim=-1) + seq_ctx.position_ids = position_ids # type: ignore[assignment] + assert position_ids.size(-1) == input_ids.size(-1) + + if multimodal_train_info: + seq_ctx.pixel_values = multimodal_train_info.get("pixel_values") + seq_ctx.image_grid_thw = _to_cpu_tensor(multimodal_train_info.get("image_grid_thw"), dtype=torch.long) + num_img_tokens = multimodal_train_info.get("num_img_tokens") + if num_img_tokens is not None: + seq_ctx.num_img_tokens = [num_img_tokens] + return seq_ctx + + +def _to_cpu_tensor(value: np.ndarray | None, *, dtype: torch.dtype | None = None) -> torch.Tensor | None: + if value is None: + return None + assert isinstance(value, np.ndarray), f"Expected np.ndarray, got {type(value)}" + return torch.as_tensor(value, dtype=dtype, device="cpu") + + def _rollout_routed_experts_storage_dtype(n_routed_experts: int) -> torch.dtype: """Choose a compact lossless dtype for stored routed-expert IDs.""" if n_routed_experts <= _UINT16_CAPACITY: @@ -149,6 +207,14 @@ class WorkerConfig(BaseModel): updated weights to rollout workers. Defaults to 0.5. seed (int | None): Training worker random seed. When None, the RL trainer seed is used. Defaults to None. + pack_strategy (Literal["legacy", "greedy", "balance", "native"]): Index-only + packing strategy used by the controller-side packer. ``legacy`` exactly + reproduces the previous greedy pack, interleaved DP allocation and + optimizer-step grouping. Defaults to "legacy". + pack_extra_buffer_size (int): Candidate buffer used by greedy chunk packing; + zero disables gap filling. Defaults to 1000. + pack_seed (int | None): Random seed for length-grouped packing strategies such + as "balance". Defaults to None. offload_rollout_routed_experts (bool): Keep rollout routed experts on CPU and move each layer's slice to device during forward. Defaults to False. offload_old_logprobs (bool): Keep actor old log probabilities on CPU between @@ -188,6 +254,10 @@ class WorkerConfig(BaseModel): log_dir: str | Path | None = None update_weight_bucket_size_in_gb: float = 0.5 # 512MB seed: None | int = None # if None, use RLTrainer seed + pack_strategy: Literal["legacy", "greedy", "balance", "native"] = "legacy" + pack_extra_buffer_size: int = Field(default=1000, ge=0) + pack_seed: int | None = None + distillation_config: DistillationConfig | None = None profile_step: list[int] | int | None = None # 1-based global RL train_step ids to profile. profile_time: bool = True profile_memory: bool = False @@ -232,14 +302,36 @@ def build( ) -class WorkerInputItem(TypedDict): +class TrainBatchAttr(TypedDict): + """Attributes of one training batch, derived once by the controller from + the lightweight metadata written by AgentLoop canonicalization. + + The controller scans the batch meta and passes this single structure across the + Ray boundary so workers materialize packs uniformly: cross-rank collectives stay + symmetric because every worker of one batch sees the same attributes. + + rollout_idx (int): Global rollout (train step) id used for logging. + use_3d_position_ids (bool): Whether the batch uses 3D MRoPE position ids; padding + packs mirror this layout. + pack_loss_keys (list[str]): Batch-level optional loss inputs (e.g. + ``rollout_logprobs``) that must exist on every pack, including all-padding + packs, so cross-rank metric collectives stay uniform. + has_routed_experts (bool): Whether the batch carries routed experts; padding + positions then receive dummy expert ids. + """ + + rollout_idx: int + use_3d_position_ids: bool + pack_loss_keys: list[str] + has_routed_experts: bool + + +class _TrainItem(TypedDict): + """CPU sample input and loss tensors, before any training context is + built.""" + seq_ctx: SequenceContext - shifted_labels: torch.LongTensor - advantages: torch.Tensor - rollout_logprobs: torch.Tensor | None - teacher_logprobs: NotRequired[torch.Tensor | None] - target_token_ids: NotRequired[torch.Tensor | None] - teacher_indices: NotRequired[torch.Tensor | None] + loss_inputs: dict[str, torch.Tensor | None] class WorkerTrainLogItem(TypedDict, total=False): @@ -259,6 +351,8 @@ class WorkerLogItem(TypedDict): rollout_is_metrics: NotRequired[dict[str, float]] train_metrics: list[WorkerTrainLogItem] sft_train_metrics: NotRequired[dict[str, float]] + packing_conversion_time_s: NotRequired[float] + packing_pack_time_s: NotRequired[float] class TrainingWorker(SingleAcceleratorWorker): @@ -316,6 +410,8 @@ def __init__( chunk_size=worker_cfg.loss_cfg.chunk_size, ) + self._distillation = DistillationTrainerAdapter(worker_cfg.distillation_config) + self._optimizer_steps = worker_cfg.optimizer_steps profile_step = worker_cfg.profile_step if isinstance(profile_step, int): @@ -669,28 +765,308 @@ def _maybe_profiling(self, global_train_step: int, phase: str): stack.enter_context(profiling_memory(memory_dir)) yield + def _convert_rollout_items_to_train_items( + self, + rollout_items: list[RolloutState], + advantages: Sequence[float | list[float]], + batch_attr: TrainBatchAttr, + ) -> list[_TrainItem]: + """Convert unshifted states into CPU items in unchanged sample order. + + Apply input_ids[:-1], labels[1:] and logprobs[1:] once. Broadcast scalar advantages or use aligned token + targets; mask unsupervised positions. Prepare position, multimodal and teacher inputs without building a loss + context or applying the plan. Final pack contexts are built by packing. + """ + return [ + self._convert_one_rollout_state(state, advantage, batch_attr["use_3d_position_ids"]) + for state, advantage in zip(rollout_items, advantages) + ] + + def _convert_one_rollout_state( + self, + state: RolloutState, + advantage_val: float | list[float], + use_3d_position_ids: bool, + ) -> _TrainItem: + input_ids = cast(list[int], state.input_ids) + labels = cast(list[int], state.labels) + shifted_labels = labels[1:] + + input_ids_t = cast(torch.LongTensor, torch.tensor(input_ids[:-1], dtype=torch.int64).unsqueeze(0)) + shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) + + rollout_logprobs: torch.Tensor | None = None + if state.logprobs is not None: + rollout_logprobs = torch.tensor(state.logprobs[1:], dtype=torch.float32).unsqueeze(0) + + if isinstance(advantage_val, list): + if len(advantage_val) != len(shifted_labels): + raise ValueError( + f"Token advantage length mismatch for rollout_id={state.rollout_id}: " + f"{len(advantage_val)} vs shifted_labels={len(shifted_labels)}" + ) + token_advantages = [value if label != -100 else 0.0 for value, label in zip(advantage_val, shifted_labels)] + else: + token_advantages = [advantage_val if label != -100 else 0.0 for label in shifted_labels] + advantages = torch.tensor([token_advantages], dtype=torch.float32) + + position_ids = state.position_ids + if use_3d_position_ids and (position_ids is None or position_ids.ndim != 3): + seq_len = input_ids_t.size(-1) + text_position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1) + # Mixed batches use three-axis MRoPE position IDs. Repeat the normal text + # positions on every axis so text samples have the same rank as the VLM + # samples when their sequence contexts are packed together. + position_ids = np.broadcast_to(text_position_ids, (3, 1, seq_len)).copy() + multimodal_train_info = cast(dict | None, state.mm_info) + seq_ctx = get_train_seq_ctx(input_ids_t, position_ids, multimodal_train_info) + seq_ctx.rollout_routed_experts = state.routed_experts # type: ignore[assignment] # n,layer*expert + + loss_inputs: dict[str, torch.Tensor | None] = { + "shifted_labels": shifted_labels_t, + "advantages": advantages, + "rollout_logprobs": rollout_logprobs, + } + loss_inputs.update(self._distillation.rollout_teacher_targets(state, shifted_labels=shifted_labels)) + + return {"seq_ctx": seq_ctx, "loss_inputs": loss_inputs} + @ray_method - def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLogItem: + def fit( + self, + rollout_item_refs: list[ray.ObjectRef], + advantages: Sequence[float | list[float]], + pack_plan: DPRankPackIndices, + batch_attr: TrainBatchAttr, + ) -> WorkerLogItem: + """Train one rollout batch from the controller's packed per-DP plan. + + Args: + rollout_item_refs (list[ray.ObjectRef]): Nested per-DP object refs; element + zero wraps this DP rank's ``RolloutState`` list and is shared by all + data replicas of the rank. + advantages: Scalar or shifted-token advantages per rollout state, aligned with + the dereferenced state list. + pack_plan (DPRankPackIndices): Local ``[optimizer_step][pack][sample]`` plan + whose indices address the dereferenced state list. + batch_attr (TrainBatchAttr): Batch-level attributes derived by the + controller; uniform across all workers of the batch. + + Returns: + WorkerLogItem: Training metrics of one rollout batch. + """ # NOTE: sglang会清除logger handle, 重新创建 self.logger = get_logger(log_dir=self.log_dir, tag="TrainingWorker") + convert_begin = time.perf_counter() + rollout_items = ray.get(rollout_item_refs[0]) + train_items = self._convert_rollout_items_to_train_items(rollout_items, advantages, batch_attr) + conversion_time_s = time.perf_counter() - convert_begin + del rollout_items, advantages, rollout_item_refs + + pack_begin = time.perf_counter() + step_batches = self._pack_train_items(train_items, pack_plan, batch_attr) + pack_time_s = time.perf_counter() - pack_begin + del train_items, pack_plan + + worker_log_item = self._fit(step_batches, batch_attr["rollout_idx"]) + worker_log_item["packing_conversion_time_s"] = conversion_time_s + worker_log_item["packing_pack_time_s"] = pack_time_s + return worker_log_item + + def _pack_train_items( + self, + items: list[_TrainItem], + plan: DPRankPackIndices, + batch_attr: TrainBatchAttr, + ) -> list[list[tuple[SequenceContext, BaseRLLossContext]]]: + """Materialize local [optimizer_step][pack][sample_index] plans on CPU. + + Indices address items. Concatenate and pad sequence/loss tensors to pack_max_length, including empty all- + padding packs. Batch and teacher configuration define required extra fields on every pack. Build one loss + context per final pack; do not shift again or recompute advantages. Return [optimizer_step][pack] + (SequenceContext, loss context) pairs. Device transfer and SP splitting happen later in _fit. + """ + template_keys = list(batch_attr["pack_loss_keys"]) + for key in self._configured_teacher_loss_keys(): + if key not in template_keys: + template_keys.append(key) + + step_batches: list[list[tuple[SequenceContext, BaseRLLossContext]]] = [] + for step_plan in plan: + step_batches.append( + [ + self._pack_one_batch( + [items[index] for index in pack], + template_keys=template_keys, + use_3d_position_ids=batch_attr["use_3d_position_ids"], + has_routed_experts=batch_attr["has_routed_experts"], + ) + for pack in step_plan + ] + ) + return step_batches + + def _configured_teacher_loss_keys(self) -> list[str]: + distillation_config = self.config.distillation_config + if distillation_config is None: + return [] + keys: list[str] = [] + if distillation_config.rollout_teachers: + keys.append("teacher_logprobs") + if not distillation_config.loss_config.uses_sampled_token_targets: + keys.append("target_token_ids") + if distillation_config.train_teachers: + keys.append("teacher_indices") + return keys + + def _pack_one_batch( + self, + pack_items: list[_TrainItem], + template_keys: list[str], + use_3d_position_ids: bool, + has_routed_experts: bool, + ) -> tuple[SequenceContext, BaseRLLossContext]: + real_tokens = 0 + for item in pack_items: + input_ids = item["seq_ctx"].input_ids + assert input_ids is not None, "Packed seq_ctx.input_ids must not be None" + real_tokens += input_ids.shape[1] + padding_len = self.config.pack_max_length - real_tokens + assert padding_len >= 0, f"Pack tokens {real_tokens} exceed pack_max_length {self.config.pack_max_length}" + + if pack_items: + loss_inputs_list = [item["loss_inputs"] for item in pack_items] + loss_inputs: dict[str, torch.Tensor | None] = { + "shifted_labels": torch.cat( + [cast(torch.Tensor, inputs["shifted_labels"]) for inputs in loss_inputs_list], dim=1 + ), + "advantages": torch.cat( + [cast(torch.Tensor, inputs["advantages"]) for inputs in loss_inputs_list], dim=1 + ), + } + for key in template_keys: + values = [inputs.get(key) for inputs in loss_inputs_list] + present = [value is not None for value in values] + if not all(present): + raise ValueError(f"Required pack field {key} must be present on every sample: {present}") + loss_inputs[key] = torch.cat([cast(torch.Tensor, value) for value in values], dim=1) + if padding_len > 0: + for key, tensor in loss_inputs.items(): + assert tensor is not None + loss_inputs[key] = torch.cat([tensor, self._pack_padding_loss_input(key, padding_len)], dim=1) + else: + # Scheduling-only padding slot: the whole pack materializes as padding. + loss_inputs = { + "shifted_labels": self._pack_padding_loss_input("shifted_labels", padding_len), + "advantages": self._pack_padding_loss_input("advantages", padding_len), + } + for key in template_keys: + loss_inputs[key] = self._pack_padding_loss_input(key, padding_len) + + seq_ctx_list = [item["seq_ctx"] for item in pack_items] + if padding_len > 0: + seq_ctx_list.append( + self._pack_padding_seq_ctx( + padding_len, + use_3d_position_ids=use_3d_position_ids, + has_routed_experts=has_routed_experts, + ) + ) + seq_ctx = SequenceContext.cat(seq_ctx_list) + + assert seq_ctx.input_ids is not None + seq_len = seq_ctx.input_ids.shape[1] + for key, tensor in loss_inputs.items(): + tensor_len = None if tensor is None else tensor.shape[1] + assert tensor_len == seq_len, f"{key}: {tensor_len} vs {seq_len}" + + # Build the loss context on CPU; `_fit` transfers it to the accelerator and + # applies the optional sequence-parallel split, mirroring the seq_ctx lifecycle. + loss_ctx = self.config.loss_cfg.build(data=loss_inputs, device="cpu") + assert loss_ctx is not None + return seq_ctx, loss_ctx + + def _pack_padding_loss_input(self, key: str, padding_len: int) -> torch.Tensor: + if key == "shifted_labels": + return torch.full((1, padding_len), -100, dtype=torch.long) + if key == "advantages": + return torch.zeros(1, padding_len, dtype=torch.float32) + if key == "rollout_logprobs": + return torch.zeros(1, padding_len, dtype=torch.float32) + if key == "teacher_indices": + return torch.full((1, padding_len), -1, dtype=torch.long) + if key in ("teacher_logprobs", "target_token_ids"): + distillation_config = self.config.distillation_config + assert distillation_config is not None, f"Padding {key} requires distillation_config" + if key == "teacher_logprobs" and distillation_config.loss_config.uses_sampled_token_targets: + return torch.zeros(1, padding_len, dtype=torch.float32) + top_k = distillation_config.loss_config.top_k + assert top_k is not None, f"Padding {key} requires loss_config.top_k" + if key == "teacher_logprobs": + return torch.zeros(1, padding_len, top_k, dtype=torch.float32) + return torch.zeros(1, padding_len, top_k, dtype=torch.long) + raise ValueError(f"Unknown pack loss input key: {key}") + + def _pack_padding_seq_ctx( + self, + padding_len: int, + use_3d_position_ids: bool, + has_routed_experts: bool, + ) -> SequenceContext: + # Reduce the attn calculation time by using multiple short sequence packs + pad_tokens = tuple(torch.zeros(1, 1024, dtype=torch.long) for _ in range(padding_len // 1024)) + if padding_len % 1024 > 0: + pad_tokens = pad_tokens + (torch.zeros(1, padding_len % 1024, dtype=torch.long),) + pad_seq_ctx = SequenceContext.from_input_ids(pad_tokens, device="cpu") # type: ignore[arg-type] + pad_seq_ctx.num_padding = padding_len + if use_3d_position_ids: + _position_ids_list = [] + for pad_token in pad_tokens: + _position_ids = torch.arange(pad_token.size(-1)).view(1, 1, -1).expand(3, 1, -1) + _position_ids_list.append(_position_ids) + pad_seq_ctx.position_ids = torch.cat(_position_ids_list, dim=-1) # type: ignore[assignment] + if has_routed_experts: + language_cfg = ( + self.config.model_cfg.text_config + if isinstance(self.config.model_cfg, BaseComposeConfig) + else self.config.model_cfg + ) + pad_seq_ctx.rollout_routed_experts = torch.randint( + low=0, high=language_cfg.n_routed_experts, size=(padding_len, 1, 1) + ) + return pad_seq_ctx + + def _fit( + self, + step_batches: list[list[tuple[SequenceContext, BaseRLLossContext]]], + rollout_idx: int, + ) -> WorkerLogItem: loss_cfg: BaseRLLossConfig = self.config.loss_cfg - num_batches = len(data_batches) - iters_per_step = math.ceil(num_batches / self._optimizer_steps) - if num_batches < self._optimizer_steps: + num_optimizer_steps = len(step_batches) + assert num_optimizer_steps <= self._optimizer_steps, ( + f"Pack plan optimizer steps {num_optimizer_steps} exceed configured " + f"optimizer_steps {self._optimizer_steps}" + ) + if num_optimizer_steps < self._optimizer_steps: self.logger.info( - f"Optimizer only step once because num_batches {num_batches} < optimizer_steps {self._optimizer_steps}." + f"Optimizer only steps {num_optimizer_steps} times because pack plan steps " + f"{num_optimizer_steps} < optimizer_steps {self._optimizer_steps}." ) - - # Update seq_ctx: pixel_values, rollout_routed_experts - # Init loss_ctx: shifted_labels, advantages, rollout_logprobs + packs_per_step = [len(step_data_batches) for step_data_batches in step_batches] + flat_data_batches = [data for step_data_batches in step_batches for data in step_data_batches] + del step_batches + + # Update seq_ctx: resolve pixel_values and rollout_routed_experts, then move to the + # accelerator and split for sequence parallelism. loss_ctx follows the same lifecycle: + # it was built on CPU during packing and is only transferred and split here. MTP + # contexts must be built on the full (pre-split) labels of the split seq_ctx. seq_ctx_list: list[SequenceContext] = [] loss_ctx_list: list[BaseRLLossContext] = [] mtp_loss_ctx_list: list[list[MTPLossContext]] = [] prepare_inputs_begin = time.perf_counter() - for data in data_batches: + for data in flat_data_batches: + seq_ctx, loss_ctx = data # update seq_ctx - seq_ctx = data["seq_ctx"] - teacher_indices = data.get("teacher_indices") pixel_values = seq_ctx.pixel_values if pixel_values is not None: if not isinstance(pixel_values, np.ndarray): @@ -709,30 +1085,13 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo self._add_rollout_routed_experts(seq_ctx, rollout_routed_experts) seq_ctx.offload_rollout_routed_experts = self.config.offload_rollout_routed_experts - seq_ctx = data["seq_ctx"].to(DEVICE) + seq_ctx = seq_ctx.to(DEVICE) if self.sp_mesh.size() > 1: seq_ctx = seq_ctx.split(self.sp_mesh) - # init loss_ctx - shifted_labels = data["shifted_labels"].to(DEVICE) - advantages = data["advantages"].to(DEVICE) - rollout_logprobs = data.get("rollout_logprobs", None) - rollout_logprobs = rollout_logprobs.to(DEVICE) if rollout_logprobs is not None else None - loss_ctx = loss_cfg.build( - data={ - "shifted_labels": shifted_labels, - "advantages": advantages, - "rollout_logprobs": rollout_logprobs, - "teacher_logprobs": data.get("teacher_logprobs", None), - "target_token_ids": data.get("target_token_ids", None), - "teacher_indices": teacher_indices, - }, - sp_mesh=self.sp_mesh, - ) - - seq_ctx_list.append(seq_ctx) - assert loss_ctx is not None - loss_ctx_list.append(loss_ctx) + # update loss_ctx + loss_kwargs = loss_ctx.loss_kwargs + loss_kwargs.to(DEVICE) if self.mtp_config is not None: mtp_loss_ctxs_per_batch: list[MTPLossContext] = [] for mtp_idx in range(self.mtp_config.num_layers): @@ -743,21 +1102,26 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo ) mtp_ctx = mtp_loss_cfg.build( data={ - "shifted_labels": shifted_labels, + "shifted_labels": loss_kwargs.shifted_labels, "seq_ctx": seq_ctx, - "logprobs": rollout_logprobs, + "logprobs": loss_kwargs.rollout_logprobs, }, sp_mesh=self.sp_mesh, ) if mtp_ctx is not None: mtp_loss_ctxs_per_batch.append(mtp_ctx) mtp_loss_ctx_list.append(mtp_loss_ctxs_per_batch) + if self.sp_mesh.size() > 1: + loss_kwargs.sp_split(self.sp_mesh) + + seq_ctx_list.append(seq_ctx) + loss_ctx_list.append(loss_ctx) self.logger.debug( f"Rank{self.rank} Rollout {rollout_idx} prepare_inputs elapsed=" f"{time.perf_counter() - prepare_inputs_begin:.4f}s" ) - del data_batches + del flat_data_batches # When sp_mesh.size() > 1, get the sp_split shifted_labels and rollout_logprobs shifted_labels_list = [loss_ctx.loss_kwargs.shifted_labels for loss_ctx in loss_ctx_list] @@ -862,25 +1226,25 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo if teacher_timings is not None: worker_log_item.update({"teacher_timings": teacher_timings.to_dict()}) - # compute batched loss context + # compute batched loss context along the plan step boundaries batched_loss_ctx_list: list[BaseRLLossContext] = [] batched_mtp_loss_ctx_list: list[list[MTPLossContext]] = [] LossContext = loss_cfg.loss_ctx_cls - for i in range(0, len(loss_ctx_list), iters_per_step): - batches_loss_ctx = loss_ctx_list[i : i + iters_per_step] + start_idx = 0 + for packs_in_step in packs_per_step: + step_loss_ctxs = loss_ctx_list[start_idx : start_idx + packs_in_step] batched_loss_ctx_list.extend( - LossContext.build_batches(batches_loss_ctx) # type: ignore[arg-type] + LossContext.build_batches(step_loss_ctxs) # type: ignore[arg-type] ) if self.mtp_config is not None: - batches_seq_ctx = seq_ctx_list[i : i + iters_per_step] - cu_seq_lens_list = [seq_ctx.cu_seq_lens_q for seq_ctx in batches_seq_ctx] + step_seq_ctxs = seq_ctx_list[start_idx : start_idx + packs_in_step] + cu_seq_lens_list = [seq_ctx.cu_seq_lens_q for seq_ctx in step_seq_ctxs] # mtp_loss_ctx_list: list[list[MTPLossContext]], outer=batch, inner=mtp_depth num_mtp_depths = len(mtp_loss_ctx_list[0]) if mtp_loss_ctx_list else 0 for mtp_idx in range(num_mtp_depths): depth_mtp_loss_ctxs: list[LMHeadLossContext] = [ - mtp_loss_ctx_list[j][mtp_idx] - for j in range(i, min(i + iters_per_step, len(mtp_loss_ctx_list))) + mtp_loss_ctx_list[j][mtp_idx] for j in range(start_idx, start_idx + packs_in_step) ] batched_mtp_depth_ctxs = cast( list[MTPLossContext], @@ -892,17 +1256,19 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo ) # Append each depth's batched ctx to the corresponding batch index for batch_offset, mtp_ctx in enumerate(batched_mtp_depth_ctxs): - global_batch_idx = i + batch_offset + global_batch_idx = start_idx + batch_offset if global_batch_idx >= len(batched_mtp_loss_ctx_list): batched_mtp_loss_ctx_list.append([mtp_ctx]) else: batched_mtp_loss_ctx_list[global_batch_idx].append(mtp_ctx) + start_idx += packs_in_step - # train optimizer steps - for i in range(0, len(seq_ctx_list), iters_per_step): + # train optimizer steps along the plan step boundaries + start_idx = 0 + for step_idx, packs_in_step in enumerate(packs_per_step): global_train_step = self._global_train_step + 1 - batches_seq_ctx = seq_ctx_list[i : i + iters_per_step] - batches_loss_ctx = batched_loss_ctx_list[i : i + iters_per_step] + batches_seq_ctx = seq_ctx_list[start_idx : start_idx + packs_in_step] + batches_loss_ctx = batched_loss_ctx_list[start_idx : start_idx + packs_in_step] engine_input = [ ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) @@ -910,7 +1276,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo ] if self.mtp_config is not None: - batches_mtp_loss_ctxs = batched_mtp_loss_ctx_list[i : i + iters_per_step] + batches_mtp_loss_ctxs = batched_mtp_loss_ctx_list[start_idx : start_idx + packs_in_step] engine_input = [ ModelItem( seq_ctx=seq_ctx, @@ -936,7 +1302,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo ) self.logger.debug( f"Rank{self.rank} Rollout {rollout_idx} GlobalStep {global_train_step} " - f"train_step[{i}].engine_train_step elapsed={time.perf_counter() - train_step_begin:.4f}s" + f"train_step[{step_idx}].engine_train_step elapsed={time.perf_counter() - train_step_begin:.4f}s" ) mtp_grad_metrics: dict[str, float] = {} if self._mtp_trainable_parameters: @@ -989,9 +1355,10 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo for key, value in train_log_item.items() if not key.startswith("reduced_train_policy_") and key != "max_ratio" ) - log_str = f"Rank{self.rank} Rollout {rollout_idx} Step {i}: " + log_str + log_str = f"Rank{self.rank} Rollout {rollout_idx} Step {step_idx}: " + log_str self.logger.info(log_str) self._global_train_step = global_train_step + start_idx += packs_in_step self._rollout_step += 1 if self._sft_dataloader is not None and self._rollout_step % self._rollout_steps_per_sft == 0: @@ -1125,6 +1492,11 @@ def get_data_replicate_size(self) -> int: # sp will affect the data replicate size in worker return self._engine.data_replicate_size * self.sp_mesh.size() + @ray_method + def get_worker_cfg(self) -> WorkerConfig: + """Expose the worker config for controller-side pack planning.""" + return self.config + @ray_method def get_model_cfg(self): model_cfg = self._engine.model_cfg diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index f4166653de..86856a3a81 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -17,7 +17,7 @@ from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1._writer import get_writer -from xtuner.v1.data_proto.rl_data import RolloutState, is_valid_for_training +from xtuner.v1.data_proto.rl_data import RolloutState, validate_train_group from xtuner.v1.patch import patch_default_save_plan from xtuner.v1.rl.advantage import BaseAdvantageConfig, GRPOAdvantageConfig from xtuner.v1.rl.agent_loop_manager import ( @@ -39,6 +39,7 @@ _snapshot_nested_objectrefs, ) from xtuner.v1.rl.rollout.controller import RolloutControllerProxy +from xtuner.v1.rl.rollout.trace_store import release_and_discard_rollout_groups from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.rollout.worker_registry import WorkerLifecycleState from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace @@ -562,6 +563,7 @@ class BaseRLTrainer: eval_agent_loop_manager: AgentLoopManager _debug_train_files: dict[int, Path] _distillation: DistillationTrainerAdapter + _distillation_config: DistillationConfig | None def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: str) -> None: distillation_config = cfg.distillation_config @@ -570,6 +572,7 @@ def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: distillation_config = distillation_config.resolve_teacher_endpoints(endpoint_map) self._distillation = DistillationTrainerAdapter(distillation_config) + self._distillation_config = distillation_config check_fa3() self._init_work_dir_and_meta(cfg, meta_path) @@ -675,6 +678,7 @@ def _init_train_worker_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> cfg.train_worker_cfg.load_from = cfg.load_from cfg.train_worker_cfg.log_dir = log_dir cfg.train_worker_cfg.train_teacher_manager_config = self._distillation.train_teacher_manager_config + cfg.train_worker_cfg.distillation_config = self._distillation_config self._train_worker_cfg = cfg.train_worker_cfg def _init_rollout_config(self, cfg: BaseRLTrainerConfig, log_dir: Path) -> None: @@ -922,6 +926,63 @@ def _release_trace_sessions_after_train_batch(self, train_batch: list[list[Rollo """Release training traces when no concurrent rollout owner exists.""" self._release_all_trace_sessions() + def _prepare_train_batch( + self, + train_batch: list[list[RolloutState]], + train_step: int, + produce_result: ProduceBatchResult | None = None, + ) -> tuple[int, int]: + """Drop invalid groups before training; preserve the labels from batch + retrieval. + + Evaluation never enters this path. Staleness is already handled by take_train_batch, and debug replay preserves + its saved label masks. + + ``ProduceBatchResult`` production counters and raw rewards describe the producer window (including groups that + may not be in this retrieved training batch); ``task_batch_sizes`` records requested buffer counts. Keep those + metrics unchanged when this method filters invalid training groups. The dropped groups are reported separately + through the returned counts and ``response/invalid_train_groups`` / ``response/invalid_train_samples``. + """ + valid_groups = [] + rejected_groups = [] + rejected_samples = 0 + for group in train_batch: + task_name = (group[0].task_name if group else None) or "train" + try: + validate_train_group(group, task_name=task_name) + except ValueError as exc: + self.logger.warning( + f"Discard invalid training group: task={task_name!r}, " + f"rollout_ids={[state.rollout_id for state in group]}, reason={exc}" + ) + rejected_groups.append(group) + rejected_samples += len(group) + continue + + valid_groups.append(group) + + if rejected_groups: + asyncio_run(release_and_discard_rollout_groups(rejected_groups)) + train_batch[:] = valid_groups + if produce_result is not None: + produce_result.rollout_states[:] = valid_groups + if produce_result is not None and produce_result.task_results is not None: + kept_states = {id(state) for group in valid_groups for state in group} + for result in produce_result.task_results.values(): + result.rollout_states[:] = [ + group for group in result.rollout_states if group and id(group[0]) in kept_states + ] + self.logger.info( + f"Training field check: kept_groups={len(valid_groups)}, " + f"discarded_groups={len(rejected_groups)}, discarded_samples={rejected_samples}" + ) + if not valid_groups: + raise ValueError( + f"No valid training data remains after checking training fields: train_step={train_step}, " + f"discarded_groups={len(rejected_groups)}, discarded_samples={rejected_samples}" + ) + return len(rejected_groups), rejected_samples + def _train_one_batch( self, train_batch: list[list[RolloutState]], @@ -932,7 +993,9 @@ def _train_one_batch( onload_train_before_train: bool = False, raw_rewards_sum: float = 0.0, raw_rewards_count: int = 0, + produce_result: ProduceBatchResult | None = None, ) -> TrainInfo: + dropped_groups, dropped_samples = self._prepare_train_batch(train_batch, train_step, produce_result) train_sample_count = sum(len(group) for group in train_batch) self.logger.info(f"generate {train_sample_count} samples for training") @@ -963,6 +1026,8 @@ def _train_one_batch( raw_rewards_count=raw_rewards_count, ) + data_info["invalid_train_groups"] = dropped_groups + data_info["invalid_train_samples"] = dropped_samples self._release_trace_sessions_after_train_batch(train_batch) return { @@ -1180,8 +1245,6 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P valid_groups = [] for group in data_groups: - if not is_valid_for_training(group, self.logger): - continue valid_groups.append(group) for data in group: reward = data.reward.get("score") if data.reward is not None else None @@ -1230,7 +1293,7 @@ def _save_trajectories(self, data_groups: list[list[RolloutState]], save_path: P "label": ground_truth, "response": response, "reward": reward, - "prompt_len": data.num_tokens, + "prompt_len": data.extra_fields.get("train_prompt_length", data.num_tokens), "response_len": len(response_ids), "reward_payload": data.reward, "agent": { @@ -1295,7 +1358,7 @@ def _save_eval_trajectories(self, data_groups: list[list[RolloutState]], save_pa "status": data.status.value if hasattr(data.status, "value") else str(data.status), "prompt": data.message, "response": data.response or "", - "prompt_len": data.num_tokens, + "prompt_len": data.extra_fields.get("train_prompt_length", data.num_tokens), "response_len": len(response_ids), "label": ground_truth, "reward": reward, @@ -1522,6 +1585,7 @@ def _fit(self): onload_train_before_train=True, raw_rewards_sum=produce_result.raw_rewards_sum, raw_rewards_count=produce_result.raw_rewards_count, + produce_result=produce_result, ) else: self._save_debug_rollout_batch(train_batch, train_step) @@ -1822,6 +1886,7 @@ async def _fit(self): step_timer_dict, raw_rewards_sum=produce_result.raw_rewards_sum, raw_rewards_count=produce_result.raw_rewards_count, + produce_result=produce_result, ) sync_model_step = train_step