Skip to content

[WIP] Refactor prepare_train_data to reduce centralized data transfer - #2084

Open
YanhuiDua wants to merge 3 commits into
InternLM:mainfrom
YanhuiDua:refactor_prepare_train_data
Open

YanhuiDua wants to merge 3 commits into
InternLM:mainfrom
YanhuiDua:refactor_prepare_train_data

Conversation

@YanhuiDua

@YanhuiDua YanhuiDua commented Sep 11, 2026

Copy link
Copy Markdown
Collaborator

核心思想:

agent_loop 数据生产结束后,先调用 prepare_training_artifacts 准备训练格式,然后将数据直接放到共享存储中,只传递 metadataRLTrainerTrainController

后续开发:
重构pack策略,pack 时仅依赖 meta 信息:#2055

待优化项:

  • agent loop actor 每个rollout worker绑定一个,缓解中心化的 agent loop 放到共享存储的 IO 压力
  • 数据放到共享存储时可以有更好的序列化/反序列化方式
  • 接入 transfer queue 作为共享存储

…t batches

- Add release_storage/best_effort ownership options to RolloutMetadata.to_rollout_state
  and discard_rollout_state; extract release_rollout_metadata_storage for outer-ref release.
- Restore complete states from Object Store in the trainer via _restore_rollout_batch and
  release them deterministically in finally blocks for train/eval/debug paths.
- Consume agent-loop prepared next-token-aligned input_ids/labels in _prepare_train_data
  and read rewards/lengths/tool_turns from metadata.
- Validate ReplayBuffer.put input is RolloutMetadata with storage for completed samples.
- Apply the token shift in localhost/sandbox loop preparation; release storage on sampler
  rerollout; generalize deterministic sorting and advantage response_len to metadata.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant