diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b975cf4af..f26ac5049 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -40,12 +40,10 @@ /unirl/train/ @CjhHa1 @celve @haonan3 /unirl/train/backend/ @celve @zzhuoxin1508 @CjhHa1 /unirl/train/stack/ @leviking98z-rgb @haonan3 -/unirl/train/refl/ @celve @Ideny42 /unirl/train/sft/ @haonan3 @xshrz /unirl/train/unified_model_stack.py @Ideny42 @zzhuoxin1508 /unirl/trainer/ @celve @Ideny42 @CjhHa1 /unirl/train_sft.py @haonan3 @xshrz -/unirl/train_refl.py @celve @Ideny42 /unirl/train_pe.py @Jayce-Ping @Ideny42 /unirl/train_unified_model.py @Ideny42 @zzhuoxin1508 diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 0d011a22f..4c4e376db 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -50,6 +50,19 @@ environment, not the code. Reward and actor share one Python process, so recipe `requirements.txt` files may only ADD packages, never re-pin the core stack. +## Deliberately not ported from the legacy core path + +Cut on the tier's minimalism rule (features return only with a recipe that +needs them; see the placement comment in `trainer.py`): + +- **Periodic eval** (`eval_interval` / `eval_sample` / `eval_cfg_text_scale` / + multi-reward `eval_rewards` suites, from #202) — the training reward curve is + this package's success signal; bring eval back via `unirl.trainer.eval_suites` + when a recipe needs checkpoint selection. +- **`reward_fraction` cross-slab reward placement** — colocated-only here. +- `param_checksum` debug probe, `activation_checkpoint_vae` knob (checkpointing + is unconditionally on in the grad decode paths). + ## Verification | Config | Hardware | Head | Status | diff --git a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml index e229d3435..9365a4856 100644 --- a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml +++ b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml @@ -18,7 +18,7 @@ max_grad_norm: 1.0 actor: _target_: experimental.refl.roles.ReflActorRole # Family selector — swap pipeline_target + model_config for another family, - # no code changes (mirrors ReFLPolicy's pipeline_target contract). + # no code changes (the pipeline_target contract). pipeline_target: experimental.refl.models.wan21.Wan21ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. diff --git a/experimental/refl/roles.py b/experimental/refl/roles.py index 2265695d6..d361bb0b1 100644 --- a/experimental/refl/roles.py +++ b/experimental/refl/roles.py @@ -1,6 +1,6 @@ """ReflActorRole — family-agnostic REFL/BPTT actor Remote for the refl recipe. -Mirrors :class:`unirl.train.refl.policy.ReFLPolicy` (the SD3 image-ReFL actor): +The family-agnostic REFL actor: a **config-chosen** ``Pipeline`` (``pipeline_target`` + ``model_config``, no per-family imports), FSDP-wrapped in place via ``FSDPBackend``, driven by three driver RPCs per step under the distributed ``enable_grad()`` context:: @@ -100,7 +100,7 @@ def __init__( def initialize(self) -> None: torch.cuda.set_device(self.device) # Default PG over the actor role's workers (env:// from Remote.setup's - # dist_env); FSDP2 fully_shard wraps over it. Same order as ReFLPolicy. + # dist_env); FSDP2 fully_shard wraps over it. Phase-0-validated order. if self.rank_info is not None and int(self.rank_info.world_size) > 1 and not dist.is_initialized(): dist.init_process_group(backend="nccl") diff --git a/experimental/refl/trainer.py b/experimental/refl/trainer.py index 3217fb3d8..a915d7add 100644 --- a/experimental/refl/trainer.py +++ b/experimental/refl/trainer.py @@ -1,7 +1,6 @@ """REFLTrainer — recipe driver for WAN REFL/BPTT (video reward backprop). -The video sibling of :class:`unirl.trainer.refl.RewardBackpropTrainer`: two -roles, always — a :class:`experimental.refl.roles.ReflActorRole` (FSDP WAN + +Two roles, always — a :class:`experimental.refl.roles.ReflActorRole` (FSDP WAN + grad BPTT sampling + optimizer) and a frozen differentiable video reward (:class:`unirl.reward.service.RewardService`), colocated on the same worker slab so decoded video never leaves the GPU. Each step runs, under the diff --git a/unirl/models/sd3/pipeline.py b/unirl/models/sd3/pipeline.py index ed00b0f99..26710acae 100644 --- a/unirl/models/sd3/pipeline.py +++ b/unirl/models/sd3/pipeline.py @@ -136,7 +136,8 @@ def build_conditions( ) -> SD3Conditions: """Encode prompts (+ optional CFG negatives) into ``SD3Conditions``. - Shared by :meth:`generate` and the ReFL draft path (``draft_generate``). + Shared by :meth:`generate` and grad-sampling adaptations (e.g. the + experimental ReFL pipeline). Applies SD3's empty-negative default (diffusers parity) when CFG is on and no negative was supplied — see the rationale quoted in :meth:`generate`. """ diff --git a/unirl/trainer/eval_suites.py b/unirl/trainer/eval_suites.py index 331ac319c..89c210ddf 100644 --- a/unirl/trainer/eval_suites.py +++ b/unirl/trainer/eval_suites.py @@ -25,7 +25,7 @@ Placement: :func:`build_eval_suites` must be called inside the SAME placement context that created the trainer's training reward — each suite reward becomes a sibling remote there, so where the trainer has a ``reward_fraction`` slab -(DiffusionTrainer, ReFL) ALL eval rewards share that dedicated-GPU slab, and +(DiffusionTrainer) ALL eval rewards share that dedicated-GPU slab, and elsewhere (PE, UnifiedModel) they colocate with the training reward. Data: an own-set suite instantiates its own driver-side data source — the @@ -34,10 +34,9 @@ format the trainer's data source reads (txt / JSONL / JSON manifests with metadata) works per suite. -Scoring uses the trainer's own reward interface: composed/rollout trainers call -``suite.reward.score_and_attach``, ReFL calls ``score_differentiable`` — a -suite's backend must support whichever its trainer uses (the same contract as -the training reward). +Scoring uses the trainer's own reward interface (every current consumer calls +``suite.reward.score_and_attach``) — a suite's backend must support whichever +its trainer uses (the same contract as the training reward). """ from __future__ import annotations