Skip to content

refactor(train): share per-micro anchor preparation between train stacks - #495

Open
chi030303 wants to merge 2 commits into
Tencent-Hunyuan:mainfrom
chi030303:refactor/share-prepare-segment-anchors
Open

chi030303 wants to merge 2 commits into
Tencent-Hunyuan:mainfrom
chi030303:refactor/share-prepare-segment-anchors

Conversation

@chi030303

@chi030303 chi030303 commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

Summary

TrainStack and UnifiedModelTrainStack both froze π_old by slicing planned micros, calling algorithm.prepare_segment, rejecting a missing anchor_fields tensor, then concatenating in plan order. That stitch now lives in one helper next to _validate_anchor_contract. Each stack still builds its own micro_slices (Plan vs _optimizer_step_slices). old_logp_source=rollout, packed concat order, and the single-micro / non-recompute fast paths are unchanged.

Related Issue

Closes #439

Test Plan

  • python -m py_compile unirl/train/stack/base.py unirl/train/unified_model_stack.py — pass
  • Uncommitted packed-varlen smoke (not added under tests/): multiple micros with non-uniform packed lengths, token-budget planner concat order, num_updates_per_batch>1, missing declared field, old_logp_source=rollout whole-shard, TrainStack and UnifiedModelTrainStack — 6/6 passed
  • Live train: Not run; reason: refactor-only, no recipe or algorithm change

Compatibility / Risk

No config, checkpoint, data format, or public API changes. Unified missing-field errors now say "a micro produced None" instead of "a micro-slice produced None"; the stack prefix is unchanged.

Reviewer Notes

Proposed on #439 before opening this PR. No overlapping open PR touches this stitch (#476 edits eval_track only). AI-assisted implementation; I reviewed the diff. Follow-up commit cd2cbc2 is ruff-format only (error string on one line).

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

Extract the slice/prepare/cat stitch so TrainStack and UnifiedModelTrainStack
cannot drift as the anchor contract evolves. Each stack still builds its own plan.
@github-actions github-actions Bot added the need review Ready and waiting for review label Sep 20, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

need review Ready and waiting for review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

refactor(train): share per-micro anchor preparation between train stacks

1 participant