[Refactor] Move train data preparation into AgentLoop and TrainingController by YanhuiDua · Pull Request #2102 · InternLM/xtuner

YanhuiDua · 2026-09-20T11:58:13Z

Summary

Training data construction is split by ownership: per-sample token fields are now built at generation time inside AgentLoop, while step-level conversion (validation, advantage, shift, tensorization, seq_ctx/teacher fields, data_info) moves into TrainingController. BaseRLTrainer._prepare_train_data is deleted.

AgentLoop.canonicalize_train_fields (base + localhost/sandbox overrides) builds the unified full-sequence input_ids/labels/logprobs convention; semantic holes (tool/env tokens) are baked into labels as -100 by the loops themselves.
RolloutState drops response_mask; labels becomes the only supervision carrier. New agent_loop_type records the producing loop class name (type(self).__name__), and AGENTIC_AGENT_LOOP_TYPES discriminates agentic full-sequence samples (None/unregistered types fall back to prompt+response semantics).
TrainingController.fit accepts list[list[RolloutState]] and absorbs group validation, session-clustered advantages, the one-position shift, tensorization, seq_ctx/teacher fields and data_info statistics.
calculate_group_effective_response_masks now bakes token staleness into labels in place (clearing is monotone, so replay-buffer expiry checks and the train-batch bake converge to identical labels); the agentic exclusion condition now keys on agent_loop_type instead of input_ids/labels presence, which would have silently disabled token staleness after canonicalization.

Behavior changes

prompt+response samples whose field construction fails are marked FAILED at generation time (their group is skipped) instead of asserting at training time.
Missing reward with task_adv_weight == 0 trains with zero advantage; with weight > 0 it still raises (crash semantics preserved, location moved).
The prepare_data timing key disappears; conversion cost is now inside the training timer.
prompt_len/response_len stats keep exact denominators for reasoning samples; only mixed-form edge cases may drift by +/-1 (log metrics only).

ProduceBatchResult impact: leftover_*/raw_rewards*/produced_* counts unchanged; generation-time canonicalize failures can flip individual samples to FAILED, which converts to group-level skips via validation (affects rollout_states content, not the accounting).
RoutedExperts impact: not affected; routed_experts is still attached to seq_ctx by the controller, ownership and release paths unchanged.
Ray concurrency impact: not affected; no agent-loop actor methods, decorators, or concurrency groups changed.

Test plan

tests/rl/test_prepare_train_data.py rewritten as contract tests for canonicalize_train_fields (base + both agentic overrides) and TrainingController._convert_rollout_groups, including the packed-alignment regression (advantage used to be one position longer than input_ids).
Staleness, replay-buffer, producer, multi-task manager and trainer tests updated to the labels-only contract; added regression cases: repeated staleness bakes converge, reset_rollout_response clears canonical train fields, canonicalized prompt+response states remain staleness-eligible.
pre-commit (ruff, ruff-format, docformatter, pyupgrade, mypy strict) passes on all touched files; unit tests run in CI (pytest tests/rl).

YanhuiDua · 2026-09-21T06:24:54Z

@claude review

github-actions · 2026-09-21T06:53:40Z

Claude:

Summary

本次重构将训练数据构造按归属拆分：per-sample 全序列 token 字段（input_ids/labels/logprobs）在生成期由 AgentLoop.canonicalize_train_fields（基类 + localhost/sandbox/gsm8k/verl 各自覆写）构造，step 级转换（组校验、session 聚类 advantage、shift、张量化、seq_ctx/teacher 字段、data_info）下沉到 TrainingController；RolloutState.response_mask 被删除，labels 成为唯一监督载体，token staleness 直接烙进 labels。抽象与线性业务流程组织良好，未发现 Warning 及以上问题。

核心原理实现与单测

Loop 侧 canonicalize：基类为 prompt+response 型样本装配全序列并对失败样本置 FAILED（隔离到单样本、不上抛 producer）；agentic 覆写只做统一约定校验。核心行为由 tests/rl/test_prepare_train_data.py 的 TestAgentLoopCanonicalizeTrainFields/TestAgenticLoopCanonicalizeTrainFields 经 public canonicalize_train_fields 覆盖（含 Tensor flatten、tokenizer 兜底、VLM train_prompt_ids、失败隔离），mock 仅限项目外依赖（lagent、tokenizer）。
Controller 侧转换：_rollout_groups_to_colate_items 以清晰的 Phase 1-4 线性流程组织，advantage/shift/张量对齐由 TestConvertRolloutGroups 覆盖，并含 packed 对齐回归（advantage 曾比 input_ids 长 1）；test_rl_colocate_trainer_integration.py 经 public TrainingController.fit(rollout_groups, ...) 覆盖端到端真实路径。
Staleness 烙制：calculate_group_effective_response_masks 原位清 -100（单调收敛），agentic 全序列样本首次纳入 token 级 staleness，由 test_multi_task_agent_loop_manager.py 与 test_staleness_policy.py 经 public produce_batch/calculate_group_effective_response_masks 覆盖（含全序列逐 token 烙制与同步路径零操作）。

其他 Issues

未发现 Warning 及以上问题。

