Skip to content

refactor: split unirl core from self-contained recipe packages#188

Draft
celve wants to merge 1 commit into
Tencent-Hunyuan:mainfrom
celve:LIN-532/main
Draft

refactor: split unirl core from self-contained recipe packages#188
celve wants to merge 1 commit into
Tencent-Hunyuan:mainfrom
celve:LIN-532/main

Conversation

@celve

@celve celve commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

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>Trainer classes (unirl/trainer/), the six thin Hydra entrypoints
(unirl/train_*.py), and the run configs in a separate top-level examples/ tree. This
change makes unirl/ a pure component library and moves each training task into a
self-contained package under a new top-level recipes/ — trainer, entrypoint, and
its configs together (the verl recipe/ convention).

  • recipes/<task>/ = trainer.py (the <Domain>Trainer merged with its @hydra.main
    main()), a 2-line __main__.py shim, and configs/. Tasks: diffusion, ar,
    async_ar, pe, refl (promoted to first-class), unified_model.
  • Launch becomes python -m recipes.<task> --config-name=<config>, with config-names
    dropping the now-redundant domain prefix (e.g. sd3/sd3_trainside).
  • BaseTrainer + build_sampling_dict move to unirl/train/base_trainer.py (the shared
    framework base recipes subclass). Nothing in unirl/ imports a recipe.
  • examples/ is dissolved into each recipe's configs/; the shell launchers move to
    scripts/. Packaging (pyproject/setup.py) ships recipes + config package-data,
    and scripts/check_recipe_targets.py scans recipes/.

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.

  • Static: python -m compileall -q recipes unirl scripts;
    python scripts/check_recipe_targets.py (1710 _target_ paths resolve);
    bash -n scripts/*.sh.
  • Compose-check — all 6 recipes pass (rc=0) (imports the full stack + resolves the
    packaged config_path and prefix-less config-names):
    python -m recipes.diffusion     --config-name=sd3/sd3_trainside                    --cfg job
    python -m recipes.ar            --config-name=qwen_vl_grpo_geo3k_mc_4x8            --cfg job
    python -m recipes.async_ar      --config-name=qwen3_grpo_4b_base_dapo_sglang_async --cfg job
    python -m recipes.pe            --config-name=pe_trainside_pickscore              --cfg job
    python -m recipes.refl          --config-name=refl_sd3                            --cfg job
    python -m recipes.unified_model --config-name=hi3_vllmomni                        --cfg job
    
  • Training smoke — real diffusion-RL loop:
    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.

ruff was unavailable in the test env — run pre-commit run --all-files before merge for
the final lint/format pass.

Compatibility / Risk

Breaking layout change — the user-facing launch surface changes:

  • Launch: python -m unirl.train_<domain>python -m recipes.<task>.
  • Config-names drop the domain prefix: --config-name=diffusion/sd3/sd3_trainside
    --config-name=sd3/sd3_trainside.
  • Launchers move: bash examples/run_experiment_*.shbash scripts/…;
    ENTRY=train_arENTRY=ar.
  • examples/, unirl/trainer/, and unirl/train_*.py are removed.
  • No change to model/checkpoint/data formats or the training math — trainer class
    logic is unchanged and each main() is a verbatim move of the old entrypoint's
    config→kwargs mapping.
  • Hydra note: packaged recipes need configs/__init__.py (included) so config_path
    resolves when run via python -m — this was surfaced and fixed via the compose-check.

Reviewer Notes

  • Draft — opening for discussion of the split before finalizing.
  • The diffusion smoke exercises the reference loop; the other five recipes are
    compose-validated (same BaseTrainer + merged-main() pattern), not smoke-run, since
    each needs a different large model.
  • async_ar is a variant sharing ARTrainer; refl (ReFL / DRaFT-K) is promoted to a
    first-class recipe.
  • Possible follow-up (not in this PR): collapse each main()'s config→kwargs mapping into
    a Trainer.from_cfg(cfg) classmethod.

Checklist

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

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.
@github-actions github-actions Bot added the wip Draft / work in progress label Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

wip Draft / work in progress

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant