refactor: split unirl core from self-contained recipe packages#188
Draft
celve wants to merge 1 commit into
Draft
refactor: split unirl core from self-contained recipe packages#188celve wants to merge 1 commit into
celve wants to merge 1 commit into
Conversation
Move the per-domain training orchestration out of the `unirl/` library into a
top-level `recipes/` package — one self-contained package per task — so `unirl/`
becomes a pure component library and each recipe bundles its trainer, entrypoint,
and configs together. Launch is now `python -m recipes.<task>`.
- recipes/<task>/{trainer.py, __main__.py, configs/} for diffusion, ar, async_ar,
pe, refl, unified_model. Each trainer.py merges the former unirl/trainer/<task>.py
with its unirl/train_<task>.py entrypoint under one @hydra.main main(); __main__.py
is a 2-line shim so `python -m recipes.<task>` runs it. refl is promoted to a
first-class recipe.
- BaseTrainer + build_sampling_dict move to unirl/train/base_trainer.py — the shared
framework base a recipe subclasses.
- All configs relocate into each recipe's configs/ and config-names drop the now
redundant domain prefix (e.g. diffusion/sd3/sd3_trainside -> sd3/sd3_trainside).
The 3 misfiled diffusion configs are filed under their model subdir, the 2 async
configs split into async_ar/, and the stray make_v2v_smoke_dataset.py moves to
scripts/.
- Launchers move to scripts/; examples/ and unirl/trainer/ are removed.
- Packaging (pyproject/setup.py) ships `recipes` + its config package-data,
check_recipe_targets.py scans recipes/, and ruff isort knows `recipes` first-party.
- Sweep docs/tutorials/config headers/launchers: `python -m unirl.train_X` ->
`python -m recipes.X`, examples/ paths -> recipes/<task>/configs/, and the
recipe-vs-config vocabulary in the top-level and recipes READMEs.
Statically verified: compileall, check_recipe_targets (1554 _target_ paths resolve),
launcher `bash -n`, and default config_name resolution. Hydra compose-check and a
ruff/pre-commit pass are pending a dependency environment.
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Proposal (draft): split the repo into a library core and self-contained recipe packages.
Today the per-domain training orchestration lives inside the
unirl/library — the<Domain>Trainerclasses (unirl/trainer/), the six thin Hydra entrypoints(
unirl/train_*.py), and the run configs in a separate top-levelexamples/tree. Thischange makes
unirl/a pure component library and moves each training task into aself-contained package under a new top-level
recipes/— trainer, entrypoint, andits configs together (the verl
recipe/convention).recipes/<task>/=trainer.py(the<Domain>Trainermerged with its@hydra.mainmain()), a 2-line__main__.pyshim, andconfigs/. Tasks:diffusion,ar,async_ar,pe,refl(promoted to first-class),unified_model.python -m recipes.<task> --config-name=<config>, with config-namesdropping the now-redundant domain prefix (e.g.
sd3/sd3_trainside).BaseTrainer+build_sampling_dictmove tounirl/train/base_trainer.py(the sharedframework base recipes subclass). Nothing in
unirl/imports a recipe.examples/is dissolved into each recipe'sconfigs/; the shell launchers move toscripts/. Packaging (pyproject/setup.py) shipsrecipes+ config package-data,and
scripts/check_recipe_targets.pyscansrecipes/.Opening as a draft to gather feedback on the core/recipe split before finalizing.
Related Issue
N/A (internal tracking: LIN-532).
Test Plan
Verified on 8×H20 (TaiJi), branch
LIN-532/main, SD3.5-medium.python -m compileall -q recipes unirl scripts;python scripts/check_recipe_targets.py(1710_target_paths resolve);bash -n scripts/*.sh.packaged
config_pathand prefix-less config-names):python -m recipes.diffusion --config-name=sd3/sd3_trainside num_devices=8 +num_rollouts=10→ 10/10 rollouts, clean exit (rc=0), reward trend ~0.75 → ~0.83 (FlowGRPO improving the
SD3 policy); all progress logged from
recipes.diffusion.trainer.ruffwas unavailable in the test env — runpre-commit run --all-filesbefore merge forthe final lint/format pass.
Compatibility / Risk
Breaking layout change — the user-facing launch surface changes:
python -m unirl.train_<domain>→python -m recipes.<task>.--config-name=diffusion/sd3/sd3_trainside→--config-name=sd3/sd3_trainside.bash examples/run_experiment_*.sh→bash scripts/…;ENTRY=train_ar→ENTRY=ar.examples/,unirl/trainer/, andunirl/train_*.pyare removed.logic is unchanged and each
main()is a verbatim move of the old entrypoint'sconfig→kwargs mapping.
configs/__init__.py(included) soconfig_pathresolves when run via
python -m— this was surfaced and fixed via the compose-check.Reviewer Notes
compose-validated (same
BaseTrainer+ merged-main()pattern), not smoke-run, sinceeach needs a different large model.
async_aris a variant sharingARTrainer;refl(ReFL / DRaFT-K) is promoted to afirst-class recipe.
main()'s config→kwargs mapping intoa
Trainer.from_cfg(cfg)classmethod.Checklist