Skip to content

refactor(trainer): extract the sync-loop template into BaseTrainer - #282

Closed
haonan3 wants to merge 1 commit into
mainfrom
refactor/base-trainer-sync-template
Closed

refactor(trainer): extract the sync-loop template into BaseTrainer#282
haonan3 wants to merge 1 commit into
mainfrom
refactor/base-trainer-sync-template

Conversation

@haonan3

@haonan3 haonan3 commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

Summary

Extracts the synchronous training loop into a BaseTrainer template method. Nine trainers hand-copy the same outer skeleton today (~700–900 duplicated lines repo-wide, measured); the diffusion and AR train() bodies differ in exactly four lines. BaseTrainer.train() now owns the loop — load checkpoint → data fast-forward → wandb → baseline eval → per-rollout build / sync-gate / train_step / log / periodic eval / save → finally finish-wandb + on_train_end — with four hooks for the observed divergences: _wandb_extra, evaluate_baseline, evaluate_periodic, on_train_end (AR routes _shutdown_runtime through it, preserving its nested-finally ordering).

DiffusionTrainer and ARTrainer drop their train() bodies (line-equivalent to the template, including the resume force-sync semantics and per-module progress-logger names). PETrainer / UnifiedModelTrainer / SFTTrainer keep their own loops and migrate in follow-ups; async/agentic loops are intentionally out of scope (producer/consumer shapes, not this template). Independent of #210 (based on main).

Test Plan

Static (CPU): DiffusionTrainer.train is BaseTrainer.train and ARTrainer.train is BaseTrainer.train assert-verified; ARTrainer.on_train_end override confirmed; compose-check diffusion/sd3/sd3_trainside and ar/qwen3_grpo_4b_base_dapo_sglang rc=0; pre-commit run --all-files green.

Pending GPU parity before merge:

  • diffusion smoke (sd3_trainside, ~10 rollouts) — identical logs/cadence vs main, incl. a resume-from-checkpoint step (force-sync path);
  • AR smoke (qwen3 grpo, ~5 rollouts) with eval_interval>0 — baseline + periodic eval fire at the same steps; teardown runs on both clean exit and Ctrl-C.

Compatibility / Risk

No config or entrypoint changes; train(**kwargs) signature is unchanged for both migrated trainers. Behavioral risk is confined to the four hook seams — hence the explicit parity checklist above. Not-yet-migrated trainers are untouched.

AI-assisted; reviewed and directed by the maintainer.

Nine trainers each hand-copy the same outer loop (~700-900 duplicated
lines repo-wide); the diffusion and AR bodies differ in exactly four
lines. BaseTrainer now owns the loop as a template method (load ckpt →
data fast-forward → wandb → baseline eval → per-rollout build/sync-gate/
train_step/log/eval/save → finally finish-wandb + on_train_end), with
four hooks for the observed divergences: _wandb_extra,
evaluate_baseline, evaluate_periodic, on_train_end (AR routes
_shutdown_runtime through it, preserving its nested-finally ordering).

DiffusionTrainer and ARTrainer drop their train() bodies (the template
is line-equivalent, incl. the resume force-sync and the per-module
progress logger name); PETrainer / UnifiedModelTrainer / SFTTrainer keep
their own loops and migrate in follow-ups. Async/agentic loops are
intentionally out of scope — their producer/consumer shapes are not this
template.
@haonan3

haonan3 commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator Author

Closing by triage, not by verdict on the idea. The duplication is real and has grown since this was cut — ar / diffusion / pe / unified_model (and experimental refl) all hand-roll the same load→wandb→loop→save→finish skeleton. But the trainer surface is in an active evolution burst (AR PPO/GAE #259, value head #256, hi3 it2i #186), a template landed now adds hook-friction to every in-flight change, and this draft only migrated 2 of the 5 candidates with no GPU parity smoke.

Re-derive when the AR-PPO wave settles, scoped to all five sync trainers in one pass — the branch (refactor/base-trainer-sync-template) stays as reference.

@haonan3 haonan3 closed this Aug 2, 2026
@github-actions github-actions Bot removed the wip Draft / work in progress label Aug 2, 2026
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