…troller - AgentLoop.canonicalize_train_fields (base + localhost/sandbox overrides) builds the unified full-sequence train fields (input_ids/labels/logprobs) at generation time; semantic holes are baked into labels by the loops. - RolloutState drops response_mask; labels become the only supervision carrier. agent_loop_type records the producing loop class name and AGENTIC_AGENT_LOOP_TYPES discriminates agentic full-sequence samples. - TrainingController.fit accepts list[list[RolloutState]] and absorbs validation, session-clustered advantages, shift/tensorization, seq_ctx, teacher fields and data_info; BaseRLTrainer._prepare_train_data is gone. - calculate_group_effective_response_masks now bakes token staleness into labels in place (monotone, convergent) and its agentic exclusion uses agent_loop_type instead of the input_ids/labels presence heuristic.

…ut conversion - response_ids now denotes the contiguous suffix of input_ids after the prompt (env/tool tokens included) in every loop; localhost/sandbox export input_ids[len(prompt_ids):] so response_model_steps stay aligned with the full response region. - Token staleness baking becomes branch-free: effective mask = semantic mask (labels != -100 on the suffix) * per-token staleness mask; the zero-prompt suffix state becomes eligible and agent_loop_type plus AGENTIC_AGENT_LOOP_TYPES are removed. - data_info stats: prompt_len reports the original prompt length and response_len the supervised (LLM-generated) token count; env/tool injected tokens count in neither. - _rollout_groups_to_colate_items is phase-split into session reward clustering, group advantage estimation, per-state ColateItem conversion and data_info summarization.

YanhuiDua · 2026-09-22T08:39:04Z

Following work: #2055

PengchengShi00 · 2026-09-22T09:00:08Z

            continue
        for group in batch_by_task.get(task.task_name, []):
-            effective_masks = calculate_group_effective_response_masks(
+            calculate_group_effective_response_masks(


这里调用了但没有接收以及后续的操作

已修复，改为返回 effective_masks 并且显式赋值

PengchengShi00 · 2026-09-22T09:28:12Z

+        )
        rollout_state.response_ids = final_response_ids
-        rollout_state.response_mask = final_response_mask
        rollout_state.logprobs = final_logprobs


这里没有对logprobs进行拼接，但后面要求assert相等了

Suggested change

rollout_state.logprobs = final_logprobs

rollout_state.logprobs = [0.0] * len(prompt_ids) + final_logprobs

PengchengShi00 · 2026-09-22T10:28:01Z

+        semantic_mask = output.response_mask if output.response_mask is not None else [1] * len(response_ids)
+        rollout_state.prompt_ids = prompt_ids
+        rollout_state.response_ids = response_ids
        rollout_state.logprobs = output.response_logprobs


这个少一个前置拼接0

PengchengShi00 · 2026-09-22T11:20:53Z

+        # Unified response_ids convention: the contiguous suffix of input_ids after the prompt
+        # (env-injected tokens included), aligned with response_model_steps for per-token
+        # staleness; also the token count used by rollout throughput logging.
+        prompt_len = len(rollout_state.prompt_ids or [])


这里用prompt_ids这个字段切分在多轮的时候会有问题吗

文本多轮应该也没问题，多轮的赋值过程如下：
第一轮：token = prompt_id
第二轮：token = token + response_ids
...

所以不会修改原始的prompt_id

另外多模的输入这里会有问题，我增加了显式的判断agent_in_localhost_loop不支持多模输入的判断

…ing, multimodal assert

YanhuiDua force-pushed the refactor/prepare-train-data-into-agent-loop branch from 7d6b5a9 to 4bb5b80 Compare September 20, 2026 12:09

YanhuiDua added 2 commits September 21, 2026 20:13

YanhuiDua force-pushed the refactor/prepare-train-data-into-agent-loop branch from 7c29de9 to 066bd24 Compare September 21, 2026 12:13

PengchengShi00 reviewed Sep 22, 2026

View reviewed changes

[Test] Align tests with labels-based staleness and sequence conventions

d92d059

YanhuiDua force-pushed the refactor/prepare-train-data-into-agent-loop branch from fbf2b05 to d92d059 Compare September 23, 2026 03:35

[Fix] Apply review fixes: pure staleness masks, agentic logprobs padd…

92ab26e

…ing, multimodal assert

Navigation Menu

[Refactor] Move train data preparation into AgentLoop and TrainingController - #2102

[Refactor] Move train data preparation into AgentLoop and TrainingController#2102
YanhuiDua wants to merge 4 commits into
mainfrom
refactor/prepare-train-data-into-agent-loop

YanhuiDua commented Sep 20, 2026

Uh oh!

YanhuiDua commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026

Uh oh!

YanhuiDua commented Sep 22, 2026

Uh oh!

PengchengShi00 Sep 22, 2026

Uh oh!

YanhuiDua Sep 23, 2026 •

edited

Loading

Uh oh!

PengchengShi00 Sep 22, 2026

Uh oh!

YanhuiDua Sep 23, 2026

Uh oh!

PengchengShi00 Sep 22, 2026

Uh oh!

YanhuiDua Sep 23, 2026

Uh oh!

PengchengShi00 Sep 22, 2026

Uh oh!

YanhuiDua Sep 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	rollout_state.logprobs = final_logprobs
	rollout_state.logprobs = [0.0] * len(prompt_ids) + final_logprobs

Conversation

YanhuiDua commented Sep 20, 2026

Summary

Behavior changes

Test plan

Uh oh!

YanhuiDua commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026

Summary

核心原理实现与单测

其他 Issues

Uh oh!

YanhuiDua commented Sep 22, 2026

Uh oh!

PengchengShi00 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

YanhuiDua Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PengchengShi00 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

YanhuiDua Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

PengchengShi00 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

YanhuiDua Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

PengchengShi00 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

YanhuiDua Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

YanhuiDua Sep 23, 2026 •

edited

Loading