Skip to content

refactor(trainer): separate lifecycle and loop programs - #25

Draft
leviking98z-rgb wants to merge 1 commit into
refactor/execution-plan-capabilitiesfrom
refactor/trainer-loop-programs
Draft

refactor(trainer): separate lifecycle and loop programs#25
leviking98z-rgb wants to merge 1 commit into
refactor/execution-plan-capabilitiesfrom
refactor/trainer-loop-programs

Conversation

@leviking98z-rgb

Copy link
Copy Markdown
Owner

Summary

Separates trainer lifecycle and outer-loop policy from model/domain wiring.

  • Adds typed LoopSpec / LoopState, an ordered TrainerLifecycle, and three explicit program families: BatchRLProgram, AgenticRLProgram, and SFTProgram.
  • Migrates synchronous AR, diffusion, PE, and unified-model loops; buffered async AR/diffusion loops; barrier/partial/async agentic loops; and SFT cursor handling.
  • Makes LoopKind select the program centrally. Only BaseTrainer and the intentionally distinct SFTTrainer now define train().
  • Preserves resume force-sync, trainer-specific evaluation coordinates, HI3's first-sync/debug state, async quiescence before eval/save/sync, agentic carry/drop tails, SFT sidecar cursors, and primary-error precedence during cleanup.
  • Moves idempotent rollout/device-pool shutdown into the shared lifecycle and extends the CPU contract guard so copied trainer loops cannot silently return.

This is stacked on fork Draft PR #24 (refactor/execution-plan-capabilities).

Related Issue

N/A

Test Plan

  • pre-commit run --all-files — all hooks passed, including recipe targets, architecture boundaries, and framework contracts.
  • python lint/check_framework_contracts.py7 rollout engines, 6 weight syncs, 19 model pipelines, 2 train backends, 2 wire types, 4 loop programs, 13 entrypoints conform.
  • python -m compileall -q unirl/trainer unirl/config/execution.py — passed.
  • Two one-off stdlib fake-trainer harnesses, run inline with python - <<'PY' and not committed per repository policy:
    • synchronous resume/sync/eval/save ordering, cleanup-error propagation, primary-error preservation, and buffered-async quiescence — passed;
    • partial agentic, resident async agentic, and SFT cursor/evaluation/checkpoint ordering — passed.
  • GPU training was not run; this Draft needs representative sync, async, agentic, and SFT smoke runs before it is ready for upstream review.

Compatibility / Risk

  • No recipe, checkpoint, wire-data, or entrypoint changes.
  • Existing save_mode defaults remain full for async AR/agentic and auto elsewhere through DEFAULT_SAVE_MODE.
  • Successful train() completion now consistently shuts down trainer-owned rollout engines and the DevicePool; shutdown is idempotent for the AR signal-handler path.
  • Async ordering is high risk: the program performs one quiesce before coincident eval/save/sync boundaries, then preserves the original eval → save → sync order.
  • Official PR refactor(rollout): engine-level sync/async contracts for async rollout Tencent-Hunyuan/UniRL#287 replaces AsyncRolloutScheduler with engine-owned async contracts and overlaps the async trainer files. If it lands first, this branch should be rebased by adapting BatchRLProgram.run_async and AgenticRLProgram.run_async to its quiesce contract.

Reviewer Notes

  • Official Draft PR refactor(trainer): extract the sync-loop template into BaseTrainer Tencent-Hunyuan/UniRL#282 migrates only synchronous AR and diffusion into a BaseTrainer template and explicitly leaves PE, unified, SFT, async, and agentic follow-ups. This branch is the broader program-family design rather than a duplicate.
  • Official Draft PR fix(trainer): preserve async diffusion teardown cause Tencent-Hunyuan/UniRL#239 covers only async-diffusion teardown error precedence; TrainerLifecycle generalizes that rule across every program family.
  • The diff removes more than 850 lines of per-trainer implementation and adds the reusable programs, lifecycle, docs, and permanent guard; net repository change is about -80 lines. The architectural benefit is that new trainers add one-step semantics without copying an outer loop.
  • AI-assisted. The complete diff, overlap checks, and listed validation results were reviewed before publication.

Checklist

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

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

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant