diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 9c38aaeb9..0d011a22f 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` (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) | 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..bd96e729c 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,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 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"]