Skip to content

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

Closed
celve wants to merge 1 commit into
mainfrom
LIN-532/main
Closed

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

Conversation

@celve

@celve celve commented Jul 7, 2026

Copy link
Copy Markdown
Owner

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 label Jul 7, 2026
@celve

celve commented Jul 7, 2026

Copy link
Copy Markdown
Owner Author

Retargeted to upstream: Tencent-Hunyuan#188

@celve celve closed this Jul 7, 2026
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