Trainer & Training Stack

UniRL runs on one trainer architecture: a single-controller Remote / placement layer with per-domain trainers and a pluggable FSDP train stack. The earlier v1 actor-group runtime (and its conf_v1/ recipes) has been retired; everything below is the current, default path.

Entrypoints and trainers

Each domain has its own Hydra entrypoint and <Domain>Trainer (unirl/trainer/). They share construction utilities and train-side contracts, but their lifecycles and placement layouts differ:

trainer/base.py owns shared facilities such as the Ray DevicePool, checkpointing, logging, and transport cleanup. Each concrete trainer owns its workers, placement blocks, and phase ordering.

Lifecycles and placement

• Synchronous AR uses one full-pool placement. Direct rollout shares the train pipeline; a dedicated colocated engine adds weight sync and phase offload/onload.
• Synchronous diffusion supports either a colocated placement or distinct train and rollout placements selected by layout; reward_fraction > 0 creates a third reward placement.
• SFT uses one placement for bundle, pipeline, backend, supervised algorithm, stack, and track builder. Its lifecycle is records → build Part → train_track: there is no rollout, reward, advantage, or weight sync.
• Async AR and async diffusion use exactly two disjoint placement slabs: train (including reward, algorithm, stack, and sync sender) and a resident dedicated rollout engine. Their manager buffers completed Samples and publishes versioned weights at explicit barriers.
• PE uses one placement containing AR and diffusion sides. Each trainable side has its own bundle, pipeline, backend, algorithm, and TrainStack; with freeze_llm, the AR side is rollout-only. A composed rollout generates the ordered AR and diffusion Parts, then reward propagation and per-Part advantages precede sibling stack updates.
• Unified model uses one shared bundle/backend and UnifiedModelTrainStack for AR + image losses on one backbone. Its rollout may be one engine or an AR/DiT engine pair; it is not PE's sibling-stack topology.
• Agentic uses one colocated placement and requires a dedicated engine plus TensorWeightSync. It runs a barrier batch through RolloutManager, scores terminal answers, assigns one group advantage to every generated assistant Part, and concatenates those Parts for training.

A common single-stage colocated RL wiring looks like this:

<Domain>Trainer(num_devices, batch_size, ...cfg blocks...)
  with placement(pool, fraction=1.0, shared_workers=True):
    bundle    = remote(bundle_cfg)
    pipeline  = remote(pipeline_cfg, bundle=bundle)
    backend   = remote(backend_cfg, bundle=bundle)            # FSDPBackend
    rollout   = remote(rollout_cfg[, pipeline=pipeline])      # engine
    reward    = remote(reward_cfg)                            # RewardService
    algorithm = remote(algorithm_cfg, pipeline=pipeline)      # StageAlgorithm
    stack     = remote(stack_cfg, fsdp_backend=backend, algorithm=algorithm)
    sync      = remote(sync_cfg, backend=backend, rollout=rollout)   # dedicated rollout only

The single-controller layer (unirl/distributed/group/: Remote, placement, RankInfo) carries DP / TP / PP / SP / EP rank information, which is what later parallelism work (SP for long sequences, EP for MoE) plugs into. This sketch does not describe SFT, two-slab async trainers, PE, unified-model, or reward-slab diffusion. trainside recipes share the trained module (no sync); dedicated engines add a weight-sync bridge under sync.

FSDP backend and train stack

Two Remote siblings split the training work:

• FSDPBackend (unirl/train/backend/fsdp/backend.py) owns the trainable model: structural injection (LoRA / NFT / mirror EMA) before the FSDP2 wrap, plus the optimizer, LR scheduler, EMA shadow, eval-EMA swap, checkpoint, and onload/offload.
• TrainStack is implemented in unirl/train/stack/base.py and exported by unirl.train.stack. It owns loss/backward sequencing and takes handles to one FSDPBackend and one StageAlgorithm, and per rollout Part runs prepare_segment → mini-batch compute_loss_and_backward loop → optimizer_step, returning a TrainStepResult.

TrainStack is single-stage by design — one Part, no part-name dict. PE uses sibling TrainStacks, one per trainable generated Part. HunyuanImage3's mixed AR + diffusion training uses unified_model_stack.py, a multi-stage variant.

num_updates_per_batch > 1 runs several PPO mini-batch updates on one rollout shard with π_old frozen once by prepare_segment; it is gated on StageAlgorithm.supports_multi_update.

The full backend contract — components, the FSDPConfig / LoraConfig / EmaLoraConfig / EmaFullConfig / OptimizerConfig / LrSchedulerConfig schemas, and TrainTopology — is documented in the generated Train Stack README.

Config shape

The YAML files under examples/ form a domain tree, with additional model subdirectories where needed (for example examples/diffusion/sd3/sd3_trainside.yaml). Each recipe instantiates its siblings by _target_ (no Hydra config-group overrides):

bundle:
  _target_: unirl.models.sd3.bundle.SD3Bundle.from_config
  config:
    _target_: unirl.models.sd3.config.SD3PipelineConfig
    pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL}

pipeline:
  _target_: unirl.models.sd3.pipeline.SD3Pipeline
  shift: 3.0

backend:
  _target_: unirl.train.backend.fsdp.FSDPBackend
  block_class_names: ["JointTransformerBlock"]
  trainable_attr: transformer
  fsdp_cfg:
    _target_: unirl.train.configs.FSDPConfig
    fsdp_mode: full
    activation_checkpointing: false
  optimizer_cfg:
    _target_: unirl.train.backend.base.OptimizerConfig
    learning_rate: 3.0e-4
  lora_cfg:
    _target_: unirl.train.configs.LoraConfig
    rank: 32
    alpha: 64

stack:
  _target_: unirl.train.stack.TrainStack
  micro_batch_size: 1
  max_grad_norm: 1.0
  num_updates_per_batch: 2

Distributed backends

FSDPBackend is the standard path. VeOmniBackend is also available for topology-specific FSDP2, sequence-parallel, and expert-parallel recipes; current VeOmni recipes require fsdp_mode=full, mixed precision enabled, and CPU offload disabled. Use the existing *_veomni_* recipes rather than treating these constraints as generic command-line toggles.

For extension routing (which surface to build on), see Extending UniRL.