Models

Where it fits: cross-cutting — the model code that both the rollout step (pipeline.generate) and the train step (stage.replay) run, on one shared set of weights. Full map: ../README.md.

The one idea that makes UniRL's RL correct: the stage runs the rollout forward (diffuse / autoregress) and the train forward (replay) over the same bundle.transformer — the one the FSDP backend mutates in place — so new_logp and old_logp land on identical, just-updated weights.

What it is

unirl/models/ holds one self-contained subpackage per model (sd3/, qwen3/, hunyuan_image3/, pe/, …). Each implements the shared bundle / pipeline / stage / conditions contract so the same model code, on the same weights, serves both the rollout engine (generate) and the train stack (replay).

Not to be confused with the repository-root models/ directory, which only holds local checkpoint and reward-model symlinks.

Why it exists

Rollout and training must score the same weights or the GRPO ratio is meaningless — and "same weights" is subtle, because the FSDP backend mutates bundle.transformer in place (LoRA injection, FSDP sharding). The package is built so both sides reach that one mutated object by attribute — the backend wraps bundle.<trainable_attr>, the engine evals pipeline.<stage_attr>.trainable_module() — instead of holding a private copy. That is the real reason a bundle is bare modules with no lifecycle logic: a self-loaded second copy (from_config) would never see the adapters or the optimizer step, so replay would silently score stale weights and the ratio would drift off 1.0 with no error. (The uniform contract across SD3 / Qwen3 / HunyuanImage3 is the secondary payoff — one backend drives them all by attribute name.)

How it works

A model package bridges three concerns through one shared bundle:

• Bundle (bundle.py) — a pure container of weights + tokenizer/scheduler, loaded by from_config. No lifecycle logic. The biggest models (HunyuanImage3, Bagel) add a from_meta_config + materialize() path so each rank loads only its shard of the large transformer; SD3 / Qwen3 load eagerly.
• Pipeline (pipeline.py) — the endomorphic generate entrypoint (generate(sample) -> sample): it reads ancestor primitives through sample.conditioning(), reads sampling params from a pre-forked generation Part, builds typed Conditions, runs the stage, and fills that Part.
• Stages (diffusion.py / ar.py) — the trainable units. DiffusionStage exposes diffuse (rollout), replay (train), and predict_noise_at_step (DiffusionNFT); ARStage exposes autoregress and replay. Each exposes trainable_module(), returning the nn.Module the FSDP backend wraps and engines eval-scope.
• Conditions (conditions.py) — typed Batch subclasses with from_dict/to_dict, so the pipeline works in typed form internally but stores the generic Dict[str, Condition] shape on the generated Part.

The bundle is shared, not duplicated: the trainer builds it once and injects it into both sides, so replay reads exactly the weights training updates — hence pipelines offer a bundle-injected constructor alongside the self-loading from_config. σ is engine-pinned: a diffusion pipeline reads the generation Part's DiffusionSamplingParams.sigmas verbatim and raises if it's None — it doesn't compute σ at generate time. Diffusion pipelines also implement latent_shape() so the trainer can author the byte-identical x_T recipe.

Extending it: a new model is a new unirl/models/<model>/ with config.py / bundle.py / diffusion.py|ar.py / conditions.py / pipeline.py mirroring an existing one — there is no registry, models are selected purely by _target_. Read .claude/skills/development/add-model-bundle/SKILL.md first — it is the authoritative bundle / pipeline / stage / conditions contract.

Gotchas

