feat(experimental): port SD3 image ReFL and remove the legacy core ReFL path

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` (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) |

examples/diffusion/refl_sd3.yaml → 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,18 @@ policy:
   fsdp_cfg:
     _target_: unirl.train.configs.FSDPConfig
     param_dtype: bf16
+    # 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
-    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 +99,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

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"]