From fd3fcdca5bb136e6e932fe641ce0962990d802a2 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 22:59:42 +0800 Subject: [PATCH 1/3] feat(experimental): port SD3 image ReFL; remove the legacy core ReFL path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ReFL bypasses the core StageAlgorithm abstraction (no replay, no advantages) and per the tier design lives outside core. This lands the second family on the experimental.refl contract and deletes the legacy path, ending the two-implementation transition: - experimental/refl/models/sd3.py — Sd3ReflPipeline/Stage: the legacy DRaFT-K flow re-expressed through the package contract (diffuse_with_grad with the mid/final window + per-sample KL; decode_with_grad reuses the mainline grad decode, pixels in [0,1] for the core image scorers). CFG under BPTT is deliberately unsupported (guidance_scale must be 1.0, as the legacy path always ran). - experimental/refl/examples/sd3_pickscore_refl.yaml — ported from examples/diffusion/refl_sd3.yaml with hyperparameters preserved; reward comes straight from core (PickScore), demonstrating that a package config can pair core rewards without any package-local code. - Removed: unirl/train_refl.py, unirl/trainer/refl.py, unirl/train/refl/ (ReFLPolicy), unirl/models/draft.py (its only consumer), examples/diffusion/refl_sd3.yaml; trainer README rows updated. eval_suites stays (DiffusionTrainer/PETrainer use it). - pipeline_target is now the ONLY thing a config swaps between WAN video and SD3 image ReFL — the boundary is not WAN-shaped. Static: compose x3 rc=0; 2311 _target_ paths resolve; pre-commit green. Pending GPU: sd3_pickscore_refl smoke (checklist in PR). --- experimental/refl/README.md | 7 +- .../refl/examples/sd3_pickscore_refl.yaml | 72 ++--- experimental/refl/models/sd3.py | 239 ++++++++++++++++ unirl/models/draft.py | 148 ---------- unirl/train/refl/__init__.py | 12 - unirl/train/refl/policy.py | 236 ---------------- unirl/train_refl.py | 44 --- unirl/trainer/README.md | 11 +- unirl/trainer/refl.py | 256 ------------------ 9 files changed, 291 insertions(+), 734 deletions(-) rename examples/diffusion/refl_sd3.yaml => experimental/refl/examples/sd3_pickscore_refl.yaml (57%) create mode 100644 experimental/refl/models/sd3.py delete mode 100644 unirl/models/draft.py delete mode 100644 unirl/train/refl/__init__.py delete mode 100644 unirl/train/refl/policy.py delete mode 100755 unirl/train_refl.py delete mode 100644 unirl/trainer/refl.py diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 9c38aaeb9..33d2d2409 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -18,6 +18,10 @@ export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers \ DATA_PATH=/path/to/prompts.txt RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan21_t2v_videoalign_refl num_devices=8 +# SD3.5 T2I + PickScore reward (core scorer — no package-local reward) +export PRETRAINED_MODEL=/path/to/stable-diffusion-3.5-medium # or the HF default +RAY_ADDRESS=auto python -m experimental.refl.run --config-name=sd3_pickscore_refl num_devices=8 + # WAN 2.2 I2V + Face-identity reward (first frame via (image, condition) # MediaRef; face reference via per-sample metadata ref_video_path) pip install -r experimental/refl/reward/face/requirements.txt @@ -33,7 +37,7 @@ RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan22_i2v_face_re |---|---| | `trainer.py` | `REFLTrainer(BaseTrainer)` — driver: wiring + the 3-RPC train step | | `roles.py` | `ReflActorRole(Remote)` — family-agnostic actor (`pipeline_target` + `model_config`) | -| `models/` | Per-model BPTT adaptations subclassing the core pipelines (`types.py` defines the contract): `wan21.py`, `wan22.py` — mirrors `unirl/models/` (graduates into the matching model packages) | +| `models/` | Per-model BPTT adaptations subclassing the core pipelines (`types.py` defines the contract): `wan21.py`, `wan22.py`, `sd3.py` — mirrors `unirl/models/` (graduates into the matching model packages) | | `reward/` | Package-local differentiable rewards (VideoAlign, Face), each with an additive-only `requirements.txt` — mirrors `unirl/reward/` (graduates into it) | | `examples/` | Flat Hydra configs (repo-wide schema) — mirrors the top-level `examples/` (graduates into it) | @@ -54,3 +58,4 @@ core stack. | `wan21_t2v_videoalign_refl` (2-rollout smoke, full 81f/480x832 geometry) | 8xH20, fleet image | `40b3f4c9` | PASS — grads flow reward → VAE → DiT LoRA | | VideoAlign load + differentiable fwd/bwd on transformers 5.6.2 + peft 0.20 | 8xH20 (isolated venv) | `9087a671` | PASS — `grad_abs_mean=3.5e-3` | | `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | +| `sd3_pickscore_refl` (ported from the legacy core path; hyperparameters preserved) | 8xH20 | current head | pending smoke | diff --git a/examples/diffusion/refl_sd3.yaml b/experimental/refl/examples/sd3_pickscore_refl.yaml similarity index 57% rename from examples/diffusion/refl_sd3.yaml rename to experimental/refl/examples/sd3_pickscore_refl.yaml index b3f409753..a0844454d 100644 --- a/examples/diffusion/refl_sd3.yaml +++ b/experimental/refl/examples/sd3_pickscore_refl.yaml @@ -1,37 +1,31 @@ # @package _global_ -# ReFL (DRaFT-K direct differentiable-reward backprop) — SD3.5 trainside, 1x8. +# REFL SD3.5 T2I — PickScore reward (ported from the legacy core path, +# examples/diffusion/refl_sd3.yaml; hyperparameters preserved). # -# Two roles (always): ReFLPolicy (FSDP SD3 + grad DRaFT-K sampling + optimizer) -# and a frozen differentiable PickScore reward, placed via `reward_fraction` (the -# SAME unified knob as the GRPO DiffusionTrainer): reward gets its own disjoint -# tail slab, policy takes the rest. Gradient crosses the role boundary via the -# distributed enable_grad() context. Loss = -reward.mean(); NO advantages / -# replay / ratio / rollout-engine / weight-sync. -# -# Divisibility: batch_size must be divisible by BOTH role dp sizes. With -# reward_fraction=0.25 on 8 GPUs → policy dp=6, reward dp=2 → batch % 6 == 0. +# Second family on the experimental.refl contract: only pipeline_target, +# model_config, the reward backend, and the sampling geometry differ from +# the WAN configs. The reward comes straight from core +# (unirl.reward.local.pickscore) — no package-local reward needed. num_devices: 8 -batch_size: 12 # divisible by policy dp=6 and reward dp=2 +batch_size: 8 num_rollouts: 200 -reward_fraction: 0.25 # differentiable reward on its own 2 GPUs; policy on 6 -max_grad_norm: 1.0 save_interval: 0 +save_dir: ${oc.env:OUTPUT_DIR,outputs/sd3_pickscore_refl} save_mode: adapter +max_grad_norm: 1.0 -policy: - _target_: unirl.train.refl.policy.ReFLPolicy - pipeline_target: unirl.models.sd3.pipeline.SD3Pipeline # family selector — swap this + model_config for another family, no code - draft_num_steps: 1 # grad through the final step only (DRaFT-1) - reward_loss_scale: 1.0 - guidance_scale: 1.0 - num_inference_steps: 4 - height: 512 - width: 512 - seed: 42 - activation_checkpoint_vae: true +actor: + _target_: experimental.refl.roles.ReflActorRole + pipeline_target: experimental.refl.models.sd3.Sd3ReflPipeline block_class_names: ["JointTransformerBlock"] + # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. + reward_weight: 1.0 + reward_baseline: 0.0 + reward_scale: 1.0 + kl_weight: 0.0 strategy: + # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE. _target_: unirl.sde.kernels.FlowSDEStrategy model_config: _target_: unirl.models.sd3.config.SD3PipelineConfig @@ -44,14 +38,14 @@ policy: fsdp_cfg: _target_: unirl.train.configs.FSDPConfig param_dtype: bf16 + # master_dtype deliberately omitted → LoRA stays bf16 (uniform dtype); + # fp32 LoRA over a bf16 base trips torch FSDP2's uniform-dtype assert. cpu_offload: false mixed_precision: true fsdp_mode: full - reshard_after_forward: true # deeper ReFL grad graph — re-gather to save memory + reshard_after_forward: true activation_checkpointing: true - # master_dtype omitted → LoRA stays bf16 (uniform dtype). fp32 LoRA + bf16 base - # trips torch FSDP2's uniform-original-dtype assert; bf16 is portable. If small - # DRaFT grads underflow (reward-collapse), give LoRA its own FSDP group + fp32 (TODO). + use_torch_compile: false root_wrap: true optimizer_cfg: _target_: unirl.train.backend.base.OptimizerConfig @@ -101,13 +95,29 @@ data_source: data_path: datasets/pickscore/train.txt eval_data_path: datasets/pickscore/test.txt seed: 42 + shuffle: false algorithm: prompts_per_rollout: ${batch_size} +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + # 4 inference steps, grad through the final step only (DRaFT-1): + # mid_timestep = final_timestep = 3. + num_inference_steps: 4 + guidance_scale: 1.0 + height: 512 + width: 512 + eta: 0.0 + samples_per_prompt: 1 + seed: 42 + init_same_noise: false + sampler_kwargs: + mid_timestep: 3 + final_timestep: 3 + logging: report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} project_name: ${oc.env:WANDB_PROJECT,unirl-refl} - run_name: ${oc.env:WANDB_RUN_NAME,refl_sd3} - entity: ${oc.env:WANDB_ENTITY,null} - tags: ["sd3.5", "refl", "draft-k", "reward-backprop", "trainside"] + run_name: ${oc.env:WANDB_RUN_NAME,sd3_pickscore_refl} + tags: ["sd3.5", "refl", "draft-k", "pickscore"] log_media: false diff --git a/experimental/refl/models/sd3.py b/experimental/refl/models/sd3.py new file mode 100644 index 000000000..24c4f032c --- /dev/null +++ b/experimental/refl/models/sd3.py @@ -0,0 +1,239 @@ +"""SD3 image ReFL adaptation — the second family on the BPTT contract. + +Ported from the legacy core path (``unirl/trainer/refl.py`` + +``unirl/train/refl/policy.py`` + ``unirl/models/draft.py``, removed in this +change): the same DRaFT-K direct reward backprop, expressed through the +``experimental.refl`` contract so ``pipeline_target`` is the only thing a +config swaps between WAN video and SD3 image ReFL. + +Deliberate non-support: CFG under BPTT (``guidance_scale > 1``). The core +``SD3DiffusionStep`` batches both CFG branches through one forward, which +under grad would add a ``(1 - g) * d(uncond)/dθ`` term; the legacy path +always trained at ``guidance_scale == 1`` and so does this port — a wrong +config fails loudly. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any, Dict, Optional, Tuple + +import torch + +from experimental.refl.models.types import DiffuseWithGradResult +from unirl.models.sd3.conditions import SD3Conditions +from unirl.models.sd3.diffusion import SD3DiffusionStage +from unirl.models.sd3.pipeline import SD3Pipeline +from unirl.models.sd3.vae import SD3VAEDecodeStage +from unirl.train.lora import adapters_disabled +from unirl.types.primitives import Images, Texts +from unirl.types.sampling import DiffusionSamplingParams +from unirl.types.segments.latent import LatentSegment + +# Inclusive max for torch.Generator.manual_seed conventions. +MAX_TORCH_SEED = (1 << 63) - 1 + + +class Sd3ReflDiffusionStage(SD3DiffusionStage): + """SD3 diffusion stage + REFL BPTT sampling override. + + Reuses the mainline single-forward ``SD3DiffusionStep.predict_noise`` + (no CFG at ``guidance_scale == 1``); adds the grad-window loop with + optional per-step KL against the LoRA-disabled reference. + """ + + def generate_latents( + self, + batch_size: int, + latent_shape: Tuple[int, ...], + device: torch.device, + dtype: torch.dtype = torch.float32, + base_seed: Optional[int] = None, + ) -> torch.Tensor: + if base_seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(int(base_seed) % (MAX_TORCH_SEED + 1)) + return torch.randn(batch_size, *latent_shape, device=device, dtype=dtype, generator=generator) + return torch.randn(batch_size, *latent_shape, device=device, dtype=dtype) + + def diffuse_with_grad( + self, + conditions: SD3Conditions, + *, + schedule: torch.Tensor, + params: DiffusionSamplingParams, + initial_latents: Optional[torch.Tensor] = None, + ) -> DiffuseWithGradResult: + """Differentiable SD3 sampling for REFL-style BPTT training. + + Same knobs as the WAN stages (``params.sampler_kwargs``: + ``mid_timestep`` / ``final_timestep``; KL switched by the actor via + ``kl_weight``). Returns the live-grad ``z_final`` + per-sample + ``kl_loss`` ``[B]``. + """ + if float(params.guidance_scale) > 1.0: + raise ValueError( + "Sd3ReflDiffusionStage.diffuse_with_grad: CFG under BPTT is not supported " + f"(guidance_scale={params.guidance_scale}); train at guidance_scale=1.0." + ) + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Sd3ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None") + prompt_embeds = conditions.text.embeds + device = prompt_embeds.device + batch_size = int(prompt_embeds.shape[0]) + T = int(params.num_inference_steps) + schedule = schedule.to(device) + if int(schedule.shape[0]) != T + 1: + raise ValueError( + f"Sd3ReflDiffusionStage.diffuse_with_grad: schedule length {schedule.shape[0]} != T+1={T + 1}" + ) + self.strategy.init_schedule(schedule) + + latent_hw = (int(params.height) // self.vae_scale_factor, int(params.width) // self.vae_scale_factor) + latent_shape = (self.latent_channels, *latent_hw) + if initial_latents is not None: + if int(initial_latents.shape[0]) != batch_size or tuple(initial_latents.shape[1:]) != latent_shape: + raise ValueError( + f"Sd3ReflDiffusionStage.diffuse_with_grad: initial_latents shape " + f"{tuple(initial_latents.shape)} != ({batch_size}, {latent_shape})." + ) + latents = initial_latents.to(device=device, dtype=self.trajectory_dtype) + else: + latents = self.generate_latents( + batch_size=batch_size, + latent_shape=latent_shape, + device=device, + dtype=self.trajectory_dtype, + base_seed=int(params.seed) if params.seed is not None else None, + ) + + sk: Dict[str, Any] = dict(params.sampler_kwargs or {}) + mid_timestep = int(sk.get("mid_timestep", 0)) + final_timestep = int(sk.get("final_timestep", T - 1)) + kl_weight = float(sk.get("kl_weight", 0.0)) + if not (0 <= mid_timestep <= final_timestep < T): + raise ValueError( + f"Sd3ReflDiffusionStage.diffuse_with_grad: require 0 <= mid_timestep <= " + f"final_timestep < num_inference_steps, got mid={mid_timestep} " + f"final={final_timestep} T={T}." + ) + + autocast_ctx = ( + torch.autocast("cuda", self.autocast_dtype) + if device.type == "cuda" and self.autocast_dtype in (torch.float16, torch.bfloat16) + else nullcontext() + ) + sigma_max = float(schedule[1].item()) if int(schedule.shape[0]) > 1 else 0.99 + + transformer = self.model.transformer + kl_total = torch.zeros(batch_size, device=device, dtype=torch.float32) + kl_steps = 0 + + for i in range(T): + sigma = schedule[i].to(device) + sigma_next = schedule[i + 1].to(device) + grad_enabled = i >= mid_timestep + + pred_ctx = nullcontext() if grad_enabled else torch.no_grad() + with pred_ctx, autocast_ctx: + noise_pred = self.step.predict_noise( + self.model, + latents, + sigma, + conditions, + guidance_scale=1.0, + ) + noise_pred = noise_pred.float() + + if kl_weight != 0.0 and grad_enabled: + with torch.no_grad(), autocast_ctx, adapters_disabled(transformer): + ref_pred = self.step.predict_noise( + self.model, + latents, + sigma, + conditions, + guidance_scale=1.0, + ) + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((noise_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).flatten(1).mean(dim=1) + kl_total = kl_total + kl_step + kl_steps += 1 + + transition_ctx = nullcontext() if grad_enabled else torch.no_grad() + with transition_ctx: + new_latents, _, _ = self.step.forward( + strategy=self.strategy, + noise_pred=noise_pred, + sample=latents, + sigma=sigma, + sigma_next=sigma_next, + eta=float(params.eta), + sigma_max=sigma_max, + step_index=i, + ) + latents = new_latents.to(dtype=self.trajectory_dtype) + + if i >= final_timestep: + break + + kl_loss = kl_total / max(kl_steps, 1) if kl_steps > 0 else kl_total + return DiffuseWithGradResult(z_final=latents, kl_loss=kl_loss) + + +class Sd3ReflVAEDecodeStage(SD3VAEDecodeStage): + """SD3 VAE decode + the BPTT entry point. + + Reuses the mainline grad decode path (fp32 VAE + activation + checkpoint); returns pixels in ``[0, 1]`` — the range the core + differentiable image scorers (PickScore et al.) expect. + """ + + def decode_with_grad(self, z_final: torch.Tensor) -> torch.Tensor: + if z_final.ndim != 4: + raise ValueError( + f"Sd3ReflVAEDecodeStage.decode_with_grad: expected 4D z_final [B, C, H, W], got {tuple(z_final.shape)}" + ) + segment = LatentSegment(latents=z_final.unsqueeze(1)) + return self.decode(segment, grad=True, activation_checkpoint=True).pixels + + +class Sd3ReflPipeline(SD3Pipeline): + """SD3 pipeline for the REFL package (post-swaps diffusion + vae_decode).""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + old = self.diffusion + assert isinstance(old, SD3DiffusionStage), ( + f"Sd3ReflPipeline expects parent to build SD3DiffusionStage, got {type(old).__name__}" + ) + self.diffusion = Sd3ReflDiffusionStage( + model=old.model, + step=old.step, + strategy=old.strategy, + autocast_precision=old.autocast_dtype, + trajectory_precision=old.trajectory_dtype, + logprob_precision=old.logprob_dtype, + vae_scale_factor=old.vae_scale_factor, + latent_channels=old.latent_channels, + ) + self.vae_decode = Sd3ReflVAEDecodeStage(self.bundle) + + def build_refl_conditions( + self, + texts: Texts, + *, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> SD3Conditions: + """Text-only REFL conditioning (SD3 is T2I; ``images`` must be None).""" + if images is not None: + raise ValueError("Sd3ReflPipeline.build_refl_conditions: SD3 ReFL is text-to-image; images must be None.") + guidance = float(params.guidance_scale) + negative_prompt = (params.sampler_kwargs or {}).get("negative_prompt") + negatives = ( + Texts(texts=[str(negative_prompt)] * len(texts.texts)) if negative_prompt and guidance > 1.0 else None + ) + return self.build_conditions(texts, negatives=negatives, guidance_scale=guidance) + + +__all__ = ["Sd3ReflDiffusionStage", "Sd3ReflVAEDecodeStage", "Sd3ReflPipeline"] diff --git a/unirl/models/draft.py b/unirl/models/draft.py deleted file mode 100644 index 552d85976..000000000 --- a/unirl/models/draft.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Shared DRaFT-K grad sampling for ReFL (direct differentiable-reward backprop). - -Family-agnostic: drives the deterministic (eta=0), grad-windowed sampling loop -through the ``DiffusionStage`` protocol's ``predict_noise_at_step`` + the shared -``strategy.denoise`` (sde/kernels.py). Every diffusion family implements -``predict_noise_at_step`` (all packing/CFG/routing lives inside it), so this one -loop reproduces each family's ``diffuse`` at eta=0 — no per-family -``diffuse_draft_k`` needed. - -- ``draft_k_sample`` — the loop; returns a 1-step ``LatentSegment`` the existing - ``decode(grad=True)`` consumes unchanged. -- ``draft_generate`` — the family-agnostic ReFL forward (conditions → sample → - grad-decode) that ``ReFLPolicy`` calls; selects the family purely via the - ``Pipeline`` it's handed. -""" - -from __future__ import annotations - -from contextlib import nullcontext -from typing import Any, Optional - -import torch - -from unirl.sde.noise import generate_latents -from unirl.sde.runtime import get_sigma_schedule -from unirl.types.primitives import Images, Texts -from unirl.types.sampling import DiffusionSamplingParams -from unirl.types.segments import LatentSegment - - -def draft_k_sample( - stage: Any, - conditions: Any, - *, - schedule: torch.Tensor, - params: DiffusionSamplingParams, - draft_num_steps: int, - initial_latents: torch.Tensor, -) -> LatentSegment: - """Deterministic (eta=0) sampling with gradients through only the final - ``draft_num_steps`` steps (DRaFT-K). Returns a 1-step ``LatentSegment`` whose - clean latent ``x_0`` carries grad_fn into the stage's transformer params. - - ``draft_num_steps <= 0`` keeps grad through ALL steps; ``K>0`` runs the first - ``T-K`` steps under ``torch.no_grad()`` (detached) and only the last ``K`` carry - grad. Family-agnostic: uses ``stage.predict_noise_at_step`` (protocol) + - ``stage.strategy.denoise`` (shared kernel). Caller owns ``.train()`` mode + the - outer grad scope (the distributed ``enable_grad()`` context) and supplies - ``initial_latents`` (noise shapes are per-family). - """ - device = initial_latents.device - T = int(params.num_inference_steps) - if int(schedule.shape[0]) != T + 1: - raise ValueError(f"draft_k_sample: schedule length {schedule.shape[0]} != T+1={T + 1}") - schedule = schedule.to(device) - stage.strategy.init_schedule(schedule) - - # cache_enabled=False: PEFT fp32 LoRA adapters lose parameter grads under - # autocast's weight cache on the differentiable pass (see train_reward_dpgo.py). - autocast_dtype = getattr(stage, "autocast_dtype", None) - autocast_ctx = ( - torch.autocast("cuda", autocast_dtype, cache_enabled=False) - if device.type == "cuda" and autocast_dtype in (torch.float16, torch.bfloat16) - else nullcontext() - ) - sigma_max = schedule[1].float() if int(schedule.shape[0]) > 1 else torch.tensor(0.99) - - # K<=0 → grad through all steps; K>0 → grad only through the final K. - grad_start_index = max(0, T - draft_num_steps) if draft_num_steps > 0 else 0 - latents = initial_latents - - for i in range(T): - sigma = schedule[i].to(device) - sigma_next = schedule[i + 1].to(device) - if i < grad_start_index: - # Frozen prefix: deterministic, no autograd graph retained. - with torch.no_grad(), autocast_ctx: - noise_pred = stage.predict_noise_at_step(conditions, sample=latents, sigma=sigma, params=params) - new_latents = stage.strategy.denoise( - noise_pred, latents, sigma, sigma_next, eta=0.0, sigma_max=sigma_max, step_index=i - )[0] - latents = new_latents.detach() - else: - # Grad window: backprop flows through these transformer forwards. - if draft_num_steps > 0 and i == grad_start_index: - latents = latents.detach().requires_grad_(True) - with autocast_ctx: - noise_pred = stage.predict_noise_at_step(conditions, sample=latents, sigma=sigma, params=params) - new_latents = stage.strategy.denoise( - noise_pred, latents, sigma, sigma_next, eta=0.0, sigma_max=sigma_max, step_index=i - )[0] - latents = new_latents - - indices = torch.tensor([T], dtype=torch.long, device=device) - return LatentSegment(latents=latents.unsqueeze(1), sigmas=schedule, indices=indices) - - -def draft_generate( - pipeline: Any, - *, - model_config: Any, - texts: Texts, - params: DiffusionSamplingParams, - draft_num_steps: int, - negatives: Optional[Texts] = None, - activation_checkpoint: bool = False, -) -> Images: - """Family-agnostic ReFL forward: conditions → DRaFT-K sample → grad VAE decode. - - Selects the family purely via ``pipeline`` (its ``build_conditions`` / - ``diffusion`` / ``vae_decode`` stages + ``latent_shape`` classmethod). The only - output tensor is the image, carrying grad_fn into the policy weights. - """ - diffusion = pipeline.diffusion - device = pipeline.bundle.device - - # Text encoders are frozen — keep their (large, e.g. T5-XXL) forward graph out - # of the DRaFT backward; the transformer still gets grad via the latents. - with torch.no_grad(): - conditions = pipeline.build_conditions(texts, negatives=negatives, guidance_scale=float(params.guidance_scale)) - - shift = float(getattr(model_config, "shift", 3.0)) - schedule = get_sigma_schedule(int(params.num_inference_steps), shift=shift, device=device) - - per_sample_shape = type(pipeline).latent_shape(model_config=model_config, sampling_spec=params) - latents = generate_latents( - batch_size=len(texts.texts), - latent_shape=tuple(per_sample_shape), - device=device, - dtype=getattr(diffusion, "trajectory_dtype", torch.bfloat16), - init_same_noise=bool(params.init_same_noise), - samples_per_prompt=int(params.samples_per_prompt), - noise_group_ids=params.noise_group_ids, - base_seed=int(params.seed), - ) - - seg = draft_k_sample( - diffusion, - conditions, - schedule=schedule, - params=params, - draft_num_steps=draft_num_steps, - initial_latents=latents, - ) - return pipeline.vae_decode.decode(seg, grad=True, activation_checkpoint=activation_checkpoint) - - -__all__ = ["draft_k_sample", "draft_generate"] diff --git a/unirl/train/refl/__init__.py b/unirl/train/refl/__init__.py deleted file mode 100644 index 976ce946e..000000000 --- a/unirl/train/refl/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -"""ReFL (direct differentiable-reward backprop) training-side roles. - -``ReFLPolicy`` is the worker-side ``Remote`` that fuses FSDP-wrapped SD3 sampling -+ VAE decode (grad-enabled, DRaFT-K) with the optimizer, so a frozen -differentiable reward on a sibling role can backprop end-to-end onto the policy -weights via the distributed ``enable_grad()`` context. The driver orchestrator -lives in ``unirl/trainer/refl.py``. -""" - -from unirl.train.refl.policy import ReFLPolicy - -__all__ = ["ReFLPolicy"] diff --git a/unirl/train/refl/policy.py b/unirl/train/refl/policy.py deleted file mode 100644 index 84f4bb2c5..000000000 --- a/unirl/train/refl/policy.py +++ /dev/null @@ -1,236 +0,0 @@ -"""ReFLPolicy — family-agnostic ReFL policy Remote. - -Builds a **config-chosen** ``Pipeline`` (no per-family imports), FSDP-wraps its -bundle's transformer via ``FSDPBackend``, and drives grad DRaFT-K sampling + grad -VAE decode through the shared :func:`unirl.models.draft.draft_generate`. The family -is selected entirely by ``pipeline_target`` + ``model_config``; the -``loss_backward`` seed and ``optimizer_step`` are family-agnostic. - -Construction follows the Phase-0/e2e-validated order: the FSDP process group is -initialized in ``initialize()`` (after ``Remote.setup`` populated the dist env), -then the pipeline + ``FSDPBackend`` (which calls ``fully_shard``) are built over it. -""" - -from __future__ import annotations - -import logging -from typing import Any, Optional, Tuple - -import torch -import torch.distributed as dist -from hydra.utils import get_class - -from unirl.distributed.group.dispatch import Dispatch, Execute, distributed -from unirl.distributed.group.remote import Remote -from unirl.models.draft import draft_generate -from unirl.train.backend.base import LrSchedulerConfig, OptimizerConfig -from unirl.train.backend.fsdp import FSDPBackend -from unirl.train.configs import FSDPConfig, LoraConfig -from unirl.types.primitives import Images, Texts -from unirl.types.sampling import DiffusionSamplingParams - -logger = logging.getLogger(__name__) - - -class ReFLPolicy(Remote): - """Family-agnostic ReFL policy: config-chosen Pipeline + FSDP + grad DRaFT-K.""" - - def __init__( - self, - *, - pipeline_target: str, - model_config: Any, - fsdp_cfg: FSDPConfig, - optimizer_cfg: OptimizerConfig, - scheduler_cfg: LrSchedulerConfig, - lora_cfg: Optional[LoraConfig] = None, - strategy: Optional[Any] = None, - block_class_names: Tuple[str, ...] = ("JointTransformerBlock",), - draft_num_steps: int = 1, - reward_loss_scale: float = 1.0, - guidance_scale: float = 1.0, - num_inference_steps: int = 4, - height: int = 512, - width: int = 512, - seed: int = 42, - activation_checkpoint_vae: bool = True, - ) -> None: - super().__init__() - self._pipeline_target = str(pipeline_target) - self._model_config = model_config - self._fsdp_cfg = fsdp_cfg - self._optimizer_cfg = optimizer_cfg - self._scheduler_cfg = scheduler_cfg - self._lora_cfg = lora_cfg - self._strategy = strategy - self._block_class_names = tuple(block_class_names) - self.draft_num_steps = int(draft_num_steps) - self.reward_loss_scale = float(reward_loss_scale) - self.guidance_scale = float(guidance_scale) - self.num_inference_steps = int(num_inference_steps) - self.height = int(height) - self.width = int(width) - self.base_seed = int(seed) - self.activation_checkpoint_vae = bool(activation_checkpoint_vae) - - def initialize(self) -> None: - torch.cuda.set_device(self.device) - # Default PG over the policy role's workers (env:// from Remote.setup's - # dist_env); FSDP2 fully_shard (mode=full) wraps over it. Phase-0-validated. - 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") - - try: - self._model_config.device = self.device # runtime device injection - except Exception: - pass - - pipeline_cls = get_class(self._pipeline_target) - self.pipeline = pipeline_cls.from_config(self._model_config, strategy=self._strategy) - - # FSDP-wrap pipeline.bundle.transformer in place + inject LoRA + optimizer. - # The pipeline's stages reference the same bundle, so sampling uses the - # wrapped trainable transformer. - self.backend = FSDPBackend( - bundle=self.pipeline.bundle, - block_class_names=self._block_class_names, - trainable_attr="transformer", - fsdp_cfg=self._fsdp_cfg, - optimizer_cfg=self._optimizer_cfg, - scheduler_cfg=self._scheduler_cfg, - device=self.device, - rank=int(self.rank_info.rank) if self.rank_info is not None else 0, - lora_cfg=self._lora_cfg, - ) - logger.info( - "ReFLPolicy initialized: pipeline=%s draft_num_steps=%d nfe=%d guidance=%.2f res=%dx%d", - self._pipeline_target, - self.draft_num_steps, - self.num_inference_steps, - self.guidance_scale, - self.height, - self.width, - ) - - # ------------------------------------------------------------------ - # Grad chain (run under the driver's enable_grad() context) - # ------------------------------------------------------------------ - - @distributed(dispatch_mode=Dispatch.DP_SCATTER) - def sample_and_decode(self, *, prompts: Texts, rollout_id: int = 0) -> Images: - """Grad-enabled DRaFT-K sample + in-graph VAE decode via the shared, - family-agnostic ``draft_generate``. Returns ``Images`` whose pixels carry - grad_fn into the FSDP transformer params — the single cross-role tensor.""" - self.backend.model.train() - dp_rank = int(self.rank_info.dp_rank) if self.rank_info is not None else 0 - params = DiffusionSamplingParams( - num_inference_steps=self.num_inference_steps, - guidance_scale=self.guidance_scale, - height=self.height, - width=self.width, - eta=0.0, # deterministic ODE for clean DRaFT gradients - samples_per_prompt=1, - seed=self.base_seed + 1000 * int(rollout_id) + dp_rank, - init_same_noise=False, - ) - return draft_generate( - self.pipeline, - model_config=self._model_config, - texts=prompts, - params=params, - draft_num_steps=self.draft_num_steps, - activation_checkpoint=self.activation_checkpoint_vae, - ) - - @distributed(dispatch_mode=Dispatch.DP_SCATTER) - def eval_sample(self, *, prompts: Texts, rollout_id: int = 0, guidance_scale: Optional[float] = None) -> Images: - """Eval sampling: ``model.eval()`` + ``no_grad`` DRaFT-K, no autograd graph. - - The eval sibling of :meth:`sample_and_decode`. Reused outside the driver's - ``enable_grad()`` context, ``sample_and_decode`` would still run in - ``train()`` mode and build the DRaFT-K activation graph (``draft_generate`` - has no ``no_grad`` guard) only to discard it. This method fixes both: eval - mode + ``torch.no_grad()`` so the returned ``Images`` carry no grad_fn and - no graph is retained. The next ``train_step`` re-asserts ``model.train()`` - via ``sample_and_decode``, so no explicit mode restore is needed. - - ``guidance_scale`` overrides the training CFG strength for eval (the - trainer passes ``eval_cfg_text_scale``; recipes train at 1.0 = no CFG); - ``None`` falls back to the training value. - - The seed is offset from ``sample_and_decode``'s so eval does not reuse the - exact training-step params at the same ``rollout_id``. NOTE: the init latent - is currently drawn unseeded (``generate_latents`` ignores the seed for - ``init_same_noise=False`` with no ``noise_group_ids``), so eval is NOT - bit-exact reproducible — it agrees within sampling noise (~σ). Follow-up to - make it exact: thread ``noise_group_ids`` into the DRaFT noise path.""" - self.backend.model.eval() - dp_rank = int(self.rank_info.dp_rank) if self.rank_info is not None else 0 - params = DiffusionSamplingParams( - num_inference_steps=self.num_inference_steps, - guidance_scale=self.guidance_scale if guidance_scale is None else float(guidance_scale), - height=self.height, - width=self.width, - eta=0.0, # deterministic ODE eval - samples_per_prompt=1, - seed=self.base_seed + 500_000 + 1000 * int(rollout_id) + dp_rank, - init_same_noise=False, - ) - with torch.no_grad(): - return draft_generate( - self.pipeline, - model_config=self._model_config, - texts=prompts, - params=params, - draft_num_steps=self.draft_num_steps, - activation_checkpoint=False, - ) - - @distributed(dispatch_mode=Dispatch.DP_SCATTER) - def loss_backward(self, *, rewards: torch.Tensor) -> None: - """``-reward.mean()`` seed node: local backward populates ``rewards.grad``; - the empty return makes this an always-run backward node so GradContext - chains the grad up through score → sample → transformer params.""" - loss = -self.reward_loss_scale * rewards.to(self.device).float().mean() - loss.backward() - return None - - # ------------------------------------------------------------------ - # Optimizer / checkpoint (delegate to the composed FSDPBackend) - # ------------------------------------------------------------------ - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def optimizer_step(self, *, max_grad_norm: float) -> float: - return self.backend.optimizer_step(max_grad_norm=max_grad_norm) - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def zero_grad(self) -> None: - self.backend.zero_grad() - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def param_checksum(self) -> float: - """L1 sum of local trainable-param shards — a cheap weight-change probe.""" - total = 0.0 - for p in self.backend.model.parameters(): - if not p.requires_grad: - continue - t = p.detach() - if hasattr(t, "to_local"): # FSDP2 sharded DTensor - t = t.to_local() - total += float(t.float().abs().sum().item()) - return total - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def save(self, path: str, step: Optional[int] = None, mode: str = "adapter") -> None: - self.backend.save(path, step=step, mode=mode) - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def load(self, path: str) -> int: - return self.backend.load(path) - - @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) - def wait_for_checkpoint(self) -> None: - self.backend.wait_for_checkpoint() - - -__all__ = ["ReFLPolicy"] diff --git a/unirl/train_refl.py b/unirl/train_refl.py deleted file mode 100755 index cc246b150..000000000 --- a/unirl/train_refl.py +++ /dev/null @@ -1,44 +0,0 @@ -#!/usr/bin/env python -"""UniRL ReFL training entry point (Hydra-native). - -Thin wrapper around :class:`unirl.trainer.refl.RewardBackpropTrainer` — direct -differentiable-reward backprop (DRaFT-K) for SD3. Pairs with -``examples/diffusion/refl_sd3.yaml``. Like ``train_diffusion``, the Ray cluster -is started by the launcher (``ray start --head`` + ``RAY_ADDRESS=auto``). -""" - -from __future__ import annotations - -import hydra -from omegaconf import DictConfig - -from unirl.trainer.refl import RewardBackpropTrainer - - -@hydra.main(version_base=None, config_path="../examples", config_name="diffusion/refl_sd3") -def main(cfg: DictConfig) -> None: - trainer = RewardBackpropTrainer( - cfg=cfg, - batch_size=cfg.batch_size, - policy_cfg=cfg.policy, - reward_cfg=cfg.reward, - data_source_cfg=cfg.data_source, - max_grad_norm=float(cfg.get("max_grad_norm", 1.0)), - reward_fraction=float(cfg.get("reward_fraction", 0.25)), - eval_interval=int(cfg.get("eval_interval", 0)), - eval_num_prompts=int(cfg.get("eval_num_prompts", cfg.batch_size)), - eval_cfg_text_scale=float(cfg.get("eval_cfg_text_scale", 4.0)), - eval_rewards_cfg=cfg.get("eval_rewards"), - logging_cfg=cfg.get("logging"), - ) - trainer.train( - num_rollouts=int(cfg.get("num_rollouts", 100)), - save_interval=int(cfg.get("save_interval", 0)), - save_dir=cfg.get("save_dir"), - load_dir=cfg.get("load_dir"), - save_mode=str(cfg.get("save_mode", "adapter")), - ) - - -if __name__ == "__main__": - main() diff --git a/unirl/trainer/README.md b/unirl/trainer/README.md index e7c107ef6..0bd60c127 100644 --- a/unirl/trainer/README.md +++ b/unirl/trainer/README.md @@ -48,9 +48,9 @@ stay swappable by `_target_`. whole `Sample` so AR and image Parts are sharded by the same prompt trees. Agentic engines return a `List[Sample]` of variable-depth trajectories; their trainers assign each trajectory's advantage to all generated turns and - concatenate those turn Parts for training. `RewardBackpropTrainer` is the one - intentional exception: ReFL differentiates directly through decoded images and - therefore does not use rollout Samples or advantages. + concatenate those turn Parts for training. (ReFL — which differentiates + directly through decoded media and uses no rollout Samples or advantages — + lives outside core as `experimental/refl`.) The current trainer surface is: @@ -63,7 +63,6 @@ The current trainer surface is: | `AsyncDiffusionTrainer` | buffered diffusion `Sample` groups → one `TrainStack` | The same separate-slab async loop for DiT. Requires `max_inflight=1` and reaps each generation before launching the next, so the cross-slab trajectory transfer never queues behind a fresh generation. | | `PETrainer` | `ar` + `diffusion` Parts → two `TrainStack`s | Composed prompt-rewrite/image rollout; image rewards propagate to AR rewrites. `freeze_llm=true` trains and checkpoints diffusion only. | | `UnifiedModelTrainer` | whole `Sample` → one `UnifiedModelTrainStack` | AR and image losses accumulate into shared-backbone optimizer steps while prompt-tree lineage remains intact during DP scatter. | -| `RewardBackpropTrainer` | differentiable image reward → policy step | ReFL/DRaFT-K path; no rollout engine, `Sample`, GRPO advantage, replay ratio, or weight sync. | | `AgenticTrainer` / `AgenticEnvTrainer` | variable-depth `List[Sample]` → concatenated turn `Part` | Barrier multi-turn tool use. The base variant scores terminal answers; the env variant consumes per-trajectory environment returns. | | `AgenticPartialTrainer` / `AgenticEnvPartialTrainer` | freshest complete trajectory groups → concatenated turn `Part` | Colocated over-sample/commit/abort loop. `carry` is for Sample-resumable stateless tools; `drop` purges tails from stateful environments that restart episodes. | | `AsyncAgenticTrainer` / `AsyncAgenticEnvTrainer` | buffered complete trajectory groups → concatenated turn `Part` | Disaggregated train/rollout slabs, resident agentic drive, weight-version staleness control, and the same explicit `carry`/`drop` tail policy. | @@ -234,8 +233,8 @@ an evaluation and checkpoint fall on the same step, evaluation runs first. - `ARTrainer` evaluates the requested prompt set in bounded batches and reports mean reward, also exposed as avg@k accuracy for binary evaluators. `AsyncARTrainer` quiesces its resident engine first. -- `DiffusionTrainer`, `PETrainer`, `UnifiedModelTrainer`, and - `RewardBackpropTrainer` report image reward; optional `eval_rewards` suites can +- `DiffusionTrainer`, `PETrainer`, and `UnifiedModelTrainer` report image + reward; optional `eval_rewards` suites can score the same generated samples or their own prompt sets. PE scores only the diffusion/image frontier. `AsyncDiffusionTrainer` quiesces first and then scores the policy already resident in its rollout engine, without a weight sync and diff --git a/unirl/trainer/refl.py b/unirl/trainer/refl.py deleted file mode 100644 index 84742a2e7..000000000 --- a/unirl/trainer/refl.py +++ /dev/null @@ -1,256 +0,0 @@ -"""RewardBackpropTrainer — driver orchestrator for ReFL (direct reward backprop). - -Two roles, always: a :class:`ReFLPolicy` (FSDP SD3 + grad DRaFT-K sampling + -optimizer) and a frozen differentiable reward (:class:`RewardService`), placed on -disjoint device fractions. Each step runs, under the distributed -``enable_grad()`` context:: - - img = policy.sample_and_decode(prompts) # grad through final K steps + VAE - rew = reward.score_differentiable(img) # frozen reward, grad → image - policy.loss_backward(rew) # -reward.mean() seed - # ctx exit → grad lands on FSDP transformer params - policy.optimizer_step(max_grad_norm) - -No advantages / replay / ratio / segment / weight-sync — those are PG-RL concepts -ReFL does not use. Success signal: the reward curve rises. -""" - -from __future__ import annotations - -import logging -import time -from typing import Any, Dict, List, Optional, Tuple - -from hydra.utils import instantiate -from omegaconf import DictConfig - -from unirl.distributed.group.placement import placement -from unirl.distributed.tensor.grad_context import enable_grad -from unirl.distributed.tensor.ref import hydrate -from unirl.trainer.base import BaseTrainer -from unirl.trainer.eval_suites import build_eval_suites -from unirl.types.primitives import Texts -from unirl.types.sample import Sample -from unirl.utils.hydra import remote_hydra - -logger = logging.getLogger(__name__) - - -def _text_inputs(inputs: Sample) -> Texts: - """Read the text root from a data-source input ``Sample``.""" - if not isinstance(inputs, Sample): - raise TypeError(f"ReFL data source must return Sample, got {type(inputs).__name__}.") - if not inputs.parts: - raise ValueError("ReFL data-source Sample has no input Parts.") - prompts = inputs.parts[0].primitives.get("text") - if not isinstance(prompts, Texts): - raise TypeError(f"ReFL data-source root requires Texts, got {type(prompts).__name__}.") - return prompts - - -class RewardBackpropTrainer(BaseTrainer): - """ReFL trainer: policy + frozen differentiable reward, grad via enable_grad().""" - - def __init__( - self, - *, - cfg: DictConfig, - batch_size: int, - policy_cfg: DictConfig, - reward_cfg: DictConfig, - data_source_cfg: DictConfig, - max_grad_norm: float = 1.0, - reward_fraction: float = 0.25, - eval_interval: int = 0, - eval_num_prompts: int = 12, - eval_cfg_text_scale: float = 4.0, - eval_rewards_cfg: Optional[Any] = None, - logging_cfg: Optional[DictConfig] = None, - ) -> None: - super().__init__(cfg=cfg, logging_cfg=logging_cfg) - self.batch_size = int(batch_size) - self.max_grad_norm = float(max_grad_norm) - # Periodic eval on the eval set (run.eval_data_path), logged under eval/*; - # eval_interval=0 disables it. ReFL has no rollout engine / tracks: eval - # samples via ReFLPolicy.eval_sample (deterministic ODE, no grad, CFG= - # eval_cfg_text_scale — same knob/semantics as DiffusionTrainer, mapped - # onto the SD3-family guidance_scale) and scores with the differentiable - # reward. Extra eval-only rewards: unirl.trainer.eval_suites. - self.eval_interval = int(eval_interval) - self.eval_num_prompts = int(eval_num_prompts) - self.eval_cfg_text_scale = float(eval_cfg_text_scale) - self.data_source = instantiate(data_source_cfg) - - # Unified reward placement — SAME knob/semantics as DiffusionTrainer's - # ``reward_fraction``: ``> 0`` carves the (frozen, differentiable) reward - # its OWN disjoint slab (the tail of the pool), policy takes the rest; the - # DRaFT-K gradient crosses the slab boundary back into the policy via the - # distributed ``enable_grad()`` context. ``== 0`` colocates reward on the - # policy's cards (cheapest grad) at the cost of sharing its GPU memory. - reward_frac = float(reward_fraction) - if not 0.0 <= reward_frac < 1.0: - raise ValueError(f"reward_fraction must be in [0, 1), got {reward_frac}") - if reward_frac > 0.0: - with placement(self.pool, fraction=1.0 - reward_frac, shared_workers=True): - self.policy = remote_hydra(policy_cfg) - with placement(self.pool, fraction=reward_frac, shared_workers=True): - self.reward = remote_hydra(reward_cfg) - # Extra eval-only rewards (eval_rewards) ride the reward slab — - # see unirl.trainer.eval_suites. Their backends must be - # differentiable-capable like the training reward. - self._eval_suites = build_eval_suites( - eval_rewards_cfg, data_source_cfg=data_source_cfg, enabled=self.eval_interval > 0 - ) - else: - with placement(self.pool, fraction=1.0, shared_workers=True): - self.policy = remote_hydra(policy_cfg) - self.reward = remote_hydra(reward_cfg) - self._eval_suites = build_eval_suites( - eval_rewards_cfg, data_source_cfg=data_source_cfg, enabled=self.eval_interval > 0 - ) - - self.policy.initialize() - # BaseTrainer.maybe_save/load_checkpoint operate on ``self.backend``. - self.backend = self.policy - - pdp, rdp = self.policy.dp_size, self.reward.dp_size - if self.batch_size % pdp or self.batch_size % rdp: - raise ValueError(f"batch_size={self.batch_size} must be divisible by policy dp={pdp} and reward dp={rdp}") - logger.info( - "RewardBackpropTrainer ready: policy dp=%d reward dp=%d batch=%d max_grad_norm=%.2f", - pdp, - rdp, - self.batch_size, - self.max_grad_norm, - ) - - def train_step(self, prompts: Texts, *, rollout_id: int) -> Tuple[float, float, float]: - """One enable_grad() sample → score → backward → step. Returns - (mean_reward, grad_norm, step_time_s).""" - t0 = time.perf_counter() - with enable_grad(): - images = self.policy.sample_and_decode(prompts=prompts, rollout_id=rollout_id) - rewards = self.reward.score_differentiable(images.pixels, list(prompts.texts)) - # Detached value for logging (does not disturb the worker-side graph). - mean_reward = float(hydrate(rewards).float().mean().item()) - self.policy.loss_backward(rewards=rewards) - grad_norm = self.policy.optimizer_step(max_grad_norm=self.max_grad_norm) - if isinstance(grad_norm, list): # BROADCAST → one result per worker - grad_norm = grad_norm[0] - self.policy.zero_grad() - return mean_reward, float(grad_norm or 0.0), time.perf_counter() - t0 - - def evaluate(self, step: int) -> float: - """Periodic eval — mean reward(s) over the eval prompt set (no training). - - ReFL has no rollout engine or tracks, so this mirrors :meth:`train_step`'s - sample→score path (minus ``enable_grad``/backward): sample images with - :meth:`ReFLPolicy.eval_sample` (deterministic ODE, ``model.eval()`` + - ``no_grad``, CFG at ``eval_cfg_text_scale``) and score. The training - reward plus every shared-set ``eval_rewards`` suite scores the SAME - images from the default eval set (``run.eval_data_path``, - ``eval_num_prompts`` prompts); each own-set suite then runs its own - sample→score pass over its own prompts. All means land in one ``eval/*`` - row (``eval/reward`` + ``eval/``); returns ``eval/reward``. - - ``step`` keys the wandb log axis (and ``eval_sample``'s seed), mirroring - :meth:`DiffusionTrainer.evaluate` — so re-running a checkpoint via - ``num_rollouts=0 load_dir=checkpoint-k`` (→ baseline ``evaluate(k)``) evals - the restored weights at the same step. NOTE: eval is NOT bit-exact — the - init latent is drawn unseeded (see ``eval_sample``), so A vs B agree only - within sampling noise (~σ), not byte-identically. - """ - # Default pass: training reward + shared-set suites score the SAME images. - scorers = [("reward", self.reward)] + [(s.name, s.reward) for s in self._eval_suites if s.data_source is None] - metrics = self._eval_pass(self.data_source, self.eval_num_prompts, scorers, step) - for suite in self._eval_suites: - if suite.data_source is not None: - n = suite.num_prompts or self.eval_num_prompts - metrics.update(self._eval_pass(suite.data_source, n, [(suite.name, suite.reward)], step)) - logger.info( - "EVAL step %d (cfg=%.1f) %s", - step, - self.eval_cfg_text_scale, - " ".join(f"{k}={v:.4f}" for k, v in metrics.items()), - ) - self.wandb_logger.log_eval(step, metrics) - return metrics["reward"] - - def _eval_pass( - self, data_source: Any, num_prompts: int, scorers: List[Tuple[str, Any]], step: int - ) -> Dict[str, float]: - """One sample→score sweep over one eval set; returns each scorer's mean. - - Chunked by ``self.batch_size`` — both ``eval_sample`` and - ``score_differentiable`` are DP_SCATTER, and ``batch_size`` is validated - divisible by both the policy and reward dp sizes in ``__init__``; a ragged - tail (``num_prompts`` not a multiple of ``batch_size``) is floored off. - """ - eval_inputs = data_source.get_eval_samples(num_prompts) - prompts = _text_inputs(eval_inputs) - texts = list(prompts.texts) - chunk = max(1, self.batch_size) - usable = len(texts) - len(texts) % chunk or len(texts) - sums = {name: 0.0 for name, _ in scorers} - counts = {name: 0 for name, _ in scorers} - for start in range(0, usable, chunk): - sub = Texts(texts=texts[start : start + chunk]) - images = self.policy.eval_sample(prompts=sub, rollout_id=step, guidance_scale=self.eval_cfg_text_scale) - for name, reward in scorers: - r = hydrate(reward.score_differentiable(images.pixels, list(sub.texts))).float() - sums[name] += float(r.sum().item()) - counts[name] += int(r.numel()) - return {name: sums[name] / max(1, counts[name]) for name, _ in scorers} - - def train( - self, - *, - num_rollouts: int, - save_interval: int = 0, - save_dir: Optional[str] = None, - load_dir: Optional[str] = None, - save_mode: str = "adapter", - ) -> None: - start = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) - self._init_wandb(num_rollouts=num_rollouts) - try: - if self.eval_interval > 0: - self.evaluate(start) # baseline eval before any training - for rollout_id in range(start, num_rollouts): - inputs = self.data_source.get_samples(self.batch_size) - prompts = _text_inputs(inputs) - mean_reward, grad_norm, dt = self.train_step(prompts, rollout_id=rollout_id) - logger.info( - "rollout %d/%d reward=%.4f grad_norm=%.4f %.1fs", - rollout_id + 1, - num_rollouts, - mean_reward, - grad_norm, - dt, - ) - self.wandb_logger.log_step( - rollout_id + 1, - { - "rollout/mean_reward": mean_reward, - "train/loss": -mean_reward, - "train/grad_norm": grad_norm, - "perf/step_time_s": dt, - }, - prefix="", - ) - # eval(k) BEFORE save(checkpoint-k) at the same step, so a - # resumed checkpoint re-runs the same eval (A/B consistency). - if self.eval_interval > 0 and (rollout_id + 1) % self.eval_interval == 0: - self.evaluate(rollout_id + 1) - self.maybe_save_checkpoint( - rollout_id, - num_rollouts, - save_interval=save_interval, - save_dir=save_dir, - save_mode=save_mode, - ) - finally: - self._finish_wandb() - - -__all__ = ["RewardBackpropTrainer"] From 9edcab004e15ad9495878844c9d089ad7a6e40ba Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 12:17:16 +0800 Subject: [PATCH 2/3] fix(experimental/refl): sd3 example trains with the fp32 LoRA master Same correction as the wan configs on #210: the bf16-master variant was an environment workaround for the fleet image's torch 2.7.1 (FSDP2 uniform-dtype assert over all params); the pinned torch checks trainables only, and the historical SD3 ReFL curves were produced with the fp32 master. --- experimental/refl/examples/sd3_pickscore_refl.yaml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/experimental/refl/examples/sd3_pickscore_refl.yaml b/experimental/refl/examples/sd3_pickscore_refl.yaml index a0844454d..bd96e729c 100644 --- a/experimental/refl/examples/sd3_pickscore_refl.yaml +++ b/experimental/refl/examples/sd3_pickscore_refl.yaml @@ -38,8 +38,12 @@ actor: fsdp_cfg: _target_: unirl.train.configs.FSDPConfig param_dtype: bf16 - # master_dtype deliberately omitted → LoRA stays bf16 (uniform dtype); - # fp32 LoRA over a bf16 base trips torch FSDP2's uniform-dtype assert. + # fp32 LoRA master over the bf16 base — the setting the historical ReFL + # curves (#120 era) were trained with; a bf16 master rounds away most + # AdamW steps at lora_A's magnitude. Needs the pinned torch (>=2.11): + # older FSDP2 asserts uniform dtype over ALL params in a group; the + # pinned family checks trainables only. + master_dtype: fp32 cpu_offload: false mixed_precision: true fsdp_mode: full From c122c04ad07d266aa809c8e84fdf5239d5530c90 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 12:53:31 +0800 Subject: [PATCH 3/3] docs(experimental/refl): record the sd3 effect-alignment result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 200-rollout run on the pinned stack reproduces the legacy core ReFL curve (#120): reward first-10 mean 0.743 -> last-10 0.903 vs the historical 0.757 -> 0.899 — effect-level validation for the port that this change ships while deleting the legacy path. --- experimental/refl/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 33d2d2409..0d011a22f 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -58,4 +58,4 @@ core stack. | `wan21_t2v_videoalign_refl` (2-rollout smoke, full 81f/480x832 geometry) | 8xH20, fleet image | `40b3f4c9` | PASS — grads flow reward → VAE → DiT LoRA | | VideoAlign load + differentiable fwd/bwd on transformers 5.6.2 + peft 0.20 | 8xH20 (isolated venv) | `9087a671` | PASS — `grad_abs_mean=3.5e-3` | | `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | -| `sd3_pickscore_refl` (ported from the legacy core path; hyperparameters preserved) | 8xH20 | current head | pending smoke | +| `sd3_pickscore_refl` (200 rollouts, ported from the legacy core path; hyperparameters preserved) | 8xH20 (torch 2.11 + transformers 5.6.2 + peft 0.20, fp32 LoRA master) | `9edcab00` | PASS — reward first-10 0.743 → last-10 0.903, matching the legacy-path curve (#120: 0.757 → 0.899) |