• trainable_module() is on the stage, not the bundle. The backend reaches the model via getattr(bundle, trainable_attr); engines via getattr(pipeline, stage_attr).trainable_module(). Keep both pointing at the same module.
• Bundles must stay pure containers — no LoRA, FSDP, autocast, or weight-sync logic; those are train-side lifecycle concerns.
• Share one bundle. For colocate/trainside runs the pipeline must be built from the injected bundle (not from_config), or replay reads a stale second copy of the weights.
• σ is engine-pinned — a diffusion pipeline reads the generation Part's sigmas and raises if it's None; never build the σ tensor inside generate.
• CFG empty-negative differs per model (SD3 "", Qwen-Image " ") — use the model's canonical upstream value or the rollout/replay ratio drifts off 1.0.
• HunyuanVideo-1.5 prompt-template whitespace is tokenizer state. Keep PROMPT_TEMPLATE_SYSTEM_MESSAGE byte-identical to upstream because mllm_crop_start=108 is tied to that exact prefix length; collapsing its indentation drops user-prompt tokens.
• Work that needs real storage goes through types/post_materialize.py. A bundle or a structural injector (LoRA / NFT / mirror) may run while the module is still on the meta device, where writing a tensor is a no-op. Register such work with defer_after_materialize; the train backend drains it with apply_deferred_ops after the weight load. The pair describes model-construction state, not training policy, which is why it lives under models/types/ and not under train/.
• Name-keyed matching around materialization compares canonical names. Deferred ops (and materialize()) run post-wrap, and checkpoint_wrapper (activation checkpointing) interposes a _checkpoint_wrapped_module. segment in named_parameters() names while state_dict() keys stay canonical (the wrapper's state-dict hooks strip it). Any match between module params and checkpoint keys goes through canonical_param_name, exported beside the deferred-op pair in types/post_materialize.py.

Conversation composition

unirl/models/types/conversations.py holds the transpose helpers with more than one model consumer: build_text_messages (qwen3, qwen3_5), build_vision_messages (qwen_vl, qwen3_5), and the two composition rules both use — system_prefix and group_consecutive_roles. The sglang wire render imports those two rather than keeping private copies, so the rules have one definition.

• One system rule: the config system_instruction is a default; an explicit system turn in the data wins. A path that wants "data is the sole source of truth" leaves system_instruction unset in the recipe — not a code branch.
• A render with one model consumer stays in that model's package and imports the shared rules (Qwen3-Omni's URI-media render lives in unirl/models/qwen3_omni/media.py, next to the media contract it enforces). It moves into the shared layer when a second consumer actually exists.
• Model chat stages own encoding — chat template, processor, Bagel splits — never the composition rules.

Support matrix

Model and algorithm support are independent dimensions (see ../algorithms/README.md); this table is the model dimension. A row is runnable only when the package and a recipe under examples/ exist for the listed entrypoint and rollout engine — a package existing by itself is not evidence of end-to-end support. The Model column lists the pretrained_model_ckpt_path defaults of the listed recipes, not marketing aliases; where a recipe has no Hub default, it names the env var to set and the checkpoint it expects, and the row is tagged checkpoint-path-required. Engines: trainside = in-process rollout on the training weights; sglang / sglang_diffusion / vllm / vllm_omni / fastvideo = external engines with a UniRL adapter (see ../rollout/engine/README.md).

Status legend

Restriction tags

• engine-specific — only the listed engines have a recipe; other engines have no adapter or no recipe and are unverified.
• scorer-dependent — the recipe's reward does not run from the base install alone. It needs a remote scorer served by unirl-reward-service through RemoteRewardBackend (e.g. editreward, wise, geneval, the hpsv2 / pickscore / hpsv3 / … bundle), a judge process in its own venv through ManagedScorerProcessBackend (editscore), an external judge endpoint (llm_judge → $JUDGE_URL), or a local scorer with its own install or checkpoint: hpsv3 / hpsv3pp (hpsv3 package), ocr ([eval] extra), videoalign (VIDEOALIGN_CKPT), imagebind (git install, CC-BY-NC-SA). Local scorers that only pull Hugging Face weights through transformers (pickscore, clip, clap, videopickscore, videoclipdelta) do not count. Registry: ../reward/local/registry.py.
• multi-node — the recipe's launch topology is NxM with N ≥ 2 nodes.
• LoRA-only — every shipped recipe injects LoRA; full fine-tuning has no recipe.
• checkpoint-path-required — the recipe has no Hugging Face default; set PRETRAINED_MODEL (or the named env var).

Image diffusion

Video diffusion

Unified and multimodal generators

Autoregressive LLM / VLM / omni

Composed pipelines and workflows

Not supported (no package): the sglang_diffusion adapter registry also exposes FluxAdapter (plain FLUX) and MochiAdapter, but there is no unirl/models/ package and no recipe for either, so they are not model rows. unirl/models/types/ is the shared contract layer, not a model.

Adding a row. A new model is a new package following .claude/skills/development/add-model-bundle/SKILL.md plus at least one recipe per examples/README.md#adding-or-editing-a-recipe; add the row here and the short form in the root README.md only once both exist, and list only the engines a recipe actually exercises.

Source: unirl/models/README.md — edit the README next to the code, then run npm run sync:readmes from docs/.