From 4b552f0010871da552e22c51d0719784fa5d306d Mon Sep 17 00:00:00 2001 From: yohunawu Date: Tue, 14 Jul 2026 17:34:47 +0800 Subject: [PATCH 01/24] new feature: support wan series refl recipes. --- pyproject.toml | 11 + recipes/__init__.py | 0 recipes/refl_wan/__init__.py | 0 .../configs/wan21_t2v_videoalign_refl.yaml | 161 +++ recipes/refl_wan/configs/wan22_face_refl.yaml | 143 ++ recipes/refl_wan/models/__init__.py | 0 recipes/refl_wan/models/wan21.py | 370 +++++ recipes/refl_wan/models/wan22.py | 344 +++++ recipes/refl_wan/rewards/__init__.py | 1 + recipes/refl_wan/rewards/face/__init__.py | 5 + recipes/refl_wan/rewards/face/face_tools.py | 439 ++++++ recipes/refl_wan/rewards/face/scorer.py | 308 +++++ .../refl_wan/rewards/videoalign/__init__.py | 17 + .../rewards/videoalign/model/__init__.py | 42 + .../rewards/videoalign/model/checkpoint.py | 165 +++ .../rewards/videoalign/model/configs.py | 155 +++ .../rewards/videoalign/model/factory.py | 190 +++ .../videoalign/model/prompt_template.py | 188 +++ .../rewards/videoalign/model/reward_model.py | 285 ++++ recipes/refl_wan/rewards/videoalign/scorer.py | 239 ++++ .../refl_wan/rewards/videoalign/wrapper.py | 374 +++++ recipes/refl_wan/roles.py | 224 +++ recipes/refl_wan/run.py | 19 + recipes/refl_wan/start_wan21_t2v.sh | 29 + recipes/refl_wan/start_wan22_i2v.sh | 13 + recipes/refl_wan/trainer.py | 125 ++ scripts/check_recipe_targets.py | 4 +- unirl/data/data_source.py | 7 +- unirl/distributed/group/worker.py | 11 +- unirl/models/types/diffusion.py | 26 +- unirl/models/wan21/bundle.py | 15 +- unirl/models/wan21/image_encode.py | 14 - unirl/models/wan21/vae.py | 101 +- unirl/models/wan21/wan_video_vae.py | 1203 +++++++++++++++++ unirl/reward/base.py | 12 +- unirl/train/backend/base_backend.py | 16 +- unirl/train/configs.py | 3 + unirl/train/lora.py | 46 +- unirl/train/optim.py | 25 +- unirl/trainer/base_role.py | 117 ++ unirl/trainer/trainer.py | 416 ++++++ 41 files changed, 5772 insertions(+), 91 deletions(-) create mode 100644 recipes/__init__.py create mode 100644 recipes/refl_wan/__init__.py create mode 100644 recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml create mode 100644 recipes/refl_wan/configs/wan22_face_refl.yaml create mode 100644 recipes/refl_wan/models/__init__.py create mode 100644 recipes/refl_wan/models/wan21.py create mode 100644 recipes/refl_wan/models/wan22.py create mode 100644 recipes/refl_wan/rewards/__init__.py create mode 100644 recipes/refl_wan/rewards/face/__init__.py create mode 100644 recipes/refl_wan/rewards/face/face_tools.py create mode 100644 recipes/refl_wan/rewards/face/scorer.py create mode 100644 recipes/refl_wan/rewards/videoalign/__init__.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/__init__.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/checkpoint.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/configs.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/factory.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/prompt_template.py create mode 100644 recipes/refl_wan/rewards/videoalign/model/reward_model.py create mode 100644 recipes/refl_wan/rewards/videoalign/scorer.py create mode 100644 recipes/refl_wan/rewards/videoalign/wrapper.py create mode 100644 recipes/refl_wan/roles.py create mode 100644 recipes/refl_wan/run.py create mode 100644 recipes/refl_wan/start_wan21_t2v.sh create mode 100644 recipes/refl_wan/start_wan22_i2v.sh create mode 100644 recipes/refl_wan/trainer.py create mode 100644 unirl/models/wan21/wan_video_vae.py create mode 100644 unirl/trainer/base_role.py create mode 100644 unirl/trainer/trainer.py diff --git a/pyproject.toml b/pyproject.toml index adc1dd358..6c7fd9d71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,17 @@ eval = [ "torchvision>=0.16", "easyocr>=1.7", ] +face_reward = [ + "onnx>=1.14", + "onnx2torch>=1.5", + "scikit-image>=0.21", + "imageio>=2.31", + "imageio-ffmpeg>=0.4", + "scipy>=1.11", +] +video_align_reward = [ + "flash-attn==2.7.0.post2", +] # VeOmni's torch-native distributed layer (FSDP2/EP parallelize), consumed # exclusively through unirl.train.backend.veomni._compat — a selective import # that never executes veomni/__init__.py or veomni/models/__init__.py, so the diff --git a/recipes/__init__.py b/recipes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/recipes/refl_wan/__init__.py b/recipes/refl_wan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml new file mode 100644 index 000000000..6cce35f73 --- /dev/null +++ b/recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml @@ -0,0 +1,161 @@ +# @package _global_ +# REFL WAN 2.1 T2V — VideoAlign (Qwen2-VL VQ/MQ/TA) reward — role config. + +num_devices: 8 +batch_size: 8 +num_rollouts: 1000 +save_interval: 100 +save_dir: ${oc.env:OUTPUT_DIR,outputs/wan21_t2v_videoalign_refl} + +roles: + - name: actor + _target_: recipes.refl_wan.roles.ReflActorRole + placement: + n_devices: ${num_devices} + + model: + _target_: unirl.models.wan21.bundle.WAN21Bundle.from_config + config: + _target_: unirl.models.wan21.config.WAN21PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + shift: 5.0 + max_sequence_length: 512 + + pipeline: + _target_: recipes.refl_wan.models.wan21.Wan21ReflPipeline + shift: 5.0 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + max_sequence_length: 512 + strategy: + # eta=0.0 below reduces FlowSDE to deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + + backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["WanTransformerBlock"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + # BPTT keeps the full 24→final grad window alive; activation + # checkpointing is mandatory for any non-trivial + # num_frames × num_inference_steps product. + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 5.0e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 64 + alpha: 128 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + # Peft path takes suffix matches → these six suffixes catch both + # attn1 (self-attention) and attn2 (cross-attention) LoRA targets + # inside every WanTransformerBlock. + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + + algorithm: ${algorithm} + + - name: reward + _target_: recipes.refl_wan.roles.ReflRewardRole + placement: + # BPTT gradients must live on the same worker — no cross-process + # autograd. VideoAlign co-locates with the actor. + share_with: actor + + backend: + _target_: recipes.refl_wan.rewards.videoalign.VideoAlignRewardScorer + base_device: cuda + config: + _target_: recipes.refl_wan.rewards.videoalign.VideoAlignSpec + reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} + device: auto + batch_size: 1 + resize_height: 336 + resize_width: 588 + micro_batch_size: 1 + reward_num_frames: 36 # 81 → 36 uniform frames before scoring + use_norm: true + w_vq: 1.0 + w_mq: 1.0 + w_ta: 1.0 + differentiable: true # REFL requires grad through the reward forward + +algorithm: + reward_weight: 0.25 + reward_baseline: 0.0 + reward_scale: 1.0 + kl_weight: 0.0 + max_grad_norm: 1.0 + sampling_params: ${sampling} + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH}} + seed: 42 + shuffle: false + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + # 25 inference steps with mid_timestep = final_timestep = 24 + # → DRaFT-1 (only the final step is differentiable). + num_inference_steps: 25 + guidance_scale: 5.0 + height: 480 + width: 832 + # (num_frames - 1) % 4 == 0 is the WAN 2.1 T2V hard constraint. + num_frames: 81 + eta: 0.0 + samples_per_prompt: 1 + seed: 42 + init_same_noise: false + sampler_kwargs: + negative_prompt: "blurry, low quality, distorted, oversaturated" + # Truncated BPTT window: no_grad until ``mid_timestep``, grad on from + # ``mid_timestep`` .. ``final_timestep`` (inclusive). 24/24/25 → only + # the final step traverses the reward backward pass (DRaFT-1 / ReFL). + mid_timestep: 24 + final_timestep: 24 + # KL branch off — outer algorithm.kl_weight also 0. Single-pass forward + # per step, no LoRA-disabled reference recompute. + kl_weight: 0.0 + +logging: + report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} + project_name: ${oc.env:WANDB_PROJECT,unirl-refl-t2v} + run_name: ${oc.env:WANDB_RUN_NAME,wan21_t2v_videoalign_refl} + tags: ["wan21", "t2v", "refl", "bptt", "videoalign"] + log_media: false + metric_prefix: refl_recipe/ diff --git a/recipes/refl_wan/configs/wan22_face_refl.yaml b/recipes/refl_wan/configs/wan22_face_refl.yaml new file mode 100644 index 000000000..43151fa3f --- /dev/null +++ b/recipes/refl_wan/configs/wan22_face_refl.yaml @@ -0,0 +1,143 @@ +# @package _global_ +# REFL WAN 2.2 I2V — Remote role config. + +num_devices: 8 +batch_size: 8 +num_rollouts: 1000 +save_interval: 100 +save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_face_refl} + +roles: + - name: actor + _target_: recipes.refl_wan.roles.ReflActorRole + placement: + n_devices: ${num_devices} + + model: + _target_: unirl.models.wan22.bundle.WAN22Bundle.from_config + config: + _target_: unirl.models.wan22.config.WAN22PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + shift: 5.0 + max_sequence_length: 512 + boundary_ratio: 0.9 + num_train_timesteps: 1000 + + pipeline: + _target_: recipes.refl_wan.models.wan22.Wan22ReflPipeline + shift: 5.0 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + max_sequence_length: 512 + strategy: + # eta=0.0 below reduces FlowSDE to deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + + backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["WanTransformerBlock"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 2.5e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 32 + alpha: 32 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + module_prefix: low_noise + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + + algorithm: ${algorithm} + + - name: reward + _target_: recipes.refl_wan.roles.ReflRewardRole + placement: + share_with: actor + + backend: + _target_: recipes.refl_wan.rewards.face.FaceRewardScorer + base_device: cuda + config: + _target_: recipes.refl_wan.rewards.face.FaceRewardSpec + model_path: ${oc.env:FACE_MODEL_PATH,/apdcephfs_gy8/share_301869871/staryding/gy6_data/ckpt/ckpt_0303/apdcephfs_wza/xiangwshen/a/antelodev2/antelodev2} + device: cuda + batch_size: 1 + image_size: 112 + ref_max_frames: 81 + ref_max_pixels: 230400 + differentiable: true + +algorithm: + reward_weight: 0.1 + reward_baseline: 0.54 + reward_scale: 0.16 + kl_weight: 1.0 + max_grad_norm: 1.0 + sampling_params: ${sampling} + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH,/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl}} + seed: 42 + shuffle: false + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + num_inference_steps: 8 + guidance_scale: 1.0 + guidance_scale_2: 1.0 + height: 352 + width: 640 + num_frames: 81 + eta: 0.0 + samples_per_prompt: 1 + seed: 42 + init_same_noise: false + sampler_kwargs: + mid_timestep: 4 + final_timestep: 7 + kl_weight: 1.0 + +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,wan22_face_refl} + tags: ["wan22", "i2v", "refl", "face", "recipe"] + log_media: false + metric_prefix: refl_recipe/ diff --git a/recipes/refl_wan/models/__init__.py b/recipes/refl_wan/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/recipes/refl_wan/models/wan21.py b/recipes/refl_wan/models/wan21.py new file mode 100644 index 000000000..ef37241a9 --- /dev/null +++ b/recipes/refl_wan/models/wan21.py @@ -0,0 +1,370 @@ +"""Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. + +Mirrors ``recipes.refl_wan.models.wan22`` but targets the WAN 2.1 T2V +single-DiT stack. The REFL-specific pieces live here: + +- :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch + CFG predictor used only by REFL BPTT. +- :class:`Wan21ReflDiffusionStage` — subclass of the mainline + :class:`WAN21DiffusionStage` with an extra :meth:`diffuse_with_grad` + method implementing single-branch CFG + BPTT windowing (mid/final + timestep) + optional per-step KL against the LoRA-disabled reference. +- :class:`Wan21ReflPipeline` — subclass of the mainline + :class:`WAN21Pipeline` that post-swaps ``self.diffusion`` for the REFL + variant above, keeping every other stage (text/vae/scheduler) intact. + +The mainline WAN 2.1 code path is unchanged. REFL-specific single-branch +CFG prediction lives in recipe-local :class:`Wan21ReflDiffusionStep`. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any, Dict, Optional, Tuple + +import torch + +from unirl.models.types.diffusion import DiffuseWithGradResult +from unirl.models.wan21.bundle import WAN21Bundle +from unirl.models.wan21.conditions import WAN21Conditions +from unirl.models.wan21.diffusion import WAN21DiffusionStage, WAN21DiffusionStep +from unirl.models.wan21.pipeline import WAN21Pipeline +from unirl.types.sampling import DiffusionSamplingParams + +# Matches the mainline module-level constant in unirl/models/wan21/diffusion.py. +# Not imported directly to keep the recipe decoupled from the mainline's +# private surface. +_WAN_TIMESTEP_SCALE: float = 1000.0 + +# Inclusive max for torch.Generator.manual_seed and torch initial_seed conventions. +MAX_TORCH_SEED = (1 << 63) - 1 + + +class Wan21ReflDiffusionStep(WAN21DiffusionStep): + """Recipe-local WAN 2.1 single-branch denoising for REFL BPTT.""" + + def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: WAN21Bundle, + sample: torch.Tensor, + sigma: torch.Tensor, + conditions: WAN21Conditions, + *, + branch: str, + ) -> torch.Tensor: + """Run exactly one CFG branch through the WAN 2.1 transformer.""" + if branch not in {"cond", "uncond"}: + raise ValueError(f"Wan21ReflDiffusionStep.predict_noise: unknown branch={branch!r}") + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan21ReflDiffusionStep.predict_noise: conditions.text.embeds is None") + + prompt_embeds = conditions.text.embeds + if branch == "cond": + branch_embeds = prompt_embeds + else: + neg = conditions.negative_text + branch_embeds = neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + + batch_size = int(sample.shape[0]) + timestep = sigma * _WAN_TIMESTEP_SCALE + if timestep.dim() == 0: + timestep = timestep.expand(batch_size) + elif int(timestep.shape[0]) != batch_size: + timestep = timestep.expand(batch_size) + + embeds_dtype = branch_embeds.dtype + sample_cast = sample.to(dtype=embeds_dtype) + + image_latent = conditions.image_latent + if image_latent is not None and image_latent.latents is not None: + sample_cat = torch.cat( + [sample_cast, image_latent.latents.to(device=sample_cast.device, dtype=embeds_dtype)], + dim=1, + ) + else: + sample_cat = sample_cast + + extra: Dict[str, Any] = {} + image_embed = conditions.image_embed + image_embeds = image_embed.embeds if image_embed is not None and image_embed.embeds is not None else None + if image_embeds is not None: + extra["encoder_hidden_states_image"] = image_embeds.to(device=sample_cast.device, dtype=embeds_dtype) + + return model.transformer( + hidden_states=sample_cat, + encoder_hidden_states=branch_embeds, + timestep=timestep, + return_dict=False, + **extra, + )[0] + + +class Wan21ReflDiffusionStage(WAN21DiffusionStage): + """WAN 2.1 T2V diffusion stage + REFL BPTT sampling override. + + Adds ``diffuse_with_grad`` that keeps the autograd graph alive on the + returned final latent, honouring the REFL-specific BPTT knobs + (``mid_timestep`` / ``final_timestep`` / ``kl_weight``) via + ``params.sampler_kwargs``. Everything else is inherited from the + mainline stage — ``diffuse`` / ``replay`` / ``predict_noise`` / + ``trainable_module`` are untouched. + """ + + 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: + """ + High-level function for generating initial latents. + + Args: + batch_size: Total number of samples + latent_shape: Shape of a single latent (C, H, W) or (C, T, H, W) + device: Device for the tensor + dtype: Data type for the tensor + + Returns: + Latent tensor [batch_size, *latent_shape] + """ + 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: WAN21Conditions, + *, + schedule: torch.Tensor, + params: DiffusionSamplingParams, + initial_latents: Optional[torch.Tensor] = None, + ) -> DiffuseWithGradResult: + """Differentiable WAN 2.1 T2V sampling for REFL-style BPTT training. + + Returns :class:`DiffuseWithGradResult` with the live-grad + ``z_final`` + scalar ``kl_loss``. + + BPTT knobs (read from ``params.sampler_kwargs``): + + - ``mid_timestep`` (int, default 0): step index at which the forward + switches from ``torch.no_grad`` to grad-enabled — implements + truncated BPTT (only the last ``T - mid_timestep`` steps + participate in backward). DRaFT-1 sets ``mid_timestep = T - 1``. + - ``final_timestep`` (int, default ``num_inference_steps - 1``): + early stop. The loop breaks once ``i >= final_timestep``. + - ``kl_weight`` (float, default 0.0): when non-zero AND the + transformer exposes ``disable_adapter_layers`` (PEFT LoRA), + per-step KL ``mean((pred - ref_pred)**2 / (2 * sigma**2))`` is + accumulated and returned in ``kl_loss``. The trainer multiplies + it by its own ``kl_weight`` at the loss site. + + ``initial_latents`` follows the same contract as :meth:`diffuse`: + when provided, used verbatim and the internal RNG path is + bypassed. + """ + if conditions.text is None or conditions.text.embeds is None: + raise ValueError( + "Wan21ReflDiffusionStage.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"Wan21ReflDiffusionStage.diffuse_with_grad: " + f"schedule length {schedule.shape[0]} != T+1={T + 1}" + ) + self.strategy.init_schedule(schedule) + + latent_shape = self._latent_shape( + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ) + if initial_latents is not None: + if int(initial_latents.shape[0]) != batch_size: + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[0]=" + f"{int(initial_latents.shape[0])} != batch_size={batch_size}." + ) + if tuple(initial_latents.shape[1:]) != tuple(latent_shape): + raise ValueError( + f"Wan21ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[1:]=" + f"{tuple(initial_latents.shape[1:])} != expected {tuple(latent_shape)} " + f"for num_frames={int(params.num_frames)}, " + f"height={int(params.height)}, width={int(params.width)}." + ) + 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), + ) + + # BPTT knobs. + sk: Dict[str, Any] = dict(getattr(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)) + + 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((), device=device, dtype=torch.float32) + kl_steps = 0 + + guidance_scale = float(params.guidance_scale) + use_cfg = guidance_scale > 1.0 + step = self.step + if not isinstance(step, Wan21ReflDiffusionStep): + raise TypeError( + f"Wan21ReflDiffusionStage.diffuse_with_grad requires Wan21ReflDiffusionStep, " + f"got {type(step).__name__}." + ) + + for i in range(T): + sigma = schedule[i].to(device) + sigma_next = schedule[i + 1].to(device) + grad_enabled = i >= mid_timestep + + if use_cfg: + # REFL parity: the conditional branch follows the BPTT grad + # window, while the uncond/negative branch is always + # stop-grad. Batched CFG would incorrectly add a + # ``(1 - guidance_scale) * d(uncond)/dθ`` term. + cond_ctx = nullcontext() if grad_enabled else torch.no_grad() + with cond_ctx, autocast_ctx: + noise_pred_cond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + with torch.no_grad(), autocast_ctx: + noise_pred_uncond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="uncond", + ) + noise_pred_cond = noise_pred_cond.float() + noise_pred_uncond = noise_pred_uncond.float() + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + kl_pred = noise_pred_cond + else: + pred_ctx = nullcontext() if grad_enabled else torch.no_grad() + with pred_ctx, autocast_ctx: + noise_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + noise_pred = noise_pred.float() + kl_pred = noise_pred + + # Per-step KL against the LoRA-disabled reference (REFL-style). + # When CFG is enabled, this intentionally computes KL on the + # conditional prediction before the stop-grad uncond CFG mix. + if ( + kl_weight != 0.0 + and grad_enabled + and hasattr(transformer, "disable_adapter_layers") + and hasattr(transformer, "enable_adapter_layers") + ): + with torch.no_grad(), autocast_ctx: + transformer.disable_adapter_layers() + try: + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) + finally: + transformer.enable_adapter_layers() + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() + kl_total = kl_total + kl_step + kl_steps += 1 + + transition_ctx = nullcontext() if grad_enabled else torch.no_grad() + with transition_ctx: + new_latents, _, _ = 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 Wan21ReflPipeline(WAN21Pipeline): + """WAN 2.1 T2V pipeline for the REFL recipe. + + Reuses the mainline :class:`WAN21Pipeline` construction (text encode / + condition build / VAE decode / schedule) and post-swaps + ``self.diffusion`` for the REFL-flavoured stage. The parent already + validates every constructor argument and wires the strategy / step / + precision policy through — we only need to rebuild the diffusion + stage with the same underlying components. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + old = self.diffusion + assert isinstance(old, WAN21DiffusionStage), ( + f"Wan21ReflPipeline expects parent to build WAN21DiffusionStage, " + f"got {type(old).__name__}" + ) + self.diffusion = Wan21ReflDiffusionStage( + model=old.model, + step=Wan21ReflDiffusionStep(), + strategy=old.strategy, + autocast_precision=old.autocast_dtype, + trajectory_precision=old.trajectory_dtype, + logprob_precision=old.logprob_dtype, + ) + + +__all__ = ["Wan21ReflDiffusionStep", "Wan21ReflDiffusionStage", "Wan21ReflPipeline"] diff --git a/recipes/refl_wan/models/wan22.py b/recipes/refl_wan/models/wan22.py new file mode 100644 index 000000000..86e38dbd9 --- /dev/null +++ b/recipes/refl_wan/models/wan22.py @@ -0,0 +1,344 @@ +"""Recipe-local WAN 2.2 step + stage + pipeline for REFL BPTT.""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any, Dict, Optional, Tuple + +import torch + +from unirl.models.types.diffusion import DiffuseWithGradResult +from unirl.models.wan21.conditions import WAN21Conditions +from unirl.models.wan22.bundle import WAN22Bundle +from unirl.models.wan22.diffusion import WAN22DiffusionStage, WAN22DiffusionStep +from unirl.models.wan22.pipeline import WAN22Pipeline +from unirl.types.sampling import DiffusionSamplingParams + +# Matches the mainline module-level constant in unirl/models/wan22/diffusion.py. +# We do not import the private symbol so the recipe stays decoupled from the +# mainline's private surface. +_WAN_TIMESTEP_SCALE: float = 1000.0 + +# Inclusive max for torch.Generator.manual_seed and torch initial_seed conventions. +MAX_TORCH_SEED = (1 << 63) - 1 + + +class Wan22ReflDiffusionStep(WAN22DiffusionStep): + """Recipe-local WAN 2.2 single-branch denoising for REFL BPTT.""" + + def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: WAN22Bundle, + sample: torch.Tensor, + sigma: torch.Tensor, + conditions: WAN21Conditions, + *, + branch: str, + use_high_noise: Optional[bool] = None, + ) -> torch.Tensor: + """Run exactly one CFG branch through the currently active WAN 2.2 DiT.""" + if branch not in {"cond", "uncond"}: + raise ValueError(f"Wan22ReflDiffusionStep.predict_noise: unknown branch={branch!r}") + if conditions.text is None or conditions.text.embeds is None: + raise ValueError("Wan22ReflDiffusionStep.predict_noise: conditions.text.embeds is None") + + prompt_embeds = conditions.text.embeds + if branch == "cond": + branch_embeds = prompt_embeds + else: + neg = conditions.negative_text + branch_embeds = neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + + if use_high_noise is None: + use_high_noise, _ = self._select_for_sigma( + sigma, + 1.0, + None, + boundary_ratio=model.boundary_ratio, + ) + + batch_size = int(sample.shape[0]) + timestep = sigma * _WAN_TIMESTEP_SCALE + if timestep.dim() == 0: + timestep = timestep.expand(batch_size) + elif int(timestep.shape[0]) != batch_size: + timestep = timestep.expand(batch_size) + + embeds_dtype = branch_embeds.dtype + sample_cast = sample.to(dtype=embeds_dtype) + + image_latent = conditions.image_latent + if image_latent is not None and image_latent.latents is not None: + sample_cat = torch.cat( + [sample_cast, image_latent.latents.to(device=sample_cast.device, dtype=embeds_dtype)], + dim=1, + ) + else: + sample_cat = sample_cast + + extra: Dict[str, Any] = {} + image_embed = conditions.image_embed + image_embeds = image_embed.embeds if image_embed is not None and image_embed.embeds is not None else None + if image_embeds is not None: + extra["encoder_hidden_states_image"] = image_embeds.to(device=sample_cast.device, dtype=embeds_dtype) + + return model.transformer( + use_high_noise=bool(use_high_noise), + hidden_states=sample_cat, + encoder_hidden_states=branch_embeds, + timestep=timestep, + return_dict=False, + **extra, + )[0] + + +class Wan22ReflDiffusionStage(WAN22DiffusionStage): + """WAN 2.2 diffusion stage + REFL BPTT sampling override.""" + + 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: + """ + High-level function for generating initial latents. + + Args: + batch_size: Total number of samples + latent_shape: Shape of a single latent (C, H, W) or (C, T, H, W) + device: Device for the tensor + dtype: Data type for the tensor + + Returns: + Latent tensor [batch_size, *latent_shape] + """ + 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: WAN21Conditions, + *, + schedule: torch.Tensor, + params: DiffusionSamplingParams, + initial_latents: Optional[torch.Tensor] = None, + ) -> DiffuseWithGradResult: + """Differentiable WAN 2.2 sampling for REFL-style BPTT training. + + Returns :class:`DiffuseWithGradResult` with the live-grad + ``z_final`` + scalar ``kl_loss``. + """ + + if conditions.text is None or conditions.text.embeds is None: + raise ValueError( + "Wan22ReflDiffusionStage.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"Wan22ReflDiffusionStage.diffuse_with_grad: " + f"schedule length {schedule.shape[0]} != T+1={T + 1}" + ) + self.strategy.init_schedule(schedule) + + latent_shape = self._latent_shape( + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ) + if initial_latents is not None: + if int(initial_latents.shape[0]) != batch_size: + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[0]=" + f"{int(initial_latents.shape[0])} != batch_size={batch_size}." + ) + if tuple(initial_latents.shape[1:]) != tuple(latent_shape): + raise ValueError( + f"Wan22ReflDiffusionStage.diffuse_with_grad: initial_latents.shape[1:]=" + f"{tuple(initial_latents.shape[1:])} != expected {tuple(latent_shape)} " + f"for num_frames={int(params.num_frames)}, " + f"height={int(params.height)}, width={int(params.width)}." + ) + 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), + ) + + + # BPTT knobs. + sk: Dict[str, Any] = dict(getattr(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)) + + 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 + + guidance_scale = float(params.guidance_scale) + guidance_scale_2 = ( + params.guidance_scale_2 if params.guidance_scale_2 is not None else self.model.guidance_scale_2 + ) + boundary_ratio = float(self.model.boundary_ratio) + step = self.step + if not isinstance(step, Wan22ReflDiffusionStep): + raise TypeError( + f"Wan22ReflDiffusionStage.diffuse_with_grad requires Wan22ReflDiffusionStep, " + f"got {type(step).__name__}." + ) + + dual = self.model.transformer + kl_total = torch.zeros((), 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) + + use_high_noise, active_guidance = step._select_for_sigma( + sigma, + guidance_scale, + guidance_scale_2, + boundary_ratio=boundary_ratio, + ) + active_sub = getattr(dual, "high_noise" if use_high_noise else "low_noise", None) + use_cfg = active_guidance > 1.0 + grad_enabled = i >= mid_timestep + + if use_cfg: + cond_ctx = nullcontext() if grad_enabled else torch.no_grad() + with cond_ctx, autocast_ctx: + noise_pred_cond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + with torch.no_grad(), autocast_ctx: + noise_pred_uncond = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="uncond", + use_high_noise=use_high_noise, + ) + noise_pred_cond = noise_pred_cond.float() + noise_pred_uncond = noise_pred_uncond.float() + noise_pred = noise_pred_uncond + active_guidance * (noise_pred_cond - noise_pred_uncond) + kl_pred = noise_pred_cond + else: + pred_ctx = nullcontext() if grad_enabled else torch.no_grad() + with pred_ctx, autocast_ctx: + noise_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + noise_pred = noise_pred.float() + kl_pred = noise_pred + + # Per-step KL against the LoRA-disabled reference, on the + # conditional prediction of the currently active branch. When CFG + # is enabled, this intentionally excludes the stop-grad uncond mix. + if kl_weight != 0.0 and grad_enabled: + if active_sub is not None and hasattr(active_sub, "disable_adapter_layers") and hasattr( + active_sub, "enable_adapter_layers" + ): + with torch.no_grad(), autocast_ctx: + active_sub.disable_adapter_layers() + try: + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + finally: + active_sub.enable_adapter_layers() + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() + kl_total = kl_total + kl_step + kl_steps += 1 + + # Keep the transition outside autocast and under the same BPTT + # window used by WAN21: grad-enabled for train steps, no-grad + # before ``mid_timestep``. + transition_ctx = nullcontext() if grad_enabled else torch.no_grad() + with transition_ctx: + new_latents, _, _ = 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 Wan22ReflPipeline(WAN22Pipeline): + """WAN 2.2 pipeline for the REFL recipe (Scheme A: post-swap ``self.diffusion``).""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + old = self.diffusion + assert isinstance(old, WAN22DiffusionStage), ( + f"Wan22ReflPipeline expects parent to build WAN22DiffusionStage, " + f"got {type(old).__name__}" + ) + self.diffusion = Wan22ReflDiffusionStage( + model=old.model, + step=Wan22ReflDiffusionStep(), + strategy=old.strategy, + autocast_precision=old.autocast_dtype, + trajectory_precision=old.trajectory_dtype, + logprob_precision=old.logprob_dtype, + ) + + +__all__ = ["Wan22ReflDiffusionStep", "Wan22ReflDiffusionStage", "Wan22ReflPipeline"] diff --git a/recipes/refl_wan/rewards/__init__.py b/recipes/refl_wan/rewards/__init__.py new file mode 100644 index 000000000..8eedd85f6 --- /dev/null +++ b/recipes/refl_wan/rewards/__init__.py @@ -0,0 +1 @@ +"""Recipe-local reward implementations.""" diff --git a/recipes/refl_wan/rewards/face/__init__.py b/recipes/refl_wan/rewards/face/__init__.py new file mode 100644 index 000000000..d4b9c0826 --- /dev/null +++ b/recipes/refl_wan/rewards/face/__init__.py @@ -0,0 +1,5 @@ +"""Recipe-local face reward for REFL.""" + +from .scorer import FaceRewardScorer, FaceRewardSpec + +__all__ = ["FaceRewardScorer", "FaceRewardSpec"] diff --git a/recipes/refl_wan/rewards/face/face_tools.py b/recipes/refl_wan/rewards/face/face_tools.py new file mode 100644 index 000000000..6a544c4a8 --- /dev/null +++ b/recipes/refl_wan/rewards/face/face_tools.py @@ -0,0 +1,439 @@ +"""A simple, flexible implementation of a face analysis tool.""" +import math +import os + +import onnx +import torch +import torch.nn.functional as F +from onnx2torch import convert +from skimage import transform as trans +from torchvision.transforms.functional import resize +import torchvision.ops as ops + + +arcface_dst = torch.tensor( + [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], + [41.5493, 92.3655], [70.7299, 92.2041]]).float() + +def distance2bbox(points, distance, max_shape=None): + """Decode distance prediction to bounding box.""" + x1 = points[:, 0] - distance[:, 0] + y1 = points[:, 1] - distance[:, 1] + x2 = points[:, 0] + distance[:, 2] + y2 = points[:, 1] + distance[:, 3] + if max_shape is not None: + x1 = x1.clamp(min=0, max=max_shape[1]) + y1 = y1.clamp(min=0, max=max_shape[0]) + x2 = x2.clamp(min=0, max=max_shape[1]) + y2 = y2.clamp(min=0, max=max_shape[0]) + return torch.stack([x1, y1, x2, y2], axis=-1) + +def distance2kps(points, distance, max_shape=None): + """Decode distance prediction to keypoints.""" + preds = [] + for i in range(0, distance.shape[1], 2): + px = points[:, i % 2] + distance[:, i] + py = points[:, i % 2 + 1] + distance[:, i + 1] + if max_shape is not None: + px = px.clamp(min=0, max=max_shape[1]) + py = py.clamp(min=0, max=max_shape[0]) + preds.append(px) + preds.append(py) + return torch.stack(preds, axis=-1) + + +def face_transform(data, center, output_size, scale, rotation, device): + def to_homogeneous(mat): + return torch.vstack([mat, torch.tensor([0., 0., 1.])]) + scale_ratio = scale + rot = float(rotation) * math.pi / 180.0 + cx = center[0] * scale_ratio + cy = center[1] * scale_ratio + + C, H, W = data.shape + + t1 = to_homogeneous(torch.tensor([ + [scale_ratio, 0, 0], + [0, scale_ratio, 0] + ])).float() + t2 = to_homogeneous(torch.tensor([ + [1, 0, -cx], + [0, 1, -cy] + ])).float() + cos_theta = math.cos(rot) + sin_theta = math.sin(rot) + t3 = to_homogeneous(torch.tensor([ + [cos_theta, -sin_theta, 0], + [sin_theta, cos_theta, 0] + ])).float() + t4 = to_homogeneous(torch.tensor([ + [1, 0, output_size / 2], + [0, 1, output_size / 2] + ])).float() + M_homogeneous = t4 @ t3 @ t2 @ t1 + M = M_homogeneous[:2, :] + T = torch.tensor([[2 / W, 0, -1], + [0, 2 / H, -1], + [0, 0, 1]]) + theta = torch.inverse(T @ M_homogeneous @ torch.inverse(T)) + theta = theta[:2, :].unsqueeze(0).to(device) + grid = F.affine_grid(theta, data.unsqueeze(0).size(), align_corners=True) + transformed = F.grid_sample(data.unsqueeze(0), grid, align_corners=True) + cropped = transformed[0] + cropped = cropped[:, :output_size, :output_size] + return cropped.unsqueeze(0), M + + +def trans_points2d(pts, M): + ones = torch.ones((pts.shape[0], 1), dtype=pts.dtype, device=pts.device) + points_hom = torch.cat([pts, ones], dim=1) + points_hom = points_hom.unsqueeze(-1) + transformed_hom = torch.matmul(M, points_hom) + transformed = transformed_hom[:, :2, :].squeeze(-1) + return transformed + + +def estimate_norm(lmk, image_size=112, mode='arcface'): + assert lmk.shape == (5, 2) + assert image_size % 112 == 0 or image_size % 128 == 0 + if image_size % 112 == 0: + ratio = float(image_size) / 112.0 + diff_x = 0 + else: + ratio = float(image_size) / 128.0 + diff_x = 8.0 * ratio + dst = arcface_dst * ratio + dst[:, 0] += diff_x + tform = trans.SimilarityTransform() + tform.estimate(lmk, dst) + M = torch.from_numpy(tform.params).float() + return M + + +def norm_crop(img, landmark, image_size=112, mode='arcface'): + """Align an image into ArcFace canonical 112x112 using a similarity transform. + + Differentiable in ``img`` — the only non-grad path is ``landmark``, which + comes from the no-grad SCRFD detector. This is what carries the REFL + gradient from cosine reward all the way back to the generated pixels. + """ + M_homogeneous = estimate_norm(landmark, image_size, mode) + C, H, W = img.shape + img = img.unsqueeze(0) + T = torch.tensor([[2 / W, 0, -1], + [0, 2 / H, -1], + [0, 0, 1]]) + T_inv = torch.inverse(T) + theta = torch.inverse(T @ M_homogeneous @ T_inv) + theta = theta[:2, :].unsqueeze(0).to(img.device) + grid = F.affine_grid(theta, img.size(), align_corners=True) + transformed = F.grid_sample(img, grid, align_corners=True) + cropped = transformed[0] + warped = cropped[:, :image_size, :image_size] + return warped + + +def invert_affine_transform(matrix): + L = matrix[..., :2] + T = matrix[..., 2:] + a, b = L[..., 0, 0], L[..., 0, 1] + c, d = L[..., 1, 0], L[..., 1, 1] + det = a * d - b * c + inv_det = 1.0 / det + inv_L = torch.stack([ + torch.stack([d * inv_det, -b * inv_det], dim=-1), + torch.stack([-c * inv_det, a * inv_det], dim=-1) + ], dim=-2) + inv_T = -torch.matmul(inv_L, T) + inv_matrix = torch.cat([inv_L, inv_T], dim=-1) + return inv_matrix + + +class Face(dict): + + def __init__(self, d=None, **kwargs): + if d is None: + d = {} + if kwargs: + d.update(**kwargs) + for k, v in d.items(): + setattr(self, k, v) + + def __setattr__(self, name, value): + if isinstance(value, (list, tuple)): + value = [self.__class__(x) + if isinstance(x, dict) else x for x in value] + elif isinstance(value, dict) and not isinstance(value, self.__class__): + value = self.__class__(value) + super(Face, self).__setattr__(name, value) + super(Face, self).__setitem__(name, value) + + __setitem__ = __setattr__ + + def __getattr__(self, name): + return None + + @property + def embedding_norm(self): + if self.embedding is None: + return None + return torch.norm(self.embedding) + + @property + def normed_embedding(self): + if self.embedding is None: + return None + return self.embedding / self.embedding_norm + + +class SCRFD: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + self.center_cache = {} + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.requires_grad_(False) + self.torch_model.to(self.device) + self.use_kps = True + self.fmc = 3 + self._num_anchors = 2 + self._feat_stride_fpn = [8, 16, 32] + self.input_size = (640, 640) + + def forward(self, det_img, threshold=0.5): + input_height = det_img.shape[2] + input_width = det_img.shape[3] + scores_list = [] + bboxes_list = [] + kpss_list = [] + net_outs = self.torch_model(det_img.float()) + + for idx, stride in enumerate(self._feat_stride_fpn): + scores = net_outs[idx].cpu() + bbox_preds = net_outs[idx + self.fmc].cpu() + bbox_preds = bbox_preds * stride + if self.use_kps: + kps_preds = net_outs[idx + self.fmc * 2].cpu() * stride + + height = input_height // stride + width = input_width // stride + K = height * width + key = (height, width, stride) + if key in self.center_cache: + anchor_centers = self.center_cache[key] + else: + rows = torch.arange(height) + cols = torch.arange(width) + grid_y, grid_x = torch.meshgrid(rows, cols, indexing='ij') + anchor_centers = torch.stack([grid_x, grid_y], dim=-1).float() + anchor_centers = (anchor_centers * stride).reshape((-1, 2)) + if self._num_anchors > 1: + anchor_centers = torch.stack([anchor_centers] * self._num_anchors, axis=1).reshape((-1, 2)) + if len(self.center_cache) < 100: + self.center_cache[key] = anchor_centers + + scores = scores.reshape(-1, 1) + bbox_preds = bbox_preds.reshape(-1, 4) + if self.use_kps: + kps_preds = kps_preds.reshape(-1, kps_preds.shape[-1]) + pos_mask = scores[:, 0] >= threshold + bboxes = distance2bbox(anchor_centers, bbox_preds) + pos_scores = scores[pos_mask] + pos_bboxes = bboxes[pos_mask] + scores_list.append(pos_scores) + bboxes_list.append(pos_bboxes) + if self.use_kps: + kpss = distance2kps(anchor_centers, kps_preds) + kpss = kpss.reshape((kpss.shape[0], -1, 2)) + pos_kpss = kpss[pos_mask] + kpss_list.append(pos_kpss) + + return scores_list, bboxes_list, kpss_list + + @torch.no_grad() + def detect(self, image, input_size=None, max_num=0, metric='default', nms_thresh=0.4, det_thresh=0.5): + assert input_size is not None or self.input_size is not None + input_size = self.input_size if input_size is None else input_size + + im_ratio = float(image.shape[1]) / image.shape[2] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + det_scale = float(new_height) / image.shape[1] + resized_img = resize(image, (new_height, new_width), antialias=False) + det_img = torch.zeros((3, input_size[1], input_size[0]), device=self.device) + det_img[:, :new_height, :new_width] = resized_img + det_img = det_img.unsqueeze(0) + scores_list, bboxes_list, kpss_list = self.forward(det_img, det_thresh) + + scores = torch.vstack(scores_list) + scores_ravel = scores.flatten() + order = torch.argsort(scores_ravel, descending=True) + bboxes = torch.vstack(bboxes_list) / det_scale + if self.use_kps: + kpss = torch.vstack(kpss_list) / det_scale + + pre_det = torch.cat((bboxes, scores), dim=1).float() + pre_det = pre_det[order] + keep = self.nms(pre_det, nms_thresh) + det = pre_det[keep, :] + + if self.use_kps: + kpss = kpss[order, :, :] + kpss = kpss[keep, :, :] + else: + kpss = None + return det, kpss + + def nms(self, dets, nms_thresh): + boxes = dets[:, :4] + scores = dets[:, 4] + keep = ops.nms(boxes, scores, iou_threshold=nms_thresh) + return keep.tolist() + + +class ArcFace: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.to(self.device) + # Frozen weights — gradient flows through inputs only. + self.torch_model.requires_grad_(False) + self.taskname = 'recognition' + self.input_size = (112, 112) + + def get(self, img, face, input_size=(112, 112)): + aimg = norm_crop(img, landmark=face.kps, image_size=self.input_size[0]) + im_ratio = float(aimg.shape[1]) / aimg.shape[2] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + resized_img = resize(aimg, (new_height, new_width), antialias=False) + face.embedding = self.get_feat(resized_img.unsqueeze(0)).flatten() + return face.embedding + + def compute_sim(self, feat1, feat2): + feat1 = feat1.ravel() + feat2 = feat2.ravel() + sim = torch.dot(feat1, feat2) / (torch.norm(feat1) * torch.norm(feat2)) + return sim + + def get_feat(self, imgs): + imgs = imgs[:, [2, 1, 0], :, :] + net_out = self.torch_model(imgs) + return net_out + + +class Landmark: + def __init__(self, model_file=None, device="cuda"): + self.model_file = model_file + self.device = device + model = onnx.load(self.model_file) + self.torch_model = convert(model) + self.torch_model.eval() + self.torch_model.to(device) + self.torch_model.requires_grad_(False) + self.lmk_dim = 2 + self.lmk_num = 106 + self.taskname = 'landmark_%dd_%d' % (self.lmk_dim, self.lmk_num) + self.input_size = (192, 192) + + def get(self, img, face, input_size=(192, 192)): + bbox = face.bbox + w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1]) + center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2 + rotate = 0 + _scale = self.input_size[0] / (max(w, h) * 1.5) + aimg, M = face_transform(img, center, self.input_size[0], _scale, rotate, img.device) + aimg = (aimg + 1) / 2 * 255. + aimg = aimg[:, [2, 1, 0], :, :] + + input_size = self.input_size if input_size is None else input_size + im_ratio = float(aimg.shape[2]) / aimg.shape[3] + model_ratio = float(input_size[1]) / input_size[0] + if im_ratio > model_ratio: + new_height = input_size[1] + new_width = int(new_height / im_ratio) + else: + new_width = input_size[0] + new_height = int(new_width * im_ratio) + det_scale = float(new_height) / aimg.shape[2] + resized_img = resize(aimg, (new_height, new_width), antialias=False) + det_img = torch.zeros((aimg.shape[0], 3, input_size[1], input_size[0]), device=self.device) + det_img[:, :, :new_height, :new_width] = resized_img + + pred = self.torch_model(det_img)[0] + pred = pred.reshape((-1, 2)) + if self.lmk_num < pred.shape[0]: + pred = pred[self.lmk_num * -1:, :] + pred[:, 0:2] += 1 + pred[:, 0:2] *= (self.input_size[0] // 2) + + IM = invert_affine_transform(M).to(img.device) + pred = trans_points2d(pred, IM) + face[self.taskname] = pred + return pred + + +class FaceAnalysis: + def __init__(self, root="~/.insightface", device="cuda"): + self.root = root + self.device = device + self.detection_root = os.path.join(root, "scrfd_10g_bnkps.onnx") + self.landmark_root = os.path.join(root, "2d106det.onnx") + self.arcface_root = os.path.join(root, "glintr100.onnx") + self.detection_model = SCRFD(self.detection_root, self.device) + self.landmark_model = Landmark(self.landmark_root, self.device) + self.arcface_model = ArcFace(self.arcface_root, self.device) + + def landmark_loss(self, id_landmark=None, gt_landmark=None, mask=None): + mask = mask.unsqueeze(-1).unsqueeze(-1) + error = torch.abs(id_landmark - gt_landmark) * mask + valid_frame_count = mask.sum() + 1e-8 + loss = error.sum() / valid_frame_count / id_landmark.shape[-2] + return loss + + def embedding_loss(self, id_embedding=None, gt_embedding=None, mask=None): + cos_sim = F.cosine_similarity(id_embedding, gt_embedding, dim=2) + cos_loss = (1 - cos_sim) * mask + valid_frame_count = mask.sum() + 1e-8 + loss = cos_loss.sum() / valid_frame_count + return loss + + def pool_embedding_loss(self, id_embedding=None, gt_embedding=None, id_mask=None): + """Pool-style cosine similarity between gen frames and any ref frame. + + Returns a scalar reward (per-call). The scorer calls this once per + sample and stacks the result into a per-sample reward tensor. + """ + id_emb_expanded = id_embedding.unsqueeze(2) + gt_emb_expanded = gt_embedding.unsqueeze(1) + gt_mask = torch.ones(gt_embedding.shape[0], gt_embedding.shape[1]).to(id_mask.device) + if gt_mask.shape[1] > 1: + gt_mask[:, 0] = 0 + is_all_zero = (gt_embedding == 0).all(dim=-1) + gt_mask[is_all_zero] = 0 + + cos_sim_all = F.cosine_similarity(id_emb_expanded, gt_emb_expanded, dim=3) + valid_mask = id_mask.unsqueeze(2) * gt_mask.unsqueeze(1) + + gt_valid_count = gt_mask.sum(dim=1) + 1e-8 + weight_matrix = valid_mask / (gt_valid_count.unsqueeze(1).unsqueeze(2) + 1e-8) + mean_similarities = (cos_sim_all * weight_matrix).sum(dim=2) + cos_loss = mean_similarities * id_mask + valid_frame_count = id_mask.sum() + 1e-8 + loss = cos_loss.sum() / valid_frame_count + return loss diff --git a/recipes/refl_wan/rewards/face/scorer.py b/recipes/refl_wan/rewards/face/scorer.py new file mode 100644 index 000000000..79d5330d7 --- /dev/null +++ b/recipes/refl_wan/rewards/face/scorer.py @@ -0,0 +1,308 @@ +"""Face similarity reward scorer (REFL-compatible / BPTT-differentiable).""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import List, Optional + +import torch + +from unirl.reward.base import BaseRewardComponentSpec +from unirl.reward.local.base import LocalRewardBackend +from unirl.reward.local.device import resolve_device +from unirl.types.reward import RewardRequest + +from .face_tools import Face, FaceAnalysis + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Reference-video loader (imageio + resize-to-cover + center_crop) +# --------------------------------------------------------------------------- + +def _load_ref_video_frames( + video_path: str, + *, + max_frames: int = 81, + max_pixels: int = 480 * 480, + height_div: int = 16, + width_div: int = 16, +) -> torch.Tensor: + """Load a reference video as ``(C, T, H, W)`` float32 in ``[-1, 1]``. + + Frame-for-frame preprocessing so the reference-side stays aligned with the + reference-loader used to train the REFL data pipeline. + """ + import imageio + import PIL.Image + import torchvision.transforms.functional as TF + + reader = imageio.get_reader(video_path) + total = int(reader.count_frames()) # type: ignore[attr-defined] + + # time_division_factor=4, remainder=1 + nf = min(max_frames, total) + while nf > 1 and nf % 4 != 1: + nf -= 1 + + first = PIL.Image.fromarray(reader.get_data(0)) + fw, fh = first.size + if fw * fh > max_pixels: + scale = (fw * fh / max_pixels) ** 0.5 + fh, fw = int(fh / scale), int(fw / scale) + target_h = fh // height_div * height_div + target_w = fw // width_div * width_div + + frames = [] + for i in range(nf): + frame = PIL.Image.fromarray(reader.get_data(i)) + iw, ih = frame.size + scale = max(target_w / iw, target_h / ih) + frame = TF.resize( + frame, + (round(ih * scale), round(iw * scale)), + interpolation=TF.InterpolationMode.BILINEAR, + ) + frame = TF.center_crop(frame, (target_h, target_w)) + t = TF.to_tensor(frame) # (C, H, W) [0, 1] + frame = TF.to_pil_image(t) # round-trip PIL to match training-time preprocessing + frames.append(TF.to_tensor(frame) * 2.0 - 1.0) + reader.close() + + # (T, C, H, W) -> (C, T, H, W) + return torch.stack(frames).permute(1, 0, 2, 3).contiguous() + + +# --------------------------------------------------------------------------- +# Reward scorer +# --------------------------------------------------------------------------- + +class FaceRewardScorer(LocalRewardBackend): + """SCRFD detection + ArcFace embedding + pool cosine similarity for REFL.""" + + canonical_model_name = "face" + input_kind = "video" + + def __init__(self, *, config: "FaceRewardSpec", base_device: str) -> None: + super().__init__( + model_name=self.canonical_model_name, + device=resolve_device(config.device, base_device), + batch_size=int(config.batch_size), + model_path=config.model_path, + image_size=config.image_size, + ref_max_frames=config.ref_max_frames, + ref_max_pixels=config.ref_max_pixels, + differentiable=config.differentiable, + ) + + def _load_model(self) -> None: + model_path = self.model_kwargs["model_path"] + self.model = FaceAnalysis(root=model_path, device=self.device) + self._image_size = int(self.model_kwargs.get("image_size", 112)) + self._ref_max_frames = int(self.model_kwargs.get("ref_max_frames", 81)) + self._ref_max_pixels = int(self.model_kwargs.get("ref_max_pixels", 480 * 480)) + self._differentiable = bool(self.model_kwargs.get("differentiable", True)) + self._ref_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + + def _compute_model_rewards(self, request: RewardRequest) -> List[float]: + raise NotImplementedError( + "FaceRewardScorer is REFL-only; use compute_rewards_differentiable()." + ) + + # ------------------------------------------------------------------ + # Per-video face embedding (recipe-local REFL implementation) + # ------------------------------------------------------------------ + + def _extract_face_embeddings( + self, + video: torch.Tensor, + *, + with_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Extract per-frame face embeddings. + + Args: + video: ``(C, T, H, W)`` float in ``[-1, 1]``. + with_grad: When True, the ArcFace forward is kept inside autograd + (used for the GENERATED video on the BPTT path); when False, + the whole computation runs under ``torch.no_grad`` (used for + the REFERENCE video, which never receives gradient). + + Returns: + embeddings: ``(1, T, 512)`` on ``self.device``. Differentiable in + ``video`` when ``with_grad=True``. + mask: ``(1, T)`` on ``self.device`` (1 = face found). + """ + fa = self.model + fa.detection_model.torch_model.to(self.device) + fa.arcface_model.torch_model.to(self.device) + + T = int(video.shape[1]) + embeddings: List[torch.Tensor] = [] + mask: List[int] = [] + zero_emb = None + + autograd_ctx = torch.enable_grad if with_grad else torch.no_grad + + for t in range(T): + frame = video[:, t].float().to(self.device) + bboxes, kpss = fa.detection_model.detect(frame) + if bboxes.shape[0] > 0: + indexed = [(i, x) for i, x in enumerate(bboxes)] + sorted_bboxes = sorted( + indexed, + key=lambda item: (item[1][2] - item[1][0]) * (item[1][3] - item[1][1]), + ) + max_index, _ = sorted_bboxes[-1] + face = Face( + bbox=bboxes[max_index][0:4], + kps=kpss[max_index], + det_score=bboxes[max_index][4], + ) + with autograd_ctx(): + emb = fa.arcface_model.get(frame, face) + embeddings.append(emb) + mask.append(1) + else: + if zero_emb is None: + zero_emb = torch.zeros(512, device=self.device) + embeddings.append(zero_emb) + mask.append(0) + + emb_stack = torch.stack(embeddings).unsqueeze(0) # (1, T, 512) + mask_tensor = torch.tensor(mask, device=self.device).unsqueeze(0).float() # (1, T) + return emb_stack, mask_tensor + + # ------------------------------------------------------------------ + # Reference-video caching + # ------------------------------------------------------------------ + + def _get_ref_embeddings(self, ref_video_path: str) -> tuple[torch.Tensor, torch.Tensor]: + cached = self._ref_cache.get(ref_video_path) + if cached is not None: + return cached + with torch.no_grad(): + video = _load_ref_video_frames( + ref_video_path, + max_frames=self._ref_max_frames, + max_pixels=self._ref_max_pixels, + ).to(self.device) + ref_emb, ref_mask = self._extract_face_embeddings(video, with_grad=False) + # Detach (defence-in-depth) before caching. + cached = (ref_emb.detach(), ref_mask.detach()) + self._ref_cache[ref_video_path] = cached + return cached + + # ------------------------------------------------------------------ + # Differentiable REFL reward entry point + # ------------------------------------------------------------------ + + def compute_rewards_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + """Score generated videos ``[B,C,T,H,W]`` against reference-video metadata.""" + if media_tensor.ndim != 5: + raise ValueError(f"FaceRewardScorer expects [B,C,T,H,W], got {tuple(media_tensor.shape)}") + B = int(media_tensor.shape[0]) + if len(prompts) != B: + raise ValueError(f"FaceRewardScorer: prompts length {len(prompts)} != batch size {B}.") + metadata = records or [None] * B + if len(metadata) != B: + raise ValueError(f"FaceRewardScorer: records length {len(metadata)} != batch size {B}.") + + rewards: List[torch.Tensor] = [] + for i in range(B): + md = metadata[i] or {} + ref_path: Optional[str] = md.get("ref_video_path") or md.get("ref_video") + if not ref_path: + raise ValueError( + f"FaceRewardScorer: sample {i} is missing metadata['ref_video_path']. " + "Make sure the JSONL row was written with absolute ref video paths." + ) + + ref_emb, ref_mask = self._get_ref_embeddings(str(ref_path)) + + gen_video = torch.clamp(media_tensor[i], -1, 1).to(self.device) + gen_emb, gen_mask = self._extract_face_embeddings( + gen_video, with_grad=self._differentiable + ) + + if int(gen_mask.sum().item()) == 0: + rewards.append(gen_video.sum() * 0.0) + continue + + score = self.model.pool_embedding_loss( + id_embedding=gen_emb, + gt_embedding=ref_emb, + id_mask=gen_mask, + ) + rewards.append(score) + + return torch.stack(rewards).float() + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def offload(self) -> None: + self._ref_cache.clear() + fa = self.model + if fa is not None: + for sub in (fa.detection_model, fa.landmark_model, fa.arcface_model): + if sub is not None and hasattr(sub.torch_model, "cpu"): + sub.torch_model.cpu() + torch.cuda.empty_cache() + + def onload(self) -> None: + fa = self.model + if fa is not None: + for sub in (fa.detection_model, fa.landmark_model, fa.arcface_model): + if sub is not None and hasattr(sub.torch_model, "to"): + sub.torch_model.to(self.device) + + def is_available(self) -> bool: + return bool(self._is_loaded) + + def dispose(self) -> None: + self.offload() + + +@dataclass +class FaceRewardSpec(BaseRewardComponentSpec): + """Typed config for :class:`FaceRewardScorer`. + + Args: + model_path: Directory containing ONNX model files + (scrfd_10g_bnkps.onnx, 2d106det.onnx, glintr100.onnx). Passed + straight to :class:`FaceAnalysis` as its model root. + device: "auto" / "cuda" / "cuda:N" — resolved against ``base_device``. + batch_size: kept for parity with sibling specs; the scorer iterates + per-sample internally because each sample needs its own ref video. + image_size: ArcFace alignment size, must be a multiple of 112 or 128. + ref_max_frames: Frame cap for ref-video decoding (default 81, matches + the REFL training-time cap). + ref_max_pixels: Pixel cap for ref-video decoding (default 480*480 — + keeps SCRFD's short-side input near its trained resolution). + differentiable: Keep autograd on the generated ArcFace forward (default + True — required for REFL). Set False to score under no_grad for + the historical GRPO/replay path. + """ + + model_path: str = "" + device: str = "auto" + batch_size: int = 1 + image_size: int = 112 + ref_max_frames: int = 81 + ref_max_pixels: int = 480 * 480 + differentiable: bool = True + + +__all__ = [ + "FaceRewardScorer", + "FaceRewardSpec", +] diff --git a/recipes/refl_wan/rewards/videoalign/__init__.py b/recipes/refl_wan/rewards/videoalign/__init__.py new file mode 100644 index 000000000..3ddfdd59e --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/__init__.py @@ -0,0 +1,17 @@ +"""Recipe-local VideoAlign reward (Qwen2-VL VQ/MQ/TA scorer). + +Self-contained port of the VideoAlign reward family for the REFL WAN recipe. +The reward model code (``Qwen2VLRewardModelBT``), prompt templates, config +dataclasses and checkpoint-loading helpers all live under +:mod:`recipes.refl_wan.rewards.videoalign.model`, matching the recipe-local +layout used by the WAN22 face reward. + +Public API +---------- +- :class:`VideoAlignRewardScorer` — REFL-compatible reward backend. +- :class:`VideoAlignSpec` — typed config dataclass. +""" + +from .scorer import VideoAlignRewardScorer, VideoAlignSpec + +__all__ = ["VideoAlignRewardScorer", "VideoAlignSpec"] diff --git a/recipes/refl_wan/rewards/videoalign/model/__init__.py b/recipes/refl_wan/rewards/videoalign/model/__init__.py new file mode 100644 index 000000000..ddd5a7684 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/__init__.py @@ -0,0 +1,42 @@ +"""Inference-only port of the VideoAlign reward model. + +This subpackage contains everything needed to *load and forward* the +VideoAlign Qwen2-VL reward checkpoint, but **none** of the training-time +code (no ``VideoVLMRewardTrainer``, no GSB CSV data loader, no Bradley-Terry +loss). The original training pipeline lives at the upstream +``TIGER-AI-Lab/VideoScore`` / ``Tencent-Hunyuan/VideoAlign`` repos; mmrl +ships a vendored snapshot for its own RL training. UniRL does not need any +of that — at REFL rollout time we only call ``forward`` on the reward model +to read out three scalars per (video, prompt) pair. + +Files +----- +- :mod:`prompt_template` — ``build_prompt`` + the four prompt variants used + by the published checkpoints. +- :mod:`configs` — ``ModelConfig`` / ``PEFTLoraConfig`` / + ``TrainingConfig`` (the *inference-relevant* subset of fields; loaded from + the checkpoint's ``model_config.json``). +- :mod:`reward_model` — ``Qwen2VLRewardModelBT`` — Qwen2-VL with an + ``rm_head`` linear projection to (VQ, MQ, TA) scalars. +- :mod:`checkpoint` — ``load_model_from_checkpoint`` (full / LoRA + branches, plus the transformers>=5 key-remap). +- :mod:`factory` — ``create_model_and_processor`` — builds the + model, processor and (optionally) wraps with PEFT LoRA. +""" + +from .checkpoint import load_model_from_checkpoint +from .configs import ModelConfig, PEFTLoraConfig, TrainingConfig +from .factory import create_model_and_processor +from .prompt_template import DIMENSION_DESCRIPTIONS, build_prompt +from .reward_model import Qwen2VLRewardModelBT + +__all__ = [ + "DIMENSION_DESCRIPTIONS", + "ModelConfig", + "PEFTLoraConfig", + "Qwen2VLRewardModelBT", + "TrainingConfig", + "build_prompt", + "create_model_and_processor", + "load_model_from_checkpoint", +] diff --git a/recipes/refl_wan/rewards/videoalign/model/checkpoint.py b/recipes/refl_wan/rewards/videoalign/model/checkpoint.py new file mode 100644 index 000000000..c201690e5 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/checkpoint.py @@ -0,0 +1,165 @@ +"""Checkpoint loader for VideoAlign reward checkpoints. + +Supports both layouts produced by the upstream trainer: + +1. **Full state dict** (``save_full_model=True``) — ``checkpoint-K/model.pth``. + Loaded directly into the freshly-constructed + :class:`Qwen2VLRewardModelBT`. + +2. **LoRA split** (the default for the published VideoAlign release) — + ``checkpoint-K/adapter_model.safetensors`` (LoRA weights) plus + ``checkpoint-K/non_lora_state_dict.pth`` (rm_head + special-token + embeddings + any other ``requires_grad`` non-LoRA tensor). Both pieces + are merged into the PEFT-wrapped model. + +Includes the transformers>=5 compatibility shim +``base_model.model.model.*`` → ``base_model.model.model.language_model.*`` +when the target model uses the new Qwen2-VL submodule layout but the +checkpoint was saved against the old one. The remap is a no-op when both +sides agree. +""" + +from __future__ import annotations + +import glob +import logging +import os +from typing import Dict, Optional, Tuple + +import safetensors.torch +import torch + +logger = logging.getLogger(__name__) + + +def _insert_adapter_name_into_state_dict( + state_dict: Dict[str, torch.Tensor], + adapter_name: str, + parameter_prefix: str, +) -> Dict[str, torch.Tensor]: + """Rewrite raw LoRA keys to match peft's ``{module}.lora_{A,B}.{adapter}.*`` layout.""" + peft_model_state_dict: Dict[str, torch.Tensor] = {} + for key, val in state_dict.items(): + if parameter_prefix in key: + suffix = key.split(parameter_prefix)[1] + if "." in suffix: + suffix_to_replace = ".".join(suffix.split(".")[1:]) + key = key.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + else: + key = f"{key}.{adapter_name}" + peft_model_state_dict[key] = val + else: + peft_model_state_dict[key] = val + return peft_model_state_dict + + +def _pick_checkpoint_dir(checkpoint_dir: str, checkpoint_step: Optional[int]) -> str: + """Return the absolute path of the ``checkpoint-`` subdir to load. + + ``checkpoint_step is None`` or ``-1`` selects the largest step found. + A specific step falls back to the latest with a warning if missing. + """ + candidates = glob.glob(os.path.join(checkpoint_dir, "checkpoint-*")) + if not candidates: + raise FileNotFoundError( + f"No 'checkpoint-*' subdirectory under {checkpoint_dir!r}. " + "Make sure you point ``reward_model_path`` at the directory that " + "contains ``model_config.json`` alongside one or more ``checkpoint-K`` dirs." + ) + candidates.sort(key=lambda x: int(x.split("-")[-1]), reverse=True) + + if checkpoint_step is None or checkpoint_step == -1: + chosen = candidates[0] + logger.info("Using latest VideoAlign checkpoint: %s", chosen) + return chosen + + explicit = os.path.join(checkpoint_dir, f"checkpoint-{checkpoint_step}") + if explicit in candidates: + logger.info("Using requested VideoAlign checkpoint: %s", explicit) + return explicit + + chosen = candidates[0] + logger.warning( + "Requested VideoAlign checkpoint-%s not found; falling back to latest %s", + checkpoint_step, chosen, + ) + return chosen + + +def load_model_from_checkpoint( + model: torch.nn.Module, + checkpoint_dir: str, + checkpoint_step: Optional[int], +) -> Tuple[torch.nn.Module, str]: + """Load VideoAlign weights into ``model`` from a ``checkpoint-K`` subdir. + + Args: + model: an already-constructed reward model (PEFT-wrapped or not). + checkpoint_dir: parent dir containing ``checkpoint-*`` subdirs. + checkpoint_step: which step to load; ``None`` / ``-1`` → latest. + + Returns: + ``(model, step_str)`` — the same model object (mutated in place by + ``load_state_dict``) and the step number of the loaded checkpoint + as a string (e.g. ``"3000"``), for downstream logging. + """ + checkpoint_path = _pick_checkpoint_dir(checkpoint_dir, checkpoint_step) + loaded_step = checkpoint_path.split("checkpoint-")[-1].split("/")[0] + + full_ckpt = os.path.join(checkpoint_path, "model.pth") + lora_ckpt = os.path.join(checkpoint_path, "adapter_model.safetensors") + non_lora_ckpt = os.path.join(checkpoint_path, "non_lora_state_dict.pth") + + if os.path.exists(full_ckpt): + model_state_dict = torch.load(full_ckpt, map_location="cpu", weights_only=True) + # The upstream checkpoints were written against the old Qwen2VL + # submodule layout (LM under ``base_model.model.model.*``, vision + # under ``base_model.model.visual.*``). transformers>=5 moved + # language layers under + # ``base_model.model.model.language_model.*`` and the visual tower + # under ``base_model.model.model.visual.*``. Detect by inspecting + # the target model's own state_dict keys and remap when needed. + target_keys = model.state_dict().keys() + needs_remap = any( + k.startswith("base_model.model.model.language_model.") + or k.startswith("base_model.model.model.visual.") + for k in target_keys + ) + if needs_remap: + new_state_dict: Dict[str, torch.Tensor] = {} + for key, value in model_state_dict.items(): + if key.startswith("base_model.model.model"): + new_key = "base_model.model.model.language_model" + key[len("base_model.model.model"):] + new_state_dict[new_key] = value + elif key.startswith("base_model.model.visual"): + new_key = "base_model.model.model.visual" + key[len("base_model.model.visual"):] + new_state_dict[new_key] = value + else: + new_state_dict[key] = value + model.load_state_dict(new_state_dict) + else: + model.load_state_dict(model_state_dict) + else: + # LoRA branch — merge LoRA adapter + non-LoRA tensors. + if not os.path.exists(lora_ckpt) or not os.path.exists(non_lora_ckpt): + raise FileNotFoundError( + f"Neither {full_ckpt!r} nor the LoRA pair " + f"({lora_ckpt!r} + {non_lora_ckpt!r}) was found under {checkpoint_path!r}." + ) + + lora_state_dict = safetensors.torch.load_file(lora_ckpt) + non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu") + + lora_state_dict = _insert_adapter_name_into_state_dict( + lora_state_dict, adapter_name="default", parameter_prefix="lora_", + ) + + model_state_dict = model.state_dict() + model_state_dict.update(non_lora_state_dict) + model_state_dict.update(lora_state_dict) + model.load_state_dict(model_state_dict) + + return model, loaded_step + + +__all__ = ["load_model_from_checkpoint"] diff --git a/recipes/refl_wan/rewards/videoalign/model/configs.py b/recipes/refl_wan/rewards/videoalign/model/configs.py new file mode 100644 index 000000000..478bb1def --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/configs.py @@ -0,0 +1,155 @@ +"""Dataclass configs for the VideoAlign reward model. + +Inference-only subset of the originals at +``mmrl/recipes/rewards/videoalign/vendor/videoalign/utils.py``. We keep the +exact field names + default values that appear in checkpoints' +``model_config.json``, so ``ModelConfig(**dict_from_json)`` / +``PEFTLoraConfig(**dict_from_json)`` continue to round-trip every public +VideoAlign release. Training-only fields (``vision_lr``, ``merger_lr``, +``conduct_eval``, ``logging_epochs`` …) live on :class:`TrainingConfig` and +are accepted-but-ignored by the inference path — they only end up here +because ``model_config.json`` was written by the trainer. + +Why we DON'T inherit from ``transformers.TrainingArguments`` +----------------------------------------------------------- +The mmrl vendor's ``TrainingConfig`` extends ``TrainingArguments`` which +runs an ``__post_init__`` validating distributed launch fields +(``local_rank``, ``deepspeed`` …). At reward *inference* time the call site +is ``TrainingConfig(load_from_pretrained=..., bf16=..., output_dir="")`` and +the only fields we actually read downstream are ``bf16`` / ``fp16`` / +``gradient_checkpointing`` / ``disable_flash_attn2``. So we re-declare a +plain ``@dataclass`` with just those fields plus a generous ``**kwargs`` +catch-all (``__init__`` ignores unknown keys via ``__init_subclass__`` — +no, simpler: we filter in the loader). That keeps load-time light and +removes the transformers private-symbol coupling entirely. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Literal, Optional + + +@dataclass +class TrainingConfig: + """Inference-relevant slice of the trainer's :class:`TrainingArguments`. + + Only ``bf16`` / ``fp16`` / ``gradient_checkpointing`` / + ``disable_flash_attn2`` are actually consumed by the inference path + (see :func:`recipes.refl_wan.rewards.videoalign.model.factory.create_model_and_processor`). + The other fields are kept so ``TrainingConfig(**model_config_json["training_args"])`` + still works when someone wants to introspect the saved config. + """ + + # Inference-time knobs (read by ``create_model_and_processor``) + output_dir: str = "" + bf16: bool = False + fp16: bool = False + gradient_checkpointing: bool = False + disable_flash_attn2: bool = False + + # Train-only — kept as no-op defaults so __init__ accepts them. + max_length: Optional[int] = None + dataset_num_proc: Optional[int] = None + center_rewards_coefficient: Optional[float] = None + vision_lr: Optional[float] = None + merger_lr: Optional[float] = None + special_token_lr: Optional[float] = None + conduct_eval: Optional[bool] = True + load_from_pretrained: Optional[str] = None + load_from_pretrained_step: Optional[int] = None + logging_epochs: Optional[float] = None + eval_epochs: Optional[float] = None + save_epochs: Optional[float] = None + remove_unused_columns: Optional[bool] = False + save_full_model: Optional[bool] = False + + +@dataclass +class PEFTLoraConfig: + """LoRA wiring for the reward model. + + For inference we typically want ``lora_enable=True`` only if the + checkpoint was saved as a LoRA split (``adapter_model.safetensors`` + + ``non_lora_state_dict.pth``) — :func:`load_model_from_checkpoint` + auto-detects which branch to take. The other fields drive the LoRA + target-module discovery inside :func:`create_model_and_processor`. + """ + + lora_enable: bool = False + vision_lora: bool = False + lora_r: int = 16 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + lora_target_modules: Optional[List[str]] = None + lora_namespan_exclude: Optional[List[str]] = None + lora_modules_to_save: Optional[List[str]] = None + lora_task_type: str = "CAUSAL_LM" + use_rslora: bool = False + num_lora_modules: int = -1 + + def __post_init__(self) -> None: + # Mirror the upstream normalisation: a single-element list is + # flattened to a scalar so peft's ``target_modules`` accepts it. + if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1: + self.lora_target_modules = self.lora_target_modules[0] + if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1: + self.lora_namespan_exclude = self.lora_namespan_exclude[0] + + +@dataclass +class ModelConfig: + """Backbone + reward-head configuration. + + Read from ``model_config.json::model_config``. Most fields just pass + through to :meth:`Qwen2VLRewardModelBT.from_pretrained`. + """ + + model_name_or_path: Optional[str] = None + model_revision: str = "main" + + # Output dimensionality of the reward head: 3 for joint (VQ, MQ, TA) + # heads, 1 for single-attribute checkpoints. + output_dim: int = 1 + + # Whether the checkpoint introduces ``<|VQ_reward|>`` / ``<|MQ_reward|>`` + # / ``<|TA_reward|>`` special tokens (matches the + # ``detailed_special`` prompt template). When True, the reward is read + # from those token positions instead of the last token. + use_special_tokens: bool = False + + freeze_vision_tower: bool = field(default=False) + freeze_llm: bool = field(default=False) + tune_merger: bool = field(default=False) + + torch_dtype: Optional[Literal["auto", "bfloat16", "float16", "float32"]] = None + trust_remote_code: bool = False + attn_implementation: Optional[str] = None + load_in_8bit: bool = False + load_in_4bit: bool = False + bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4" + use_bnb_nested_quant: bool = False + + # Where the reward is read from the hidden-state sequence: + # ``"last"`` — last non-pad token (standard). + # ``"mean"`` — masked mean over the prompt. + # ``"special"`` — special-token positions (set automatically when + # ``use_special_tokens=True`` and the checkpoint + # declares ``additional_special_tokens``). + reward_token: Literal["last", "mean", "special"] = "last" + + # Bradley-Terry / regression flavour — train-only. Inference ignores + # this entirely; kept for round-trip compatibility with + # ``model_config.json``. + loss_type: Literal[ + "bt", "reg", "btt", "margin", "constant_margin", "scaled", "regular" + ] = "regular" + + def __post_init__(self) -> None: + if self.load_in_8bit and self.load_in_4bit: + raise ValueError( + "You can't use 8 bit and 4 bit precision at the same time" + ) + + +__all__ = ["ModelConfig", "PEFTLoraConfig", "TrainingConfig"] diff --git a/recipes/refl_wan/rewards/videoalign/model/factory.py b/recipes/refl_wan/rewards/videoalign/model/factory.py new file mode 100644 index 000000000..a2a6bcc98 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/factory.py @@ -0,0 +1,190 @@ +"""Build the VideoAlign reward model + Qwen2-VL processor. + +Inference-only counterpart to ``create_model_and_processor`` in +``mmrl/recipes/rewards/videoalign/vendor/videoalign/train_reward.py``. +Specifically: + +- Drops the ``trl.get_kbit_device_map`` / ``trl.get_quantization_config`` + dependency — we never 4-/8-bit quantise a reward model at REFL time. +- Drops the optimiser / loss / dataset side entirely. +- Keeps the optional LoRA wrapping and the optional ``<|VQ_reward|>`` / + ``<|MQ_reward|>`` / ``<|TA_reward|>`` special-token registration, since + both are needed to construct a *load-able* parameter graph for the + published checkpoints. +""" + +from __future__ import annotations + +import logging +from typing import Optional, Tuple + +import torch +from transformers import AutoProcessor + +from .configs import ModelConfig, PEFTLoraConfig, TrainingConfig +from .reward_model import Qwen2VLRewardModelBT + +logger = logging.getLogger(__name__) + + +def _find_target_linear_names( + model: torch.nn.Module, + num_lora_modules: int = -1, + lora_namespan_exclude: Optional[list] = None, +) -> list: + """Discover LoRA target modules by introspecting the model graph. + + Mirrors the upstream selection rule: every ``nn.Linear`` / ``nn.Embedding`` + whose qualified name doesn't contain any excluded namespace keyword. + ``num_lora_modules > 0`` truncates to the last N modules (the upstream + knob for cheaply LoRA-tuning only the top of the network). + """ + lora_namespan_exclude = lora_namespan_exclude or [] + if isinstance(lora_namespan_exclude, str): + lora_namespan_exclude = [lora_namespan_exclude] + + target_classes = (torch.nn.Linear, torch.nn.Embedding) + out = [] + for name, module in model.named_modules(): + if any(keyword in name for keyword in lora_namespan_exclude): + continue + if isinstance(module, target_classes): + out.append(name) + + if num_lora_modules > 0: + out = out[-num_lora_modules:] + return out + + +def create_model_and_processor( + model_config: ModelConfig, + peft_lora_config: PEFTLoraConfig, + training_args: TrainingConfig, + cache_dir: Optional[str] = None, +) -> Tuple[torch.nn.Module, AutoProcessor, object]: + """Build the Qwen2-VL reward model + matching processor. + + Args: + model_config: backbone configuration (matches the saved + ``model_config`` block in ``model_config.json``). + peft_lora_config: LoRA wiring (matches ``peft_lora_config``). + training_args: dtype / flash-attn knobs only; the dataset/optimiser + fields of :class:`TrainingConfig` are not used here. + cache_dir: optional HF cache directory. + + Returns: + ``(model, processor, peft_config)`` — same tuple shape as the + upstream helper. ``peft_config`` is ``None`` when LoRA is disabled. + """ + # Resolve the torch dtype string to a real dtype. + torch_dtype = ( + model_config.torch_dtype + if model_config.torch_dtype in ["auto", None] + else getattr(torch, model_config.torch_dtype) + ) + + # Build processor + (optional) special tokens. + processor = AutoProcessor.from_pretrained( + model_config.model_name_or_path, + padding_side="right", + cache_dir=cache_dir, + ) + + special_token_ids = None + if model_config.use_special_tokens: + special_tokens = ["<|VQ_reward|>", "<|MQ_reward|>", "<|TA_reward|>"] + processor.tokenizer.add_special_tokens( + {"additional_special_tokens": special_tokens} + ) + special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens) + + # Build the reward model. Quantisation is intentionally not supported + # here — the reward path expects full-precision (or bf16/fp16) weights. + # + # Aligned with mmrl's ``create_model_and_processor``: forward + # ``revision`` and ``use_cache`` directly into ``from_pretrained``. + # ``use_cache`` is a standard ``PretrainedConfig`` field, so HF's + # ``from_pretrained`` absorbs it into the config rather than passing + # it down to ``Qwen2VLRewardModelBT.__init__`` — no ``TypeError``. + # ``revision`` is a no-op for local-path checkpoints but matches the + # mmrl call site verbatim. + model = Qwen2VLRewardModelBT.from_pretrained( + model_config.model_name_or_path, + output_dim=model_config.output_dim, + reward_token=model_config.reward_token, + special_token_ids=special_token_ids, + torch_dtype=torch_dtype, + attn_implementation=( + "flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa" + ), + cache_dir=cache_dir, + revision=getattr(model_config, "model_revision", "main"), + use_cache=True if training_args.gradient_checkpointing else False, + ) + + # ------------------------------------------------------------------ + # transformers>=4.58/5.x compatibility shim (DISABLED). + # + # Older comment / behaviour preserved here for reference: under some + # future transformers versions, ``from_pretrained`` may stop treating + # ``use_cache`` as a known ``PretrainedConfig`` field and pass it + # through to ``Qwen2VLRewardModelBT.__init__``, which only accepts + # ``output_dim`` / ``reward_token`` / ``special_token_ids`` and would + # raise ``TypeError: unexpected keyword argument 'use_cache'``. If + # that happens, drop ``use_cache`` / ``revision`` from the + # ``from_pretrained`` call above and re-enable the post-hoc setter + # below: + # + # model.config.use_cache = bool( + # True if training_args.gradient_checkpointing else False + # ) + # ------------------------------------------------------------------ + + if model_config.use_special_tokens: + model.resize_token_embeddings(len(processor.tokenizer)) + + if training_args.bf16: + model.to(torch.bfloat16) + if training_args.fp16: + model.to(torch.float16) + + # Optional LoRA wrapping — required to *load* a LoRA-split checkpoint. + if peft_lora_config.lora_enable: + # peft is an optional dep for the non-LoRA inference path; import + # lazily so users who only ever load full-state-dict ckpts don't + # need it installed. + from peft import LoraConfig, get_peft_model + + namespan_exclude = list(peft_lora_config.lora_namespan_exclude or []) + if isinstance(peft_lora_config.lora_namespan_exclude, str): + namespan_exclude = [peft_lora_config.lora_namespan_exclude] + # Mirror upstream: when vision_lora is off, exclude the visual tower. + if not peft_lora_config.vision_lora and "visual" not in namespan_exclude: + namespan_exclude.append("visual") + + target_modules = _find_target_linear_names( + model, + num_lora_modules=peft_lora_config.num_lora_modules, + lora_namespan_exclude=namespan_exclude, + ) + peft_config = LoraConfig( + target_modules=target_modules, + r=peft_lora_config.lora_r, + lora_alpha=peft_lora_config.lora_alpha, + lora_dropout=peft_lora_config.lora_dropout, + task_type=peft_lora_config.lora_task_type, + use_rslora=peft_lora_config.use_rslora, + bias="none", + modules_to_save=peft_lora_config.lora_modules_to_save, + ) + model = get_peft_model(model, peft_config) + else: + peft_config = None + + model.config.tokenizer_padding_side = processor.tokenizer.padding_side + model.config.pad_token_id = processor.tokenizer.pad_token_id + + return model, processor, peft_config + + +__all__ = ["create_model_and_processor"] diff --git a/recipes/refl_wan/rewards/videoalign/model/prompt_template.py b/recipes/refl_wan/rewards/videoalign/model/prompt_template.py new file mode 100644 index 000000000..20b8f9d37 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/prompt_template.py @@ -0,0 +1,188 @@ +"""Prompt templates for the VideoAlign reward model. + +Verbatim port of ``mmrl/recipes/rewards/videoalign/vendor/videoalign/prompt_template.py``. +The string constants are part of the *training contract* of the published +checkpoints — changing them would silently shift the reward distribution +(the Qwen2-VL chat template was fed exactly these strings during reward-model +training). Do not edit unless you are retraining from scratch. +""" + +from __future__ import annotations + +from typing import List, Union + +VIDEOSCORE_QUERY_PROMPT = """ +Suppose you are an expert in judging and evaluating the quality of AI-generated videos, +please watch the frames of a given video and see the text prompt for generating the video, +then give scores based on its {dimension_name}, i.e., {dimension_description}. +Output a float number from 1.0 to 5.0 for this dimension, +the higher the number is, the better the video performs in that sub-score, +the lowest 1.0 means Bad, the highest 5.0 means Perfect/Real (the video is like a real video). +The text prompt used for generation is "{text_prompt}". +""" + +DIMENSION_DESCRIPTIONS = { + "VQ": [ + "visual quality", + "the quality of the video in terms of clearness, resolution, brightness, and color", + ], + "TA": [ + "text-to-video alignment", + "the alignment between the text prompt and the video content and motion", + ], + "MQ": [ + "motion quality", + "the quality of the motion in terms of consistency, smoothness, and completeness", + ], + "Overall": [ + "Overall Performance", + "the overall performance of the video in terms of visual quality, " + "text-to-video alignment, and motion quality", + ], +} + +SIMPLE_PROMPT = """ +Please evaluate the {dimension_name} of a generated video. Consider {dimension_description}. +The text prompt used for generation is "{text_prompt}". +""" + +DETAILED_PROMPT_WITH_SPECIAL_TOKEN = """ +You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. + +**Visual Quality:** +Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: +- **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. +- **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. +- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). +- **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. + +Please provide the ratings of Visual Quality: <|VQ_reward|> +END + +**Motion Quality:** +Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: +- **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). +- **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). +- **Aesthetic Quality:** The movement should be smooth and fluid. The transitions between different motions or camera angles should be seamless, and the overall dynamic feel should be visually pleasing. +- **Fusion:** Ensure that elements in motion (e.g., edges of the subject, hair, clothing) blend naturally with the background, without obvious artifacts or the feeling of cut-and-paste effects. +- **Clarity of Motion:** The video should be clear and smooth in motion. Pay attention to any areas where the video might have blurry or unsteady sections that hinder visual continuity. +- **Amplitude:** If the video is largely static or has little movement, assign a low score for motion quality. + +Please provide the ratings of Motion Quality: <|MQ_reward|> +END + +**Text Alignment:** +Assess how well the video matches the textual prompt across the following sub-dimensions: +- **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. +- **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. +- **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. + +Textual prompt - {text_prompt} +Please provide the ratings of Text Alignment: <|TA_reward|> +END +""" + +DETAILED_PROMPT = """ +You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. + +**Visual Quality:** +Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: +- **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. +- **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. +- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). +- **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. + +**Motion Quality:** +Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: +- **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). +- **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). +- **Aesthetic Quality:** The movement should be smooth and fluid. The transitions between different motions or camera angles should be seamless, and the overall dynamic feel should be visually pleasing. +- **Fusion:** Ensure that elements in motion (e.g., edges of the subject, hair, clothing) blend naturally with the background, without obvious artifacts or the feeling of cut-and-paste effects. +- **Clarity of Motion:** The video should be clear and smooth in motion. Pay attention to any areas where the video might have blurry or unsteady sections that hinder visual continuity. +- **Amplitude:** If the video is largely static or has little movement, assign a low score for motion quality. + + +**Text Alignment:** +Assess how well the video matches the textual prompt across the following sub-dimensions: +- **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. +- **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. +- **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. + +Textual prompt - {text_prompt} +Please provide the ratings of Visual Quality, Motion Quality, and Text Alignment. +""" + +SIMPLE_PROMPT_NO_PROMPT = """ +Please evaluate the {dimension_name} of a generated video. Consider {dimension_description}. +""" + + +def build_prompt( + prompt: str, + dimension: Union[str, List[str]], + template_type: str, +) -> str: + """Render the user-message text for the Qwen2-VL chat template. + + Args: + prompt: the text-to-video user prompt that produced the candidate + generation. May be ``""`` for the ``"none"`` template (then the + user message is empty — the model still reads the video frames). + dimension: ``"VQ"`` / ``"MQ"`` / ``"TA"`` / ``"Overall"`` (or a list + for multi-dim prompts). Drives the ``{dimension_name}`` / + ``{dimension_description}`` placeholders. + template_type: must match what the checkpoint was trained with. Read + this from ``model_config.json::data_config.prompt_template_type``. + One of ``"none"`` / ``"simple"`` / ``"video_score"`` / + ``"detailed_special"`` / ``"detailed"``. + + Returns: + The fully rendered user message string. + """ + if isinstance(dimension, list) and len(dimension) > 1: + dimension_name = ", ".join(DIMENSION_DESCRIPTIONS[d][0] for d in dimension) + dimension_name = f"overall performance({dimension_name})" + dimension_description = "the overall performance of the video" + else: + if isinstance(dimension, list): + dimension = dimension[0] + dimension_name = DIMENSION_DESCRIPTIONS[dimension][0] + dimension_description = DIMENSION_DESCRIPTIONS[dimension][1] + + if template_type == "none": + return prompt + elif template_type == "simple": + return SIMPLE_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + elif template_type == "video_score": + return VIDEOSCORE_QUERY_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + elif template_type == "detailed_special": + return DETAILED_PROMPT_WITH_SPECIAL_TOKEN.format(text_prompt=prompt) + elif template_type == "detailed": + return DETAILED_PROMPT.format(text_prompt=prompt) + else: + raise ValueError(f"Invalid template type: {template_type!r}") + + +__all__ = [ + "DIMENSION_DESCRIPTIONS", + "VIDEOSCORE_QUERY_PROMPT", + "SIMPLE_PROMPT", + "SIMPLE_PROMPT_NO_PROMPT", + "DETAILED_PROMPT", + "DETAILED_PROMPT_WITH_SPECIAL_TOKEN", + "build_prompt", +] diff --git a/recipes/refl_wan/rewards/videoalign/model/reward_model.py b/recipes/refl_wan/rewards/videoalign/model/reward_model.py new file mode 100644 index 000000000..654f89d83 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/model/reward_model.py @@ -0,0 +1,285 @@ +"""Qwen2-VL reward model with a 3-dim (VQ, MQ, TA) regression head. + +Verbatim port of ``Qwen2VLRewardModelBT`` from +``mmrl/recipes/rewards/videoalign/vendor/videoalign/trainer.py`` — the +*only* class from that file we actually need at inference time. Training- +specific machinery (``VideoVLMRewardTrainer``, ``PartialEmbeddingUpdateCallback``, +``compute_multi_attr_accuracy``) lived in the same file but are not part of +the inference contract and have been intentionally dropped to remove the +``transformers.trainer`` private-symbol coupling that breaks on +transformers>=5. + +The forward signature is preserved so the same checkpoint state_dict loads +cleanly (no key remapping needed beyond the standard transformers>=5 +``base_model.model.model.language_model.*`` shim handled by +:func:`recipes.refl_wan.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). +""" + +from __future__ import annotations + +from typing import Any, List, Optional + +import torch +import torch.nn as nn +from transformers import Qwen2VLForConditionalGeneration + + +def _cfg_get(config: Any, name: str) -> Any: + """Read a config attribute that may live on the top-level config or + on the nested ``text_config`` sub-config. + + transformers>=4.52 refactored ``Qwen2VLConfig`` so that fields like + ``hidden_size`` / ``image_token_id`` / ``video_token_id`` / ``pad_token_id`` + moved under ``config.text_config``. Older transformers exposed them at + the top level. This helper transparently supports both layouts and + raises a friendly error when the field is genuinely absent. + """ + if hasattr(config, name): + val = getattr(config, name) + if val is not None: + return val + text_cfg = getattr(config, "text_config", None) + if text_cfg is not None and hasattr(text_cfg, name): + return getattr(text_cfg, name) + # Fall back to whatever the top-level had (possibly None) so callers + # that tolerate ``None`` (e.g. ``pad_token_id``) keep working. + if hasattr(config, name): + return getattr(config, name) + raise AttributeError( + f"{type(config).__name__} has no attribute {name!r} " + f"(checked top-level and .text_config)." + ) + + +class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): + """Qwen2-VL backbone + ``nn.Linear`` reward head (Bradley-Terry / regression). + + Differences vs ``Qwen2VLForConditionalGeneration``: + + - Replaces the LM-head pathway with an ``rm_head: Linear(hidden, output_dim)`` + whose output is pooled per-sample (last / mean / special-token). + - ``forward`` returns ``{"logits": Tensor[B, output_dim]}`` rather than a + ``CausalLMOutputWithPast``. + + Output dim conventions used by VideoAlign checkpoints: + + - ``output_dim=3`` + ``reward_token="last"`` → joint head; final-token + logits are the (VQ, MQ, TA) scalars directly. + - ``output_dim=3`` + ``reward_token="special"`` + 3 special-token IDs → + each special token contributes its own row of the 3×3 head; we + diagonal-extract to (VQ from <|VQ_reward|>, MQ from <|MQ_reward|>, …). + - ``output_dim=1`` → single-attribute heads + (one ckpt per dimension, rarely used by the public release). + """ + + def __init__( + self, + config, + output_dim: int = 4, + reward_token: str = "last", + special_token_ids: Optional[List[int]] = None, + ) -> None: + super().__init__(config) + self.output_dim = output_dim + hidden_size = _cfg_get(config, "hidden_size") + self.rm_head = nn.Linear(hidden_size, output_dim, bias=False) + self.reward_token = reward_token + + self.special_token_ids = special_token_ids + # When special tokens are configured, the pooling mode is forced — + # otherwise the trainer-set ``reward_token`` setting on disk wins. + if self.special_token_ids is not None: + self.reward_token = "special" + + # ------------------------------------------------------------------ + # Backwards-compat shim: vision tower location + # ------------------------------------------------------------------ + # transformers <= 4.51 exposed the Qwen2-VL vision tower as + # ``self.visual`` directly on ``Qwen2VLForConditionalGeneration``. + # transformers >= 4.52 refactored it under ``self.model.visual`` + # (the LM backbone now owns the vision encoder). The forward code + # below was written against the old layout; rather than fork it, + # we expose a thin property that resolves to whichever location + # the currently-installed transformers uses. + # + # Implementation notes: + # - Read from ``self._modules`` (the underlying ``OrderedDict``) + # instead of ``getattr``/``hasattr`` to avoid recursing back into + # ``nn.Module.__getattr__`` which would re-trigger this property + # and yield a misleading "no attribute 'visual'" error. + # - We deliberately do *not* define a setter: ``super().__init__`` + # registers any ``self.visual = module`` assignment via + # ``nn.Module.__setattr__`` (which writes into ``_modules``), and + # the property here just reads that slot back out — so old-style + # checkpoints still load and the ``state_dict`` layout is + # completely unchanged. + @property + def visual(self): # type: ignore[override] + own = self._modules.get("visual", None) + if own is not None: + return own # transformers <= 4.51 layout + inner = self._modules.get("model", None) + if inner is not None: + inner_visual = getattr(inner, "visual", None) + if inner_visual is not None: + return inner_visual # transformers >= 4.52 layout + raise AttributeError( + "Qwen2VLRewardModelBT: vision tower not found on either " + "self.visual (transformers<=4.51) or self.model.visual " + "(transformers>=4.52). The installed transformers version " + "may be incompatible with this checkpoint." + ) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, # noqa: ARG002 — kept for API parity + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, # noqa: ARG002 — kept for API parity + ): + output_attentions = ( + output_attentions if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Visual + text token fusion. Identical to the upstream + # Qwen2VLForConditionalGeneration.forward except we don't compute + # an LM head — only the regression head over the final hidden + # states. + # + # Backwards-compat shim for transformers >= 4.52: the Qwen2-VL + # vision tower's forward used to return a raw ``Tensor`` (the + # flattened visual hidden states) but now returns a + # ``BaseModelOutputWithPooling`` dataclass. We accept both shapes + # by extracting ``.last_hidden_state`` when present. + def _as_tensor(visual_out): + if isinstance(visual_out, torch.Tensor): + return visual_out + # transformers >= 4.52 returns a ModelOutput-like dataclass; + # the visual hidden states live on ``last_hidden_state``. + t = getattr(visual_out, "last_hidden_state", None) + if t is not None: + return t + # Last resort: some forks return a tuple where the first + # element is the hidden-states tensor. + if isinstance(visual_out, (tuple, list)) and len(visual_out) > 0: + return visual_out[0] + raise TypeError( + "Qwen2-VL vision tower returned an unsupported type: " + f"{type(visual_out).__name__}. Expected Tensor or " + "BaseModelOutputWithPooling (transformers>=4.52)." + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.get_dtype()) + image_embeds = _as_tensor( + self.visual(pixel_values, grid_thw=image_grid_thw) + ) + image_token_id = _cfg_get(self.config, "image_token_id") + image_mask = (input_ids == image_token_id).unsqueeze(-1).expand_as(inputs_embeds) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) + video_embeds = _as_tensor( + self.visual(pixel_values_videos, grid_thw=video_grid_thw) + ) + video_token_id = _cfg_get(self.config, "video_token_id") + video_mask = (input_ids == video_token_id).unsqueeze(-1).expand_as(inputs_embeds) + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] # (B, L, D) + logits = self.rm_head(hidden_states) # (B, L, output_dim) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + # Locate per-sample sequence length so we can pool the right token. + try: + pad_token_id = _cfg_get(self.config, "pad_token_id") + except AttributeError: + pad_token_id = None + if pad_token_id is None and batch_size != 1: + raise ValueError( + "Cannot handle batch sizes > 1 if no padding token is defined." + ) + if pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # Find the first pad token; previous token is the last real + # token. ``%`` keeps the index in range for ONNX export. + sequence_lengths = torch.eq(input_ids, pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + if self.reward_token == "last": + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + elif self.reward_token == "mean": + valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1) + pooled_logits = torch.stack( + [logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)] + ) + elif self.reward_token == "special": + special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool) + for special_token_id in self.special_token_ids: + special_token_mask = special_token_mask | (input_ids == special_token_id) + pooled_logits = logits[special_token_mask, ...] + # Each sample has exactly 3 special tokens (VQ/MQ/TA), each + # producing an ``output_dim``-row of logits. Reshape → + # (B, 3, output_dim); when output_dim=3 we keep the diagonal + # (matching the special-token training contract). + pooled_logits = pooled_logits.view(batch_size, 3, -1) + if self.output_dim == 3: + pooled_logits = pooled_logits.diagonal(dim1=1, dim2=2) + pooled_logits = pooled_logits.view(batch_size, -1) + else: + raise ValueError(f"Invalid reward_token: {self.reward_token!r}") + + return {"logits": pooled_logits} + + +__all__ = ["Qwen2VLRewardModelBT"] diff --git a/recipes/refl_wan/rewards/videoalign/scorer.py b/recipes/refl_wan/rewards/videoalign/scorer.py new file mode 100644 index 000000000..180ee4a76 --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/scorer.py @@ -0,0 +1,239 @@ +"""VideoAlign reward scorer (REFL-compatible / BPTT-differentiable). + +Wraps :class:`VideoRewardWrapper` (Qwen2-VL-based reward model producing +three scalar scores per (video, prompt) pair: VQ / MQ / TA) and exposes a +recipe-local differentiable REFL entry point. + +Reward = ``w_vq * VQ + w_mq * MQ + w_ta * TA`` (defaults to 1 / 1 / 1). + +Gradient flow +------------- +The Qwen2-VL vision encoder is differentiable w.r.t. the input pixels when +the *fast* image processor is used (the wrapper force-installs +``Qwen2VLImageProcessorFast`` on construction). The generated video arrives +via ``compute_rewards_differentiable`` as ``[B, C, T, H, W]`` float in +``[-1, 1]`` with a live ``grad_fn`` (BPTT path); we forward into the wrapper under +``torch.enable_grad`` so the linear combination of VQ/MQ/TA traces back +through the vision tower into the diffusion graph. + +Self-containment +---------------- +This scorer no longer requires the sibling ``mmrl`` repo on disk. The +Qwen2-VL reward backbone, prompt template, checkpoint loader and +inference wrapper all live under +:mod:`recipes.refl_wan.rewards.videoalign.model` / :mod:`...wrapper`. The +``mmrl_repo_root`` Spec field has been removed; ``MMRL_REPO_ROOT`` env +var is now irrelevant. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import List, Optional + +import torch + +from unirl.reward.base import BaseRewardComponentSpec +from unirl.reward.local.base import LocalRewardBackend +from unirl.reward.local.device import resolve_device +from unirl.types.reward import RewardRequest + +from .wrapper import VideoRewardWrapper + +logger = logging.getLogger(__name__) + + +class VideoAlignRewardScorer(LocalRewardBackend): + """Qwen2-VL VideoAlign reward — VQ + MQ + TA → scalar per sample.""" + + canonical_model_name = "videoalign" + input_kind = "video" + + def __init__(self, *, config: "VideoAlignSpec", base_device: str) -> None: + super().__init__( + model_name=self.canonical_model_name, + device=resolve_device(config.device, base_device), + batch_size=int(config.batch_size), + reward_model_path=config.reward_model_path, + resize_height=config.resize_height, + resize_width=config.resize_width, + micro_batch_size=config.micro_batch_size, + reward_num_frames=config.reward_num_frames, + use_norm=config.use_norm, + w_vq=config.w_vq, + w_mq=config.w_mq, + w_ta=config.w_ta, + differentiable=config.differentiable, + ) + + # ------------------------------------------------------------------ + # Model loading + # ------------------------------------------------------------------ + + def _load_model(self) -> None: + reward_model_path = str(self.model_kwargs["reward_model_path"]) + if not reward_model_path: + raise ValueError( + "VideoAlignRewardScorer: ``reward_model_path`` must be set " + "(the directory containing ``model_config.json`` + the " + "``checkpoint-*`` subdir)." + ) + + self._w_vq = float(self.model_kwargs["w_vq"]) + self._w_mq = float(self.model_kwargs["w_mq"]) + self._w_ta = float(self.model_kwargs["w_ta"]) + self._use_norm = bool(self.model_kwargs["use_norm"]) + self._reward_num_frames = int(self.model_kwargs["reward_num_frames"]) + self._differentiable = bool(self.model_kwargs["differentiable"]) + + logger.info( + "VideoAlignRewardScorer: loading VideoRewardWrapper from %s", + reward_model_path, + ) + self.model = VideoRewardWrapper( + checkpoint_dir=reward_model_path, + device=self.device, + dtype=torch.bfloat16, + use_norm=self._use_norm, + resize_height=int(self.model_kwargs["resize_height"]), + resize_width=int(self.model_kwargs["resize_width"]), + micro_batch_size=int(self.model_kwargs["micro_batch_size"]), + ) + + def _compute_model_rewards(self, request: RewardRequest) -> List[float]: + raise NotImplementedError( + "VideoAlignRewardScorer is REFL-only; use compute_rewards_differentiable()." + ) + + # ------------------------------------------------------------------ + # Differentiable REFL reward entry point + # ------------------------------------------------------------------ + + def compute_rewards_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + del records + if media_tensor.ndim != 5: + raise ValueError(f"VideoAlignRewardScorer expects [B,C,T,H,W], got {tuple(media_tensor.shape)}") + if len(prompts) != int(media_tensor.shape[0]): + raise ValueError( + f"VideoAlignRewardScorer: prompts length {len(prompts)} " + f"!= batch size {int(media_tensor.shape[0])}." + ) + + # The wrapper expects per-sample [T, C, H, W] in [-1, 1]. + # NOTE: ``.clamp(-1.0, 1.0)`` mirrors mmrl's ``role.score`` — VAE + # decode can produce slightly out-of-range pixels (e.g. -1.02 / + # 1.03), and ``_pixels_neg1_to_255`` only clamps the [0, 1] + # midpoint afterwards, so values can still spill above 255 or + # below 0 without this safeguard. Required for numeric parity + # with the mmrl baseline. + per_sample_videos: List[torch.Tensor] = [] + for v in media_tensor: + v = ( + v.to(self.device) + .permute(1, 0, 2, 3) + .clamp(-1.0, 1.0) + .contiguous() + ) # → (T, C, H, W) + per_sample_videos.append(v) + + if self._reward_num_frames > 0: + ds: List[torch.Tensor] = [] + for v in per_sample_videos: + if v.shape[0] > self._reward_num_frames: + idx = torch.linspace( + 0, v.shape[0] - 1, self._reward_num_frames, + device=v.device, + ).long() + v = v[idx] + ds.append(v) + per_sample_videos = ds + + autograd_ctx = torch.enable_grad if self._differentiable else torch.no_grad + with autograd_ctx(): + scores = self.model.forward_scores( + per_sample_videos, prompts, use_norm=self._use_norm, + ) + + reward = ( + self._w_vq * scores["VQ"] + + self._w_mq * scores["MQ"] + + self._w_ta * scores["TA"] + ) + return reward.float() + + # ------------------------------------------------------------------ + # Lifecycle hooks (CPU offload between rollouts to free VRAM) + # ------------------------------------------------------------------ + + def offload(self) -> None: + if self.model is not None and getattr(self.model, "model", None) is not None: + self.model.model.cpu() + torch.cuda.empty_cache() + + def onload(self) -> None: + if self.model is not None and getattr(self.model, "model", None) is not None: + self.model.model.to(self.device) + + def is_available(self) -> bool: + return bool(self._is_loaded) + + def dispose(self) -> None: + self.offload() + + +# --------------------------------------------------------------------------- +# Spec +# --------------------------------------------------------------------------- + +@dataclass +class VideoAlignSpec(BaseRewardComponentSpec): + """Typed config for :class:`VideoAlignRewardScorer`. + + Args: + reward_model_path: Directory containing ``model_config.json`` and + the ``checkpoint-*`` subdir (with ``model.pth`` or LoRA split). + Required. + device: ``"auto"`` / ``"cuda"`` / ``"cuda:N"`` — resolved against + ``base_device``. + batch_size: kept for parity with sibling specs. + resize_height / resize_width: bicubic target before the Qwen2-VL + vision encoder. Defaults (336 × 588) match the published + checkpoints. + micro_batch_size: max samples per reward forward (peak-VRAM knob). + reward_num_frames: temporal downsample to this many uniformly + spaced frames before scoring; ``<= 0`` disables. + use_norm: z-score normalise each dimension using the means / stds + stored under ``inference_config`` in ``model_config.json``. + w_vq, w_mq, w_ta: linear combination weights into the final scalar. + differentiable: keep autograd on the reward forward (default True — + required for REFL). Set False for historical GRPO / replay. + """ + + reward_model_path: str = "" + + device: str = "auto" + batch_size: int = 1 + + resize_height: int = 336 + resize_width: int = 588 + micro_batch_size: int = 1 + reward_num_frames: int = 36 + + use_norm: bool = True + w_vq: float = 1.0 + w_mq: float = 1.0 + w_ta: float = 1.0 + + differentiable: bool = True + + +__all__ = [ + "VideoAlignRewardScorer", + "VideoAlignSpec", +] diff --git a/recipes/refl_wan/rewards/videoalign/wrapper.py b/recipes/refl_wan/rewards/videoalign/wrapper.py new file mode 100644 index 000000000..a85fbc0bc --- /dev/null +++ b/recipes/refl_wan/rewards/videoalign/wrapper.py @@ -0,0 +1,374 @@ +"""Differentiable forward over a loaded VideoAlign reward model. + +Self-contained re-implementation — does NOT touch ``sys.path`` . The Qwen2-VL reward +backbone, its prompt template and the checkpoint loader all live under +:mod:`recipes.refl_wan.rewards.videoalign.model`. + +Public API +---------- +``VideoRewardWrapper(checkpoint_dir, device, ...) -> .forward_scores(...)`` + +The ``forward_scores`` signature is preserved verbatim from the mmrl +wrapper so the same callsite — ``self.model.forward_scores(per_sample_videos, +prompts, use_norm=...)`` — keeps returning ``{"VQ": Tensor[B], "MQ": +Tensor[B], "TA": Tensor[B], "Overall": Tensor[B]}`` with autograd-live +graphs when the input videos have ``grad_fn`` set. + +Gradient flow +------------- +The Qwen2-VL vision encoder is fully differentiable w.r.t. its pixel input +when the *fast* image processor (``Qwen2VLImageProcessorFast``) is used — +the slow PIL-based variant routes through ``numpy`` and silently cuts the +graph. We force-swap to the fast processor on construction; this is the +single line that makes REFL gradients work end-to-end. +""" + +from __future__ import annotations + +# NOTE: ``inspect`` was previously used to filter ``batch`` against the +# reward backbone's ``forward`` signature in order to drop the +# ``mm_token_type_ids`` kwarg that transformers>=4.58/5.x injects via +# ``Qwen2VLProcessor``. We have aligned the runtime back to transformers +# 4.54 (matching the mmrl baseline), where that kwarg is not produced, so +# the filter is no longer necessary — and under PEFT it can spuriously +# strip ``pixel_values_videos`` / ``video_grid_thw`` when ``base.forward`` +# turns out to be ``LoraModel.forward(*args, **kwargs)``. The import is +# kept commented for future re-enablement once we move past 4.58. +# import inspect +import json +import logging +import os +from collections.abc import Mapping +from typing import Dict, List, Optional + +import torch +import torchvision.transforms.functional as TF +from torchvision.transforms import InterpolationMode +from transformers import AutoImageProcessor + +from .model import ( + ModelConfig, + PEFTLoraConfig, + TrainingConfig, + build_prompt, + create_model_and_processor, + load_model_from_checkpoint, +) + +logger = logging.getLogger(__name__) + + +def _load_configs_from_json(config_path: str): + """Parse ``model_config.json`` into the four typed configs. + + Drops trainer-only fields that may carry absolute filesystem paths + (they would break on a different machine and are never read at + inference). + """ + with open(config_path, "r", encoding="utf-8") as f: + config_dict = json.load(f) + + data_config = dict(config_dict["data_config"]) + data_config.pop("meta_data", None) + data_config.pop("data_dir", None) + + return ( + data_config, + config_dict["model_config"], + config_dict["peft_lora_config"], + config_dict.get("inference_config", None), + ) + + +class VideoRewardWrapper: + """Frozen VideoAlign reward model with a differentiable forward. + + Constructor responsibilities: + + 1. Parse ``model_config.json`` into typed configs. + 2. Build the Qwen2-VL reward model + processor via + :func:`create_model_and_processor`. + 3. Load weights from ``checkpoint-K`` via + :func:`load_model_from_checkpoint`. + 4. Replace the slow image processor with the autograd-friendly fast one. + 5. Move to the requested device + dtype, ``eval()`` + freeze + parameters (RL gradients flow *through* the activations, not into + the reward weights). + """ + + def __init__( + self, + checkpoint_dir: str, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + use_norm: bool = True, + resize_height: int = 336, + resize_width: int = 588, + micro_batch_size: int = 1, + ) -> None: + self.device = device + self.dtype = dtype + self.use_norm = use_norm + self.resize_height = resize_height + self.resize_width = resize_width + self.micro_batch_size = max(1, int(micro_batch_size)) + + config_path = os.path.join(checkpoint_dir, "model_config.json") + if not os.path.exists(config_path): + raise FileNotFoundError( + f"VideoRewardWrapper: expected ``model_config.json`` at " + f"{config_path!r}. Make sure ``reward_model_path`` points at " + "the directory containing both ``model_config.json`` and " + "one or more ``checkpoint-K`` subdirs (the layout produced " + "by the upstream VideoAlign trainer)." + ) + + data_config_dict, model_config_dict, peft_lora_config_dict, inference_config = \ + _load_configs_from_json(config_path) + + # We only need two fields out of the data_config block — the + # template type and the eval-dim list. Stash them directly without + # constructing the full dataclass. + self.prompt_template_type = data_config_dict.get("prompt_template_type", "none") + self.eval_dim = data_config_dict.get("eval_dim", "VQ") + self.inference_config = inference_config + self.build_prompt = build_prompt + + model_config = ModelConfig(**model_config_dict) + peft_lora_config = PEFTLoraConfig(**peft_lora_config_dict) + training_args = TrainingConfig( + load_from_pretrained=checkpoint_dir, + load_from_pretrained_step=-1, + gradient_checkpointing=False, + disable_flash_attn2=False, + bf16=(dtype == torch.bfloat16), + fp16=(dtype == torch.float16), + output_dir="", + ) + + model, processor, _ = create_model_and_processor( + model_config=model_config, + peft_lora_config=peft_lora_config, + training_args=training_args, + ) + model, _ = load_model_from_checkpoint(model, checkpoint_dir, -1) + + # CRITICAL: force the *fast* image processor. The slow variant + # converts inputs to numpy/PIL and (a) rejects float pixels outside + # [0, 1] via ``_rescale_for_pil_conversion``, and (b) silently + # severs autograd even when you push the values past the range + # check. The fast variant operates on torch tensors end-to-end via + # torchvision.transforms.v2.functional.resize, so gradients flow + # from pixels through the processor into the vision encoder. + fast_ip = AutoImageProcessor.from_pretrained( + model_config.model_name_or_path, use_fast=True + ) + processor.image_processor = fast_ip + + model.to(self.device) + model.eval() + model.requires_grad_(False) + + self.model = model + self.processor = processor + self.data_config = data_config_dict # raw dict, kept for debugging + + logger.info( + "VideoRewardWrapper loaded: ckpt=%s device=%s dtype=%s " + "resize=%dx%d micro_bs=%d use_norm=%s template=%s", + checkpoint_dir, self.device, self.dtype, self.resize_height, + self.resize_width, self.micro_batch_size, self.use_norm, + self.prompt_template_type, + ) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _prepare_input(self, data): + if isinstance(data, Mapping): + return type(data)({k: self._prepare_input(v) for k, v in data.items()}) + if isinstance(data, (tuple, list)): + return type(data)(self._prepare_input(v) for v in data) + if isinstance(data, torch.Tensor): + return data.to(device=self.device) + return data + + def _prepare_inputs(self, inputs): + inputs = self._prepare_input(inputs) + if len(inputs) == 0: + raise ValueError("VideoReward inputs must not be empty.") + return inputs + + @staticmethod + def _pixels_neg1_to_255(video: torch.Tensor) -> torch.Tensor: + """[-1, 1] float → [0, 255] float — stays differentiable.""" + return ((video * 0.5 + 0.5).clamp(0, 1) * 255.0).float() + + def _resize_for_reward(self, video_tchw: torch.Tensor) -> torch.Tensor: + return TF.resize( # type: ignore[no-any-return] + video_tchw, + [self.resize_height, self.resize_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + + # ------------------------------------------------------------------ + # Public interface + # ------------------------------------------------------------------ + + def prepare_batch_from_frames( + self, + video_tensors: List[torch.Tensor], + prompts: List[str], + ): + """Build a Qwen2-VL processor batch from ``[T,C,H,W]`` videos + prompts. + + Auto-transposes ``[C,T,H,W]`` callers when ``C==3``. + """ + chat_data = [ + [ + { + "role": "user", + "content": [ + {"type": "video", "video": "file://dummy_path"}, + { + "type": "text", + "text": self.build_prompt( + prompt, + self.eval_dim, + self.prompt_template_type, + ), + }, + ], + } + ] + for prompt in prompts + ] + + processed: List[torch.Tensor] = [] + for video in video_tensors: + if video.dim() != 4: + raise ValueError( + f"Expected video tensor shape [T,C,H,W], got {tuple(video.shape)}" + ) + # Auto-transpose if caller passed [C,T,H,W] with T > 3. + if video.shape[0] == 3 and video.shape[1] > 3: + video = video.permute(1, 0, 2, 3) + video = self._pixels_neg1_to_255(video) + video = self._resize_for_reward(video) + processed.append(video) + + batch = self.processor( + text=self.processor.apply_chat_template( + chat_data, tokenize=False, add_generation_prompt=True, + ), + images=None, + videos=processed, + padding=True, + return_tensors="pt", + videos_kwargs={"do_rescale": True}, + ) + return self._prepare_inputs(batch) + + def _norm( + self, + vq: torch.Tensor, + mq: torch.Tensor, + ta: torch.Tensor, + ): + if self.inference_config is None: + return vq, mq, ta + vq = (vq - self.inference_config["VQ_mean"]) / self.inference_config["VQ_std"] + mq = (mq - self.inference_config["MQ_mean"]) / self.inference_config["MQ_std"] + ta = (ta - self.inference_config["TA_mean"]) / self.inference_config["TA_std"] + return vq, mq, ta + + def forward_scores( + self, + video_tensors: List[torch.Tensor], + prompts: List[str], + use_norm: Optional[bool] = None, + ) -> Dict[str, torch.Tensor]: + """Differentiable scoring → per-sample (VQ, MQ, TA, Overall) scalars. + + Args: + video_tensors: list of ``(T, C, H, W)`` tensors, pixels in + ``[-1, 1]``. May carry ``grad_fn`` (REFL path). + prompts: per-sample user prompt text. + use_norm: override ``self.use_norm`` for this call. + + Returns: + ``{"VQ": Tensor[B], "MQ": Tensor[B], "TA": Tensor[B], "Overall": Tensor[B]}``. + Gradients flow back into ``video_tensors`` when they require grad. + """ + if len(video_tensors) != len(prompts): + raise ValueError( + "video_tensors and prompts must have the same batch size." + ) + use_norm = self.use_norm if use_norm is None else use_norm + + all_vq, all_mq, all_ta = [], [], [] + for start in range(0, len(video_tensors), self.micro_batch_size): + end = start + self.micro_batch_size + batch = self.prepare_batch_from_frames( + video_tensors[start:end], prompts[start:end] + ) + # Aligned with mmrl's wrapper: pass the processor batch + # straight through to the model. transformers 4.54 (current + # pinned version) does not inject ``mm_token_type_ids``, so + # no filtering is required. + logits = self.model(**batch, return_dict=True)["logits"] # (B, 3) + + # ---------------------------------------------------------- + # transformers>=4.58/5.x compatibility shim (DISABLED). + # + # Newer ``Qwen2VLProcessor`` versions inject ``mm_token_type_ids`` + # (text=0 / image=1 / video=2) into the batch dict for the + # rewritten ``get_rope_index``. ``Qwen2VLRewardModelBT`` was + # authored against transformers 4.45 and its ``forward`` does + # NOT list this kwarg, so ``self.model(**batch)`` would blow + # up with ``TypeError: ... got an unexpected keyword argument + # 'mm_token_type_ids'`` once upgraded. + # + # The previous implementation used ``inspect.signature`` on + # the unwrapped base forward to filter the batch, but PEFT + # wraps the model as ``PeftModel -> LoraModel -> Qwen2VLRewardModelBT`` + # and depending on the PEFT version ``get_base_model()`` may + # return ``LoraModel`` whose forward signature is + # ``(*args, **kwargs)`` — the resulting filter would silently + # drop ``pixel_values_videos`` / ``video_grid_thw``, leaving + # the reward model "blind" to the video. + # + # When we re-upgrade past 4.58, the safer replacement is an + # explicit black-list (NOT signature inspection): + # + # _DROP = {"mm_token_type_ids"} + # filtered_batch = {k: v for k, v in batch.items() if k not in _DROP} + # logits = self.model(**filtered_batch, return_dict=True)["logits"] + # + # Original code preserved below for reference: + # + # base = ( + # self.model.get_base_model() + # if hasattr(self.model, "get_base_model") + # else self.model + # ) + # allowed = set(inspect.signature(base.forward).parameters) + # filtered_batch = {k: v for k, v in batch.items() if k in allowed} + # logits = self.model(**filtered_batch, return_dict=True)["logits"] + # ---------------------------------------------------------- + vq, mq, ta = logits[:, 0], logits[:, 1], logits[:, 2] + if use_norm: + vq, mq, ta = self._norm(vq, mq, ta) + all_vq.append(vq) + all_mq.append(mq) + all_ta.append(ta) + + vq = torch.cat(all_vq, dim=0) + mq = torch.cat(all_mq, dim=0) + ta = torch.cat(all_ta, dim=0) + return {"VQ": vq, "MQ": mq, "TA": ta, "Overall": vq + mq + ta} + + +__all__ = ["VideoRewardWrapper"] diff --git a/recipes/refl_wan/roles.py b/recipes/refl_wan/roles.py new file mode 100644 index 000000000..4ad6dbe35 --- /dev/null +++ b/recipes/refl_wan/roles.py @@ -0,0 +1,224 @@ +"""Remote roles for the WAN22 REFL recipe.""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from typing import Any, List, Optional + +import torch +from hydra.utils import instantiate +from omegaconf import OmegaConf + +from unirl.distributed.group.dispatch import distributed +from unirl.distributed.tensor.batch import Batch, concat_field, shared_field +from unirl.reward.base import DifferentiableReward +from unirl.sde.runtime import get_sigma_schedule +from unirl.trainer.base_role import Role +from unirl.types.primitives import Images, Texts +from unirl.types.rollout_req import RolloutReq + + +@dataclass +class REFLGenerated(Batch): + """Generated BPTT payload carrying live-grad decoded pixels and KL loss.""" + + decoded: torch.Tensor = concat_field(default_factory=lambda: torch.empty(0)) + kl_loss: torch.Tensor = shared_field(default_factory=lambda: torch.zeros(1)) + + +@dataclass +class REFLLossMetrics(Batch): + """Per-DP-shard REFL scalar metrics.""" + + loss: List[float] = concat_field(default_factory=list) + reward_loss: List[float] = concat_field(default_factory=list) + kl_loss: List[float] = concat_field(default_factory=list) + reward_mean: List[float] = concat_field(default_factory=list) + + +def _maybe_instantiate(value: Any) -> Any: + if OmegaConf.is_config(value) and value.get("_target_") is not None: + return instantiate(value) + return value + + +class ReflActorRole(Role): + """Actor role: bundle + pipeline + backend + REFL BPTT logic.""" + + bundle: Any + pipeline: Any + backend: Any + algo_cfg: Any + sampling_params: Any + + def initialize(self) -> None: + super().initialize() + self.algo_cfg = self.cfg.algorithm + self.sampling_params = _maybe_instantiate(self.algo_cfg.get("sampling_params")) + + @distributed + def generate_samples(self, req: RolloutReq) -> REFLGenerated: + """Run live-grad diffusion sampling and VAE decode.""" + stage = self.pipeline.diffusion + decode_stage = self.pipeline.vae_decode + if not hasattr(stage, "diffuse_with_grad"): + raise RuntimeError("ReflActorRole: pipeline.diffusion lacks diffuse_with_grad(...).") + if not hasattr(decode_stage, "decode_with_grad"): + raise RuntimeError("ReflActorRole: pipeline.vae_decode lacks decode_with_grad(...).") + + texts = req.primitives.get("text") if req.primitives else None + if not isinstance(texts, Texts): + raise TypeError( + f"ReflActorRole.generate_samples: req.primitives['text'] must be Texts, " + f"got {type(texts).__name__ if texts is not None else 'None'}" + ) + negatives_raw = req.primitives.get("negative_text") if req.primitives else None + negatives = negatives_raw if isinstance(negatives_raw, Texts) else None + if negatives is not None and len(negatives.texts) != len(texts.texts): + raise ValueError( + f"ReflActorRole.generate_samples: negative_text length {len(negatives.texts)} " + f"!= text length {len(texts.texts)}" + ) + + params = self.sampling_params + primary_g = float(getattr(params, "guidance_scale", 1.0)) + secondary_g = getattr(params, "guidance_scale_2", None) + effective_guidance = max(primary_g, float(secondary_g)) if secondary_g is not None else primary_g + conditions = self.pipeline.build_conditions(texts, negatives=negatives, guidance_scale=effective_guidance) + + images_prim = req.primitives.get("image") if req.primitives else None + if images_prim is not None: + if not isinstance(images_prim, Images): + raise TypeError( + f"ReflActorRole.generate_samples: req.primitives['image'] must be Images, " + f"got {type(images_prim).__name__}" + ) + if int(images_prim.pixels.shape[0]) != len(texts.texts): + raise ValueError( + f"ReflActorRole.generate_samples: image count {images_prim.pixels.shape[0]} " + f"!= text count {len(texts.texts)}" + ) + from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage + from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage + + image_latent_cond = WAN21ImageLatentEncodeStage( + self.pipeline.bundle, + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ).encode(images_prim) + image_embed_cond = ( + WAN21CLIPVisionEncodeStage(self.pipeline.bundle).encode(images_prim) + if getattr(self.pipeline.bundle, "uses_clip_vision", False) + else None + ) + if image_latent_cond is not None or image_embed_cond is not None: + conditions = dataclasses.replace( + conditions, + image_latent=image_latent_cond, + image_embed=image_embed_cond, + ) + + device = getattr(getattr(self.pipeline, "bundle", None), "device", None) + schedule = get_sigma_schedule( + int(params.num_inference_steps), + shift=float(getattr(self.pipeline, "shift", 5.0)), + device=device, + ) + + train_model = getattr(self.backend, "model", None) + if train_model is not None and hasattr(train_model, "train"): + train_model.train() + self.backend.zero_grad() + + if bool(getattr(params, "init_same_noise", False)) and not getattr(params, "noise_group_ids", None): + params = dataclasses.replace(params, noise_group_ids=list(req.group_ids)) + + result = stage.diffuse_with_grad( + conditions, + schedule=schedule, + params=params, + ) + kl_loss = result.kl_loss + pixels = decode_stage.decode_with_grad(result.z_final) + return REFLGenerated(decoded=pixels, kl_loss=kl_loss.unsqueeze(0) if kl_loss.ndim == 0 else kl_loss) + + @distributed + def forward_backward_loss( + self, + *, + rewards: torch.Tensor, + kl_loss: Optional[torch.Tensor] = None, + ) -> REFLLossMetrics: + """Assemble REFL loss and run backward on the actor graph.""" + algo = self.algo_cfg + reward_weight = float(algo.get("reward_weight", 1.0)) + reward_baseline = float(algo.get("reward_baseline", 0.0)) + reward_scale = float(algo.get("reward_scale", 1.0)) + kl_weight = float(algo.get("kl_weight", 0.0)) + + reward = rewards.to(dtype=torch.bfloat16) + reward_loss = (-(reward - reward_baseline) / reward_scale * reward_weight).mean() + if kl_loss is not None and kl_weight != 0.0: + kl_term = kl_weight * kl_loss.squeeze() + else: + kl_term = torch.zeros((), device=reward.device, dtype=reward_loss.dtype) + loss = reward_loss + kl_term + loss.backward() + return REFLLossMetrics( + loss=[float(loss.detach().item())], + reward_loss=[float(reward_loss.detach().item())], + kl_loss=[float(kl_term.detach().item())], + reward_mean=[float(reward.detach().mean().item())], + ) + + +class ReflRewardRole(Role): + """Generic local differentiable reward role for REFL recipes.""" + + backend: Any + + @property + def preferred_input_kind(self) -> str: + kind = str(getattr(self.backend, "preferred_input_kind", "") or "").strip().lower() + if kind not in {"image", "video", "text"}: + raise ValueError(f"Reward backend must expose preferred_input_kind image/video/text, got {kind!r}.") + return kind + + @distributed + def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: + """Score live-grad decoded media through the official differentiable reward path.""" + decoded = generated.decoded + if not isinstance(decoded, torch.Tensor): + raise TypeError( + f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " + f"got {type(decoded).__name__}." + ) + texts = req.primitives.get("text") if req.primitives else None + if not isinstance(texts, Texts): + raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") + + if not isinstance(self.backend, DifferentiableReward): + raise TypeError( + f"{type(self.backend).__name__} must implement compute_rewards_differentiable " + "for REFL reward backprop." + ) + + records = list(req.metadata) if req.metadata else None + kind = self.preferred_input_kind + if kind == "image" and decoded.ndim != 4: + raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") + if kind == "video" and decoded.ndim != 5: + raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") + if kind not in {"image", "video"}: + raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") + + return self.backend.compute_rewards_differentiable( + decoded, + list(texts.texts), + records=records, + ) + + +__all__ = ["ReflActorRole", "ReflRewardRole"] diff --git a/recipes/refl_wan/run.py b/recipes/refl_wan/run.py new file mode 100644 index 000000000..673f2f06b --- /dev/null +++ b/recipes/refl_wan/run.py @@ -0,0 +1,19 @@ +#!/usr/bin/env python +"""Hydra entry point for the recipes.refl_wan REFL recipe.""" + +from __future__ import annotations + +import hydra +from omegaconf import DictConfig + +from recipes.refl_wan.trainer import REFLTrainer + + +@hydra.main(version_base=None, config_path="configs", config_name="wan22_face_refl") +def main(cfg: DictConfig) -> None: + trainer = REFLTrainer(cfg=cfg) + trainer.train() + + +if __name__ == "__main__": + main() diff --git a/recipes/refl_wan/start_wan21_t2v.sh b/recipes/refl_wan/start_wan21_t2v.sh new file mode 100644 index 000000000..dfaa20f2b --- /dev/null +++ b/recipes/refl_wan/start_wan21_t2v.sh @@ -0,0 +1,29 @@ +set -u + +# === Data === +# Newline-delimited prompts .txt (or {"prompt": ...} .jsonl). VideoAlign is +# text-conditioned; no reference video / first-frame needed. +export DATA_PATH=${DATA_PATH:-/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/filtered_720x1280_prompts.txt} +export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} + +# === Models === +# WAN 2.1 T2V 1.3B base checkpoint +export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/Wan2.1-T2V-1.3B-Diffusers} +# VideoAlign Qwen2-VL reward checkpoint +export VIDEOALIGN_MODEL_PATH=${VIDEOALIGN_MODEL_PATH:-/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/VideoReward} + +# === Output / Logging === +export OUTPUT_DIR=${OUTPUT_DIR:-./outputs/wan21_t2v_videoalign_refl} +export REPORT_TO_WANDB=${REPORT_TO_WANDB:-true} +export WANDB_PROJECT=${WANDB_PROJECT:-unirl-refl} +export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan21_t2v_videoalign_refl_recipe_opt} + +mkdir -p "${OUTPUT_DIR}" logs + +LOG_FILE="logs/wan21_t2v_videoalign_refl_$(date +%Y%m%d_%H%M%S).log" +echo "=== launching wan21 t2v videoalign refl, log → ${LOG_FILE} ===" + +RAY_ADDRESS=auto python -u -m recipes.refl_wan.run \ + --config-name=wan21_t2v_videoalign_refl \ + num_devices=8 \ + 2>&1 | tee "${LOG_FILE}" diff --git a/recipes/refl_wan/start_wan22_i2v.sh b/recipes/refl_wan/start_wan22_i2v.sh new file mode 100644 index 000000000..aeba8403f --- /dev/null +++ b/recipes/refl_wan/start_wan22_i2v.sh @@ -0,0 +1,13 @@ +export PRETRAINED_MODEL="/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/Wan2.2-I2V-A14B-Diffusers" +export DATA_PATH="/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl" +export EVAL_DATA_PATH="${DATA_PATH}" +export FACE_MODEL_PATH="/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/antelodev2" +export OUTPUT_DIR="outputs/wan22_face_refl_recipe_opt" + +export REPORT_TO_WANDB=true +export WANDB_PROJECT="unirl-refl" +export WANDB_RUN_NAME="wan22_face_refl_recipe_opt" + +RAY_ADDRESS=auto python -m recipes.refl_wan.run \ + num_devices=8 \ + 2>&1 | tee ./wan22_i2v_refl.log \ No newline at end of file diff --git a/recipes/refl_wan/trainer.py b/recipes/refl_wan/trainer.py new file mode 100644 index 000000000..2787163cc --- /dev/null +++ b/recipes/refl_wan/trainer.py @@ -0,0 +1,125 @@ +"""REFLTrainer — recipe-local trainer for WAN22 REFL/BPTT.""" + +from __future__ import annotations + +import time +from collections.abc import Mapping +from typing import Any, Dict + +from unirl.distributed.tensor.grad_context import enable_grad +from unirl.trainer.trainer import Trainer +from unirl.types.prompts import RolloutInputs +from unirl.types.rollout_req import RolloutReq +from unirl.types.sampling import total_samples_per_prompt + + +class REFLTrainer(Trainer): + """REFL / BPTT recipe trainer: role-driven 3-RPC train step.""" + + def validate_config(self) -> None: + super().validate_config() + if "actor" not in self.roles: + raise ValueError("REFLTrainer requires a role named 'actor'.") + if "reward" not in self.roles: + raise ValueError("REFLTrainer requires a role named 'reward'.") + + actor_spec = self._role_specs_by_name["actor"] + if actor_spec.raw_cfg.get("algorithm") is None: + raise ValueError("REFLTrainer requires roles[name=actor].algorithm: ${algorithm}.") + algo_cfg = self.cfg.get("algorithm") + if algo_cfg is None: + raise ValueError("REFLTrainer requires top-level cfg.algorithm.") + if hasattr(algo_cfg, "get") and algo_cfg.get("sampling_params") is None: + raise ValueError("REFLTrainer requires cfg.algorithm.sampling_params: ${sampling}.") + + reward_spec = self._role_specs_by_name["reward"] + backend_cfg = reward_spec.raw_cfg.get("backend") + backend_target = str(backend_cfg.get("_target_") or "") if backend_cfg is not None and hasattr(backend_cfg, "get") else "" + if backend_target.endswith("RemoteRewardBackend"): + raise ValueError( + "REFLTrainer requires a local differentiable reward backend; " + f"got roles[name=reward].backend._target_={backend_target!r}." + ) + + rollout_section = self.cfg.get("rollout", None) + if rollout_section is not None: + rollout_target = str(rollout_section.get("_target_", "")) if hasattr(rollout_section, "get") else "" + if rollout_target and not rollout_target.endswith("TrainsideRolloutEngine"): + raise ValueError( + "REFLTrainer: a rollout section is present but is not trainside. " + f"Got rollout._target_={rollout_target!r}." + ) + + def build_req(self, inputs: RolloutInputs, rollout_id: int) -> RolloutReq: + """Build one RolloutReq from data-source samples.""" + inputs = inputs.expand(total_samples_per_prompt(self.sampling_params)) + primitives: Dict[str, Any] = dict(inputs.primitives) + + diff_params = self.sampling_params.get("diffusion") + guidance_scale = float(getattr(diff_params, "guidance_scale", 1.0)) if diff_params is not None else 1.0 + sampler_kwargs = getattr(diff_params, "sampler_kwargs", {}) if diff_params is not None else {} + negative_prompt = sampler_kwargs.get("negative_prompt") if isinstance(sampler_kwargs, Mapping) else None + if negative_prompt is None and diff_params is not None: + negative_prompt = getattr(diff_params, "negative_prompt", None) + if negative_prompt is not None and guidance_scale > 1.0 and "negative_text" not in primitives: + texts = primitives.get("text") + if not hasattr(texts, "texts"): + raise TypeError( + "REFLTrainer.build_req: sampling negative_prompt requires " + "req.primitives['text'] to expose a `texts` field." + ) + text_list = getattr(texts, "texts") + text_cls: Any = type(texts) + primitives["negative_text"] = text_cls(texts=[str(negative_prompt)] * len(text_list)) + + return RolloutReq( + sample_ids=list(inputs.sample_ids), + group_ids=list(inputs.group_ids), + primitives=primitives, + request_conditions={}, + stage_config={}, + sampling_params=dict(self.sampling_params), + metadata=list(inputs.metadata) if inputs.metadata else [], + init_noise_group_ids=[], + init_noise_latent_shape=None, + ) + + def train_step(self, req: RolloutReq, *, training_progress: float = 0.0, rollout_id: int = 0) -> Dict[str, Any]: + """One REFL step: actor generate → reward score → actor backward → actor step.""" + t0 = time.perf_counter() + with enable_grad(): + gen = self.actor.generate_samples(req) + rewards = self.reward.score_differentiable(req=req, generated=gen) + loss_metrics = self.actor.forward_backward_loss( + rewards=rewards, + kl_loss=gen.kl_loss, + ) + step_result = self.actor.step() + + metrics: Dict[str, Any] = { + "loss": self._mean(getattr(loss_metrics, "loss", None)), + "reward_loss": self._mean(getattr(loss_metrics, "reward_loss", None)), + "kl_loss": self._mean(getattr(loss_metrics, "kl_loss", None)), + "reward_mean": self._mean(getattr(loss_metrics, "reward_mean", None)), + "grad_norm": self._mean(step_result.metrics.get("grad_norm")) if step_result.metrics else 0.0, + "lr": self._mean(step_result.metrics.get("lr")) if step_result.metrics else 0.0, + "step_time_s": time.perf_counter() - t0, + "training_progress": float(training_progress), + } + return metrics + + @staticmethod + def _mean(field: Any) -> float: + try: + if hasattr(field, "tolist"): + field = field.tolist() + if isinstance(field, (list, tuple)) and field: + return float(sum(float(x) for x in field) / len(field)) + if isinstance(field, (int, float)): + return float(field) + return 0.0 + except Exception: + return 0.0 + + +__all__ = ["REFLTrainer"] diff --git a/scripts/check_recipe_targets.py b/scripts/check_recipe_targets.py index ceb14ba3f..176ba4097 100755 --- a/scripts/check_recipe_targets.py +++ b/scripts/check_recipe_targets.py @@ -23,11 +23,11 @@ ROOT = Path(__file__).resolve().parents[1] # YAML trees that hold recipes / stage configs with ``_target_`` entries. -SCAN_DIRS = ["examples", "CPPO", "DRPO", "FlowDPPO", "unirl"] +SCAN_DIRS = ["examples", "recipes", "CPPO", "DRPO", "FlowDPPO", "unirl"] # Vendored / sub-project trees kept byte-pristine (mirror .pre-commit-config exclude). SKIP_PARTS = {".git", "vendor"} -_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?(unirl\.[A-Za-z0-9_.]+)['"]?\s*$""") +_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?((?:unirl|recipes)\.[A-Za-z0-9_.]+)['"]?\s*$""") @lru_cache(maxsize=None) diff --git a/unirl/data/data_source.py b/unirl/data/data_source.py index e2fa1a826..05d3ec87d 100644 --- a/unirl/data/data_source.py +++ b/unirl/data/data_source.py @@ -277,12 +277,14 @@ def __init__(self, args): args: Hydra ``cfg`` (DictConfig) with: - run.data_path: Path to data file (JSON, JSONL, or TXT) - run.seed: Random seed + - run.shuffle: Whether to shuffle the training prompts (default: True) - algorithm.prompts_per_rollout: Batch size """ self.args = args self.data_path = args.run.data_path self.eval_data_path = args.run.eval_data_path self.seed = args.run.seed + self.shuffle = bool(getattr(args.run, "shuffle", True)) self.prompts_per_rollout = int(args.algorithm.prompts_per_rollout) self.drop_last = True @@ -368,12 +370,13 @@ def _create_dataloader(self) -> None: f"prompts_per_rollout={self.prompts_per_rollout})." ) + should_shuffle = sampler is None and self.shuffle self._dataloader = DataLoader( self.train_dataset, batch_size=self.prompts_per_rollout, sampler=sampler, - shuffle=(sampler is None), # Only shuffle if not using custom sampler - generator=self._shuffle_generator if sampler is None else None, + shuffle=should_shuffle, # Only shuffle if not using custom sampler + generator=self._shuffle_generator if should_shuffle else None, num_workers=0, # Keep simple for Ray collate_fn=self._collate_text, drop_last=True, diff --git a/unirl/distributed/group/worker.py b/unirl/distributed/group/worker.py index 860ba4bc9..8e35940f7 100644 --- a/unirl/distributed/group/worker.py +++ b/unirl/distributed/group/worker.py @@ -306,9 +306,14 @@ def resolve(o): resolved_kwargs = map_tree(kwargs, resolve) if grad_mode: - # get_batch returns detached views/copies, so resolved tensors are - # fresh objects that don't alias store contents — mark them directly. - tensors = collect_leaves(resolved_args, Tensor) + collect_leaves(tuple(resolved_kwargs.values()), Tensor) + # Cross-RPC autograd can only propagate gradients back to controller-side + # TensorRef inputs recorded by Handle as input_metas. + tensors = [fetched[str(i)] for i in range(len(in_metas))] + if len(tensors) != len(in_metas): + raise RuntimeError( + f"Worker.call grad input mismatch for {method_name} call_id={call_id}: " + f"saved {len(tensors)} tensors for {len(in_metas)} TensorRef inputs" + ) for t in tensors: t.requires_grad_(True) t.retain_grad() diff --git a/unirl/models/types/diffusion.py b/unirl/models/types/diffusion.py index 114978141..e632ecc0d 100644 --- a/unirl/models/types/diffusion.py +++ b/unirl/models/types/diffusion.py @@ -14,6 +14,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional, Protocol, Tuple, TypeVar, runtime_checkable import torch @@ -151,5 +152,28 @@ def predict_noise_at_step( """ ... + # ------------------------------------------------------------------ + # BPTT path (REFL): generate-and-train in a single forward. + # ------------------------------------------------------------------ -__all__ = ["DiffusionStage", "DiffusionStep", "ReplayResult"] + def diffuse_with_grad( + self, + conditions: C, + *, + schedule: torch.Tensor, + params: object, + initial_latents: Optional[torch.Tensor] = None, + ) -> "DiffuseWithGradResult": + """Differentiable sampling: ``C → (z_final, kl_loss)``.""" + ... + + +@dataclass +class DiffuseWithGradResult: + """Output of :meth:`DiffusionStage.diffuse_with_grad`.""" + + z_final: torch.Tensor + kl_loss: torch.Tensor + + +__all__ = ["DiffusionStage", "DiffusionStep", "DiffuseWithGradResult", "ReplayResult"] diff --git a/unirl/models/wan21/bundle.py b/unirl/models/wan21/bundle.py index 2a5c29ccf..d7520cd39 100644 --- a/unirl/models/wan21/bundle.py +++ b/unirl/models/wan21/bundle.py @@ -81,14 +81,13 @@ def uses_clip_vision(self) -> bool: def from_config(cls, config: WAN21PipelineConfig) -> "WAN21Bundle": """Load all WAN 2.1 components from a HuggingFace checkpoint.""" try: - from diffusers import AutoencoderKLWan, WanTransformer3DModel + from diffusers import WanTransformer3DModel except ImportError: # Fallback for older diffusers: ``AutoModel`` does dynamic # dispatch on the checkpoint config. Matches the fallback in # legacy ``models/wan21.py``. from diffusers import AutoModel - AutoencoderKLWan = AutoModel WanTransformer3DModel = AutoModel try: from transformers import AutoTokenizer, UMT5EncoderModel @@ -136,7 +135,17 @@ def from_config(cls, config: WAN21PipelineConfig) -> "WAN21Bundle": vae: Optional[nn.Module] = None if config.load_vae: - vae = AutoencoderKLWan.from_pretrained(vae_path, subfolder="vae", torch_dtype=vae_dtype).to(device).eval() + from .wan_video_vae import WanVideoVAE + + vae = ( + WanVideoVAE.load_from_diffusers( + vae_path, + use_nested_grad_checkpoint=True, + use_act_grad_only_conv=True, + ) + .to(device=device, dtype=vae_dtype) + .eval() + ) vae.requires_grad_(False) text_encoder = ( diff --git a/unirl/models/wan21/image_encode.py b/unirl/models/wan21/image_encode.py index 9a591e660..1767aebde 100644 --- a/unirl/models/wan21/image_encode.py +++ b/unirl/models/wan21/image_encode.py @@ -109,20 +109,6 @@ def encode(self, p: Images) -> ImageLatentCondition: latent_condition = latent_condition.to(device=device, dtype=dtype) - # Per-channel normalization — strict inverse of ``wan21/vae.py:78-87`` - # decode: ``latent_decoded = stored * std + mean``. - vae_config = vae.config - latents_mean = getattr(vae_config, "latents_mean", None) - latents_std = getattr(vae_config, "latents_std", None) - if latents_mean is not None and latents_std is not None: - z_dim = int(getattr(vae_config, "z_dim", latent_condition.shape[1])) - mean = torch.tensor(latents_mean, device=device, dtype=dtype).view(1, z_dim, 1, 1, 1) - std = torch.tensor(latents_std, device=device, dtype=dtype).view(1, z_dim, 1, 1, 1) - latent_condition = (latent_condition - mean) / std - else: - scaling_factor = float(getattr(vae_config, "scaling_factor", 1.0)) - latent_condition = latent_condition * scaling_factor - # 4-channel first-frame mask (mirrors diffusers # ``pipeline_wan_i2v.py:468-479``): 1.0 at first pixel-time slot, # 0.0 elsewhere; then view-reshape + transpose so the temporal diff --git a/unirl/models/wan21/vae.py b/unirl/models/wan21/vae.py index 49e65238e..749fb4455 100644 --- a/unirl/models/wan21/vae.py +++ b/unirl/models/wan21/vae.py @@ -1,31 +1,24 @@ """WAN21VAEDecodeStage — LatentSegment → Videos via 3D VAE decode. +Framework-level ``DecodeStage`` implementation (no REFL-only knobs). Both +GRPO and REFL recipes route their final ``LatentSegment`` through this stage +— the memory / autograd optimizations in ``WanVideoVAE`` (below) are what +make BPTT decode feasible, so pushing this into a recipe would fork the +decode contract across recipes. Kept in the framework. + Implements ``DecodeStage[LatentSegment, Videos]``. Reads the final stored position from ``LatentSegment.latents[:, -1]`` (the clean latent at -``T``, which ``WAN21DiffusionStage`` always stores), denormalizes using -either per-channel ``latents_mean`` / ``latents_std`` (when the VAE -config carries them, as recent diffusers ``AutoencoderKLWan`` does) or -the scalar ``scaling_factor`` fallback, runs VAE decode in fp32 (bf16 is -unsupported by most VAE implementations), and packs the 5D -``[B, C, T, H, W]`` output into a varlen-batched ``Videos`` primitive. - -**Why per-channel mean/std support (vs the legacy scalar-only path):** -diffusers' canonical Wan VAE ships with ``latents_mean`` / -``latents_std`` arrays in its config; using only ``scaling_factor`` -yields off-distribution decodes. Legacy ``models/wan21.py::decode_latents`` -uses the scalar fallback alone — a known latent-norm bug. The new path -follows the diffusers spec. +``T``, which ``WAN21DiffusionStage`` always stores), then dispatches to +the bundle's ``WanVideoVAE.decode`` . All un-normalization +(per-channel ``mean`` / ``std``), spatial tiling, nested gradient +checkpointing inside the decoder, and ``Conv3dActGradOnly`` are owned +by ``WanVideoVAE`` itself; this stage is just the +``LatentSegment → Videos`` adapter. VAE encode for I2V's image-condition latent lives in -``WAN21ImageLatentEncodeStage`` (``image_encode.py``) — it shares the -same per-channel norm helper path but as a sibling stage so the decode -side stays focused on its single job. There is intentionally no -generic ``WAN21VAEEncodeStage`` here; the I2V encode is the only -encode path we need. - -Decode math derived from diffusers' ``AutoencoderKLWan`` reference -and ``unirl/models/wan21.py:325-342, 450-462`` (do NOT import -legacy code). +``WAN21ImageLatentEncodeStage`` (``image_encode.py``) — it goes through +the same ``WanVideoVAE.encode(x).latent_dist.mode()`` diffusers-style +contract; there is intentionally no generic ``WAN21VAEEncodeStage``. """ from __future__ import annotations @@ -83,44 +76,42 @@ def decode(self, s: LatentSegment, *, grad: bool = False, activation_checkpoint: f"[B, C, T_lat, H_lat, W_lat], got {tuple(clean.shape)}" ) - vae = self.bundle.vae - vae_config = vae.config - latents_mean = getattr(vae_config, "latents_mean", None) - latents_std = getattr(vae_config, "latents_std", None) - scaling_factor = getattr(vae_config, "scaling_factor", 1.0) - - device = clean.device - - def _decode(lat: torch.Tensor) -> torch.Tensor: - latents_f32 = lat.to(dtype=torch.float32) - if latents_mean is not None and latents_std is not None: - # diffusers Wan VAE spec: latent ↦ latent * std + mean - # (un-normalize). Reshape to [1, C, 1, 1, 1] to broadcast - # over the (B, T, H, W) axes. - z_dim = int(getattr(vae_config, "z_dim", lat.shape[1])) - mean = torch.tensor(latents_mean, device=device, dtype=torch.float32).view(1, z_dim, 1, 1, 1) - std = torch.tensor(latents_std, device=device, dtype=torch.float32).view(1, z_dim, 1, 1, 1) - latents_f32 = latents_f32 * std + mean - else: - latents_f32 = latents_f32 / float(scaling_factor) - - return vae.to(torch.float32).decode(latents_f32).sample - - with nullcontext() if grad else torch.no_grad(): - if grad and activation_checkpoint and clean.requires_grad: - from torch.utils.checkpoint import checkpoint - - decoded = checkpoint(_decode, clean, use_reentrant=False) - else: - decoded = _decode(clean) + with torch.no_grad(): + decoded = self._vae_decode(clean) # Decoded layout is [B, C, T_dec, H_dec, W_dec] in [-1, 1]. # Normalize to [0, 1] and clamp before packing. decoded = ((decoded + 1.0) / 2.0).clamp(0.0, 1.0) - # Pack into the varlen ``Videos`` primitive: ``Video.frames`` is - # ``[T, C, H, W]`` so we permute each sample from (C, T, H, W) → - # (T, C, H, W) and let ``Videos.from_list`` concat along T. + return self._pack_videos(decoded) + + # ------------------------------------------------------------------ + # BPTT path (REFL): differentiable decode of a live grad latent. + # ------------------------------------------------------------------ + + def decode_with_grad(self, z_final: torch.Tensor) -> torch.Tensor: + """Differentiable VAE decode: ``z_final → pixels`` with autograd alive.""" + if z_final.ndim != 5: + raise ValueError( + f"WAN21VAEDecodeStage.decode_with_grad: expected 5D z_final " + f"[B, C, T_lat, H_lat, W_lat], got {tuple(z_final.shape)}" + ) + return self._vae_decode(z_final) + + # ------------------------------------------------------------------ + # Shared decode kernel (used by both `decode` and `decode_with_grad`). + # ------------------------------------------------------------------ + + def _vae_decode(self, clean: torch.Tensor) -> torch.Tensor: + """Run ``WanVideoVAE.decode`` and return pixels in the VAE's native ``[-1, 1]``.""" + vae = self.bundle.vae + vae_dtype = vae.dtype + latents = clean.to(dtype=vae_dtype) + decoded = vae.decode(latents, device=latents.device, tiled=True) + return decoded + + def _pack_videos(self, decoded: torch.Tensor) -> Videos: + """Pack a dense ``[B, C, T_dec, H_dec, W_dec]`` pixel tensor into ``Videos``.""" videos = [Video(frames=decoded[i].permute(1, 0, 2, 3).contiguous()) for i in range(int(decoded.shape[0]))] return Videos.from_list(videos) diff --git a/unirl/models/wan21/wan_video_vae.py b/unirl/models/wan21/wan_video_vae.py new file mode 100644 index 000000000..9de86cc1d --- /dev/null +++ b/unirl/models/wan21/wan_video_vae.py @@ -0,0 +1,1203 @@ +"""WanVideoVAE — Wan 2.x Video VAE with training optimizations.""" + +from __future__ import annotations + +import os +from collections import OrderedDict + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange, repeat +from torch.nn.modules.utils import _triple + +CACHE_T = 2 + + +# ── Memory-efficient Conv layers ── + +class Conv3dActGradOnlyFunction(torch.autograd.Function): + """Conv3d that only computes input gradient, not weight gradient. + + For frozen VAE: weights never update, so no need to save activations + for weight grad computation. Saves significant GPU memory. + """ + + @staticmethod + def forward(ctx, input, weight, bias, stride, padding, dilation, groups): + ctx.save_for_backward(weight) + ctx.input_info = (input.size(), input.stride(), input.dtype) + ctx.stride = stride + ctx.padding = padding + ctx.dilation = dilation + ctx.groups = groups + return F.conv3d(input, weight, bias, stride, padding, dilation, groups) + + @staticmethod + def backward(ctx, grad_output): + weight, = ctx.saved_tensors + grad_input = None + if ctx.needs_input_grad[0]: + input = torch.empty_strided( + ctx.input_info[0], ctx.input_info[1], + dtype=ctx.input_info[2], device=grad_output.device) + grad_input = torch.ops.aten.convolution_backward( + grad_output, input, weight, None, + _triple(ctx.stride), _triple(ctx.padding), _triple(ctx.dilation), + False, [0, 0, 0], ctx.groups, (True, False, False))[0] + return grad_input, None, None, None, None, None, None + + +class Conv2dActGradOnlyFunction(torch.autograd.Function): + """Conv2d that only computes input gradient.""" + + @staticmethod + def forward(ctx, input, weight, bias, stride, padding, dilation, groups): + ctx.save_for_backward(weight) + ctx.input_shape = input.shape + ctx.stride = stride + ctx.padding = padding + ctx.dilation = dilation + ctx.groups = groups + return F.conv2d(input, weight, bias, stride, padding, dilation, groups) + + @staticmethod + def backward(ctx, grad_output): + weight, = ctx.saved_tensors + grad_input = None + if ctx.needs_input_grad[0]: + from torch.nn.grad import conv2d_input + grad_input = conv2d_input( + ctx.input_shape, weight, grad_output, + ctx.stride, ctx.padding, ctx.dilation, ctx.groups) + return grad_input, None, None, None, None, None, None + + +class Conv2dActGradOnly(nn.Conv2d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.weight.requires_grad = False + if self.bias is not None: + self.bias.requires_grad = False + + def forward(self, x): + return Conv2dActGradOnlyFunction.apply( + x, self.weight, self.bias, + self.stride, self.padding, self.dilation, self.groups) + + +class CausalConv3d(nn.Conv3d): + """Causal 3D convolution with temporal padding.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + self.time_kernel_size = self.kernel_size[0] + + def forward_pad(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return x + + def forward(self, x, cache_x=None): + x = self.forward_pad(x, cache_x) + return super().forward(x) + + +class CausalConv3dActGradOnly(CausalConv3d): + """CausalConv3d with act-grad-only optimization.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.weight.requires_grad = False + if self.bias is not None: + self.bias.requires_grad = False + + def forward(self, x, cache_x=None): + x = self.forward_pad(x, cache_x) + return Conv3dActGradOnlyFunction.apply( + x, self.weight, self.bias, + self.stride, self.padding, self.dilation, self.groups) + + +# ── Building blocks ── + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + self.channel_first = channel_first + self.scale = dim ** 0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else -1) + ) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', + 'downsample2d', 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + causal_conv3d = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + if hasattr(self, 'decoder') and getattr(self, 'decoder'): + setattr(causal_conv3d, 'decoder', True) + self.time_conv = causal_conv3d + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + causal_conv3d = CausalConv3d(dim, dim, (3, 1, 1), + stride=(2, 1, 1), padding=(0, 0, 0)) + if hasattr(self, 'decoder') and getattr(self, 'decoder'): + setattr(causal_conv3d, 'decoder', True) + self.time_conv = causal_conv3d + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if hasattr(self, 'decoder') and getattr(self, 'decoder'): + x1 = x[:, :, 1:] + x1 = self.time_conv(x1) + x1 = x1.reshape(b, 2, c, t - 1, h, w) + x1 = torch.stack((x1[:, 0], x1[:, 1]), 3) + x1 = x1.reshape(b, c, (t - 1) * 2, h, w) + x = torch.cat([x[:, :, 0:1], x1], dim=2) + elif feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != 'Rep': + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x], dim=2) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0], x[:, 1]), 3) + x = x.reshape(b, c, t * 2, h, w) + + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if hasattr(self, 'decoder') and getattr(self, 'decoder'): + x = self.time_conv(x) + elif feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + if feat_cache is None or feat_idx is None: + return x + return x, feat_cache, feat_idx + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 5: + return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, r=patch_size) + if x.dim() == 4: + return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", + q=patch_size, r=patch_size) + raise ValueError(f"Invalid input shape: {x.shape}") + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 5: + return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, r=patch_size) + if x.dim() == 4: + return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", + q=patch_size, r=patch_size) + return x + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + causal_conv3d_block1 = CausalConv3d(in_dim, out_dim, 3, padding=1) + causal_conv3d_block2 = CausalConv3d(out_dim, out_dim, 3, padding=1) + causal_conv3d_shortcut = CausalConv3d(in_dim, out_dim, 1) + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + causal_conv3d_block1, + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + causal_conv3d_block2) + self.shortcut = causal_conv3d_shortcut if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + h = self.shortcut(x) + for layer in self.residual: + if hasattr(self, 'decoder') and getattr(self, 'decoder'): + x = layer(x) + elif check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + if feat_cache is None or feat_idx is None: + return x + h + return x + h, feat_cache, feat_idx + + +class AttentionBlock(nn.Module): + """Causal self-attention with a single head.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +# ── Encoder / Decoder ── + +class Encoder3d(nn.Module): + def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], + num_res_blocks=2, attn_scales=[], + temperal_downsample=[True, True, False], dropout=0.0): + super().__init__() + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x, feat_cache, feat_idx + + +class Decoder3d(nn.Module): + def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], + num_res_blocks=2, attn_scales=[], + temperal_upsample=[False, True, True], dropout=0.0, + use_nested_grad_checkpoint=True): + super().__init__() + self.use_nested_grad_checkpoint = use_nested_grad_checkpoint + + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + setattr(self.conv1, "decoder", True) + + res_block_1 = ResidualBlock(dims[0], dims[0], dropout) + attn_block_2 = AttentionBlock(dims[0]) + res_block_3 = ResidualBlock(dims[0], dims[0], dropout) + for m in (res_block_1, attn_block_2, res_block_3): + setattr(m, 'decoder', True) + self.middle = nn.Sequential(res_block_1, attn_block_2, res_block_3) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i in (1, 2, 3): + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + res_block = ResidualBlock(in_dim, out_dim, dropout) + setattr(res_block, 'decoder', True) + upsamples.append(res_block) + if scale in attn_scales: + attn_block = AttentionBlock(out_dim) + setattr(attn_block, 'decoder', True) + upsamples.append(attn_block) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + resample = Resample(out_dim, mode=mode) + setattr(resample, 'decoder', True) + upsamples.append(resample) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + causal_conv3d = CausalConv3d(out_dim, 3, 3, padding=1) + setattr(causal_conv3d, 'decoder', True) + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), causal_conv3d) + + self.unsample_splits = [ + len(self.upsamples) // 15 * 10, + len(self.upsamples) // 15 * 13, + ] + + def forward(self, x): + def custom_forward1(x): + x = self.conv1(x) + for layer in self.middle: + x = layer(x) + for layer in self.upsamples[:self.unsample_splits[0]]: + x = layer(x) + return x + + def custom_forward2(x): + for layer in self.upsamples[self.unsample_splits[0]:self.unsample_splits[1]]: + x = layer(x) + return x + + def custom_forward3(x): + for layer in self.upsamples[self.unsample_splits[1]:]: + x = layer(x) + for layer in self.head: + x = layer(x) + return x + + if self.use_nested_grad_checkpoint: + x = torch.utils.checkpoint.checkpoint(custom_forward1, x, use_reentrant=True) + x = torch.utils.checkpoint.checkpoint(custom_forward2, x, use_reentrant=True) + x = torch.utils.checkpoint.checkpoint(custom_forward3, x, use_reentrant=True) + else: + x = custom_forward1(x) + x = custom_forward2(x) + x = custom_forward3(x) + return x + + +def count_conv3d(model): + return sum(1 for m in model.modules() if isinstance(m, CausalConv3d)) + + +# ── VideoVAE_ (core model) ── + +class VideoVAE_(nn.Module): + def __init__(self, dim=96, z_dim=16, dim_mult=[1, 2, 4, 4], + num_res_blocks=2, attn_scales=[], + temperal_downsample=[False, True, True], dropout=0.0, + use_nested_grad_checkpoint=True): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + self.use_nested_grad_checkpoint = use_nested_grad_checkpoint + + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout, + use_nested_grad_checkpoint) + + def encode(self, x, scale): + self.clear_cache() + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out, self._enc_feat_map, self._enc_conv_idx = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_, self._enc_feat_map, self._enc_conv_idx = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.conv2), z, use_reentrant=False) + out = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.decoder), x, use_reentrant=True) + return out + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +# ── WanVideoVAE (top-level wrapper) ── + +def _replace_conv_with_act_grad_only(model): + """Replace Conv layers with act-grad-only versions (post-construction). + + Avoids global monkey-patching of nn.Conv2d / CausalConv3d. + """ + for name, module in model.named_modules(): + if isinstance(module, CausalConv3d) and not isinstance(module, CausalConv3dActGradOnly): + parent_name, child_name = name.rsplit('.', 1) if '.' in name else ('', name) + parent = model if parent_name == '' else dict(model.named_modules())[parent_name] + new_module = CausalConv3dActGradOnly.__new__(CausalConv3dActGradOnly) + nn.Conv3d.__init__(new_module, module.in_channels, module.out_channels, + module.kernel_size, module.stride, (0, 0, 0), + module.dilation, module.groups, module.bias is not None) + new_module._padding = module._padding + new_module.padding = module.padding + new_module.time_kernel_size = module.time_kernel_size + new_module.weight = module.weight + new_module.weight.requires_grad = False + if module.bias is not None: + new_module.bias = module.bias + new_module.bias.requires_grad = False + setattr(parent, child_name, new_module) + + elif isinstance(module, nn.Conv2d) and not isinstance(module, Conv2dActGradOnly): + parent_name, child_name = name.rsplit('.', 1) if '.' in name else ('', name) + parent = model if parent_name == '' else dict(model.named_modules())[parent_name] + new_module = Conv2dActGradOnly( + module.in_channels, module.out_channels, module.kernel_size, + module.stride, module.padding, module.dilation, + module.groups, module.bias is not None) + new_module.weight = module.weight + new_module.weight.requires_grad = False + if module.bias is not None: + new_module.bias = module.bias + new_module.bias.requires_grad = False + setattr(parent, child_name, new_module) + + +def _get_task_range(total, world_size, rank): + """Evenly divide tasks across ranks.""" + per_rank = (total + world_size - 1) // world_size + start = rank * per_rank + end = min(start + per_rank, total) + return start, end + + +class _SimpleConfig: + """Minimal config object for diffusers pipeline compatibility. + + WanPipeline.__init__ reads vae.config.scale_factor_temporal/spatial. + """ + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +class _DeterministicLatentDist: + """Drop-in replacement for diffusers' ``DiagonalGaussianDistribution``. + + UniRL's I2V image-condition encode stage calls + ``vae.encode(x).latent_dist.mode()`` (see ``image_encode.py:100``) + to get the deterministic mode of the encoded Gaussian. ROLL's + ``VideoVAE_.encode`` only returns the (already-normalized) mean + directly, so ``mode()`` and ``sample()`` here are identical. + """ + + def __init__(self, latent: torch.Tensor) -> None: + self._latent = latent + + def mode(self) -> torch.Tensor: + return self._latent + + def sample(self, *_args, **_kwargs) -> torch.Tensor: + return self._latent + + +class _EncoderOutput: + """Drop-in replacement for diffusers' ``AutoencoderKLOutput``.""" + + def __init__(self, latent: torch.Tensor) -> None: + self.latent_dist = _DeterministicLatentDist(latent) + + +class WanVideoVAE(nn.Module): + """Wan Video VAE with training optimizations. + + Wraps VideoVAE_ with: + - Built-in latent normalization (mean/std) + - Tiled encode/decode for large resolutions + - Tiled parallel decode across SP ranks + - Optional act-grad-only conv optimization + """ + + def __init__(self, z_dim=16, use_nested_grad_checkpoint=True, + use_act_grad_only_conv=True): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + self.upsampling_factor = 8 + self.z_dim = z_dim + + # diffusers-compatible config for WanPipeline.__init__ compatibility. + # WanPipeline reads vae.config.scale_factor_temporal/spatial in __init__. + # ``scaling_factor`` is exposed only as a legacy fallback — the real + # un-normalization always uses ``latents_mean`` / ``latents_std``. + self.config = _SimpleConfig( + scale_factor_temporal=4, + scale_factor_spatial=8, + latents_mean=mean, + latents_std=[1.0 / s for s in std], + z_dim=z_dim, + scaling_factor=1.0, + ) + + # Build model with standard conv, then optionally replace with act-grad-only + self.model = VideoVAE_( + z_dim=z_dim, + use_nested_grad_checkpoint=use_nested_grad_checkpoint, + ).eval().requires_grad_(False) + + if use_act_grad_only_conv: + _replace_conv_with_act_grad_only(self.model) + + # ── Properties / dtype shim ── + + @property + def dtype(self) -> torch.dtype: + """Parameter dtype, matching the diffusers ``vae.dtype`` convention. + + UniRL's ``image_encode.py:94`` reads ``vae.dtype`` to cast the + VAE input. Expose it as a property over the first parameter so it + stays in sync if the module is later ``.to(dtype=...)``-cast. + """ + return next(self.parameters()).dtype + + # ── Encode / Decode ── + + def single_encode(self, video, device): + video = video.to(device) + return self.model.encode(video, self.scale) + + def single_decode(self, hidden_state, device): + hidden_state = hidden_state.to(device) + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + def encode(self, videos, device=None, tiled=False, + tile_size=(34, 34), tile_stride=(18, 16)): + """Encode videos to latent space. + + Two calling conventions are supported (so this class can drop into + both UniRL's batched-encode path and diffusers' single-tensor + ``vae.encode(x).latent_dist.mode()`` contract used by the I2V + image-condition encode stage): + + 1. **Batched / list**: ``videos`` is a list of ``(C, T, H, W)`` + tensors or a 5D ``(B, C, T, H, W)`` tensor, ``device`` must be + passed; returns a stacked ``(B, z_dim, T_lat, H_lat, W_lat)`` + latent tensor directly. + 2. **Diffusers-compatible**: ``videos`` is a 5D + ``(B, C, T, H, W)`` tensor, ``device`` is omitted; returns an + ``_EncoderOutput`` whose ``.latent_dist.mode()`` / + ``.latent_dist.sample()`` both return the (deterministic) + normalized latent. The VAE is non-stochastic in this port + , so ``mode()`` and ``sample()`` agree. + """ + # Diffusers-compatible path: single 5D tensor, infer device. + if device is None and isinstance(videos, torch.Tensor) and videos.dim() == 5: + target_device = videos.device + latents = self._encode_batched(videos, target_device, tiled, tile_size, tile_stride) + return _EncoderOutput(latents) + if device is None: + raise ValueError( + "WanVideoVAE.encode: ``device`` is required for the " + "batched/list calling convention." + ) + return self._encode_batched(videos, device, tiled, tile_size, tile_stride) + + def _encode_batched(self, videos, device, tiled, tile_size, tile_stride): + if isinstance(videos, torch.Tensor) and videos.dim() == 5: + videos = [videos[i] for i in range(videos.shape[0])] + + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size_px = (tile_size[0] * self.upsampling_factor, + tile_size[1] * self.upsampling_factor) + tile_stride_px = (tile_stride[0] * self.upsampling_factor, + tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, device, tile_size_px, tile_stride_px) + else: + hidden_state = self.single_encode(video, device) + hidden_states.append(hidden_state.squeeze(0)) + return torch.stack(hidden_states) + + def decode(self, hidden_states, device=None, tiled=True, sp_group=None, + tile_size=(34, 34), tile_stride=(18, 16)): + """Decode latents to video. + + ``device`` defaults to the latent's own device (so callers + upstream of UniRL's decode stage can ``vae.decode(z)`` without + threading the device through). ``tiled=True`` is the default + because uniform-tile + per-tile checkpointing is the only path + that survives BPTT memory pressure at 480×832×81f resolution. + """ + if device is None: + device = hidden_states.device + if tiled: + if sp_group is not None: + video = self.tiled_parallel_decode( + hidden_states, device, tile_size, tile_stride, sp_group) + else: + video = self.tiled_decode( + hidden_states, device, tile_size, tile_stride) + else: + video = self.single_decode(hidden_states, device) + return video + + # ── Tiled operations ── + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if not left_bound: + x[:border_width] = (torch.arange(border_width) + 1) / border_width + if not right_bound: + x[-border_width:] = torch.flip( + (torch.arange(border_width) + 1) / border_width, dims=(0,)) + return x + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + mask = torch.stack([h, w]).min(dim=0).values + return rearrange(mask, "H W -> 1 1 1 H W") + + def _make_tile_tasks(self, H, W, size_h, size_w, stride_h, stride_w): + tasks = [] + for h in range(0, H, stride_h): + if h - stride_h >= 0 and h - stride_h + size_h >= H: + continue + for w in range(0, W, stride_w): + if w - stride_w >= 0 and w - stride_w + size_w >= W: + continue + tasks.append((h, h + size_h, w, w + size_w)) + return tasks + + def tiled_decode(self, hidden_states, device, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = T * 4 - 3 + weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, + W * self.upsampling_factor), + dtype=hidden_states.dtype, device=device) + values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, + W * self.upsampling_factor), + dtype=hidden_states.dtype, device=device) + + for h, h_, w, w_ in tasks: + batch = hidden_states[:, :, :, h:h_, w:w_].to(device) + batch = self.model.decode(batch, self.scale).to(device) + mask = self.build_mask( + batch, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=((size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype, device=device) + + th = h * self.upsampling_factor + tw = w * self.upsampling_factor + values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += batch * mask + weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += mask + + values = values / weight + return values.clamp_(-1, 1) + + def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, + sp_group): + """Tiled decode with SP parallelism — each rank decodes a subset of tiles. + + Args: + sp_group: torch.distributed ProcessGroup for sequence parallelism. + """ + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tile_img_h = size_h * self.upsampling_factor + tile_img_w = size_w * self.upsampling_factor + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = T * 4 - 3 + weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, + W * self.upsampling_factor), + dtype=hidden_states.dtype, device=device) + values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, + W * self.upsampling_factor), + dtype=hidden_states.dtype, device=device) + + world_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + task_start, task_end = _get_task_range(len(tasks), world_size, sp_rank) + rank_tasks = tasks[task_start:task_end] + + all_decoded = [] + for h, h_, w, w_ in rank_tasks: + batch = hidden_states[:, :, :, h:h_, w:w_].to(device) + batch = self.model.decode(batch, self.scale).to(device) + # Pad to uniform tile size for all_gather + padding = (0, tile_img_w - batch.shape[-1], + 0, tile_img_h - batch.shape[-2]) + if padding != (0, 0, 0, 0): + batch = F.pad(batch, padding) + all_decoded.append(batch) + + # Stack and all_gather across SP ranks + local_stack = torch.stack(all_decoded, dim=0) # (num_local_tasks, B, C, T, H, W) + gathered = [torch.empty_like(local_stack) for _ in range(world_size)] + # Pad to max local tasks across ranks + max_tasks = max((task_end - task_start) + for s, e in [_get_task_range(len(tasks), world_size, r) + for r in range(world_size)] + for task_start, task_end in [(s, e)]) + if local_stack.shape[0] < max_tasks: + pad_n = max_tasks - local_stack.shape[0] + local_stack = torch.cat([ + local_stack, + torch.zeros(pad_n, *local_stack.shape[1:], + dtype=local_stack.dtype, device=local_stack.device) + ], dim=0) + + gathered = [torch.empty_like(local_stack) for _ in range(world_size)] + dist.all_gather(gathered, local_stack.contiguous(), group=sp_group) + all_decoded_global = torch.cat(gathered, dim=0) + + # Reassemble tiles (using global task order) + for i, (h, h_, w, w_) in enumerate(tasks): + if i >= all_decoded_global.shape[0]: + break + latent_h = hidden_states[:, :, :, h:h_, w:w_].shape[-2] + latent_w = hidden_states[:, :, :, h:h_, w:w_].shape[-1] + img_h = latent_h * self.upsampling_factor + img_w = latent_w * self.upsampling_factor + decoded_tile = all_decoded_global[i][:, :, :, :img_h, :img_w] + + mask = self.build_mask( + decoded_tile, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=((size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype, device=device) + + th = h * self.upsampling_factor + tw = w * self.upsampling_factor + values[:, :, :, th:th + decoded_tile.shape[3], + tw:tw + decoded_tile.shape[4]] += decoded_tile * mask + weight[:, :, :, th:th + decoded_tile.shape[3], + tw:tw + decoded_tile.shape[4]] += mask + + values = values / weight + return values.clamp_(-1, 1) + + def tiled_encode(self, video, device, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) + + out_T = (T + 3) // 4 + data_device = "cpu" + weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, + W // self.upsampling_factor), + dtype=video.dtype, device=data_device) + values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, + W // self.upsampling_factor), + dtype=video.dtype, device=data_device) + + for h, h_, w, w_ in tasks: + batch = video[:, :, :, h:h_, w:w_].to(device) + batch = self.model.encode(batch, self.scale).to(data_device) + mask = self.build_mask( + batch, + is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), + border_width=((size_h - stride_h) // self.upsampling_factor, + (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype, device=data_device) + + th = h // self.upsampling_factor + tw = w // self.upsampling_factor + values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += batch * mask + weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += mask + + return values / weight + + # ── HF Diffusers weight loading ── + + @classmethod + def load_from_diffusers(cls, pretrained_path, **kwargs): + """Load WanVideoVAE from HuggingFace diffusers format. + + Args: + pretrained_path: Path to model directory containing vae/ subfolder, + or direct path to vae directory. + """ + vae_dir = pretrained_path + if os.path.isdir(os.path.join(pretrained_path, "vae")): + vae_dir = os.path.join(pretrained_path, "vae") + + safetensors_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") + if os.path.exists(safetensors_path): + from safetensors.torch import load_file + hf_sd = load_file(safetensors_path) + else: + bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") + hf_sd = torch.load(bin_path, map_location="cpu") + + vae = cls(**kwargs) + converted_sd = convert_diffusers_state_dict(hf_sd) + vae.load_state_dict(converted_sd, strict=True) + return vae + + +# ── State dict converter ── + +def convert_diffusers_state_dict(hf_sd: dict) -> OrderedDict: + """Convert HuggingFace AutoencoderKLWan state dict to WanVideoVAE format. + + The two models have identical architecture (194 params, same shapes) but + different key naming. This function handles the mapping including the + mid_block reorder: HF [attn, res0, res1] → ROLL [res0, attn, res1]. + """ + new_sd = OrderedDict() + + # ResidualBlock internal mapping + # HF: norm1.gamma, conv1.{w,b}, norm2.gamma, conv2.{w,b}, conv_shortcut.{w,b} + # ROLL: residual.{0.gamma, 2.{w,b}, 3.gamma, 6.{w,b}}, shortcut.{w,b} + resblock_map = { + "norm1.gamma": "residual.0.gamma", + "conv1.weight": "residual.2.weight", + "conv1.bias": "residual.2.bias", + "norm2.gamma": "residual.3.gamma", + "conv2.weight": "residual.6.weight", + "conv2.bias": "residual.6.bias", + "conv_shortcut.weight": "shortcut.weight", + "conv_shortcut.bias": "shortcut.bias", + } + + # AttentionBlock internal mapping (keys match directly) + # HF: norm.gamma, to_qkv.{w,b}, proj.{w,b} + # ROLL: norm.gamma, to_qkv.{w,b}, proj.{w,b} + + for hf_key, tensor in hf_sd.items(): + roll_key = _convert_single_key(hf_key, resblock_map) + new_sd[roll_key] = tensor + + return new_sd + + +def _convert_single_key(hf_key: str, resblock_map: dict) -> str: + """Convert a single HF key to ROLL key.""" + + # quant_conv / post_quant_conv + if hf_key.startswith("quant_conv."): + return "model.conv1." + hf_key[len("quant_conv."):] + if hf_key.startswith("post_quant_conv."): + return "model.conv2." + hf_key[len("post_quant_conv."):] + + # Encoder + if hf_key.startswith("encoder."): + return "model.encoder." + _convert_encoder_key( + hf_key[len("encoder."):], resblock_map) + + # Decoder + if hf_key.startswith("decoder."): + return "model.decoder." + _convert_decoder_key( + hf_key[len("decoder."):], resblock_map) + + raise ValueError(f"Unknown HF key: {hf_key}") + + +def _convert_encoder_key(key: str, rb_map: dict) -> str: + # conv_in → conv1 + if key.startswith("conv_in."): + return "conv1." + key[len("conv_in."):] + + # norm_out → head.0 + if key.startswith("norm_out."): + return "head.0." + key[len("norm_out."):] + + # conv_out → head.2 + if key.startswith("conv_out."): + return "head.2." + key[len("conv_out."):] + + # down_blocks.{i}.{suffix} + if key.startswith("down_blocks."): + rest = key[len("down_blocks."):] + dot = rest.index(".") + block_idx = rest[:dot] + suffix = rest[dot + 1:] + return "downsamples." + block_idx + "." + _convert_resblock_suffix(suffix, rb_map) + + # mid_block: HF [attn, res0, res1] → ROLL [res0(middle.0), attn(middle.1), res1(middle.2)] + if key.startswith("mid_block."): + rest = key[len("mid_block."):] + if rest.startswith("resnets.0."): + suffix = rest[len("resnets.0."):] + return "middle.0." + _convert_resblock_suffix(suffix, rb_map) + if rest.startswith("attentions.0."): + suffix = rest[len("attentions.0."):] + return "middle.1." + suffix + if rest.startswith("resnets.1."): + suffix = rest[len("resnets.1."):] + return "middle.2." + _convert_resblock_suffix(suffix, rb_map) + + raise ValueError(f"Unknown encoder key: {key}") + + +def _convert_decoder_key(key: str, rb_map: dict) -> str: + # conv_in → conv1 + if key.startswith("conv_in."): + return "conv1." + key[len("conv_in."):] + + # norm_out → head.0 + if key.startswith("norm_out."): + return "head.0." + key[len("norm_out."):] + + # conv_out → head.2 + if key.startswith("conv_out."): + return "head.2." + key[len("conv_out."):] + + # mid_block: same reorder as encoder + if key.startswith("mid_block."): + rest = key[len("mid_block."):] + if rest.startswith("resnets.0."): + suffix = rest[len("resnets.0."):] + return "middle.0." + _convert_resblock_suffix(suffix, rb_map) + if rest.startswith("attentions.0."): + suffix = rest[len("attentions.0."):] + return "middle.1." + suffix + if rest.startswith("resnets.1."): + suffix = rest[len("resnets.1."):] + return "middle.2." + _convert_resblock_suffix(suffix, rb_map) + + # up_blocks.{block_i}.resnets.{res_j}.{suffix} → upsamples.{flat_idx}.{suffix} + # up_blocks.{block_i}.upsamplers.0.{suffix} → upsamples.{flat_idx}.{suffix} + if key.startswith("up_blocks."): + return _convert_upblock_key(key[len("up_blocks."):], rb_map) + + raise ValueError(f"Unknown decoder key: {key}") + + +def _convert_upblock_key(key: str, rb_map: dict) -> str: + """Convert decoder up_blocks flat key to ROLL upsamples flat index. + + HF decoder structure (Wan 2.1 VideoVAE_): + - up_blocks.0: 3 resnets + 1 upsampler → upsamples indices 0,1,2,3 + - up_blocks.1: 3 resnets + 1 upsampler → upsamples indices 4,5,6,7 + - up_blocks.2: 3 resnets + 1 upsampler → upsamples indices 8,9,10,11 + - up_blocks.3: 3 resnets (no upsampler) → upsamples indices 12,13,14 + + ROLL's Decoder3d has 15 sequential upsamples: + - indices 0-2: ResBlocks (up_blocks.0.resnets.0-2) + - index 3: Resample (up_blocks.0.upsamplers.0) + - indices 4-6: ResBlocks (up_blocks.1.resnets.0-2) + - index 7: Resample (up_blocks.1.upsamplers.0) + - indices 8-10: ResBlocks (up_blocks.2.resnets.0-2) + - index 11: Resample (up_blocks.2.upsamplers.0) + - indices 12-14: ResBlocks (up_blocks.3.resnets.0-2) + """ + dot = key.index(".") + block_idx = int(key[:dot]) + rest = key[dot + 1:] + + # Base offset: each block contributes 4 items (3 resnets + 1 upsampler) + # except the last block which has 3 items (3 resnets, no upsampler) + base = block_idx * 4 # works for blocks 0,1,2; block 3 won't have upsampler + + if rest.startswith("resnets."): + rest2 = rest[len("resnets."):] + dot2 = rest2.index(".") + res_idx = int(rest2[:dot2]) + suffix = rest2[dot2 + 1:] + flat_idx = base + res_idx + return f"upsamples.{flat_idx}." + _convert_resblock_suffix(suffix, rb_map) + + if rest.startswith("upsamplers.0."): + suffix = rest[len("upsamplers.0."):] + flat_idx = base + 3 # upsampler comes after 3 resnets + return f"upsamples.{flat_idx}." + suffix + + raise ValueError(f"Unknown up_block key: {key}") + + +def _convert_resblock_suffix(suffix: str, rb_map: dict) -> str: + """Convert a ResidualBlock suffix from HF to ROLL format.""" + for hf_pattern, roll_pattern in rb_map.items(): + if suffix == hf_pattern or suffix.startswith(hf_pattern): + return suffix.replace(hf_pattern, roll_pattern, 1) + # Pass through (e.g. resample.*, time_conv.*) + return suffix diff --git a/unirl/reward/base.py b/unirl/reward/base.py index a6c8f0148..1daca48f4 100644 --- a/unirl/reward/base.py +++ b/unirl/reward/base.py @@ -81,12 +81,16 @@ class DifferentiableReward(Protocol): def compute_rewards_differentiable( self, - images_tensor: "torch.Tensor", + media_tensor: "torch.Tensor", prompts: List[str], - records: Optional[List[dict]] = None, + records: Optional[List[dict[str, object]]] = None, ) -> "torch.Tensor": - """Score a grad-carrying image tensor ``[B, C, H, W]`` in ``[0, 1]`` → - ``[B]`` reward tensor with ``grad_fn`` intact (no ``no_grad``/``.item()``).""" + """Score grad-carrying image ``[B,C,H,W]`` or video ``[B,C,T,H,W]`` media. + + Returns a ``[B]`` reward tensor with ``grad_fn`` intact. ``records`` may + carry per-sample metadata such as ``ref_video_path`` for recipe-local + video rewards. + """ ... diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index bad32147f..9fec1e242 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -50,7 +50,7 @@ ) from unirl.train.configs import EmaFullConfig, EmaLoraConfig, FSDPConfig, LoraConfig from unirl.train.ema import EMA, Shadow, inject_mirror, inject_nft, make_decay_fn -from unirl.train.lora import inject_lora +from unirl.train.lora import inject_lora, resolve_target_modules_pattern from unirl.train.optim import build_lr_scheduler, build_optimizer if TYPE_CHECKING: @@ -208,6 +208,7 @@ def _inject_structural( rank=lora_cfg.rank, alpha=lora_cfg.alpha, target_modules=lora_cfg.target_modules, + module_prefix=lora_cfg.module_prefix, exclude_modules=lora_cfg.exclude_modules, dropout=lora_cfg.dropout, bias=lora_cfg.bias, @@ -273,11 +274,22 @@ def _finalize_construction( # Checkpointed for export tooling: the LoRA fold needs scaling = # alpha / rank, and alpha is not derivable from the weights. active_lora = lora_cfg or ema_lora_cfg + recorded_target_modules = None + if active_lora is not None: + recorded_target_modules = active_lora.target_modules + if lora_cfg is not None: + recorded_target_modules, _ = resolve_target_modules_pattern( + target_modules=lora_cfg.target_modules, + module_prefix=lora_cfg.module_prefix, + ) self._lora_meta = ( { "rank": active_lora.rank, "alpha": active_lora.alpha, - "target_modules": active_lora.target_modules, + # Store the selector actually passed to PEFT. For prefixed LoRA + # this is a regex, so exported adapters keep the same subtree + # restriction instead of matching equivalent suffixes elsewhere. + "target_modules": recorded_target_modules, "exclude_modules": active_lora.exclude_modules, "dropout": active_lora.dropout, "bias": active_lora.bias, diff --git a/unirl/train/configs.py b/unirl/train/configs.py index b96dd9efa..d2efcaadd 100644 --- a/unirl/train/configs.py +++ b/unirl/train/configs.py @@ -19,6 +19,9 @@ class LoraConfig: # Like ``target_modules``, tuples match exact names/suffixes and strings are # regular expressions. This lets a regex exclude an entire frozen sub-tower. exclude_modules: Any = None + # Restrict sequence-style target suffixes to a named model subtree. + # This cannot be combined with regex or ``all-linear`` target strings. + module_prefix: str = "" dropout: float = 0.0 bias: str = "none" task_type: str = "FEATURE_EXTRACTION" diff --git a/unirl/train/lora.py b/unirl/train/lora.py index a7e6796e6..7b1579b9e 100644 --- a/unirl/train/lora.py +++ b/unirl/train/lora.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +import re from contextlib import contextmanager from functools import partial from typing import Iterator, Optional, Sequence, Union @@ -43,12 +44,46 @@ def normalize_optional_module_selection( return None if modules is None else normalize_module_selection(modules) +def resolve_target_modules_pattern( + *, + target_modules: ModuleSelection, + module_prefix: str = "", +) -> tuple[PeftModuleSelection, str]: + """Resolve PEFT targets, optionally restricting suffixes to one subtree.""" + normalized = normalize_module_selection(target_modules) + if not module_prefix: + logged = normalized if isinstance(normalized, str) else tuple(normalized) + return normalized, str(logged) + + if isinstance(normalized, str): + raise ValueError( + "resolve_target_modules_pattern: module_prefix cannot be combined " + "with a regex or 'all-linear' target_modules string; provide an " + "explicit sequence of target-module suffixes." + ) + if not normalized: + raise ValueError( + "resolve_target_modules_pattern: module_prefix is set but " + "target_modules is empty; provide at least one target-module suffix." + ) + normalized_prefix = str(module_prefix).strip(".") + if not normalized_prefix: + raise ValueError( + "resolve_target_modules_pattern: module_prefix must contain a model subtree name, not only dots." + ) + prefix_re = re.escape(normalized_prefix) + leaves_re = "|".join(re.escape(str(target)) for target in normalized) + pattern = rf"^{prefix_re}\.(?:.*\.)?(?:{leaves_re})$" + return pattern, pattern + + def inject_lora( model: nn.Module, *, rank: int, alpha: int, target_modules: ModuleSelection, + module_prefix: str = "", exclude_modules: Optional[ModuleSelection] = None, dropout: float = 0.0, bias: str = "none", @@ -58,11 +93,16 @@ def inject_lora( """Inject a single LoRA adapter. No Shadow, no EMA.""" from peft import LoraConfig, inject_adapter_in_model + peft_target_modules, log_target = resolve_target_modules_pattern( + target_modules=target_modules, + module_prefix=module_prefix, + ) + peft_cfg = LoraConfig( r=int(rank), lora_alpha=int(alpha), lora_dropout=float(dropout), - target_modules=normalize_module_selection(target_modules), + target_modules=peft_target_modules, exclude_modules=normalize_optional_module_selection(exclude_modules), bias=str(bias), task_type=str(task_type), @@ -71,7 +111,6 @@ def inject_lora( if _current_rank() == 0: n_trainable = sum(1 for p in model.parameters() if p.requires_grad) - logged_targets = target_modules if isinstance(target_modules, str) else tuple(target_modules) logged_exclusions = ( exclude_modules if isinstance(exclude_modules, str) or exclude_modules is None else tuple(exclude_modules) ) @@ -80,7 +119,7 @@ def inject_lora( adapter_name, rank, alpha, - logged_targets, + log_target, logged_exclusions, n_trainable, ) @@ -136,4 +175,5 @@ def _current_rank() -> int: "inject_lora", "normalize_module_selection", "normalize_optional_module_selection", + "resolve_target_modules_pattern", ] diff --git a/unirl/train/optim.py b/unirl/train/optim.py index 3f20310de..26243013a 100644 --- a/unirl/train/optim.py +++ b/unirl/train/optim.py @@ -131,8 +131,8 @@ def build_lr_scheduler( Supports the same backend-override path as :func:`build_optimizer`. Returns ``None`` if ``config.type`` is not one of the supported values - (``constant`` / ``linear`` / ``cosine``) and the backend did not provide - an override. + (``constant`` / ``linear`` / ``linear_warmup`` / ``cosine``) and the backend + did not provide an override. """ del actor if backend is not None: @@ -166,6 +166,27 @@ def lr_lambda(step: int) -> float: return max(0.0, 1.0 - (step - warmup_steps) / (total_steps - warmup_steps)) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + if scheduler_type == "linear_warmup": + constant = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=1.0, + end_factor=1.0, + total_iters=1, + ) + if warmup_steps <= 0: + return constant + warmup = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=1e-8, + end_factor=1.0, + total_iters=warmup_steps, + ) + return torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[warmup, constant], + milestones=[warmup_steps], + ) if scheduler_type == "cosine": diff --git a/unirl/trainer/base_role.py b/unirl/trainer/base_role.py new file mode 100644 index 000000000..872f251be --- /dev/null +++ b/unirl/trainer/base_role.py @@ -0,0 +1,117 @@ +"""Base Role abstraction for trainer-managed Remote roles.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Optional + +import torch +from hydra.utils import instantiate + +from unirl.distributed.group.dispatch import Dispatch, distributed +from unirl.distributed.group.remote import Remote + + +@dataclass(frozen=True) +class RoleStepResult: + """Generic result of one role-local optimizer step.""" + + metrics: Mapping[str, object] + grad_norm: float + lr: float + + +class Role(Remote): + """A trainer role that runs as a UniRL Remote inside a Worker. + + ``Role`` stores the role config and initializes common role-local + components (model/runtime bundle / pipeline / backend) as ordinary + Worker-local objects. Recipe roles may override ``initialize`` for + recipe-specific config, but should call ``super().initialize()`` first. + """ + + def __init__(self, cfg: Any) -> None: + super().__init__() + self.cfg = cfg + + def _cfg_get(self, key: str) -> Any: + if hasattr(self.cfg, "get"): + return self.cfg.get(key) + return getattr(self.cfg, key, None) + + def initialize(self) -> None: + """Initialize common role-local components after Worker setup.""" + model_cfg = self._cfg_get("model") + bundle_cfg = model_cfg if model_cfg is not None else self._cfg_get("bundle") + if bundle_cfg is not None: + self.bundle = instantiate(bundle_cfg) + + pipeline_cfg = self._cfg_get("pipeline") + if pipeline_cfg is not None: + if hasattr(self, "bundle"): + self.pipeline = instantiate(pipeline_cfg, bundle=self.bundle) + else: + self.pipeline = instantiate(pipeline_cfg) + + backend_cfg = self._cfg_get("backend") + if backend_cfg is not None: + if hasattr(self, "bundle"): + if self.device is None or self.rank_info is None: + raise RuntimeError( + f"{type(self).__name__}.initialize called before Remote.setup injected device/rank_info." + ) + self.backend = instantiate( + backend_cfg, + bundle=self.bundle, + device=torch.device(self.device), + rank=int(self.rank_info.rank), + ) + else: + self.backend = instantiate(backend_cfg) + + def _max_grad_norm(self) -> float: + algo_cfg = getattr(self, "algo_cfg", None) + if algo_cfg is None: + algo_cfg = getattr(self, "algorithm_cfg", None) + if algo_cfg is not None and hasattr(algo_cfg, "get"): + return float(algo_cfg.get("max_grad_norm", 1.0)) + if algo_cfg is not None and hasattr(algo_cfg, "max_grad_norm"): + return float(algo_cfg.max_grad_norm) + return float(getattr(self, "max_grad_norm", 1.0)) + + @distributed + def step(self) -> RoleStepResult: + """Clip gradients and run one backend optimizer step.""" + if not hasattr(self, "backend") or not hasattr(self.backend, "optimizer_step"): + raise RuntimeError(f"{type(self).__name__}.step requires a backend with optimizer_step(...).") + grad_norm = float(self.backend.optimizer_step(max_grad_norm=self._max_grad_norm())) + lr = 0.0 + try: + sched = getattr(self.backend, "scheduler", None) + if sched is not None: + last = sched.get_last_lr() + lr = float(last[0]) if last else 0.0 + except Exception: + lr = 0.0 + return RoleStepResult( + metrics={"grad_norm": grad_norm, "lr": lr}, + grad_norm=grad_norm, + lr=lr, + ) + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def save_checkpoint(self, path: str, step: Optional[int] = None, mode: str = "auto") -> None: + """Save backend checkpoint when the role backend supports it.""" + if hasattr(self, "backend") and hasattr(self.backend, "save"): + self.backend.save(path, step=step, mode=mode) + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def load_checkpoint(self, path: str) -> int: + """Load backend checkpoint when the role backend supports it.""" + if hasattr(self, "backend") and hasattr(self.backend, "load"): + return int(self.backend.load(path) or 0) + return 0 + + +__all__ = ["Role", "RoleStepResult"] diff --git a/unirl/trainer/trainer.py b/unirl/trainer/trainer.py new file mode 100644 index 000000000..27c486728 --- /dev/null +++ b/unirl/trainer/trainer.py @@ -0,0 +1,416 @@ +"""Trainer: config-driven Remote role orchestration.""" + +from __future__ import annotations + +import json +import logging +import os +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional + +from hydra.utils import get_method +from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict + +from unirl.distributed.group.remote import Remote +from unirl.trainer.base import BaseTrainer, build_sampling_dict + +logger = logging.getLogger(__name__) + +_RESERVED_ROLE_NAMES = { + "cfg", + "pool", + "roles", + "role_specs", + "role_device_ids", + "role_slot_ids", + "data_source", + "sampling_params", + "wandb_logger", +} + + +@dataclass(frozen=True) +class RoleSpec: + """Driver-side parsed role declaration.""" + + name: str + target: str + placement: DictConfig + raw_cfg: DictConfig + index: int + + +class Trainer(BaseTrainer): + """Base trainer that creates recipe roles as UniRL Remote handles. + + ``Trainer`` owns the generic role lifecycle: + parse ``cfg.roles`` → topo-sort placement dependencies → create Remote + handles with explicit ``device_ids`` / ``slot_id`` → initialize roles. + Recipe trainers only implement ``build_req`` and ``train_step``. + """ + + def __init__(self, *, cfg: DictConfig, logging_cfg: Optional[DictConfig] = None) -> None: + self.cfg = cfg + self.role_specs: List[RoleSpec] = self.parse_role_specs(cfg) + self._role_specs_by_name: Dict[str, RoleSpec] = {spec.name: spec for spec in self.role_specs} + self._sorted_role_specs: List[RoleSpec] = self.topological_roles() + self._simulated_slot_ids: Dict[str, int] = self._simulate_role_slots(self._sorted_role_specs) + self._ensure_workers_per_device(cfg, self._simulated_slot_ids) + + super().__init__(cfg=cfg, logging_cfg=logging_cfg if logging_cfg is not None else cfg.get("logging")) + + self.batch_size = int(cfg.get("batch_size", 1)) + self.data_source = self.instantiate_data_source(cfg) + self.sampling_params = build_sampling_dict(cfg.sampling) if cfg.get("sampling") is not None else {} + self.roles: Dict[str, Any] = {} + self.role_device_ids: Dict[str, List[int]] = {} + self.role_slot_ids: Dict[str, int] = {} + self._next_colocate_slot = 1 + self._roles_initialized = False + + self.setup_roles() + self.initialize_roles() + self.validate_config() + + # ------------------------------------------------------------------ + # Config parsing and placement planning. + # ------------------------------------------------------------------ + + def parse_role_specs(self, cfg: DictConfig) -> List[RoleSpec]: + roles_cfg = cfg.get("roles") + if roles_cfg is None: + raise ValueError("Trainer requires cfg.roles as a list of role declarations.") + if not isinstance(roles_cfg, (list, tuple, ListConfig)): + raise TypeError(f"cfg.roles must be a list, got {type(roles_cfg).__name__}.") + + specs: List[RoleSpec] = [] + seen: set[str] = set() + for idx, role_cfg in enumerate(roles_cfg): + if not OmegaConf.is_config(role_cfg): + role_cfg = OmegaConf.create(role_cfg) + name = str(role_cfg.get("name") or "").strip() + if not name: + raise ValueError(f"roles[{idx}] is missing required field `name`.") + if name in seen: + raise ValueError(f"Duplicate role name {name!r} in cfg.roles.") + if name in _RESERVED_ROLE_NAMES or hasattr(self.__class__, name): + raise ValueError(f"Role name {name!r} is reserved; choose a different role name.") + target = str(role_cfg.get("_target_") or "").strip() + if not target: + raise ValueError(f"roles[{idx}] ({name}) is missing required field `_target_`.") + placement = role_cfg.get("placement") or OmegaConf.create({}) + if not OmegaConf.is_config(placement): + placement = OmegaConf.create(placement) + if placement.get("share_with") and placement.get("colocate_with"): + raise ValueError(f"Role {name!r}: placement cannot set both share_with and colocate_with.") + specs.append(RoleSpec(name=name, target=target, placement=placement, raw_cfg=role_cfg, index=idx)) + seen.add(name) + return specs + + def topological_roles(self) -> List[RoleSpec]: + specs = self.role_specs + by_name = {spec.name: spec for spec in specs} + visiting: set[str] = set() + visited: set[str] = set() + ordered: List[RoleSpec] = [] + + def parent_of(spec: RoleSpec) -> Optional[str]: + parent = spec.placement.get("share_with") or spec.placement.get("colocate_with") + return str(parent) if parent else None + + def visit(spec: RoleSpec) -> None: + if spec.name in visited: + return + if spec.name in visiting: + raise ValueError(f"Cycle detected in role placement dependencies at role {spec.name!r}.") + visiting.add(spec.name) + parent = parent_of(spec) + if parent: + if parent not in by_name: + raise ValueError(f"Role {spec.name!r} placement references unknown role {parent!r}.") + visit(by_name[parent]) + visiting.remove(spec.name) + visited.add(spec.name) + ordered.append(spec) + + for spec in specs: + visit(spec) + return ordered + + def _simulate_role_slots(self, specs: Iterable[RoleSpec]) -> Dict[str, int]: + slots: Dict[str, int] = {} + next_colocate_slot = 1 + for spec in specs: + if spec.placement.get("share_with"): + slots[spec.name] = slots[str(spec.placement.share_with)] + elif spec.placement.get("colocate_with"): + slots[spec.name] = next_colocate_slot + next_colocate_slot += 1 + else: + slots[spec.name] = 0 + return slots + + def _ensure_workers_per_device(self, cfg: DictConfig, slots: Dict[str, int]) -> None: + required = max(slots.values(), default=0) + 1 + current = int(cfg.get("workers_per_device", 1)) + if required <= current: + return + transport_kind = str(cfg.get("transport_kind", "colocate_store")) + if transport_kind in ("colocate_store", "colocate"): + raise ValueError( + "placement.colocate_with requires multiple worker slots per GPU, but " + f"transport_kind={transport_kind!r} only supports workers_per_device=1. " + "Set transport_kind='gpu_store' (or another multi-slot-capable transport) " + f"and workers_per_device>={required}." + ) + with open_dict(cfg): + cfg.workers_per_device = required + + # ------------------------------------------------------------------ + # Role creation. + # ------------------------------------------------------------------ + + def setup_roles(self) -> None: + self.roles = {} + self.role_device_ids = {} + self.role_slot_ids = {} + self._next_colocate_slot = 1 + + for spec in self._sorted_role_specs: + device_ids, slot_id = self.resolve_role_placement(spec) + handle = self.create_remote_role(spec, device_ids=device_ids, slot_id=slot_id) + self.roles[spec.name] = handle + self.role_device_ids[spec.name] = device_ids + self.role_slot_ids[spec.name] = slot_id + setattr(self, spec.name, handle) + + def resolve_role_placement(self, spec: RoleSpec) -> tuple[List[int], int]: + p = spec.placement or OmegaConf.create({}) + if p.get("share_with"): + parent = str(p.share_with) + parent_ids = self.role_device_ids[parent] + n = int(p.get("n_devices") or len(parent_ids)) + return self.step_subset(parent_ids, n), self.role_slot_ids[parent] + if p.get("colocate_with"): + parent = str(p.colocate_with) + parent_ids = self.role_device_ids[parent] + n = int(p.get("n_devices") or len(parent_ids)) + return self.step_subset(parent_ids, n), self.allocate_colocate_slot(parent) + + if p.get("device_ids") is not None: + return [int(d) for d in list(p.device_ids)], 0 + n_devices = int(p.get("n_devices") or spec.raw_cfg.get("n_devices") or self.cfg.num_devices) + if n_devices <= 0: + raise ValueError(f"Role {spec.name!r}: placement.n_devices must be positive, got {n_devices}.") + return self.pool.allocate(n_devices), 0 + + @staticmethod + def step_subset(device_ids: List[int], n_devices: int) -> List[int]: + if n_devices <= 0: + raise ValueError(f"n_devices must be positive, got {n_devices}.") + if n_devices > len(device_ids): + raise ValueError(f"Cannot take {n_devices} devices from parent device slab {device_ids}.") + if n_devices == len(device_ids): + return list(device_ids) + if n_devices == 1: + return [device_ids[0]] + last = len(device_ids) - 1 + return [device_ids[round(i * last / (n_devices - 1))] for i in range(n_devices)] + + def allocate_colocate_slot(self, parent: str) -> int: + del parent # The current implementation allocates globally unique colocate slots. + slot = max(self._next_colocate_slot, max(self.role_slot_ids.values(), default=0) + 1) + self._next_colocate_slot = slot + 1 + return slot + + def resolve_role_cls(self, target: str) -> type: + role_cls = get_method(target) + if not isinstance(role_cls, type) or not issubclass(role_cls, Remote): + raise TypeError(f"Role target {target!r} must resolve to a Remote subclass, got {role_cls!r}.") + return role_cls + + def prepare_role_cfg(self, spec: RoleSpec) -> DictConfig: + container = OmegaConf.to_container(spec.raw_cfg, resolve=True) + if not isinstance(container, dict): + raise TypeError(f"Role {spec.name!r} config must resolve to a mapping.") + for key in ("name", "_target_", "placement"): + container.pop(key, None) + return OmegaConf.create(container) + + def create_remote_role(self, spec: RoleSpec, *, device_ids: List[int], slot_id: int): + return self.pool.create_remote( + self.resolve_role_cls(spec.target), + device_ids=device_ids, + slot_id=slot_id, + role_name=spec.name, + init_kwargs={"cfg": self.prepare_role_cfg(spec)}, + ) + + def initialize_roles(self) -> None: + if self._roles_initialized: + return + for spec in self._sorted_role_specs: + self.roles[spec.name].initialize() + self._roles_initialized = True + + # ------------------------------------------------------------------ + # Data, validation, training loop. + # ------------------------------------------------------------------ + + def instantiate_data_source(self, cfg: DictConfig): + if cfg.get("data_source") is None: + return None + from hydra.utils import instantiate + + return instantiate(cfg.data_source) + + def validate_config(self) -> None: + if self.data_source is None: + raise ValueError("Trainer requires cfg.data_source.") + if not self.sampling_params: + raise ValueError("Trainer requires cfg.sampling.") + + def build_req(self, inputs: Any, rollout_id: int) -> Any: + raise NotImplementedError + + def train_step(self, req: Any, *, training_progress: float = 0.0, rollout_id: int = 0) -> Dict[str, Any]: + raise NotImplementedError + + def train( + self, + *, + num_rollouts: Optional[int] = None, + save_interval: Optional[int] = None, + save_dir: Optional[str] = None, + load_dir: Optional[str] = None, + save_mode: Optional[str] = None, + ) -> None: + num_rollouts = int(num_rollouts if num_rollouts is not None else self.cfg.get("num_rollouts", 100)) + save_interval = int(save_interval if save_interval is not None else self.cfg.get("save_interval", 0)) + save_dir = save_dir if save_dir is not None else self.cfg.get("save_dir") + load_dir = load_dir if load_dir is not None else self.cfg.get("load_dir") + save_mode = save_mode if save_mode is not None else self.cfg.get("save_mode", "auto") + + start_rollout = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) + data_source = self.data_source + if data_source is None: + raise ValueError("Trainer requires cfg.data_source.") + for _ in range(start_rollout): + data_source.get_samples(self.batch_size) + self._init_wandb(num_rollouts=num_rollouts) + try: + for rollout_id in range(start_rollout, num_rollouts): + training_progress = rollout_id / max(1, num_rollouts - 1) + inputs = data_source.get_samples(self.batch_size) + req = self.build_req(inputs, rollout_id) + metrics = self.train_step(req, training_progress=training_progress, rollout_id=rollout_id) + self.log_metrics(metrics, rollout_id=rollout_id, num_rollouts=num_rollouts) + self.maybe_save_checkpoint( + rollout_id, + num_rollouts, + save_interval=save_interval, + save_dir=save_dir, + save_mode=save_mode, + ) + finally: + self._finish_wandb() + + fit = train + + def log_metrics(self, metrics: Dict[str, Any], *, rollout_id: int, num_rollouts: int) -> None: + trainer_name = self.__class__.__name__.removesuffix("Trainer") or self.__class__.__name__ + metric_text = " ".join( + f"{key}={self._format_metric_for_log(value)}" + for key, value in metrics.items() + ) + if metric_text: + logger.info("%s rollout %d/%d %s", trainer_name, rollout_id + 1, num_rollouts, metric_text) + else: + logger.info("%s rollout %d/%d", trainer_name, rollout_id + 1, num_rollouts) + + wb = self.wandb_logger + if wb is not None and wb.initialized: + wb.log_step( + step=rollout_id + 1, + metrics={k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))}, + prefix=str((self.logging_cfg or {}).get("metric_prefix", "train/")), + ) + + @staticmethod + def _format_metric_for_log(value: Any) -> str: + if isinstance(value, bool): + return str(value) + if isinstance(value, int): + return str(value) + if isinstance(value, float): + abs_value = abs(value) + if value != 0.0 and (abs_value < 1e-3 or abs_value >= 1e4): + return f"{value:.6e}" + return f"{value:.4f}" + return repr(value) + + # ------------------------------------------------------------------ + # Role-aware checkpointing. + # ------------------------------------------------------------------ + + def checkpoint_roles(self, path: str, *, step: int, mode: str) -> None: + for name, role in self.roles.items(): + if hasattr(role, "save_checkpoint"): + role_path = os.path.join(path, name) + role.save_checkpoint(role_path, step=step, mode=mode) + + def load_checkpoint_roles(self, path: str) -> int: + starts: List[int] = [] + for name, role in self.roles.items(): + if hasattr(role, "load_checkpoint"): + role_path = os.path.join(path, name) + if os.path.exists(role_path): + result = role.load_checkpoint(role_path) + if isinstance(result, list): + result = result[0] if result else 0 + starts.append(int(result or 0)) + return max(starts, default=0) + + def maybe_save_checkpoint( + self, + rollout_id: int, + num_rollouts: int, + *, + save_interval: int, + save_dir: Optional[str], + save_mode: str = "auto", + ) -> None: + if save_interval <= 0: + return + step = rollout_id + 1 + if step % save_interval != 0 and step < num_rollouts: + return + base_dir = os.path.abspath(save_dir) if save_dir else os.path.join(os.getcwd(), "checkpoints") + path = os.path.join(base_dir, f"checkpoint-{step}") + os.makedirs(path, exist_ok=True) + logger.info("Saving role checkpoint at rollout %d/%d -> %s", step, num_rollouts, path) + self.checkpoint_roles(path, step=step, mode=save_mode) + with open(os.path.join(path, "trainer_state.json"), "w") as f: + json.dump({"wandb_run_id": self.wandb_logger.run_id, "optimizer_step": self.wandb_logger.optimizer_step}, f) + + def maybe_load_checkpoint(self, load_dir: Optional[str], *, num_rollouts: Optional[int] = None) -> int: + if not load_dir: + return 0 + load_dir = os.path.abspath(load_dir) + logger.info("Loading role checkpoint from %s", load_dir) + start = self.load_checkpoint_roles(load_dir) + state_path = os.path.join(load_dir, "trainer_state.json") + if os.path.exists(state_path): + with open(state_path) as f: + self._resume_state = json.load(f) + logger.info("Checkpoint restored; resuming at rollout %d", start) + if num_rollouts is not None and start >= num_rollouts: + logger.warning( + "Checkpoint step %d >= num_rollouts %d — nothing left to train.", + start, + num_rollouts, + ) + return start + + +__all__ = ["Trainer", "RoleSpec"] From 72762c76c847b81ee88cfd93f9620f86c4a97aac Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Tue, 14 Jul 2026 23:57:26 +0800 Subject: [PATCH 02/24] Refactor REFL reward integration and trainer cleanup - Remove trainer-side config validation and simplify mean calculation with np.mean. - Remove unused fit=train handling from trainer. - Rename REFL recipe namespace from refl_wan to refl. - Rename base_role.py to role.py for the REFL role implementation. - Move max_grad_norm handling out of the base role and let reward backends own it. - Add a generic reward role under unirl/reward that directly extends the reward service. - Remove reward package pip installation from pyproject configuration. - Remove hardcoded local paths from recipe configs and shell scripts. - Add requirements.txt files for face and videoalign reward backends. --- pyproject.toml | 11 --- recipes/{refl_wan => refl}/__init__.py | 0 .../configs/wan21_t2v_videoalign_refl.yaml | 12 +-- .../configs/wan22_face_refl.yaml | 16 +-- recipes/{refl_wan => refl}/models/__init__.py | 0 recipes/{refl_wan => refl}/models/wan21.py | 2 +- recipes/{refl_wan => refl}/models/wan22.py | 0 .../{refl_wan => refl}/rewards/__init__.py | 0 .../rewards/face/__init__.py | 0 .../rewards/face/face_tools.py | 0 recipes/refl/rewards/face/requirements.txt | 6 ++ .../{refl_wan => refl}/rewards/face/scorer.py | 0 .../rewards/videoalign/__init__.py | 2 +- .../rewards/videoalign/model/__init__.py | 0 .../rewards/videoalign/model/checkpoint.py | 0 .../rewards/videoalign/model/configs.py | 2 +- .../rewards/videoalign/model/factory.py | 0 .../videoalign/model/prompt_template.py | 0 .../rewards/videoalign/model/reward_model.py | 2 +- .../refl/rewards/videoalign/requirements.txt | 11 +++ .../rewards/videoalign/scorer.py | 2 +- .../rewards/videoalign/wrapper.py | 2 +- recipes/{refl_wan => refl}/roles.py | 52 +--------- recipes/{refl_wan => refl}/run.py | 4 +- .../scripts}/start_wan21_t2v.sh | 11 ++- recipes/refl/scripts/start_wan22_i2v.sh | 20 ++++ recipes/{refl_wan => refl}/trainer.py | 61 ++---------- recipes/refl_wan/start_wan22_i2v.sh | 13 --- unirl/reward/role.py | 99 +++++++++++++++++++ unirl/trainer/{base_role.py => role.py} | 12 +-- unirl/trainer/trainer.py | 2 - 31 files changed, 176 insertions(+), 166 deletions(-) rename recipes/{refl_wan => refl}/__init__.py (100%) rename recipes/{refl_wan => refl}/configs/wan21_t2v_videoalign_refl.yaml (93%) rename recipes/{refl_wan => refl}/configs/wan22_face_refl.yaml (83%) rename recipes/{refl_wan => refl}/models/__init__.py (100%) rename recipes/{refl_wan => refl}/models/wan21.py (99%) rename recipes/{refl_wan => refl}/models/wan22.py (100%) rename recipes/{refl_wan => refl}/rewards/__init__.py (100%) rename recipes/{refl_wan => refl}/rewards/face/__init__.py (100%) rename recipes/{refl_wan => refl}/rewards/face/face_tools.py (100%) create mode 100644 recipes/refl/rewards/face/requirements.txt rename recipes/{refl_wan => refl}/rewards/face/scorer.py (100%) rename recipes/{refl_wan => refl}/rewards/videoalign/__init__.py (88%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/__init__.py (100%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/checkpoint.py (100%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/configs.py (98%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/factory.py (100%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/prompt_template.py (100%) rename recipes/{refl_wan => refl}/rewards/videoalign/model/reward_model.py (99%) create mode 100644 recipes/refl/rewards/videoalign/requirements.txt rename recipes/{refl_wan => refl}/rewards/videoalign/scorer.py (99%) rename recipes/{refl_wan => refl}/rewards/videoalign/wrapper.py (99%) rename recipes/{refl_wan => refl}/roles.py (77%) rename recipes/{refl_wan => refl}/run.py (74%) rename recipes/{refl_wan => refl/scripts}/start_wan21_t2v.sh (66%) mode change 100644 => 100755 create mode 100755 recipes/refl/scripts/start_wan22_i2v.sh rename recipes/{refl_wan => refl}/trainer.py (53%) delete mode 100644 recipes/refl_wan/start_wan22_i2v.sh create mode 100644 unirl/reward/role.py rename unirl/trainer/{base_role.py => role.py} (88%) diff --git a/pyproject.toml b/pyproject.toml index 6c7fd9d71..adc1dd358 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,17 +80,6 @@ eval = [ "torchvision>=0.16", "easyocr>=1.7", ] -face_reward = [ - "onnx>=1.14", - "onnx2torch>=1.5", - "scikit-image>=0.21", - "imageio>=2.31", - "imageio-ffmpeg>=0.4", - "scipy>=1.11", -] -video_align_reward = [ - "flash-attn==2.7.0.post2", -] # VeOmni's torch-native distributed layer (FSDP2/EP parallelize), consumed # exclusively through unirl.train.backend.veomni._compat — a selective import # that never executes veomni/__init__.py or veomni/models/__init__.py, so the diff --git a/recipes/refl_wan/__init__.py b/recipes/refl/__init__.py similarity index 100% rename from recipes/refl_wan/__init__.py rename to recipes/refl/__init__.py diff --git a/recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml similarity index 93% rename from recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml rename to recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 6cce35f73..e3c601e7f 100644 --- a/recipes/refl_wan/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -9,7 +9,7 @@ save_dir: ${oc.env:OUTPUT_DIR,outputs/wan21_t2v_videoalign_refl} roles: - name: actor - _target_: recipes.refl_wan.roles.ReflActorRole + _target_: recipes.refl.roles.ReflActorRole placement: n_devices: ${num_devices} @@ -23,7 +23,7 @@ roles: max_sequence_length: 512 pipeline: - _target_: recipes.refl_wan.models.wan21.Wan21ReflPipeline + _target_: recipes.refl.models.wan21.Wan21ReflPipeline shift: 5.0 autocast_precision: bf16 trajectory_precision: bf16 @@ -84,19 +84,19 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.refl_wan.roles.ReflRewardRole + _target_: unirl.reward.role.RewardRole placement: # BPTT gradients must live on the same worker — no cross-process # autograd. VideoAlign co-locates with the actor. share_with: actor backend: - _target_: recipes.refl_wan.rewards.videoalign.VideoAlignRewardScorer + _target_: recipes.refl.rewards.videoalign.VideoAlignRewardScorer base_device: cuda config: - _target_: recipes.refl_wan.rewards.videoalign.VideoAlignSpec + _target_: recipes.refl.rewards.videoalign.VideoAlignSpec reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} - device: auto + device: cuda batch_size: 1 resize_height: 336 resize_width: 588 diff --git a/recipes/refl_wan/configs/wan22_face_refl.yaml b/recipes/refl/configs/wan22_face_refl.yaml similarity index 83% rename from recipes/refl_wan/configs/wan22_face_refl.yaml rename to recipes/refl/configs/wan22_face_refl.yaml index 43151fa3f..42dc2d5de 100644 --- a/recipes/refl_wan/configs/wan22_face_refl.yaml +++ b/recipes/refl/configs/wan22_face_refl.yaml @@ -9,7 +9,7 @@ save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_face_refl} roles: - name: actor - _target_: recipes.refl_wan.roles.ReflActorRole + _target_: recipes.refl.roles.ReflActorRole placement: n_devices: ${num_devices} @@ -25,7 +25,7 @@ roles: num_train_timesteps: 1000 pipeline: - _target_: recipes.refl_wan.models.wan22.Wan22ReflPipeline + _target_: recipes.refl.models.wan22.Wan22ReflPipeline shift: 5.0 autocast_precision: bf16 trajectory_precision: bf16 @@ -81,16 +81,16 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.refl_wan.roles.ReflRewardRole + _target_: unirl.reward.role.RewardRole placement: share_with: actor backend: - _target_: recipes.refl_wan.rewards.face.FaceRewardScorer + _target_: recipes.refl.rewards.face.FaceRewardScorer base_device: cuda config: - _target_: recipes.refl_wan.rewards.face.FaceRewardSpec - model_path: ${oc.env:FACE_MODEL_PATH,/apdcephfs_gy8/share_301869871/staryding/gy6_data/ckpt/ckpt_0303/apdcephfs_wza/xiangwshen/a/antelodev2/antelodev2} + _target_: recipes.refl.rewards.face.FaceRewardSpec + model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} device: cuda batch_size: 1 image_size: 112 @@ -110,8 +110,8 @@ data_source: _target_: unirl.data.data_source.MultimodalRLDataSource args: run: - data_path: ${oc.env:DATA_PATH,/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl} - eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl}} + data_path: ${oc.env:DATA_PATH,/path/to/wan22_face_refl_prompts.jsonl} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/path/to/wan22_face_refl_prompts.jsonl}} seed: 42 shuffle: false algorithm: diff --git a/recipes/refl_wan/models/__init__.py b/recipes/refl/models/__init__.py similarity index 100% rename from recipes/refl_wan/models/__init__.py rename to recipes/refl/models/__init__.py diff --git a/recipes/refl_wan/models/wan21.py b/recipes/refl/models/wan21.py similarity index 99% rename from recipes/refl_wan/models/wan21.py rename to recipes/refl/models/wan21.py index ef37241a9..238fabc85 100644 --- a/recipes/refl_wan/models/wan21.py +++ b/recipes/refl/models/wan21.py @@ -1,6 +1,6 @@ """Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. -Mirrors ``recipes.refl_wan.models.wan22`` but targets the WAN 2.1 T2V +Mirrors ``recipes.refl.models.wan22`` but targets the WAN 2.1 T2V single-DiT stack. The REFL-specific pieces live here: - :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch diff --git a/recipes/refl_wan/models/wan22.py b/recipes/refl/models/wan22.py similarity index 100% rename from recipes/refl_wan/models/wan22.py rename to recipes/refl/models/wan22.py diff --git a/recipes/refl_wan/rewards/__init__.py b/recipes/refl/rewards/__init__.py similarity index 100% rename from recipes/refl_wan/rewards/__init__.py rename to recipes/refl/rewards/__init__.py diff --git a/recipes/refl_wan/rewards/face/__init__.py b/recipes/refl/rewards/face/__init__.py similarity index 100% rename from recipes/refl_wan/rewards/face/__init__.py rename to recipes/refl/rewards/face/__init__.py diff --git a/recipes/refl_wan/rewards/face/face_tools.py b/recipes/refl/rewards/face/face_tools.py similarity index 100% rename from recipes/refl_wan/rewards/face/face_tools.py rename to recipes/refl/rewards/face/face_tools.py diff --git a/recipes/refl/rewards/face/requirements.txt b/recipes/refl/rewards/face/requirements.txt new file mode 100644 index 000000000..b5421c1be --- /dev/null +++ b/recipes/refl/rewards/face/requirements.txt @@ -0,0 +1,6 @@ +onnx>=1.14 +onnx2torch>=1.5 +scikit-image>=0.21 +imageio>=2.31 +imageio-ffmpeg>=0.4 +scipy>=1.11 diff --git a/recipes/refl_wan/rewards/face/scorer.py b/recipes/refl/rewards/face/scorer.py similarity index 100% rename from recipes/refl_wan/rewards/face/scorer.py rename to recipes/refl/rewards/face/scorer.py diff --git a/recipes/refl_wan/rewards/videoalign/__init__.py b/recipes/refl/rewards/videoalign/__init__.py similarity index 88% rename from recipes/refl_wan/rewards/videoalign/__init__.py rename to recipes/refl/rewards/videoalign/__init__.py index 3ddfdd59e..90ee502ad 100644 --- a/recipes/refl_wan/rewards/videoalign/__init__.py +++ b/recipes/refl/rewards/videoalign/__init__.py @@ -3,7 +3,7 @@ Self-contained port of the VideoAlign reward family for the REFL WAN recipe. The reward model code (``Qwen2VLRewardModelBT``), prompt templates, config dataclasses and checkpoint-loading helpers all live under -:mod:`recipes.refl_wan.rewards.videoalign.model`, matching the recipe-local +:mod:`recipes.refl.rewards.videoalign.model`, matching the recipe-local layout used by the WAN22 face reward. Public API diff --git a/recipes/refl_wan/rewards/videoalign/model/__init__.py b/recipes/refl/rewards/videoalign/model/__init__.py similarity index 100% rename from recipes/refl_wan/rewards/videoalign/model/__init__.py rename to recipes/refl/rewards/videoalign/model/__init__.py diff --git a/recipes/refl_wan/rewards/videoalign/model/checkpoint.py b/recipes/refl/rewards/videoalign/model/checkpoint.py similarity index 100% rename from recipes/refl_wan/rewards/videoalign/model/checkpoint.py rename to recipes/refl/rewards/videoalign/model/checkpoint.py diff --git a/recipes/refl_wan/rewards/videoalign/model/configs.py b/recipes/refl/rewards/videoalign/model/configs.py similarity index 98% rename from recipes/refl_wan/rewards/videoalign/model/configs.py rename to recipes/refl/rewards/videoalign/model/configs.py index 478bb1def..c22750c85 100644 --- a/recipes/refl_wan/rewards/videoalign/model/configs.py +++ b/recipes/refl/rewards/videoalign/model/configs.py @@ -36,7 +36,7 @@ class TrainingConfig: Only ``bf16`` / ``fp16`` / ``gradient_checkpointing`` / ``disable_flash_attn2`` are actually consumed by the inference path - (see :func:`recipes.refl_wan.rewards.videoalign.model.factory.create_model_and_processor`). + (see :func:`recipes.refl.rewards.videoalign.model.factory.create_model_and_processor`). The other fields are kept so ``TrainingConfig(**model_config_json["training_args"])`` still works when someone wants to introspect the saved config. """ diff --git a/recipes/refl_wan/rewards/videoalign/model/factory.py b/recipes/refl/rewards/videoalign/model/factory.py similarity index 100% rename from recipes/refl_wan/rewards/videoalign/model/factory.py rename to recipes/refl/rewards/videoalign/model/factory.py diff --git a/recipes/refl_wan/rewards/videoalign/model/prompt_template.py b/recipes/refl/rewards/videoalign/model/prompt_template.py similarity index 100% rename from recipes/refl_wan/rewards/videoalign/model/prompt_template.py rename to recipes/refl/rewards/videoalign/model/prompt_template.py diff --git a/recipes/refl_wan/rewards/videoalign/model/reward_model.py b/recipes/refl/rewards/videoalign/model/reward_model.py similarity index 99% rename from recipes/refl_wan/rewards/videoalign/model/reward_model.py rename to recipes/refl/rewards/videoalign/model/reward_model.py index 654f89d83..4276b671b 100644 --- a/recipes/refl_wan/rewards/videoalign/model/reward_model.py +++ b/recipes/refl/rewards/videoalign/model/reward_model.py @@ -12,7 +12,7 @@ The forward signature is preserved so the same checkpoint state_dict loads cleanly (no key remapping needed beyond the standard transformers>=5 ``base_model.model.model.language_model.*`` shim handled by -:func:`recipes.refl_wan.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). +:func:`recipes.refl.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). """ from __future__ import annotations diff --git a/recipes/refl/rewards/videoalign/requirements.txt b/recipes/refl/rewards/videoalign/requirements.txt new file mode 100644 index 000000000..c66c054a7 --- /dev/null +++ b/recipes/refl/rewards/videoalign/requirements.txt @@ -0,0 +1,11 @@ +transformers==4.45.2 +peft==0.10.0 +trl==0.8.6 +accelerate==0.34.0 +decord==0.6.0 +opencv-python-headless +safetensors +huggingface_hub +einops +flash-attn==2.5.8 +setuptools<70 \ No newline at end of file diff --git a/recipes/refl_wan/rewards/videoalign/scorer.py b/recipes/refl/rewards/videoalign/scorer.py similarity index 99% rename from recipes/refl_wan/rewards/videoalign/scorer.py rename to recipes/refl/rewards/videoalign/scorer.py index 180ee4a76..39f5394b4 100644 --- a/recipes/refl_wan/rewards/videoalign/scorer.py +++ b/recipes/refl/rewards/videoalign/scorer.py @@ -21,7 +21,7 @@ This scorer no longer requires the sibling ``mmrl`` repo on disk. The Qwen2-VL reward backbone, prompt template, checkpoint loader and inference wrapper all live under -:mod:`recipes.refl_wan.rewards.videoalign.model` / :mod:`...wrapper`. The +:mod:`recipes.refl.rewards.videoalign.model` / :mod:`...wrapper`. The ``mmrl_repo_root`` Spec field has been removed; ``MMRL_REPO_ROOT`` env var is now irrelevant. """ diff --git a/recipes/refl_wan/rewards/videoalign/wrapper.py b/recipes/refl/rewards/videoalign/wrapper.py similarity index 99% rename from recipes/refl_wan/rewards/videoalign/wrapper.py rename to recipes/refl/rewards/videoalign/wrapper.py index a85fbc0bc..25e5713ee 100644 --- a/recipes/refl_wan/rewards/videoalign/wrapper.py +++ b/recipes/refl/rewards/videoalign/wrapper.py @@ -2,7 +2,7 @@ Self-contained re-implementation — does NOT touch ``sys.path`` . The Qwen2-VL reward backbone, its prompt template and the checkpoint loader all live under -:mod:`recipes.refl_wan.rewards.videoalign.model`. +:mod:`recipes.refl.rewards.videoalign.model`. Public API ---------- diff --git a/recipes/refl_wan/roles.py b/recipes/refl/roles.py similarity index 77% rename from recipes/refl_wan/roles.py rename to recipes/refl/roles.py index 4ad6dbe35..a1fca7979 100644 --- a/recipes/refl_wan/roles.py +++ b/recipes/refl/roles.py @@ -13,8 +13,9 @@ from unirl.distributed.group.dispatch import distributed from unirl.distributed.tensor.batch import Batch, concat_field, shared_field from unirl.reward.base import DifferentiableReward +from unirl.reward.role import RewardRole from unirl.sde.runtime import get_sigma_schedule -from unirl.trainer.base_role import Role +from unirl.trainer.role import Role from unirl.types.primitives import Images, Texts from unirl.types.rollout_req import RolloutReq @@ -174,51 +175,4 @@ def forward_backward_loss( ) -class ReflRewardRole(Role): - """Generic local differentiable reward role for REFL recipes.""" - - backend: Any - - @property - def preferred_input_kind(self) -> str: - kind = str(getattr(self.backend, "preferred_input_kind", "") or "").strip().lower() - if kind not in {"image", "video", "text"}: - raise ValueError(f"Reward backend must expose preferred_input_kind image/video/text, got {kind!r}.") - return kind - - @distributed - def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: - """Score live-grad decoded media through the official differentiable reward path.""" - decoded = generated.decoded - if not isinstance(decoded, torch.Tensor): - raise TypeError( - f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " - f"got {type(decoded).__name__}." - ) - texts = req.primitives.get("text") if req.primitives else None - if not isinstance(texts, Texts): - raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") - - if not isinstance(self.backend, DifferentiableReward): - raise TypeError( - f"{type(self.backend).__name__} must implement compute_rewards_differentiable " - "for REFL reward backprop." - ) - - records = list(req.metadata) if req.metadata else None - kind = self.preferred_input_kind - if kind == "image" and decoded.ndim != 4: - raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") - if kind == "video" and decoded.ndim != 5: - raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") - if kind not in {"image", "video"}: - raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") - - return self.backend.compute_rewards_differentiable( - decoded, - list(texts.texts), - records=records, - ) - - -__all__ = ["ReflActorRole", "ReflRewardRole"] +__all__ = ["ReflActorRole"] diff --git a/recipes/refl_wan/run.py b/recipes/refl/run.py similarity index 74% rename from recipes/refl_wan/run.py rename to recipes/refl/run.py index 673f2f06b..83854f7f2 100644 --- a/recipes/refl_wan/run.py +++ b/recipes/refl/run.py @@ -1,12 +1,12 @@ #!/usr/bin/env python -"""Hydra entry point for the recipes.refl_wan REFL recipe.""" +"""Hydra entry point for the recipes.refl REFL recipe.""" from __future__ import annotations import hydra from omegaconf import DictConfig -from recipes.refl_wan.trainer import REFLTrainer +from recipes.refl.trainer import REFLTrainer @hydra.main(version_base=None, config_path="configs", config_name="wan22_face_refl") diff --git a/recipes/refl_wan/start_wan21_t2v.sh b/recipes/refl/scripts/start_wan21_t2v.sh old mode 100644 new mode 100755 similarity index 66% rename from recipes/refl_wan/start_wan21_t2v.sh rename to recipes/refl/scripts/start_wan21_t2v.sh index dfaa20f2b..96046457e --- a/recipes/refl_wan/start_wan21_t2v.sh +++ b/recipes/refl/scripts/start_wan21_t2v.sh @@ -1,16 +1,19 @@ set -u +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "${SCRIPT_DIR}/../../.." + # === Data === # Newline-delimited prompts .txt (or {"prompt": ...} .jsonl). VideoAlign is # text-conditioned; no reference video / first-frame needed. -export DATA_PATH=${DATA_PATH:-/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/filtered_720x1280_prompts.txt} +export DATA_PATH=${DATA_PATH:-/path/to/wan21_prompts.txt} export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} # === Models === # WAN 2.1 T2V 1.3B base checkpoint -export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/Wan2.1-T2V-1.3B-Diffusers} +export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/path/to/Wan2.1-T2V-1.3B-Diffusers} # VideoAlign Qwen2-VL reward checkpoint -export VIDEOALIGN_MODEL_PATH=${VIDEOALIGN_MODEL_PATH:-/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/VideoReward} +export VIDEOALIGN_MODEL_PATH=${VIDEOALIGN_MODEL_PATH:-/path/to/VideoReward} # === Output / Logging === export OUTPUT_DIR=${OUTPUT_DIR:-./outputs/wan21_t2v_videoalign_refl} @@ -23,7 +26,7 @@ mkdir -p "${OUTPUT_DIR}" logs LOG_FILE="logs/wan21_t2v_videoalign_refl_$(date +%Y%m%d_%H%M%S).log" echo "=== launching wan21 t2v videoalign refl, log → ${LOG_FILE} ===" -RAY_ADDRESS=auto python -u -m recipes.refl_wan.run \ +RAY_ADDRESS=auto python -u -m recipes.refl.run \ --config-name=wan21_t2v_videoalign_refl \ num_devices=8 \ 2>&1 | tee "${LOG_FILE}" diff --git a/recipes/refl/scripts/start_wan22_i2v.sh b/recipes/refl/scripts/start_wan22_i2v.sh new file mode 100755 index 000000000..ac588bc02 --- /dev/null +++ b/recipes/refl/scripts/start_wan22_i2v.sh @@ -0,0 +1,20 @@ +set -u + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "${SCRIPT_DIR}/../../.." + +export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/path/to/Wan2.2-I2V-A14B-Diffusers} +export DATA_PATH=${DATA_PATH:-/path/to/wan22_face_refl_prompts.jsonl} +export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} +export FACE_MODEL_PATH=${FACE_MODEL_PATH:-/path/to/antelodev2} +export OUTPUT_DIR=${OUTPUT_DIR:-outputs/wan22_face_refl_recipe_opt} + +export REPORT_TO_WANDB=${REPORT_TO_WANDB:-true} +export WANDB_PROJECT=${WANDB_PROJECT:-unirl-refl} +export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan22_face_refl_recipe_opt} + +mkdir -p logs + +RAY_ADDRESS=auto python -m recipes.refl.run \ + num_devices=8 \ + 2>&1 | tee ./logs/wan22_i2v_refl.log diff --git a/recipes/refl_wan/trainer.py b/recipes/refl/trainer.py similarity index 53% rename from recipes/refl_wan/trainer.py rename to recipes/refl/trainer.py index 2787163cc..3c6b1c3d2 100644 --- a/recipes/refl_wan/trainer.py +++ b/recipes/refl/trainer.py @@ -5,6 +5,7 @@ import time from collections.abc import Mapping from typing import Any, Dict +import numpy as np from unirl.distributed.tensor.grad_context import enable_grad from unirl.trainer.trainer import Trainer @@ -16,40 +17,6 @@ class REFLTrainer(Trainer): """REFL / BPTT recipe trainer: role-driven 3-RPC train step.""" - def validate_config(self) -> None: - super().validate_config() - if "actor" not in self.roles: - raise ValueError("REFLTrainer requires a role named 'actor'.") - if "reward" not in self.roles: - raise ValueError("REFLTrainer requires a role named 'reward'.") - - actor_spec = self._role_specs_by_name["actor"] - if actor_spec.raw_cfg.get("algorithm") is None: - raise ValueError("REFLTrainer requires roles[name=actor].algorithm: ${algorithm}.") - algo_cfg = self.cfg.get("algorithm") - if algo_cfg is None: - raise ValueError("REFLTrainer requires top-level cfg.algorithm.") - if hasattr(algo_cfg, "get") and algo_cfg.get("sampling_params") is None: - raise ValueError("REFLTrainer requires cfg.algorithm.sampling_params: ${sampling}.") - - reward_spec = self._role_specs_by_name["reward"] - backend_cfg = reward_spec.raw_cfg.get("backend") - backend_target = str(backend_cfg.get("_target_") or "") if backend_cfg is not None and hasattr(backend_cfg, "get") else "" - if backend_target.endswith("RemoteRewardBackend"): - raise ValueError( - "REFLTrainer requires a local differentiable reward backend; " - f"got roles[name=reward].backend._target_={backend_target!r}." - ) - - rollout_section = self.cfg.get("rollout", None) - if rollout_section is not None: - rollout_target = str(rollout_section.get("_target_", "")) if hasattr(rollout_section, "get") else "" - if rollout_target and not rollout_target.endswith("TrainsideRolloutEngine"): - raise ValueError( - "REFLTrainer: a rollout section is present but is not trainside. " - f"Got rollout._target_={rollout_target!r}." - ) - def build_req(self, inputs: RolloutInputs, rollout_id: int) -> RolloutReq: """Build one RolloutReq from data-source samples.""" inputs = inputs.expand(total_samples_per_prompt(self.sampling_params)) @@ -97,29 +64,15 @@ def train_step(self, req: RolloutReq, *, training_progress: float = 0.0, rollout step_result = self.actor.step() metrics: Dict[str, Any] = { - "loss": self._mean(getattr(loss_metrics, "loss", None)), - "reward_loss": self._mean(getattr(loss_metrics, "reward_loss", None)), - "kl_loss": self._mean(getattr(loss_metrics, "kl_loss", None)), - "reward_mean": self._mean(getattr(loss_metrics, "reward_mean", None)), - "grad_norm": self._mean(step_result.metrics.get("grad_norm")) if step_result.metrics else 0.0, - "lr": self._mean(step_result.metrics.get("lr")) if step_result.metrics else 0.0, + "loss": np.mean(loss_metrics.loss), + "reward_loss": np.mean(loss_metrics.reward_loss), + "kl_loss": np.mean(loss_metrics.kl_loss), + "reward_mean": np.mean(loss_metrics.reward_mean), + "grad_norm": np.mean(step_result.metrics.get("grad_norm")) if step_result.metrics else 0.0, + "lr": np.mean(step_result.metrics.get("lr")) if step_result.metrics else 0.0, "step_time_s": time.perf_counter() - t0, "training_progress": float(training_progress), } return metrics - @staticmethod - def _mean(field: Any) -> float: - try: - if hasattr(field, "tolist"): - field = field.tolist() - if isinstance(field, (list, tuple)) and field: - return float(sum(float(x) for x in field) / len(field)) - if isinstance(field, (int, float)): - return float(field) - return 0.0 - except Exception: - return 0.0 - - __all__ = ["REFLTrainer"] diff --git a/recipes/refl_wan/start_wan22_i2v.sh b/recipes/refl_wan/start_wan22_i2v.sh deleted file mode 100644 index aeba8403f..000000000 --- a/recipes/refl_wan/start_wan22_i2v.sh +++ /dev/null @@ -1,13 +0,0 @@ -export PRETRAINED_MODEL="/apdcephfs_gy8/share_301869871/ysunlin/mmrl/.model_cache/Wan2.2-I2V-A14B-Diffusers" -export DATA_PATH="/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/wan22_face_refl_prompts.jsonl" -export EVAL_DATA_PATH="${DATA_PATH}" -export FACE_MODEL_PATH="/apdcephfs_gy8/share_301869871/yohunawu/gy6/yohunawu/refl_wan_data/antelodev2" -export OUTPUT_DIR="outputs/wan22_face_refl_recipe_opt" - -export REPORT_TO_WANDB=true -export WANDB_PROJECT="unirl-refl" -export WANDB_RUN_NAME="wan22_face_refl_recipe_opt" - -RAY_ADDRESS=auto python -m recipes.refl_wan.run \ - num_devices=8 \ - 2>&1 | tee ./wan22_i2v_refl.log \ No newline at end of file diff --git a/unirl/reward/role.py b/unirl/reward/role.py new file mode 100644 index 000000000..c3219deaf --- /dev/null +++ b/unirl/reward/role.py @@ -0,0 +1,99 @@ +"""Reward role abstraction for trainer-managed reward remotes.""" + +from __future__ import annotations + +import logging +import torch +from typing import Any + +from hydra.utils import instantiate + +from unirl.distributed.group.remote import Remote +from unirl.distributed.group.dispatch import distributed +from unirl.reward.service import RewardService +from unirl.types.rollout_req import RolloutReq +from unirl.types.primitives import Texts +from unirl.reward.base import DifferentiableReward + +logger = logging.getLogger(__name__) + + +class RewardRole(RewardService): + """A trainer role that owns one reward backend. + + ``RewardService`` normally receives an already-instantiated backend in its + constructor. Trainer-managed roles are constructed with a role config first, + then ``initialize()`` runs after ``Remote.setup`` injects the worker device. + This class bridges those lifecycles while preserving ``RewardService``'s + scoring and offload/onload/dispose methods. + """ + + def __init__(self, cfg: Any) -> None: + Remote.__init__(self) + self.cfg = cfg + + def _cfg_get(self, key: str, default: Any = None) -> Any: + if hasattr(self.cfg, "get"): + return self.cfg.get(key, default) + return getattr(self.cfg, key, default) + + def initialize(self) -> None: + if self.device is None: + raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") + + backend_cfg = self._cfg_get("backend") + if backend_cfg is None: + raise ValueError(f"{type(self).__name__} requires cfg.backend.") + + self.backend = instantiate(backend_cfg, base_device=self.device) + self.truncated_reward = str(self._cfg_get("truncated_reward", "zero")) + self.overlong_buffer_len = int(self._cfg_get("overlong_buffer_len", 4096)) + self.overlong_penalty_factor = float(self._cfg_get("overlong_penalty_factor", 1.0)) + if self.truncated_reward not in ("zero", "keep", "soft"): + raise ValueError(f"truncated_reward must be zero|keep|soft, got {self.truncated_reward!r}") + + if hasattr(self.backend, "initialize"): + self.backend.initialize() + + logger.info( + "RewardRole initialized with backend=%s, truncated_reward=%s", + self.backend.get_model_name() or type(self.backend).__name__, + self.truncated_reward, + ) + + @distributed + def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: + """Score live-grad decoded media through the official differentiable reward path.""" + decoded = generated.decoded + if not isinstance(decoded, torch.Tensor): + raise TypeError( + f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " + f"got {type(decoded).__name__}." + ) + texts = req.primitives.get("text") if req.primitives else None + if not isinstance(texts, Texts): + raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") + + if not isinstance(self.backend, DifferentiableReward): + raise TypeError( + f"{type(self.backend).__name__} must implement compute_rewards_differentiable " + "for REFL reward backprop." + ) + + records = list(req.metadata) if req.metadata else None + kind = self.preferred_input_kind + if kind == "image" and decoded.ndim != 4: + raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") + if kind == "video" and decoded.ndim != 5: + raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") + if kind not in {"image", "video"}: + raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") + + return self.backend.compute_rewards_differentiable( + decoded, + list(texts.texts), + records=records, + ) + + +__all__ = ["RewardRole"] diff --git a/unirl/trainer/base_role.py b/unirl/trainer/role.py similarity index 88% rename from unirl/trainer/base_role.py rename to unirl/trainer/role.py index 872f251be..b01c672e2 100644 --- a/unirl/trainer/base_role.py +++ b/unirl/trainer/role.py @@ -70,22 +70,12 @@ def initialize(self) -> None: else: self.backend = instantiate(backend_cfg) - def _max_grad_norm(self) -> float: - algo_cfg = getattr(self, "algo_cfg", None) - if algo_cfg is None: - algo_cfg = getattr(self, "algorithm_cfg", None) - if algo_cfg is not None and hasattr(algo_cfg, "get"): - return float(algo_cfg.get("max_grad_norm", 1.0)) - if algo_cfg is not None and hasattr(algo_cfg, "max_grad_norm"): - return float(algo_cfg.max_grad_norm) - return float(getattr(self, "max_grad_norm", 1.0)) - @distributed def step(self) -> RoleStepResult: """Clip gradients and run one backend optimizer step.""" if not hasattr(self, "backend") or not hasattr(self.backend, "optimizer_step"): raise RuntimeError(f"{type(self).__name__}.step requires a backend with optimizer_step(...).") - grad_norm = float(self.backend.optimizer_step(max_grad_norm=self._max_grad_norm())) + grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(self.algo_cfg.get("max_grad_norm", 1.0)))) lr = 0.0 try: sched = getattr(self.backend, "scheduler", None) diff --git a/unirl/trainer/trainer.py b/unirl/trainer/trainer.py index 27c486728..26e32bcab 100644 --- a/unirl/trainer/trainer.py +++ b/unirl/trainer/trainer.py @@ -315,8 +315,6 @@ def train( finally: self._finish_wandb() - fit = train - def log_metrics(self, metrics: Dict[str, Any], *, rollout_id: int, num_rollouts: int) -> None: trainer_name = self.__class__.__name__.removesuffix("Trainer") or self.__class__.__name__ metric_text = " ".join( From db35f4c03e490f7baf4df14e43f95ecbd34b9ec0 Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Wed, 15 Jul 2026 12:06:10 +0800 Subject: [PATCH 03/24] rename wan2.2 i2v refl config file's name --- .../{wan22_face_refl.yaml => wan22_i2v_face_refl.yaml} | 8 ++++---- recipes/refl/run.py | 2 +- recipes/refl/scripts/start_wan22_i2v.sh | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) rename recipes/refl/configs/{wan22_face_refl.yaml => wan22_i2v_face_refl.yaml} (93%) diff --git a/recipes/refl/configs/wan22_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml similarity index 93% rename from recipes/refl/configs/wan22_face_refl.yaml rename to recipes/refl/configs/wan22_i2v_face_refl.yaml index 42dc2d5de..7031b0eb0 100644 --- a/recipes/refl/configs/wan22_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -5,7 +5,7 @@ num_devices: 8 batch_size: 8 num_rollouts: 1000 save_interval: 100 -save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_face_refl} +save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_i2v_face_refl} roles: - name: actor @@ -110,8 +110,8 @@ data_source: _target_: unirl.data.data_source.MultimodalRLDataSource args: run: - data_path: ${oc.env:DATA_PATH,/path/to/wan22_face_refl_prompts.jsonl} - eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/path/to/wan22_face_refl_prompts.jsonl}} + data_path: ${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl}} seed: 42 shuffle: false algorithm: @@ -137,7 +137,7 @@ sampling: 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,wan22_face_refl} + run_name: ${oc.env:WANDB_RUN_NAME,wan22_i2v_face_refl} tags: ["wan22", "i2v", "refl", "face", "recipe"] log_media: false metric_prefix: refl_recipe/ diff --git a/recipes/refl/run.py b/recipes/refl/run.py index 83854f7f2..de7f6129f 100644 --- a/recipes/refl/run.py +++ b/recipes/refl/run.py @@ -9,7 +9,7 @@ from recipes.refl.trainer import REFLTrainer -@hydra.main(version_base=None, config_path="configs", config_name="wan22_face_refl") +@hydra.main(version_base=None, config_path="configs", config_name="wan22_i2v_face_refl") def main(cfg: DictConfig) -> None: trainer = REFLTrainer(cfg=cfg) trainer.train() diff --git a/recipes/refl/scripts/start_wan22_i2v.sh b/recipes/refl/scripts/start_wan22_i2v.sh index ac588bc02..9e581b180 100755 --- a/recipes/refl/scripts/start_wan22_i2v.sh +++ b/recipes/refl/scripts/start_wan22_i2v.sh @@ -4,14 +4,14 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "${SCRIPT_DIR}/../../.." export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/path/to/Wan2.2-I2V-A14B-Diffusers} -export DATA_PATH=${DATA_PATH:-/path/to/wan22_face_refl_prompts.jsonl} +export DATA_PATH=${DATA_PATH:-/path/to/wan22_i2v_face_refl_prompts.jsonl} export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} export FACE_MODEL_PATH=${FACE_MODEL_PATH:-/path/to/antelodev2} -export OUTPUT_DIR=${OUTPUT_DIR:-outputs/wan22_face_refl_recipe_opt} +export OUTPUT_DIR=${OUTPUT_DIR:-outputs/wan22_i2v_face_refl_recipe_opt} export REPORT_TO_WANDB=${REPORT_TO_WANDB:-true} export WANDB_PROJECT=${WANDB_PROJECT:-unirl-refl} -export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan22_face_refl_recipe_opt} +export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan22_i2v_face_refl_recipe_opt} mkdir -p logs From c1333d514238e7207fba5851cf59fb423e747b4e Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Fri, 17 Jul 2026 17:48:03 +0800 Subject: [PATCH 04/24] fix bug: resolved the issue of being unable to enter the kl branch in the wan refl recipe. --- recipes/refl/models/wan21.py | 37 ++++++++++++++---------------------- recipes/refl/models/wan22.py | 35 ++++++++++++++-------------------- 2 files changed, 28 insertions(+), 44 deletions(-) diff --git a/recipes/refl/models/wan21.py b/recipes/refl/models/wan21.py index 238fabc85..ede2a477e 100644 --- a/recipes/refl/models/wan21.py +++ b/recipes/refl/models/wan21.py @@ -29,6 +29,7 @@ from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.diffusion import WAN21DiffusionStage, WAN21DiffusionStep from unirl.models.wan21.pipeline import WAN21Pipeline +from unirl.train.lora import adapters_disabled from unirl.types.sampling import DiffusionSamplingParams # Matches the mainline module-level constant in unirl/models/wan21/diffusion.py. @@ -168,11 +169,10 @@ def diffuse_with_grad( participate in backward). DRaFT-1 sets ``mid_timestep = T - 1``. - ``final_timestep`` (int, default ``num_inference_steps - 1``): early stop. The loop breaks once ``i >= final_timestep``. - - ``kl_weight`` (float, default 0.0): when non-zero AND the - transformer exposes ``disable_adapter_layers`` (PEFT LoRA), - per-step KL ``mean((pred - ref_pred)**2 / (2 * sigma**2))`` is - accumulated and returned in ``kl_loss``. The trainer multiplies - it by its own ``kl_weight`` at the loss site. + - ``kl_weight`` (float, default 0.0): when non-zero, per-step KL + ``mean((pred - ref_pred)**2 / (2 * sigma**2))`` is accumulated + and returned in ``kl_loss``. The trainer multiplies it by its own + ``kl_weight`` at the loss site. ``initial_latents`` follows the same contract as :meth:`diffuse`: when provided, used verbatim and the internal RNG path is @@ -295,24 +295,15 @@ def diffuse_with_grad( # Per-step KL against the LoRA-disabled reference (REFL-style). # When CFG is enabled, this intentionally computes KL on the # conditional prediction before the stop-grad uncond CFG mix. - if ( - kl_weight != 0.0 - and grad_enabled - and hasattr(transformer, "disable_adapter_layers") - and hasattr(transformer, "enable_adapter_layers") - ): - with torch.no_grad(), autocast_ctx: - transformer.disable_adapter_layers() - try: - ref_pred = step.predict_noise( - self.model, - latents, - sigma, - conditions, - branch="cond", - ) - finally: - transformer.enable_adapter_layers() + if kl_weight != 0.0 and grad_enabled: + with torch.no_grad(), autocast_ctx, adapters_disabled(transformer): + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + ) sigma_f32 = sigma.to(dtype=torch.float32) kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() kl_total = kl_total + kl_step diff --git a/recipes/refl/models/wan22.py b/recipes/refl/models/wan22.py index 86e38dbd9..6ce837af3 100644 --- a/recipes/refl/models/wan22.py +++ b/recipes/refl/models/wan22.py @@ -12,6 +12,7 @@ from unirl.models.wan22.bundle import WAN22Bundle from unirl.models.wan22.diffusion import WAN22DiffusionStage, WAN22DiffusionStep from unirl.models.wan22.pipeline import WAN22Pipeline +from unirl.train.lora import adapters_disabled from unirl.types.sampling import DiffusionSamplingParams # Matches the mainline module-level constant in unirl/models/wan22/diffusion.py. @@ -230,7 +231,6 @@ def diffuse_with_grad( guidance_scale_2, boundary_ratio=boundary_ratio, ) - active_sub = getattr(dual, "high_noise" if use_high_noise else "low_noise", None) use_cfg = active_guidance > 1.0 grad_enabled = i >= mid_timestep @@ -276,26 +276,19 @@ def diffuse_with_grad( # conditional prediction of the currently active branch. When CFG # is enabled, this intentionally excludes the stop-grad uncond mix. if kl_weight != 0.0 and grad_enabled: - if active_sub is not None and hasattr(active_sub, "disable_adapter_layers") and hasattr( - active_sub, "enable_adapter_layers" - ): - with torch.no_grad(), autocast_ctx: - active_sub.disable_adapter_layers() - try: - ref_pred = step.predict_noise( - self.model, - latents, - sigma, - conditions, - branch="cond", - use_high_noise=use_high_noise, - ) - finally: - active_sub.enable_adapter_layers() - sigma_f32 = sigma.to(dtype=torch.float32) - kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() - kl_total = kl_total + kl_step - kl_steps += 1 + with torch.no_grad(), autocast_ctx, adapters_disabled(dual): + ref_pred = step.predict_noise( + self.model, + latents, + sigma, + conditions, + branch="cond", + use_high_noise=use_high_noise, + ) + sigma_f32 = sigma.to(dtype=torch.float32) + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() + kl_total = kl_total + kl_step + kl_steps += 1 # Keep the transition outside autocast and under the same BPTT # window used by WAN21: grad-enabled for train steps, no-grad From ced3c37c95754ad76fd6619dc1e11a98e2ff5245 Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Mon, 20 Jul 2026 11:16:38 +0800 Subject: [PATCH 05/24] fix bug: support batched WAN VAE tiled decode instead of fixing the batch size to 1 --- unirl/models/wan21/wan_video_vae.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/unirl/models/wan21/wan_video_vae.py b/unirl/models/wan21/wan_video_vae.py index 9de86cc1d..29a570b71 100644 --- a/unirl/models/wan21/wan_video_vae.py +++ b/unirl/models/wan21/wan_video_vae.py @@ -847,16 +847,16 @@ def _make_tile_tasks(self, H, W, size_h, size_w, stride_h, stride_w): return tasks def tiled_decode(self, hidden_states, device, tile_size, tile_stride): - _, _, T, H, W = hidden_states.shape + B, _, T, H, W = hidden_states.shape size_h, size_w = tile_size stride_h, stride_w = tile_stride tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) out_T = T * 4 - 3 - weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, + weight = torch.zeros((B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) - values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, + values = torch.zeros((B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) @@ -872,8 +872,10 @@ def tiled_decode(self, hidden_states, device, tile_size, tile_stride): th = h * self.upsampling_factor tw = w * self.upsampling_factor - values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += batch * mask - weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += mask + value_slice = values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] + weight_slice = weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] + value_slice += batch * mask + weight_slice += mask.expand_as(weight_slice) values = values / weight return values.clamp_(-1, 1) @@ -885,7 +887,7 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, Args: sp_group: torch.distributed ProcessGroup for sequence parallelism. """ - _, _, T, H, W = hidden_states.shape + B, _, T, H, W = hidden_states.shape size_h, size_w = tile_size stride_h, stride_w = tile_stride tile_img_h = size_h * self.upsampling_factor @@ -893,10 +895,10 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) out_T = T * 4 - 3 - weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, + weight = torch.zeros((B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) - values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, + values = torch.zeros((B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) @@ -955,10 +957,12 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, th = h * self.upsampling_factor tw = w * self.upsampling_factor - values[:, :, :, th:th + decoded_tile.shape[3], - tw:tw + decoded_tile.shape[4]] += decoded_tile * mask - weight[:, :, :, th:th + decoded_tile.shape[3], - tw:tw + decoded_tile.shape[4]] += mask + value_slice = values[:, :, :, th:th + decoded_tile.shape[3], + tw:tw + decoded_tile.shape[4]] + weight_slice = weight[:, :, :, th:th + decoded_tile.shape[3], + tw:tw + decoded_tile.shape[4]] + value_slice += decoded_tile * mask + weight_slice += mask.expand_as(weight_slice) values = values / weight return values.clamp_(-1, 1) From 25d4ee77e36d2e08ca6701696378b19b7dac6ae5 Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Tue, 21 Jul 2026 14:14:24 +0800 Subject: [PATCH 06/24] refactor: move base role from unirl core to recipes --- recipes/common/__init__.py | 1 + {unirl/trainer => recipes/common}/role.py | 4 +- .../configs/wan21_t2v_videoalign_refl.yaml | 2 +- recipes/refl/configs/wan22_i2v_face_refl.yaml | 2 +- recipes/refl/roles.py | 72 +++++++++++++- unirl/reward/role.py | 99 ------------------- 6 files changed, 74 insertions(+), 106 deletions(-) create mode 100644 recipes/common/__init__.py rename {unirl/trainer => recipes/common}/role.py (96%) delete mode 100644 unirl/reward/role.py diff --git a/recipes/common/__init__.py b/recipes/common/__init__.py new file mode 100644 index 000000000..d80502c7d --- /dev/null +++ b/recipes/common/__init__.py @@ -0,0 +1 @@ +"""Shared helpers for recipe-level implementations.""" diff --git a/unirl/trainer/role.py b/recipes/common/role.py similarity index 96% rename from unirl/trainer/role.py rename to recipes/common/role.py index b01c672e2..6bce18140 100644 --- a/unirl/trainer/role.py +++ b/recipes/common/role.py @@ -1,4 +1,4 @@ -"""Base Role abstraction for trainer-managed Remote roles.""" +"""Recipe-level Role abstraction for trainer-managed Remote roles.""" from __future__ import annotations @@ -23,7 +23,7 @@ class RoleStepResult: class Role(Remote): - """A trainer role that runs as a UniRL Remote inside a Worker. + """Recipe-level role that runs as a UniRL Remote inside a Worker. ``Role`` stores the role config and initializes common role-local components (model/runtime bundle / pipeline / backend) as ordinary diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index e3c601e7f..c72c61fb3 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -84,7 +84,7 @@ roles: algorithm: ${algorithm} - name: reward - _target_: unirl.reward.role.RewardRole + _target_: recipes.refl.roles.ReflRewardRole placement: # BPTT gradients must live on the same worker — no cross-process # autograd. VideoAlign co-locates with the actor. diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index 7031b0eb0..0cc7b84a9 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -81,7 +81,7 @@ roles: algorithm: ${algorithm} - name: reward - _target_: unirl.reward.role.RewardRole + _target_: recipes.refl.roles.ReflRewardRole placement: share_with: actor diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index a1fca7979..572fb3a53 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -3,6 +3,7 @@ from __future__ import annotations import dataclasses +import logging from dataclasses import dataclass from typing import Any, List, Optional @@ -10,15 +11,18 @@ from hydra.utils import instantiate from omegaconf import OmegaConf +from recipes.common.role import Role from unirl.distributed.group.dispatch import distributed +from unirl.distributed.group.remote import Remote from unirl.distributed.tensor.batch import Batch, concat_field, shared_field from unirl.reward.base import DifferentiableReward -from unirl.reward.role import RewardRole +from unirl.reward.service import RewardService from unirl.sde.runtime import get_sigma_schedule -from unirl.trainer.role import Role from unirl.types.primitives import Images, Texts from unirl.types.rollout_req import RolloutReq +logger = logging.getLogger(__name__) + @dataclass class REFLGenerated(Batch): @@ -175,4 +179,66 @@ def forward_backward_loss( ) -__all__ = ["ReflActorRole"] +class ReflRewardRole(RewardService): + """Recipe-owned reward role for REFL.""" + + def __init__(self, cfg: Any) -> None: + Remote.__init__(self) + self.cfg = cfg + + def _cfg_get(self, key: str, default: Any = None) -> Any: + if hasattr(self.cfg, "get"): + return self.cfg.get(key, default) + return getattr(self.cfg, key, default) + + def initialize(self) -> None: + if self.device is None: + raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") + + backend_cfg = self._cfg_get("backend") + if backend_cfg is None: + raise ValueError(f"{type(self).__name__} requires cfg.backend.") + + self.backend = instantiate(backend_cfg, base_device=self.device) + + logger.info( + "ReflRewardRole initialized with backend=%s", + self.backend.get_model_name() or type(self.backend).__name__, + ) + + @distributed + def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: + """Score live-grad decoded media through the official differentiable reward path.""" + decoded = generated.decoded + if not isinstance(decoded, torch.Tensor): + raise TypeError( + f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " + f"got {type(decoded).__name__}." + ) + texts = req.primitives.get("text") if req.primitives else None + if not isinstance(texts, Texts): + raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") + + if not isinstance(self.backend, DifferentiableReward): + raise TypeError( + f"{type(self.backend).__name__} must implement compute_rewards_differentiable " + "for REFL reward backprop." + ) + + records = list(req.metadata) if req.metadata else None + kind = self.preferred_input_kind + if kind == "image" and decoded.ndim != 4: + raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") + if kind == "video" and decoded.ndim != 5: + raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") + if kind not in {"image", "video"}: + raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") + + return self.backend.compute_rewards_differentiable( + decoded, + list(texts.texts), + records=records, + ) + + +__all__ = ["ReflActorRole", "ReflRewardRole"] diff --git a/unirl/reward/role.py b/unirl/reward/role.py deleted file mode 100644 index c3219deaf..000000000 --- a/unirl/reward/role.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Reward role abstraction for trainer-managed reward remotes.""" - -from __future__ import annotations - -import logging -import torch -from typing import Any - -from hydra.utils import instantiate - -from unirl.distributed.group.remote import Remote -from unirl.distributed.group.dispatch import distributed -from unirl.reward.service import RewardService -from unirl.types.rollout_req import RolloutReq -from unirl.types.primitives import Texts -from unirl.reward.base import DifferentiableReward - -logger = logging.getLogger(__name__) - - -class RewardRole(RewardService): - """A trainer role that owns one reward backend. - - ``RewardService`` normally receives an already-instantiated backend in its - constructor. Trainer-managed roles are constructed with a role config first, - then ``initialize()`` runs after ``Remote.setup`` injects the worker device. - This class bridges those lifecycles while preserving ``RewardService``'s - scoring and offload/onload/dispose methods. - """ - - def __init__(self, cfg: Any) -> None: - Remote.__init__(self) - self.cfg = cfg - - def _cfg_get(self, key: str, default: Any = None) -> Any: - if hasattr(self.cfg, "get"): - return self.cfg.get(key, default) - return getattr(self.cfg, key, default) - - def initialize(self) -> None: - if self.device is None: - raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") - - backend_cfg = self._cfg_get("backend") - if backend_cfg is None: - raise ValueError(f"{type(self).__name__} requires cfg.backend.") - - self.backend = instantiate(backend_cfg, base_device=self.device) - self.truncated_reward = str(self._cfg_get("truncated_reward", "zero")) - self.overlong_buffer_len = int(self._cfg_get("overlong_buffer_len", 4096)) - self.overlong_penalty_factor = float(self._cfg_get("overlong_penalty_factor", 1.0)) - if self.truncated_reward not in ("zero", "keep", "soft"): - raise ValueError(f"truncated_reward must be zero|keep|soft, got {self.truncated_reward!r}") - - if hasattr(self.backend, "initialize"): - self.backend.initialize() - - logger.info( - "RewardRole initialized with backend=%s, truncated_reward=%s", - self.backend.get_model_name() or type(self.backend).__name__, - self.truncated_reward, - ) - - @distributed - def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: - """Score live-grad decoded media through the official differentiable reward path.""" - decoded = generated.decoded - if not isinstance(decoded, torch.Tensor): - raise TypeError( - f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " - f"got {type(decoded).__name__}." - ) - texts = req.primitives.get("text") if req.primitives else None - if not isinstance(texts, Texts): - raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") - - if not isinstance(self.backend, DifferentiableReward): - raise TypeError( - f"{type(self.backend).__name__} must implement compute_rewards_differentiable " - "for REFL reward backprop." - ) - - records = list(req.metadata) if req.metadata else None - kind = self.preferred_input_kind - if kind == "image" and decoded.ndim != 4: - raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") - if kind == "video" and decoded.ndim != 5: - raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") - if kind not in {"image", "video"}: - raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") - - return self.backend.compute_rewards_differentiable( - decoded, - list(texts.texts), - records=records, - ) - - -__all__ = ["RewardRole"] From dbc48893d5688dd1ba4a42bb5f07f679571ea6a3 Mon Sep 17 00:00:00 2001 From: Yohuna <378072862@qq.com> Date: Tue, 21 Jul 2026 15:27:24 +0800 Subject: [PATCH 07/24] refactor: move reward role class to recipes/common/roles.py --- recipes/common/{role.py => roles.py} | 73 ++++++++++++++++++- .../configs/wan21_t2v_videoalign_refl.yaml | 2 +- recipes/refl/configs/wan22_i2v_face_refl.yaml | 2 +- recipes/refl/roles.py | 72 +----------------- 4 files changed, 75 insertions(+), 74 deletions(-) rename recipes/common/{role.py => roles.py} (57%) diff --git a/recipes/common/role.py b/recipes/common/roles.py similarity index 57% rename from recipes/common/role.py rename to recipes/common/roles.py index 6bce18140..a6752a282 100644 --- a/recipes/common/role.py +++ b/recipes/common/roles.py @@ -1,7 +1,8 @@ -"""Recipe-level Role abstraction for trainer-managed Remote roles.""" +"""Recipe-level Role abstractions for trainer-managed Remote roles.""" from __future__ import annotations +import logging from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Optional @@ -11,6 +12,12 @@ from unirl.distributed.group.dispatch import Dispatch, distributed from unirl.distributed.group.remote import Remote +from unirl.reward.base import DifferentiableReward +from unirl.reward.service import RewardService +from unirl.types.primitives import Texts +from unirl.types.rollout_req import RolloutReq + +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -104,4 +111,66 @@ def load_checkpoint(self, path: str) -> int: return 0 -__all__ = ["Role", "RoleStepResult"] +class RewardRole(RewardService): + """Common differentiable reward role for recipe-local reward backends.""" + + def __init__(self, cfg: Any) -> None: + Remote.__init__(self) + self.cfg = cfg + + def _cfg_get(self, key: str, default: Any = None) -> Any: + if hasattr(self.cfg, "get"): + return self.cfg.get(key, default) + return getattr(self.cfg, key, default) + + def initialize(self) -> None: + if self.device is None: + raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") + + backend_cfg = self._cfg_get("backend") + if backend_cfg is None: + raise ValueError(f"{type(self).__name__} requires cfg.backend.") + + self.backend = instantiate(backend_cfg, base_device=self.device) + + logger.info( + "RewardRole initialized with backend=%s", + self.backend.get_model_name() or type(self.backend).__name__, + ) + + @distributed + def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: + """Score live-grad decoded media through a differentiable reward backend.""" + decoded = generated.decoded + if not isinstance(decoded, torch.Tensor): + raise TypeError( + f"RewardRole.score_differentiable: generated.decoded must be Tensor, " + f"got {type(decoded).__name__}." + ) + texts = req.primitives.get("text") if req.primitives else None + if not isinstance(texts, Texts): + raise TypeError("RewardRole.score_differentiable: req.primitives['text'] must be Texts.") + + if not isinstance(self.backend, DifferentiableReward): + raise TypeError( + f"{type(self.backend).__name__} must implement compute_rewards_differentiable " + "for reward backprop." + ) + + records = list(req.metadata) if req.metadata else None + kind = self.preferred_input_kind + if kind == "image" and decoded.ndim != 4: + raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") + if kind == "video" and decoded.ndim != 5: + raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") + if kind not in {"image", "video"}: + raise ValueError(f"RewardRole.score_differentiable: unsupported input_kind={kind!r}.") + + return self.backend.compute_rewards_differentiable( + decoded, + list(texts.texts), + records=records, + ) + + +__all__ = ["Role", "RoleStepResult", "RewardRole"] diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index c72c61fb3..93df78817 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -84,7 +84,7 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.refl.roles.ReflRewardRole + _target_: recipes.common.roles.RewardRole placement: # BPTT gradients must live on the same worker — no cross-process # autograd. VideoAlign co-locates with the actor. diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index 0cc7b84a9..bccc76c46 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -81,7 +81,7 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.refl.roles.ReflRewardRole + _target_: recipes.common.roles.RewardRole placement: share_with: actor diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index 572fb3a53..929462e04 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -3,7 +3,6 @@ from __future__ import annotations import dataclasses -import logging from dataclasses import dataclass from typing import Any, List, Optional @@ -11,18 +10,13 @@ from hydra.utils import instantiate from omegaconf import OmegaConf -from recipes.common.role import Role +from recipes.common.roles import Role from unirl.distributed.group.dispatch import distributed -from unirl.distributed.group.remote import Remote from unirl.distributed.tensor.batch import Batch, concat_field, shared_field -from unirl.reward.base import DifferentiableReward -from unirl.reward.service import RewardService from unirl.sde.runtime import get_sigma_schedule from unirl.types.primitives import Images, Texts from unirl.types.rollout_req import RolloutReq -logger = logging.getLogger(__name__) - @dataclass class REFLGenerated(Batch): @@ -179,66 +173,4 @@ def forward_backward_loss( ) -class ReflRewardRole(RewardService): - """Recipe-owned reward role for REFL.""" - - def __init__(self, cfg: Any) -> None: - Remote.__init__(self) - self.cfg = cfg - - def _cfg_get(self, key: str, default: Any = None) -> Any: - if hasattr(self.cfg, "get"): - return self.cfg.get(key, default) - return getattr(self.cfg, key, default) - - def initialize(self) -> None: - if self.device is None: - raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") - - backend_cfg = self._cfg_get("backend") - if backend_cfg is None: - raise ValueError(f"{type(self).__name__} requires cfg.backend.") - - self.backend = instantiate(backend_cfg, base_device=self.device) - - logger.info( - "ReflRewardRole initialized with backend=%s", - self.backend.get_model_name() or type(self.backend).__name__, - ) - - @distributed - def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: - """Score live-grad decoded media through the official differentiable reward path.""" - decoded = generated.decoded - if not isinstance(decoded, torch.Tensor): - raise TypeError( - f"ReflRewardRole.score_differentiable: generated.decoded must be Tensor, " - f"got {type(decoded).__name__}." - ) - texts = req.primitives.get("text") if req.primitives else None - if not isinstance(texts, Texts): - raise TypeError("ReflRewardRole.score_differentiable: req.primitives['text'] must be Texts.") - - if not isinstance(self.backend, DifferentiableReward): - raise TypeError( - f"{type(self.backend).__name__} must implement compute_rewards_differentiable " - "for REFL reward backprop." - ) - - records = list(req.metadata) if req.metadata else None - kind = self.preferred_input_kind - if kind == "image" and decoded.ndim != 4: - raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") - if kind == "video" and decoded.ndim != 5: - raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") - if kind not in {"image", "video"}: - raise ValueError(f"ReflRewardRole.score_differentiable: unsupported input_kind={kind!r}.") - - return self.backend.compute_rewards_differentiable( - decoded, - list(texts.texts), - records=records, - ) - - -__all__ = ["ReflActorRole", "ReflRewardRole"] +__all__ = ["ReflActorRole"] From 1749f2be040b800c6e7a36f13195916b278f6ab6 Mon Sep 17 00:00:00 2001 From: YSunLIN Date: Mon, 27 Jul 2026 15:20:20 +0800 Subject: [PATCH 08/24] refactor: drop Role/RewardRole, move Trainer to recipes/common - ReflActorRole inherits Remote directly: inline initialize (build bundle/pipeline/backend) + step/save_checkpoint/load_checkpoint, and move RoleStepResult here; delete recipes/common/roles.py - reward role now uses unirl.reward.service.RewardService directly; Trainer.create_remote_role special-cases it so the worker builds the backend, other roles keep the cfg-driven path - RewardService.score_differentiable takes (media_tensor, prompts, records) and forwards records to the backend; update the recipe trainer and the SD3 refl callers to the new signature - move Trainer base from unirl/trainer/trainer.py to recipes/common/trainer.py - configs: reward _target_ -> RewardService; fix the placement comment (cross-process autograd is supported via GradContext RPC) --- recipes/common/roles.py | 176 ------------------ {unirl/trainer => recipes/common}/trainer.py | 14 +- .../configs/wan21_t2v_videoalign_refl.yaml | 7 +- recipes/refl/configs/wan22_i2v_face_refl.yaml | 2 +- recipes/refl/roles.py | 64 ++++++- recipes/refl/trainer.py | 6 +- unirl/reward/service.py | 30 +-- unirl/trainer/refl.py | 4 +- 8 files changed, 100 insertions(+), 203 deletions(-) delete mode 100644 recipes/common/roles.py rename {unirl/trainer => recipes/common}/trainer.py (96%) diff --git a/recipes/common/roles.py b/recipes/common/roles.py deleted file mode 100644 index a6752a282..000000000 --- a/recipes/common/roles.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Recipe-level Role abstractions for trainer-managed Remote roles.""" - -from __future__ import annotations - -import logging -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any, Optional - -import torch -from hydra.utils import instantiate - -from unirl.distributed.group.dispatch import Dispatch, distributed -from unirl.distributed.group.remote import Remote -from unirl.reward.base import DifferentiableReward -from unirl.reward.service import RewardService -from unirl.types.primitives import Texts -from unirl.types.rollout_req import RolloutReq - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class RoleStepResult: - """Generic result of one role-local optimizer step.""" - - metrics: Mapping[str, object] - grad_norm: float - lr: float - - -class Role(Remote): - """Recipe-level role that runs as a UniRL Remote inside a Worker. - - ``Role`` stores the role config and initializes common role-local - components (model/runtime bundle / pipeline / backend) as ordinary - Worker-local objects. Recipe roles may override ``initialize`` for - recipe-specific config, but should call ``super().initialize()`` first. - """ - - def __init__(self, cfg: Any) -> None: - super().__init__() - self.cfg = cfg - - def _cfg_get(self, key: str) -> Any: - if hasattr(self.cfg, "get"): - return self.cfg.get(key) - return getattr(self.cfg, key, None) - - def initialize(self) -> None: - """Initialize common role-local components after Worker setup.""" - model_cfg = self._cfg_get("model") - bundle_cfg = model_cfg if model_cfg is not None else self._cfg_get("bundle") - if bundle_cfg is not None: - self.bundle = instantiate(bundle_cfg) - - pipeline_cfg = self._cfg_get("pipeline") - if pipeline_cfg is not None: - if hasattr(self, "bundle"): - self.pipeline = instantiate(pipeline_cfg, bundle=self.bundle) - else: - self.pipeline = instantiate(pipeline_cfg) - - backend_cfg = self._cfg_get("backend") - if backend_cfg is not None: - if hasattr(self, "bundle"): - if self.device is None or self.rank_info is None: - raise RuntimeError( - f"{type(self).__name__}.initialize called before Remote.setup injected device/rank_info." - ) - self.backend = instantiate( - backend_cfg, - bundle=self.bundle, - device=torch.device(self.device), - rank=int(self.rank_info.rank), - ) - else: - self.backend = instantiate(backend_cfg) - - @distributed - def step(self) -> RoleStepResult: - """Clip gradients and run one backend optimizer step.""" - if not hasattr(self, "backend") or not hasattr(self.backend, "optimizer_step"): - raise RuntimeError(f"{type(self).__name__}.step requires a backend with optimizer_step(...).") - grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(self.algo_cfg.get("max_grad_norm", 1.0)))) - lr = 0.0 - try: - sched = getattr(self.backend, "scheduler", None) - if sched is not None: - last = sched.get_last_lr() - lr = float(last[0]) if last else 0.0 - except Exception: - lr = 0.0 - return RoleStepResult( - metrics={"grad_norm": grad_norm, "lr": lr}, - grad_norm=grad_norm, - lr=lr, - ) - - @distributed(dispatch_mode=Dispatch.BROADCAST) - def save_checkpoint(self, path: str, step: Optional[int] = None, mode: str = "auto") -> None: - """Save backend checkpoint when the role backend supports it.""" - if hasattr(self, "backend") and hasattr(self.backend, "save"): - self.backend.save(path, step=step, mode=mode) - - @distributed(dispatch_mode=Dispatch.BROADCAST) - def load_checkpoint(self, path: str) -> int: - """Load backend checkpoint when the role backend supports it.""" - if hasattr(self, "backend") and hasattr(self.backend, "load"): - return int(self.backend.load(path) or 0) - return 0 - - -class RewardRole(RewardService): - """Common differentiable reward role for recipe-local reward backends.""" - - def __init__(self, cfg: Any) -> None: - Remote.__init__(self) - self.cfg = cfg - - def _cfg_get(self, key: str, default: Any = None) -> Any: - if hasattr(self.cfg, "get"): - return self.cfg.get(key, default) - return getattr(self.cfg, key, default) - - def initialize(self) -> None: - if self.device is None: - raise RuntimeError(f"{type(self).__name__}.initialize called before Remote.setup injected device.") - - backend_cfg = self._cfg_get("backend") - if backend_cfg is None: - raise ValueError(f"{type(self).__name__} requires cfg.backend.") - - self.backend = instantiate(backend_cfg, base_device=self.device) - - logger.info( - "RewardRole initialized with backend=%s", - self.backend.get_model_name() or type(self.backend).__name__, - ) - - @distributed - def score_differentiable(self, *, req: RolloutReq, generated: Any) -> torch.Tensor: - """Score live-grad decoded media through a differentiable reward backend.""" - decoded = generated.decoded - if not isinstance(decoded, torch.Tensor): - raise TypeError( - f"RewardRole.score_differentiable: generated.decoded must be Tensor, " - f"got {type(decoded).__name__}." - ) - texts = req.primitives.get("text") if req.primitives else None - if not isinstance(texts, Texts): - raise TypeError("RewardRole.score_differentiable: req.primitives['text'] must be Texts.") - - if not isinstance(self.backend, DifferentiableReward): - raise TypeError( - f"{type(self.backend).__name__} must implement compute_rewards_differentiable " - "for reward backprop." - ) - - records = list(req.metadata) if req.metadata else None - kind = self.preferred_input_kind - if kind == "image" and decoded.ndim != 4: - raise ValueError(f"image backend expects [B,C,H,W], got {tuple(decoded.shape)}") - if kind == "video" and decoded.ndim != 5: - raise ValueError(f"video backend expects [B,C,T,H,W], got {tuple(decoded.shape)}") - if kind not in {"image", "video"}: - raise ValueError(f"RewardRole.score_differentiable: unsupported input_kind={kind!r}.") - - return self.backend.compute_rewards_differentiable( - decoded, - list(texts.texts), - records=records, - ) - - -__all__ = ["Role", "RoleStepResult", "RewardRole"] diff --git a/unirl/trainer/trainer.py b/recipes/common/trainer.py similarity index 96% rename from unirl/trainer/trainer.py rename to recipes/common/trainer.py index 26e32bcab..5364524cf 100644 --- a/unirl/trainer/trainer.py +++ b/recipes/common/trainer.py @@ -12,6 +12,7 @@ from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict from unirl.distributed.group.remote import Remote +from unirl.reward.service import RewardService from unirl.trainer.base import BaseTrainer, build_sampling_dict logger = logging.getLogger(__name__) @@ -238,12 +239,21 @@ def prepare_role_cfg(self, spec: RoleSpec) -> DictConfig: return OmegaConf.create(container) def create_remote_role(self, spec: RoleSpec, *, device_ids: List[int], slot_id: int): + role_cls = self.resolve_role_cls(spec.target) + if issubclass(role_cls, RewardService): + # RewardService takes a materialized ``backend`` (not a ``cfg`` blob): + # pass the role's own fields as plain-dict kwargs so the worker's + # ``_resolve_init_kwargs`` walker instantiates the nested ``_target_`` + # backend in its own CUDA context. + init_kwargs = OmegaConf.to_container(self.prepare_role_cfg(spec), resolve=True) + else: + init_kwargs = {"cfg": self.prepare_role_cfg(spec)} return self.pool.create_remote( - self.resolve_role_cls(spec.target), + role_cls, device_ids=device_ids, slot_id=slot_id, role_name=spec.name, - init_kwargs={"cfg": self.prepare_role_cfg(spec)}, + init_kwargs=init_kwargs, ) def initialize_roles(self) -> None: diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 93df78817..7ba583585 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -84,10 +84,11 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.common.roles.RewardRole + _target_: unirl.reward.service.RewardService placement: - # BPTT gradients must live on the same worker — no cross-process - # autograd. VideoAlign co-locates with the actor. + # Co-locate reward with the actor to avoid shipping decoded video + # across processes. NOT an autograd requirement — GradContext + # backprops rewards across workers via RPC — purely an efficiency choice. share_with: actor backend: diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index bccc76c46..8ca767858 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -81,7 +81,7 @@ roles: algorithm: ${algorithm} - name: reward - _target_: recipes.common.roles.RewardRole + _target_: unirl.reward.service.RewardService placement: share_with: actor diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index 929462e04..ee48355de 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -3,6 +3,7 @@ from __future__ import annotations import dataclasses +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, List, Optional @@ -10,8 +11,8 @@ from hydra.utils import instantiate from omegaconf import OmegaConf -from recipes.common.roles import Role -from unirl.distributed.group.dispatch import distributed +from unirl.distributed.group.dispatch import Dispatch, distributed +from unirl.distributed.group.remote import Remote from unirl.distributed.tensor.batch import Batch, concat_field, shared_field from unirl.sde.runtime import get_sigma_schedule from unirl.types.primitives import Images, Texts @@ -42,7 +43,16 @@ def _maybe_instantiate(value: Any) -> Any: return value -class ReflActorRole(Role): +@dataclass(frozen=True) +class RoleStepResult: + """Generic result of one role-local optimizer step.""" + + metrics: Mapping[str, object] + grad_norm: float + lr: float + + +class ReflActorRole(Remote): """Actor role: bundle + pipeline + backend + REFL BPTT logic.""" bundle: Any @@ -51,11 +61,55 @@ class ReflActorRole(Role): algo_cfg: Any sampling_params: Any + def __init__(self, cfg: Any) -> None: + super().__init__() + self.cfg = cfg + def initialize(self) -> None: - super().initialize() + self.bundle = instantiate(self.cfg.get("model")) + self.pipeline = instantiate(self.cfg.get("pipeline"), bundle=self.bundle) + self.backend = instantiate( + self.cfg.get("backend"), + bundle=self.bundle, + device=torch.device(self.device), + rank=int(self.rank_info.rank), + ) self.algo_cfg = self.cfg.algorithm self.sampling_params = _maybe_instantiate(self.algo_cfg.get("sampling_params")) + @distributed + def step(self) -> RoleStepResult: + """Clip gradients and run one backend optimizer step.""" + if not hasattr(self, "backend") or not hasattr(self.backend, "optimizer_step"): + raise RuntimeError(f"{type(self).__name__}.step requires a backend with optimizer_step(...).") + grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(self.algo_cfg.get("max_grad_norm", 1.0)))) + lr = 0.0 + try: + sched = getattr(self.backend, "scheduler", None) + if sched is not None: + last = sched.get_last_lr() + lr = float(last[0]) if last else 0.0 + except Exception: + lr = 0.0 + return RoleStepResult( + metrics={"grad_norm": grad_norm, "lr": lr}, + grad_norm=grad_norm, + lr=lr, + ) + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def save_checkpoint(self, path: str, step: Optional[int] = None, mode: str = "auto") -> None: + """Save backend checkpoint when the role backend supports it.""" + if hasattr(self, "backend") and hasattr(self.backend, "save"): + self.backend.save(path, step=step, mode=mode) + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def load_checkpoint(self, path: str) -> int: + """Load backend checkpoint when the role backend supports it.""" + if hasattr(self, "backend") and hasattr(self.backend, "load"): + return int(self.backend.load(path) or 0) + return 0 + @distributed def generate_samples(self, req: RolloutReq) -> REFLGenerated: """Run live-grad diffusion sampling and VAE decode.""" @@ -173,4 +227,4 @@ def forward_backward_loss( ) -__all__ = ["ReflActorRole"] +__all__ = ["ReflActorRole", "RoleStepResult"] diff --git a/recipes/refl/trainer.py b/recipes/refl/trainer.py index 3c6b1c3d2..1b15bb91f 100644 --- a/recipes/refl/trainer.py +++ b/recipes/refl/trainer.py @@ -7,8 +7,8 @@ from typing import Any, Dict import numpy as np +from recipes.common.trainer import Trainer from unirl.distributed.tensor.grad_context import enable_grad -from unirl.trainer.trainer import Trainer from unirl.types.prompts import RolloutInputs from unirl.types.rollout_req import RolloutReq from unirl.types.sampling import total_samples_per_prompt @@ -54,9 +54,11 @@ def build_req(self, inputs: RolloutInputs, rollout_id: int) -> RolloutReq: def train_step(self, req: RolloutReq, *, training_progress: float = 0.0, rollout_id: int = 0) -> Dict[str, Any]: """One REFL step: actor generate → reward score → actor backward → actor step.""" t0 = time.perf_counter() + prompts = list(req.primitives["text"].texts) + records = list(req.metadata) if req.metadata else None with enable_grad(): gen = self.actor.generate_samples(req) - rewards = self.reward.score_differentiable(req=req, generated=gen) + rewards = self.reward.score_differentiable(gen.decoded, prompts, records) loss_metrics = self.actor.forward_backward_loss( rewards=rewards, kl_loss=gen.kl_loss, diff --git a/unirl/reward/service.py b/unirl/reward/service.py index d8e39544f..94fe6836d 100644 --- a/unirl/reward/service.py +++ b/unirl/reward/service.py @@ -15,7 +15,6 @@ from unirl.distributed.group.dispatch import Dispatch, distributed from unirl.distributed.group.remote import Remote -from unirl.types.primitives import Images, Texts from unirl.types.reward import RewardRequest, RewardResponse from unirl.types.rollout_req import PrimitiveValue, RolloutReq from unirl.types.rollout_resp import RolloutTrack, _track_with_field @@ -167,16 +166,23 @@ def compute_rewards(self, request: RewardRequest) -> RewardResponse: return self.backend.compute_rewards(request) @distributed(dispatch_mode=Dispatch.DP_SCATTER) - def score_differentiable(self, *, images: Images, prompts: Texts) -> torch.Tensor: - """ReFL scoring: score grad-carrying ``images`` (pixels ``[B, C, H, W]`` in - ``[0, 1]``) against ``prompts`` and return a ``[B]`` reward tensor with - ``grad_fn`` intact. - - Deliberately bypasses :meth:`score_and_attach` / ``RewardRequest.images`` - (those go through ``tensor_frame_to_pil`` + ``torch.tensor(...)``, which - detach). Under ``enable_grad()`` the framework marks ``images.pixels`` as a - grad leaf and chains the returned reward's grad back to it. The backend must - satisfy the :class:`~unirl.reward.base.DifferentiableReward` Protocol. + def score_differentiable( + self, + media_tensor: torch.Tensor, + prompts: List[str], + records: Optional[List[dict]] = None, + ) -> torch.Tensor: + """ReFL scoring: score grad-carrying ``media_tensor`` (image ``[B, C, H, W]`` + or video ``[B, C, T, H, W]``) against ``prompts`` and return a ``[B]`` reward + tensor with ``grad_fn`` intact. + + Deliberately bypasses :meth:`score_and_attach` / ``RewardRequest`` media + conversion (those go through ``tensor_frame_to_pil`` + ``torch.tensor(...)``, + which detach). Under ``enable_grad()`` the framework marks ``media_tensor`` as + a grad leaf and chains the returned reward's grad back to it. ``records`` + carries optional per-sample metadata (e.g. ``ref_video_path`` for recipe-local + video rewards). The backend must satisfy the + :class:`~unirl.reward.base.DifferentiableReward` Protocol. """ if not isinstance(self.backend, DifferentiableReward): raise TypeError( @@ -184,7 +190,7 @@ def score_differentiable(self, *, images: Images, prompts: Texts) -> torch.Tenso f"{type(self.backend).__name__} is not a DifferentiableReward — ReFL " f"needs a differentiable in-process reward (e.g. pickscore/clip/hpsv2)." ) - return self.backend.compute_rewards_differentiable(images.pixels, list(prompts.texts)) + return self.backend.compute_rewards_differentiable(media_tensor, list(prompts), records=records) @distributed(dispatch_mode=Dispatch.DP_SCATTER) def score_and_attach(self, *, req: RolloutReq, track: RolloutTrack) -> RolloutTrack: diff --git a/unirl/trainer/refl.py b/unirl/trainer/refl.py index d27e8ccf2..d9c64c044 100644 --- a/unirl/trainer/refl.py +++ b/unirl/trainer/refl.py @@ -117,7 +117,7 @@ def train_step(self, prompts: Texts, *, rollout_id: int) -> Tuple[float, float, 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=images, prompts=prompts) + 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) @@ -186,7 +186,7 @@ def _eval_pass( 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=images, prompts=sub)).float() + 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} From e3c6b940dc4ba1e1929595f4aea1276080288a70 Mon Sep 17 00:00:00 2001 From: leviking98z-rgb Date: Mon, 27 Jul 2026 21:52:05 +0800 Subject: [PATCH 09/24] Merge main and resolve ReFL recipe compatibility --- recipes/common/trainer.py | 5 +- recipes/refl/models/wan21.py | 19 +- recipes/refl/models/wan22.py | 20 +- recipes/refl/rewards/face/face_tools.py | 72 +-- recipes/refl/rewards/face/requirements.txt | 4 +- recipes/refl/rewards/face/scorer.py | 16 +- .../rewards/videoalign/model/checkpoint.py | 14 +- .../refl/rewards/videoalign/model/configs.py | 8 +- .../refl/rewards/videoalign/model/factory.py | 8 +- .../videoalign/model/prompt_template.py | 23 +- .../rewards/videoalign/model/reward_model.py | 33 +- .../refl/rewards/videoalign/requirements.txt | 14 +- recipes/refl/rewards/videoalign/scorer.py | 29 +- recipes/refl/rewards/videoalign/wrapper.py | 37 +- recipes/refl/run.py | 0 recipes/refl/trainer.py | 4 +- unirl/models/wan21/vae.py | 2 +- unirl/models/wan21/wan_video_vae.py | 553 ++++++++++-------- unirl/train/optim.py | 2 +- 19 files changed, 443 insertions(+), 420 deletions(-) mode change 100644 => 100755 recipes/refl/run.py diff --git a/recipes/common/trainer.py b/recipes/common/trainer.py index 5364524cf..1a020bb69 100644 --- a/recipes/common/trainer.py +++ b/recipes/common/trainer.py @@ -327,10 +327,7 @@ def train( def log_metrics(self, metrics: Dict[str, Any], *, rollout_id: int, num_rollouts: int) -> None: trainer_name = self.__class__.__name__.removesuffix("Trainer") or self.__class__.__name__ - metric_text = " ".join( - f"{key}={self._format_metric_for_log(value)}" - for key, value in metrics.items() - ) + metric_text = " ".join(f"{key}={self._format_metric_for_log(value)}" for key, value in metrics.items()) if metric_text: logger.info("%s rollout %d/%d %s", trainer_name, rollout_id + 1, num_rollouts, metric_text) else: diff --git a/recipes/refl/models/wan21.py b/recipes/refl/models/wan21.py index ede2a477e..d5fdc6834 100644 --- a/recipes/refl/models/wan21.py +++ b/recipes/refl/models/wan21.py @@ -64,7 +64,9 @@ def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] branch_embeds = prompt_embeds else: neg = conditions.negative_text - branch_embeds = neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + branch_embeds = ( + neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + ) batch_size = int(sample.shape[0]) timestep = sigma * _WAN_TIMESTEP_SCALE @@ -179,9 +181,7 @@ def diffuse_with_grad( bypassed. """ if conditions.text is None or conditions.text.embeds is None: - raise ValueError( - "Wan21ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None" - ) + raise ValueError("Wan21ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None") prompt_embeds = conditions.text.embeds device = prompt_embeds.device batch_size = int(prompt_embeds.shape[0]) @@ -189,8 +189,7 @@ def diffuse_with_grad( schedule = schedule.to(device) if int(schedule.shape[0]) != T + 1: raise ValueError( - f"Wan21ReflDiffusionStage.diffuse_with_grad: " - f"schedule length {schedule.shape[0]} != T+1={T + 1}" + f"Wan21ReflDiffusionStage.diffuse_with_grad: schedule length {schedule.shape[0]} != T+1={T + 1}" ) self.strategy.init_schedule(schedule) @@ -244,8 +243,7 @@ def diffuse_with_grad( step = self.step if not isinstance(step, Wan21ReflDiffusionStep): raise TypeError( - f"Wan21ReflDiffusionStage.diffuse_with_grad requires Wan21ReflDiffusionStep, " - f"got {type(step).__name__}." + f"Wan21ReflDiffusionStage.diffuse_with_grad requires Wan21ReflDiffusionStep, got {type(step).__name__}." ) for i in range(T): @@ -305,7 +303,7 @@ def diffuse_with_grad( branch="cond", ) sigma_f32 = sigma.to(dtype=torch.float32) - kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).mean() kl_total = kl_total + kl_step kl_steps += 1 @@ -345,8 +343,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) old = self.diffusion assert isinstance(old, WAN21DiffusionStage), ( - f"Wan21ReflPipeline expects parent to build WAN21DiffusionStage, " - f"got {type(old).__name__}" + f"Wan21ReflPipeline expects parent to build WAN21DiffusionStage, got {type(old).__name__}" ) self.diffusion = Wan21ReflDiffusionStage( model=old.model, diff --git a/recipes/refl/models/wan22.py b/recipes/refl/models/wan22.py index 6ce837af3..1c671a8f8 100644 --- a/recipes/refl/models/wan22.py +++ b/recipes/refl/models/wan22.py @@ -48,7 +48,9 @@ def predict_noise( # pyright: ignore[reportIncompatibleMethodOverride] branch_embeds = prompt_embeds else: neg = conditions.negative_text - branch_embeds = neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + branch_embeds = ( + neg.embeds if neg is not None and neg.embeds is not None else torch.zeros_like(prompt_embeds) + ) if use_high_noise is None: use_high_noise, _ = self._select_for_sigma( @@ -148,9 +150,7 @@ def diffuse_with_grad( """ if conditions.text is None or conditions.text.embeds is None: - raise ValueError( - "Wan22ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None" - ) + raise ValueError("Wan22ReflDiffusionStage.diffuse_with_grad: conditions.text.embeds is None") prompt_embeds = conditions.text.embeds device = prompt_embeds.device batch_size = int(prompt_embeds.shape[0]) @@ -158,8 +158,7 @@ def diffuse_with_grad( schedule = schedule.to(device) if int(schedule.shape[0]) != T + 1: raise ValueError( - f"Wan22ReflDiffusionStage.diffuse_with_grad: " - f"schedule length {schedule.shape[0]} != T+1={T + 1}" + f"Wan22ReflDiffusionStage.diffuse_with_grad: schedule length {schedule.shape[0]} != T+1={T + 1}" ) self.strategy.init_schedule(schedule) @@ -191,7 +190,6 @@ def diffuse_with_grad( base_seed=int(params.seed), ) - # BPTT knobs. sk: Dict[str, Any] = dict(getattr(params, "sampler_kwargs", {}) or {}) mid_timestep = int(sk.get("mid_timestep", 0)) @@ -213,8 +211,7 @@ def diffuse_with_grad( step = self.step if not isinstance(step, Wan22ReflDiffusionStep): raise TypeError( - f"Wan22ReflDiffusionStage.diffuse_with_grad requires Wan22ReflDiffusionStep, " - f"got {type(step).__name__}." + f"Wan22ReflDiffusionStage.diffuse_with_grad requires Wan22ReflDiffusionStep, got {type(step).__name__}." ) dual = self.model.transformer @@ -286,7 +283,7 @@ def diffuse_with_grad( use_high_noise=use_high_noise, ) sigma_f32 = sigma.to(dtype=torch.float32) - kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32 ** 2)).mean() + kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).mean() kl_total = kl_total + kl_step kl_steps += 1 @@ -321,8 +318,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) old = self.diffusion assert isinstance(old, WAN22DiffusionStage), ( - f"Wan22ReflPipeline expects parent to build WAN22DiffusionStage, " - f"got {type(old).__name__}" + f"Wan22ReflPipeline expects parent to build WAN22DiffusionStage, got {type(old).__name__}" ) self.diffusion = Wan22ReflDiffusionStage( model=old.model, diff --git a/recipes/refl/rewards/face/face_tools.py b/recipes/refl/rewards/face/face_tools.py index 6a544c4a8..eb52d52d0 100644 --- a/recipes/refl/rewards/face/face_tools.py +++ b/recipes/refl/rewards/face/face_tools.py @@ -1,19 +1,20 @@ """A simple, flexible implementation of a face analysis tool.""" + import math import os import onnx import torch import torch.nn.functional as F +import torchvision.ops as ops from onnx2torch import convert from skimage import transform as trans from torchvision.transforms.functional import resize -import torchvision.ops as ops - arcface_dst = torch.tensor( - [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], - [41.5493, 92.3655], [70.7299, 92.2041]]).float() + [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]] +).float() + def distance2bbox(points, distance, max_shape=None): """Decode distance prediction to bounding box.""" @@ -28,6 +29,7 @@ def distance2bbox(points, distance, max_shape=None): y2 = y2.clamp(min=0, max=max_shape[0]) return torch.stack([x1, y1, x2, y2], axis=-1) + def distance2kps(points, distance, max_shape=None): """Decode distance prediction to keypoints.""" preds = [] @@ -44,7 +46,8 @@ def distance2kps(points, distance, max_shape=None): def face_transform(data, center, output_size, scale, rotation, device): def to_homogeneous(mat): - return torch.vstack([mat, torch.tensor([0., 0., 1.])]) + return torch.vstack([mat, torch.tensor([0.0, 0.0, 1.0])]) + scale_ratio = scale rot = float(rotation) * math.pi / 180.0 cx = center[0] * scale_ratio @@ -52,29 +55,15 @@ def to_homogeneous(mat): C, H, W = data.shape - t1 = to_homogeneous(torch.tensor([ - [scale_ratio, 0, 0], - [0, scale_ratio, 0] - ])).float() - t2 = to_homogeneous(torch.tensor([ - [1, 0, -cx], - [0, 1, -cy] - ])).float() + t1 = to_homogeneous(torch.tensor([[scale_ratio, 0, 0], [0, scale_ratio, 0]])).float() + t2 = to_homogeneous(torch.tensor([[1, 0, -cx], [0, 1, -cy]])).float() cos_theta = math.cos(rot) sin_theta = math.sin(rot) - t3 = to_homogeneous(torch.tensor([ - [cos_theta, -sin_theta, 0], - [sin_theta, cos_theta, 0] - ])).float() - t4 = to_homogeneous(torch.tensor([ - [1, 0, output_size / 2], - [0, 1, output_size / 2] - ])).float() + t3 = to_homogeneous(torch.tensor([[cos_theta, -sin_theta, 0], [sin_theta, cos_theta, 0]])).float() + t4 = to_homogeneous(torch.tensor([[1, 0, output_size / 2], [0, 1, output_size / 2]])).float() M_homogeneous = t4 @ t3 @ t2 @ t1 M = M_homogeneous[:2, :] - T = torch.tensor([[2 / W, 0, -1], - [0, 2 / H, -1], - [0, 0, 1]]) + T = torch.tensor([[2 / W, 0, -1], [0, 2 / H, -1], [0, 0, 1]]) theta = torch.inverse(T @ M_homogeneous @ torch.inverse(T)) theta = theta[:2, :].unsqueeze(0).to(device) grid = F.affine_grid(theta, data.unsqueeze(0).size(), align_corners=True) @@ -93,7 +82,7 @@ def trans_points2d(pts, M): return transformed -def estimate_norm(lmk, image_size=112, mode='arcface'): +def estimate_norm(lmk, image_size=112, mode="arcface"): assert lmk.shape == (5, 2) assert image_size % 112 == 0 or image_size % 128 == 0 if image_size % 112 == 0: @@ -110,7 +99,7 @@ def estimate_norm(lmk, image_size=112, mode='arcface'): return M -def norm_crop(img, landmark, image_size=112, mode='arcface'): +def norm_crop(img, landmark, image_size=112, mode="arcface"): """Align an image into ArcFace canonical 112x112 using a similarity transform. Differentiable in ``img`` — the only non-grad path is ``landmark``, which @@ -120,9 +109,7 @@ def norm_crop(img, landmark, image_size=112, mode='arcface'): M_homogeneous = estimate_norm(landmark, image_size, mode) C, H, W = img.shape img = img.unsqueeze(0) - T = torch.tensor([[2 / W, 0, -1], - [0, 2 / H, -1], - [0, 0, 1]]) + T = torch.tensor([[2 / W, 0, -1], [0, 2 / H, -1], [0, 0, 1]]) T_inv = torch.inverse(T) theta = torch.inverse(T @ M_homogeneous @ T_inv) theta = theta[:2, :].unsqueeze(0).to(img.device) @@ -140,17 +127,15 @@ def invert_affine_transform(matrix): c, d = L[..., 1, 0], L[..., 1, 1] det = a * d - b * c inv_det = 1.0 / det - inv_L = torch.stack([ - torch.stack([d * inv_det, -b * inv_det], dim=-1), - torch.stack([-c * inv_det, a * inv_det], dim=-1) - ], dim=-2) + inv_L = torch.stack( + [torch.stack([d * inv_det, -b * inv_det], dim=-1), torch.stack([-c * inv_det, a * inv_det], dim=-1)], dim=-2 + ) inv_T = -torch.matmul(inv_L, T) inv_matrix = torch.cat([inv_L, inv_T], dim=-1) return inv_matrix class Face(dict): - def __init__(self, d=None, **kwargs): if d is None: d = {} @@ -161,8 +146,7 @@ def __init__(self, d=None, **kwargs): def __setattr__(self, name, value): if isinstance(value, (list, tuple)): - value = [self.__class__(x) - if isinstance(x, dict) else x for x in value] + value = [self.__class__(x) if isinstance(x, dict) else x for x in value] elif isinstance(value, dict) and not isinstance(value, self.__class__): value = self.__class__(value) super(Face, self).__setattr__(name, value) @@ -219,14 +203,13 @@ def forward(self, det_img, threshold=0.5): height = input_height // stride width = input_width // stride - K = height * width key = (height, width, stride) if key in self.center_cache: anchor_centers = self.center_cache[key] else: rows = torch.arange(height) cols = torch.arange(width) - grid_y, grid_x = torch.meshgrid(rows, cols, indexing='ij') + grid_y, grid_x = torch.meshgrid(rows, cols, indexing="ij") anchor_centers = torch.stack([grid_x, grid_y], dim=-1).float() anchor_centers = (anchor_centers * stride).reshape((-1, 2)) if self._num_anchors > 1: @@ -253,7 +236,7 @@ def forward(self, det_img, threshold=0.5): return scores_list, bboxes_list, kpss_list @torch.no_grad() - def detect(self, image, input_size=None, max_num=0, metric='default', nms_thresh=0.4, det_thresh=0.5): + def detect(self, image, input_size=None, max_num=0, metric="default", nms_thresh=0.4, det_thresh=0.5): assert input_size is not None or self.input_size is not None input_size = self.input_size if input_size is None else input_size @@ -308,7 +291,7 @@ def __init__(self, model_file=None, device="cuda"): self.torch_model.to(self.device) # Frozen weights — gradient flows through inputs only. self.torch_model.requires_grad_(False) - self.taskname = 'recognition' + self.taskname = "recognition" self.input_size = (112, 112) def get(self, img, face, input_size=(112, 112)): @@ -348,7 +331,7 @@ def __init__(self, model_file=None, device="cuda"): self.torch_model.requires_grad_(False) self.lmk_dim = 2 self.lmk_num = 106 - self.taskname = 'landmark_%dd_%d' % (self.lmk_dim, self.lmk_num) + self.taskname = "landmark_%dd_%d" % (self.lmk_dim, self.lmk_num) self.input_size = (192, 192) def get(self, img, face, input_size=(192, 192)): @@ -358,7 +341,7 @@ def get(self, img, face, input_size=(192, 192)): rotate = 0 _scale = self.input_size[0] / (max(w, h) * 1.5) aimg, M = face_transform(img, center, self.input_size[0], _scale, rotate, img.device) - aimg = (aimg + 1) / 2 * 255. + aimg = (aimg + 1) / 2 * 255.0 aimg = aimg[:, [2, 1, 0], :, :] input_size = self.input_size if input_size is None else input_size @@ -370,7 +353,6 @@ def get(self, img, face, input_size=(192, 192)): else: new_width = input_size[0] new_height = int(new_width * im_ratio) - det_scale = float(new_height) / aimg.shape[2] resized_img = resize(aimg, (new_height, new_width), antialias=False) det_img = torch.zeros((aimg.shape[0], 3, input_size[1], input_size[0]), device=self.device) det_img[:, :, :new_height, :new_width] = resized_img @@ -378,9 +360,9 @@ def get(self, img, face, input_size=(192, 192)): pred = self.torch_model(det_img)[0] pred = pred.reshape((-1, 2)) if self.lmk_num < pred.shape[0]: - pred = pred[self.lmk_num * -1:, :] + pred = pred[self.lmk_num * -1 :, :] pred[:, 0:2] += 1 - pred[:, 0:2] *= (self.input_size[0] // 2) + pred[:, 0:2] *= self.input_size[0] // 2 IM = invert_affine_transform(M).to(img.device) pred = trans_points2d(pred, IM) diff --git a/recipes/refl/rewards/face/requirements.txt b/recipes/refl/rewards/face/requirements.txt index b5421c1be..329432f90 100644 --- a/recipes/refl/rewards/face/requirements.txt +++ b/recipes/refl/rewards/face/requirements.txt @@ -1,6 +1,6 @@ +imageio>=2.31 +imageio-ffmpeg>=0.4 onnx>=1.14 onnx2torch>=1.5 scikit-image>=0.21 -imageio>=2.31 -imageio-ffmpeg>=0.4 scipy>=1.11 diff --git a/recipes/refl/rewards/face/scorer.py b/recipes/refl/rewards/face/scorer.py index 79d5330d7..b9b04c672 100644 --- a/recipes/refl/rewards/face/scorer.py +++ b/recipes/refl/rewards/face/scorer.py @@ -22,6 +22,7 @@ # Reference-video loader (imageio + resize-to-cover + center_crop) # --------------------------------------------------------------------------- + def _load_ref_video_frames( video_path: str, *, @@ -66,8 +67,8 @@ def _load_ref_video_frames( interpolation=TF.InterpolationMode.BILINEAR, ) frame = TF.center_crop(frame, (target_h, target_w)) - t = TF.to_tensor(frame) # (C, H, W) [0, 1] - frame = TF.to_pil_image(t) # round-trip PIL to match training-time preprocessing + t = TF.to_tensor(frame) # (C, H, W) [0, 1] + frame = TF.to_pil_image(t) # round-trip PIL to match training-time preprocessing frames.append(TF.to_tensor(frame) * 2.0 - 1.0) reader.close() @@ -79,6 +80,7 @@ def _load_ref_video_frames( # Reward scorer # --------------------------------------------------------------------------- + class FaceRewardScorer(LocalRewardBackend): """SCRFD detection + ArcFace embedding + pool cosine similarity for REFL.""" @@ -107,9 +109,7 @@ def _load_model(self) -> None: self._ref_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} def _compute_model_rewards(self, request: RewardRequest) -> List[float]: - raise NotImplementedError( - "FaceRewardScorer is REFL-only; use compute_rewards_differentiable()." - ) + raise NotImplementedError("FaceRewardScorer is REFL-only; use compute_rewards_differentiable().") # ------------------------------------------------------------------ # Per-video face embedding (recipe-local REFL implementation) @@ -171,7 +171,7 @@ def _extract_face_embeddings( embeddings.append(zero_emb) mask.append(0) - emb_stack = torch.stack(embeddings).unsqueeze(0) # (1, T, 512) + emb_stack = torch.stack(embeddings).unsqueeze(0) # (1, T, 512) mask_tensor = torch.tensor(mask, device=self.device).unsqueeze(0).float() # (1, T) return emb_stack, mask_tensor @@ -228,9 +228,7 @@ def compute_rewards_differentiable( ref_emb, ref_mask = self._get_ref_embeddings(str(ref_path)) gen_video = torch.clamp(media_tensor[i], -1, 1).to(self.device) - gen_emb, gen_mask = self._extract_face_embeddings( - gen_video, with_grad=self._differentiable - ) + gen_emb, gen_mask = self._extract_face_embeddings(gen_video, with_grad=self._differentiable) if int(gen_mask.sum().item()) == 0: rewards.append(gen_video.sum() * 0.0) diff --git a/recipes/refl/rewards/videoalign/model/checkpoint.py b/recipes/refl/rewards/videoalign/model/checkpoint.py index c201690e5..0f9609860 100644 --- a/recipes/refl/rewards/videoalign/model/checkpoint.py +++ b/recipes/refl/rewards/videoalign/model/checkpoint.py @@ -81,7 +81,8 @@ def _pick_checkpoint_dir(checkpoint_dir: str, checkpoint_step: Optional[int]) -> chosen = candidates[0] logger.warning( "Requested VideoAlign checkpoint-%s not found; falling back to latest %s", - checkpoint_step, chosen, + checkpoint_step, + chosen, ) return chosen @@ -121,18 +122,17 @@ def load_model_from_checkpoint( # the target model's own state_dict keys and remap when needed. target_keys = model.state_dict().keys() needs_remap = any( - k.startswith("base_model.model.model.language_model.") - or k.startswith("base_model.model.model.visual.") + k.startswith("base_model.model.model.language_model.") or k.startswith("base_model.model.model.visual.") for k in target_keys ) if needs_remap: new_state_dict: Dict[str, torch.Tensor] = {} for key, value in model_state_dict.items(): if key.startswith("base_model.model.model"): - new_key = "base_model.model.model.language_model" + key[len("base_model.model.model"):] + new_key = "base_model.model.model.language_model" + key[len("base_model.model.model") :] new_state_dict[new_key] = value elif key.startswith("base_model.model.visual"): - new_key = "base_model.model.model.visual" + key[len("base_model.model.visual"):] + new_key = "base_model.model.model.visual" + key[len("base_model.model.visual") :] new_state_dict[new_key] = value else: new_state_dict[key] = value @@ -151,7 +151,9 @@ def load_model_from_checkpoint( non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu") lora_state_dict = _insert_adapter_name_into_state_dict( - lora_state_dict, adapter_name="default", parameter_prefix="lora_", + lora_state_dict, + adapter_name="default", + parameter_prefix="lora_", ) model_state_dict = model.state_dict() diff --git a/recipes/refl/rewards/videoalign/model/configs.py b/recipes/refl/rewards/videoalign/model/configs.py index c22750c85..e9845c437 100644 --- a/recipes/refl/rewards/videoalign/model/configs.py +++ b/recipes/refl/rewards/videoalign/model/configs.py @@ -141,15 +141,11 @@ class ModelConfig: # Bradley-Terry / regression flavour — train-only. Inference ignores # this entirely; kept for round-trip compatibility with # ``model_config.json``. - loss_type: Literal[ - "bt", "reg", "btt", "margin", "constant_margin", "scaled", "regular" - ] = "regular" + loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled", "regular"] = "regular" def __post_init__(self) -> None: if self.load_in_8bit and self.load_in_4bit: - raise ValueError( - "You can't use 8 bit and 4 bit precision at the same time" - ) + raise ValueError("You can't use 8 bit and 4 bit precision at the same time") __all__ = ["ModelConfig", "PEFTLoraConfig", "TrainingConfig"] diff --git a/recipes/refl/rewards/videoalign/model/factory.py b/recipes/refl/rewards/videoalign/model/factory.py index a2a6bcc98..d65fab275 100644 --- a/recipes/refl/rewards/videoalign/model/factory.py +++ b/recipes/refl/rewards/videoalign/model/factory.py @@ -93,9 +93,7 @@ def create_model_and_processor( special_token_ids = None if model_config.use_special_tokens: special_tokens = ["<|VQ_reward|>", "<|MQ_reward|>", "<|TA_reward|>"] - processor.tokenizer.add_special_tokens( - {"additional_special_tokens": special_tokens} - ) + processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens}) special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens) # Build the reward model. Quantisation is intentionally not supported @@ -114,9 +112,7 @@ def create_model_and_processor( reward_token=model_config.reward_token, special_token_ids=special_token_ids, torch_dtype=torch_dtype, - attn_implementation=( - "flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa" - ), + attn_implementation=("flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa"), cache_dir=cache_dir, revision=getattr(model_config, "model_revision", "main"), use_cache=True if training_args.gradient_checkpointing else False, diff --git a/recipes/refl/rewards/videoalign/model/prompt_template.py b/recipes/refl/rewards/videoalign/model/prompt_template.py index 20b8f9d37..7687af0da 100644 --- a/recipes/refl/rewards/videoalign/model/prompt_template.py +++ b/recipes/refl/rewards/videoalign/model/prompt_template.py @@ -36,8 +36,7 @@ ], "Overall": [ "Overall Performance", - "the overall performance of the video in terms of visual quality, " - "text-to-video alignment, and motion quality", + "the overall performance of the video in terms of visual quality, text-to-video alignment, and motion quality", ], } @@ -49,18 +48,18 @@ DETAILED_PROMPT_WITH_SPECIAL_TOKEN = """ You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. -**Visual Quality:** +**Visual Quality:** Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: - **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. - **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. - **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). - **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. -- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. Please provide the ratings of Visual Quality: <|VQ_reward|> END -**Motion Quality:** +**Motion Quality:** Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: - **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). - **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). @@ -72,11 +71,11 @@ Please provide the ratings of Motion Quality: <|MQ_reward|> END -**Text Alignment:** +**Text Alignment:** Assess how well the video matches the textual prompt across the following sub-dimensions: - **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. - **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. -- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. - **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. - **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. @@ -88,15 +87,15 @@ DETAILED_PROMPT = """ You are tasked with evaluating a generated video based on three distinct criteria: Visual Quality, Motion Quality, and Text Alignment. Please provide a rating from 0 to 10 for each of the three categories, with 0 being the worst and 10 being the best. Each evaluation should be independent of the others. -**Visual Quality:** +**Visual Quality:** Evaluate the overall visual quality of the video, with a focus on static factors. The following sub-dimensions should be considered: - **Reasonableness:** The video should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups. - **Clarity:** Evaluate the sharpness and visibility of the video. The image should be clear and easy to interpret, with no blurring or indistinct areas. - **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows). - **Aesthetic and Creativity:** Assess the artistic aspects of the video, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance. -- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. +- **Safety:** The video should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. -**Motion Quality:** +**Motion Quality:** Assess the dynamic aspects of the video, with a focus on dynamic factors. Consider the following sub-dimensions: - **Stability:** Evaluate the continuity and stability between frames. There should be no sudden, unnatural jumps, and the video should maintain stable attributes (e.g., no fluctuating colors, textures, or missing body parts). - **Naturalness:** The movement should align with physical laws and be realistic. For example, clothing should flow naturally with motion, and facial expressions should change appropriately (e.g., blinking, mouth movements). @@ -106,11 +105,11 @@ - **Amplitude:** If the video is largely static or has little movement, assign a low score for motion quality. -**Text Alignment:** +**Text Alignment:** Assess how well the video matches the textual prompt across the following sub-dimensions: - **Subject Relevance** Evaluate how accurately the subject(s) in the video (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior. - **Motion Relevance:** Evaluate if the dynamic actions (e.g., gestures, posture, facial expressions like talking or blinking) align with the described prompt. The motion should match the prompt in terms of type, scale, and direction. -- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. +- **Environment Relevance:** Assess whether the background and scene fit the prompt. This includes checking if real-world locations or scenes are accurately represented, though some stylistic adaptation is acceptable. - **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the video adheres to this style. - **Camera Movement Relevance:** Check if the camera movements (e.g., following the subject, focus shifts) are consistent with the expected behavior from the prompt. diff --git a/recipes/refl/rewards/videoalign/model/reward_model.py b/recipes/refl/rewards/videoalign/model/reward_model.py index 4276b671b..d0e4765f0 100644 --- a/recipes/refl/rewards/videoalign/model/reward_model.py +++ b/recipes/refl/rewards/videoalign/model/reward_model.py @@ -45,10 +45,7 @@ def _cfg_get(config: Any, name: str) -> Any: # that tolerate ``None`` (e.g. ``pad_token_id``) keep working. if hasattr(config, name): return getattr(config, name) - raise AttributeError( - f"{type(config).__name__} has no attribute {name!r} " - f"(checked top-level and .text_config)." - ) + raise AttributeError(f"{type(config).__name__} has no attribute {name!r} (checked top-level and .text_config).") class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): @@ -152,13 +149,9 @@ def forward( video_grid_thw: Optional[torch.LongTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, # noqa: ARG002 — kept for API parity ): - output_attentions = ( - output_attentions if output_attentions is not None - else self.config.output_attentions - ) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( - output_hidden_states if output_hidden_states is not None - else self.config.output_hidden_states + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict @@ -194,9 +187,7 @@ def _as_tensor(visual_out): inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.get_dtype()) - image_embeds = _as_tensor( - self.visual(pixel_values, grid_thw=image_grid_thw) - ) + image_embeds = _as_tensor(self.visual(pixel_values, grid_thw=image_grid_thw)) image_token_id = _cfg_get(self.config, "image_token_id") image_mask = (input_ids == image_token_id).unsqueeze(-1).expand_as(inputs_embeds) image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) @@ -204,9 +195,7 @@ def _as_tensor(visual_out): if pixel_values_videos is not None: pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) - video_embeds = _as_tensor( - self.visual(pixel_values_videos, grid_thw=video_grid_thw) - ) + video_embeds = _as_tensor(self.visual(pixel_values_videos, grid_thw=video_grid_thw)) video_token_id = _cfg_get(self.config, "video_token_id") video_mask = (input_ids == video_token_id).unsqueeze(-1).expand_as(inputs_embeds) video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) @@ -227,8 +216,8 @@ def _as_tensor(visual_out): return_dict=return_dict, ) - hidden_states = outputs[0] # (B, L, D) - logits = self.rm_head(hidden_states) # (B, L, output_dim) + hidden_states = outputs[0] # (B, L, D) + logits = self.rm_head(hidden_states) # (B, L, output_dim) if input_ids is not None: batch_size = input_ids.shape[0] @@ -241,9 +230,7 @@ def _as_tensor(visual_out): except AttributeError: pad_token_id = None if pad_token_id is None and batch_size != 1: - raise ValueError( - "Cannot handle batch sizes > 1 if no padding token is defined." - ) + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") if pad_token_id is None: sequence_lengths = -1 else: @@ -260,9 +247,7 @@ def _as_tensor(visual_out): pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] elif self.reward_token == "mean": valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1) - pooled_logits = torch.stack( - [logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)] - ) + pooled_logits = torch.stack([logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) elif self.reward_token == "special": special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool) for special_token_id in self.special_token_ids: diff --git a/recipes/refl/rewards/videoalign/requirements.txt b/recipes/refl/rewards/videoalign/requirements.txt index c66c054a7..458920c6b 100644 --- a/recipes/refl/rewards/videoalign/requirements.txt +++ b/recipes/refl/rewards/videoalign/requirements.txt @@ -1,11 +1,11 @@ -transformers==4.45.2 -peft==0.10.0 -trl==0.8.6 accelerate==0.34.0 decord==0.6.0 -opencv-python-headless -safetensors -huggingface_hub einops flash-attn==2.5.8 -setuptools<70 \ No newline at end of file +huggingface_hub +opencv-python-headless +peft==0.10.0 +safetensors +setuptools<70 +transformers==4.45.2 +trl==0.8.6 diff --git a/recipes/refl/rewards/videoalign/scorer.py b/recipes/refl/rewards/videoalign/scorer.py index 39f5394b4..5c01adf30 100644 --- a/recipes/refl/rewards/videoalign/scorer.py +++ b/recipes/refl/rewards/videoalign/scorer.py @@ -102,9 +102,7 @@ def _load_model(self) -> None: ) def _compute_model_rewards(self, request: RewardRequest) -> List[float]: - raise NotImplementedError( - "VideoAlignRewardScorer is REFL-only; use compute_rewards_differentiable()." - ) + raise NotImplementedError("VideoAlignRewardScorer is REFL-only; use compute_rewards_differentiable().") # ------------------------------------------------------------------ # Differentiable REFL reward entry point @@ -121,8 +119,7 @@ def compute_rewards_differentiable( raise ValueError(f"VideoAlignRewardScorer expects [B,C,T,H,W], got {tuple(media_tensor.shape)}") if len(prompts) != int(media_tensor.shape[0]): raise ValueError( - f"VideoAlignRewardScorer: prompts length {len(prompts)} " - f"!= batch size {int(media_tensor.shape[0])}." + f"VideoAlignRewardScorer: prompts length {len(prompts)} != batch size {int(media_tensor.shape[0])}." ) # The wrapper expects per-sample [T, C, H, W] in [-1, 1]. @@ -134,12 +131,7 @@ def compute_rewards_differentiable( # with the mmrl baseline. per_sample_videos: List[torch.Tensor] = [] for v in media_tensor: - v = ( - v.to(self.device) - .permute(1, 0, 2, 3) - .clamp(-1.0, 1.0) - .contiguous() - ) # → (T, C, H, W) + v = v.to(self.device).permute(1, 0, 2, 3).clamp(-1.0, 1.0).contiguous() # → (T, C, H, W) per_sample_videos.append(v) if self._reward_num_frames > 0: @@ -147,7 +139,9 @@ def compute_rewards_differentiable( for v in per_sample_videos: if v.shape[0] > self._reward_num_frames: idx = torch.linspace( - 0, v.shape[0] - 1, self._reward_num_frames, + 0, + v.shape[0] - 1, + self._reward_num_frames, device=v.device, ).long() v = v[idx] @@ -157,14 +151,12 @@ def compute_rewards_differentiable( autograd_ctx = torch.enable_grad if self._differentiable else torch.no_grad with autograd_ctx(): scores = self.model.forward_scores( - per_sample_videos, prompts, use_norm=self._use_norm, + per_sample_videos, + prompts, + use_norm=self._use_norm, ) - reward = ( - self._w_vq * scores["VQ"] - + self._w_mq * scores["MQ"] - + self._w_ta * scores["TA"] - ) + reward = self._w_vq * scores["VQ"] + self._w_mq * scores["MQ"] + self._w_ta * scores["TA"] return reward.float() # ------------------------------------------------------------------ @@ -191,6 +183,7 @@ def dispose(self) -> None: # Spec # --------------------------------------------------------------------------- + @dataclass class VideoAlignSpec(BaseRewardComponentSpec): """Typed config for :class:`VideoAlignRewardScorer`. diff --git a/recipes/refl/rewards/videoalign/wrapper.py b/recipes/refl/rewards/videoalign/wrapper.py index 25e5713ee..fe7bc2f05 100644 --- a/recipes/refl/rewards/videoalign/wrapper.py +++ b/recipes/refl/rewards/videoalign/wrapper.py @@ -123,8 +123,9 @@ def __init__( "by the upstream VideoAlign trainer)." ) - data_config_dict, model_config_dict, peft_lora_config_dict, inference_config = \ - _load_configs_from_json(config_path) + data_config_dict, model_config_dict, peft_lora_config_dict, inference_config = _load_configs_from_json( + config_path + ) # We only need two fields out of the data_config block — the # template type and the eval-dim list. Stash them directly without @@ -160,9 +161,7 @@ def __init__( # check. The fast variant operates on torch tensors end-to-end via # torchvision.transforms.v2.functional.resize, so gradients flow # from pixels through the processor into the vision encoder. - fast_ip = AutoImageProcessor.from_pretrained( - model_config.model_name_or_path, use_fast=True - ) + fast_ip = AutoImageProcessor.from_pretrained(model_config.model_name_or_path, use_fast=True) processor.image_processor = fast_ip model.to(self.device) @@ -174,10 +173,14 @@ def __init__( self.data_config = data_config_dict # raw dict, kept for debugging logger.info( - "VideoRewardWrapper loaded: ckpt=%s device=%s dtype=%s " - "resize=%dx%d micro_bs=%d use_norm=%s template=%s", - checkpoint_dir, self.device, self.dtype, self.resize_height, - self.resize_width, self.micro_batch_size, self.use_norm, + "VideoRewardWrapper loaded: ckpt=%s device=%s dtype=%s resize=%dx%d micro_bs=%d use_norm=%s template=%s", + checkpoint_dir, + self.device, + self.dtype, + self.resize_height, + self.resize_width, + self.micro_batch_size, + self.use_norm, self.prompt_template_type, ) @@ -249,9 +252,7 @@ def prepare_batch_from_frames( processed: List[torch.Tensor] = [] for video in video_tensors: if video.dim() != 4: - raise ValueError( - f"Expected video tensor shape [T,C,H,W], got {tuple(video.shape)}" - ) + raise ValueError(f"Expected video tensor shape [T,C,H,W], got {tuple(video.shape)}") # Auto-transpose if caller passed [C,T,H,W] with T > 3. if video.shape[0] == 3 and video.shape[1] > 3: video = video.permute(1, 0, 2, 3) @@ -261,7 +262,9 @@ def prepare_batch_from_frames( batch = self.processor( text=self.processor.apply_chat_template( - chat_data, tokenize=False, add_generation_prompt=True, + chat_data, + tokenize=False, + add_generation_prompt=True, ), images=None, videos=processed, @@ -303,17 +306,13 @@ def forward_scores( Gradients flow back into ``video_tensors`` when they require grad. """ if len(video_tensors) != len(prompts): - raise ValueError( - "video_tensors and prompts must have the same batch size." - ) + raise ValueError("video_tensors and prompts must have the same batch size.") use_norm = self.use_norm if use_norm is None else use_norm all_vq, all_mq, all_ta = [], [], [] for start in range(0, len(video_tensors), self.micro_batch_size): end = start + self.micro_batch_size - batch = self.prepare_batch_from_frames( - video_tensors[start:end], prompts[start:end] - ) + batch = self.prepare_batch_from_frames(video_tensors[start:end], prompts[start:end]) # Aligned with mmrl's wrapper: pass the processor batch # straight through to the model. transformers 4.54 (current # pinned version) does not inject ``mm_token_type_ids``, so diff --git a/recipes/refl/run.py b/recipes/refl/run.py old mode 100644 new mode 100755 diff --git a/recipes/refl/trainer.py b/recipes/refl/trainer.py index 1b15bb91f..50b26fe2e 100644 --- a/recipes/refl/trainer.py +++ b/recipes/refl/trainer.py @@ -5,6 +5,7 @@ import time from collections.abc import Mapping from typing import Any, Dict + import numpy as np from recipes.common.trainer import Trainer @@ -44,7 +45,7 @@ def build_req(self, inputs: RolloutInputs, rollout_id: int) -> RolloutReq: group_ids=list(inputs.group_ids), primitives=primitives, request_conditions={}, - stage_config={}, + task_config={}, sampling_params=dict(self.sampling_params), metadata=list(inputs.metadata) if inputs.metadata else [], init_noise_group_ids=[], @@ -77,4 +78,5 @@ def train_step(self, req: RolloutReq, *, training_progress: float = 0.0, rollout } return metrics + __all__ = ["REFLTrainer"] diff --git a/unirl/models/wan21/vae.py b/unirl/models/wan21/vae.py index 749fb4455..a8b8b7f45 100644 --- a/unirl/models/wan21/vae.py +++ b/unirl/models/wan21/vae.py @@ -76,7 +76,7 @@ def decode(self, s: LatentSegment, *, grad: bool = False, activation_checkpoint: f"[B, C, T_lat, H_lat, W_lat], got {tuple(clean.shape)}" ) - with torch.no_grad(): + with nullcontext() if grad else torch.no_grad(): decoded = self._vae_decode(clean) # Decoded layout is [B, C, T_dec, H_dec, W_dec] in [-1, 1]. diff --git a/unirl/models/wan21/wan_video_vae.py b/unirl/models/wan21/wan_video_vae.py index 29a570b71..22b246266 100644 --- a/unirl/models/wan21/wan_video_vae.py +++ b/unirl/models/wan21/wan_video_vae.py @@ -18,6 +18,7 @@ # ── Memory-efficient Conv layers ── + class Conv3dActGradOnlyFunction(torch.autograd.Function): """Conv3d that only computes input gradient, not weight gradient. @@ -37,16 +38,25 @@ def forward(ctx, input, weight, bias, stride, padding, dilation, groups): @staticmethod def backward(ctx, grad_output): - weight, = ctx.saved_tensors + (weight,) = ctx.saved_tensors grad_input = None if ctx.needs_input_grad[0]: input = torch.empty_strided( - ctx.input_info[0], ctx.input_info[1], - dtype=ctx.input_info[2], device=grad_output.device) + ctx.input_info[0], ctx.input_info[1], dtype=ctx.input_info[2], device=grad_output.device + ) grad_input = torch.ops.aten.convolution_backward( - grad_output, input, weight, None, - _triple(ctx.stride), _triple(ctx.padding), _triple(ctx.dilation), - False, [0, 0, 0], ctx.groups, (True, False, False))[0] + grad_output, + input, + weight, + None, + _triple(ctx.stride), + _triple(ctx.padding), + _triple(ctx.dilation), + False, + [0, 0, 0], + ctx.groups, + (True, False, False), + )[0] return grad_input, None, None, None, None, None, None @@ -65,13 +75,14 @@ def forward(ctx, input, weight, bias, stride, padding, dilation, groups): @staticmethod def backward(ctx, grad_output): - weight, = ctx.saved_tensors + (weight,) = ctx.saved_tensors grad_input = None if ctx.needs_input_grad[0]: from torch.nn.grad import conv2d_input + grad_input = conv2d_input( - ctx.input_shape, weight, grad_output, - ctx.stride, ctx.padding, ctx.dilation, ctx.groups) + ctx.input_shape, weight, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups + ) return grad_input, None, None, None, None, None, None @@ -84,8 +95,8 @@ def __init__(self, *args, **kwargs): def forward(self, x): return Conv2dActGradOnlyFunction.apply( - x, self.weight, self.bias, - self.stride, self.padding, self.dilation, self.groups) + x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) class CausalConv3d(nn.Conv3d): @@ -93,8 +104,7 @@ class CausalConv3d(nn.Conv3d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._padding = (self.padding[2], self.padding[2], self.padding[1], - self.padding[1], 2 * self.padding[0], 0) + self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) self.time_kernel_size = self.kernel_size[0] @@ -124,12 +134,13 @@ def __init__(self, *args, **kwargs): def forward(self, x, cache_x=None): x = self.forward_pad(x, cache_x) return Conv3dActGradOnlyFunction.apply( - x, self.weight, self.bias, - self.stride, self.padding, self.dilation, self.groups) + x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups + ) # ── Building blocks ── + def check_is_instance(model, module_class): if isinstance(model, module_class): return True @@ -144,14 +155,12 @@ def __init__(self, dim, channel_first=True, images=True, bias=False): broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim,) self.channel_first = channel_first - self.scale = dim ** 0.5 + self.scale = dim**0.5 self.gamma = nn.Parameter(torch.ones(shape)) - self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): - return F.normalize( - x, dim=(1 if self.channel_first else -1) - ) * self.scale * self.gamma + self.bias + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias class Upsample(nn.Upsample): @@ -161,44 +170,38 @@ def forward(self, x): class Resample(nn.Module): def __init__(self, dim, mode): - assert mode in ('none', 'upsample2d', 'upsample3d', - 'downsample2d', 'downsample3d') + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") super().__init__() self.dim = dim self.mode = mode - if mode == 'upsample2d': + if mode == "upsample2d": self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) - elif mode == 'upsample3d': + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) + elif mode == "upsample3d": self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) + ) causal_conv3d = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - setattr(causal_conv3d, 'decoder', True) + if hasattr(self, "decoder") and getattr(self, "decoder"): + setattr(causal_conv3d, "decoder", True) self.time_conv = causal_conv3d - elif mode == 'downsample2d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - elif mode == 'downsample3d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - causal_conv3d = CausalConv3d(dim, dim, (3, 1, 1), - stride=(2, 1, 1), padding=(0, 0, 0)) - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - setattr(causal_conv3d, 'decoder', True) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + causal_conv3d = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + if hasattr(self, "decoder") and getattr(self, "decoder"): + setattr(causal_conv3d, "decoder", True) self.time_conv = causal_conv3d else: self.resample = nn.Identity() def forward(self, x, feat_cache=None, feat_idx=None): b, c, t, h, w = x.size() - if self.mode == 'upsample3d': - if hasattr(self, 'decoder') and getattr(self, 'decoder'): + if self.mode == "upsample3d": + if hasattr(self, "decoder") and getattr(self, "decoder"): x1 = x[:, :, 1:] x1 = self.time_conv(x1) x1 = x1.reshape(b, 2, c, t - 1, h, w) @@ -208,17 +211,17 @@ def forward(self, x, feat_cache=None, feat_idx=None): elif feat_cache is not None: idx = feat_idx[0] if feat_cache[idx] is None: - feat_cache[idx] = 'Rep' + feat_cache[idx] = "Rep" feat_idx[0] += 1 else: cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != 'Rep': - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x], dim=2) - if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == 'Rep': + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) - if feat_cache[idx] == 'Rep': + if feat_cache[idx] == "Rep": x = self.time_conv(x) else: x = self.time_conv(x, feat_cache[idx]) @@ -229,12 +232,12 @@ def forward(self, x, feat_cache=None, feat_idx=None): x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] - x = rearrange(x, 'b c t h w -> (b t) c h w') + x = rearrange(x, "b c t h w -> (b t) c h w") x = self.resample(x) - x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) - if self.mode == 'downsample3d': - if hasattr(self, 'decoder') and getattr(self, 'decoder'): + if self.mode == "downsample3d": + if hasattr(self, "decoder") and getattr(self, "decoder"): x = self.time_conv(x) elif feat_cache is not None: idx = feat_idx[0] @@ -243,8 +246,7 @@ def forward(self, x, feat_cache=None, feat_idx=None): feat_idx[0] += 1 else: cache_x = x[:, :, -1:, :, :].clone() - x = self.time_conv( - torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -257,11 +259,9 @@ def patchify(x, patch_size): if patch_size == 1: return x if x.dim() == 5: - return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", - q=patch_size, r=patch_size) + return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) if x.dim() == 4: - return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", - q=patch_size, r=patch_size) + return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) raise ValueError(f"Invalid input shape: {x.shape}") @@ -269,11 +269,9 @@ def unpatchify(x, patch_size): if patch_size == 1: return x if x.dim() == 5: - return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", - q=patch_size, r=patch_size) + return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) if x.dim() == 4: - return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", - q=patch_size, r=patch_size) + return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) return x @@ -287,24 +285,28 @@ def __init__(self, in_dim, out_dim, dropout=0.0): causal_conv3d_block2 = CausalConv3d(out_dim, out_dim, 3, padding=1) causal_conv3d_shortcut = CausalConv3d(in_dim, out_dim, 1) self.residual = nn.Sequential( - RMS_norm(in_dim, images=False), nn.SiLU(), + RMS_norm(in_dim, images=False), + nn.SiLU(), causal_conv3d_block1, - RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), - causal_conv3d_block2) + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + causal_conv3d_block2, + ) self.shortcut = causal_conv3d_shortcut if in_dim != out_dim else nn.Identity() def forward(self, x, feat_cache=None, feat_idx=None): h = self.shortcut(x) for layer in self.residual: - if hasattr(self, 'decoder') and getattr(self, 'decoder'): + if hasattr(self, "decoder") and getattr(self, "decoder"): x = layer(x) elif check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x], dim=2) + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -330,23 +332,30 @@ def __init__(self, dim): def forward(self, x): identity = x b, c, t, h, w = x.size() - x = rearrange(x, 'b c t h w -> (b t) c h w') + x = rearrange(x, "b c t h w -> (b t) c h w") x = self.norm(x) - q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( - 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) x = F.scaled_dot_product_attention(q, k, v) x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) x = self.proj(x) - x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) return x + identity # ── Encoder / Decoder ── + class Encoder3d(nn.Module): - def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], - num_res_blocks=2, attn_scales=[], - temperal_downsample=[True, True, False], dropout=0.0): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): super().__init__() dims = [dim * u for u in [1] + dim_mult] scale = 1.0 @@ -361,28 +370,25 @@ def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], downsamples.append(AttentionBlock(out_dim)) in_dim = out_dim if i != len(dim_mult) - 1: - mode = 'downsample3d' if temperal_downsample[i] else 'downsample2d' + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" downsamples.append(Resample(out_dim, mode=mode)) scale /= 2.0 self.downsamples = nn.Sequential(*downsamples) self.middle = nn.Sequential( - ResidualBlock(out_dim, out_dim, dropout), - AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout)) + ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), ResidualBlock(out_dim, out_dim, dropout) + ) self.head = nn.Sequential( - RMS_norm(out_dim, images=False), nn.SiLU(), - CausalConv3d(out_dim, z_dim, 3, padding=1)) + RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, padding=1) + ) def forward(self, x, feat_cache=None, feat_idx=None): if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x], dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -406,9 +412,9 @@ def forward(self, x, feat_cache=None, feat_idx=None): idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x], dim=2) + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -418,10 +424,17 @@ def forward(self, x, feat_cache=None, feat_idx=None): class Decoder3d(nn.Module): - def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], - num_res_blocks=2, attn_scales=[], - temperal_upsample=[False, True, True], dropout=0.0, - use_nested_grad_checkpoint=True): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + use_nested_grad_checkpoint=True, + ): super().__init__() self.use_nested_grad_checkpoint = use_nested_grad_checkpoint @@ -435,7 +448,7 @@ def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], attn_block_2 = AttentionBlock(dims[0]) res_block_3 = ResidualBlock(dims[0], dims[0], dropout) for m in (res_block_1, attn_block_2, res_block_3): - setattr(m, 'decoder', True) + setattr(m, "decoder", True) self.middle = nn.Sequential(res_block_1, attn_block_2, res_block_3) upsamples = [] @@ -444,23 +457,23 @@ def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], in_dim = in_dim // 2 for _ in range(num_res_blocks + 1): res_block = ResidualBlock(in_dim, out_dim, dropout) - setattr(res_block, 'decoder', True) + setattr(res_block, "decoder", True) upsamples.append(res_block) if scale in attn_scales: attn_block = AttentionBlock(out_dim) - setattr(attn_block, 'decoder', True) + setattr(attn_block, "decoder", True) upsamples.append(attn_block) in_dim = out_dim if i != len(dim_mult) - 1: - mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + mode = "upsample3d" if temperal_upsample[i] else "upsample2d" resample = Resample(out_dim, mode=mode) - setattr(resample, 'decoder', True) + setattr(resample, "decoder", True) upsamples.append(resample) scale *= 2.0 self.upsamples = nn.Sequential(*upsamples) causal_conv3d = CausalConv3d(out_dim, 3, 3, padding=1) - setattr(causal_conv3d, 'decoder', True) + setattr(causal_conv3d, "decoder", True) self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), causal_conv3d) self.unsample_splits = [ @@ -473,17 +486,17 @@ def custom_forward1(x): x = self.conv1(x) for layer in self.middle: x = layer(x) - for layer in self.upsamples[:self.unsample_splits[0]]: + for layer in self.upsamples[: self.unsample_splits[0]]: x = layer(x) return x def custom_forward2(x): - for layer in self.upsamples[self.unsample_splits[0]:self.unsample_splits[1]]: + for layer in self.upsamples[self.unsample_splits[0] : self.unsample_splits[1]]: x = layer(x) return x def custom_forward3(x): - for layer in self.upsamples[self.unsample_splits[1]:]: + for layer in self.upsamples[self.unsample_splits[1] :]: x = layer(x) for layer in self.head: x = layer(x) @@ -506,11 +519,19 @@ def count_conv3d(model): # ── VideoVAE_ (core model) ── + class VideoVAE_(nn.Module): - def __init__(self, dim=96, z_dim=16, dim_mult=[1, 2, 4, 4], - num_res_blocks=2, attn_scales=[], - temperal_downsample=[False, True, True], dropout=0.0, - use_nested_grad_checkpoint=True): + def __init__( + self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + use_nested_grad_checkpoint=True, + ): super().__init__() self.dim = dim self.z_dim = z_dim @@ -520,13 +541,21 @@ def __init__(self, dim=96, z_dim=16, dim_mult=[1, 2, 4, 4], self.temperal_upsample = temperal_downsample[::-1] self.use_nested_grad_checkpoint = use_nested_grad_checkpoint - self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout) + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) self.conv2 = CausalConv3d(z_dim, z_dim, 1) - self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, - attn_scales, self.temperal_upsample, dropout, - use_nested_grad_checkpoint) + self.decoder = Decoder3d( + dim, + z_dim, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_upsample, + dropout, + use_nested_grad_checkpoint, + ) def encode(self, x, scale): self.clear_cache() @@ -536,19 +565,20 @@ def encode(self, x, scale): self._enc_conv_idx = [0] if i == 0: out, self._enc_feat_map, self._enc_conv_idx = self.encoder( - x[:, :, :1, :, :], feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) else: out_, self._enc_feat_map, self._enc_conv_idx = self.encoder( - x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) out = torch.cat([out, out_], 2) mu, log_var = self.conv1(out).chunk(2, dim=1) if isinstance(scale[0], torch.Tensor): scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] - mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( - 1, self.z_dim, 1, 1, 1) + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) else: scale = scale.to(dtype=mu.dtype, device=mu.device) mu = (mu - scale[0]) * scale[1] @@ -557,8 +587,7 @@ def encode(self, x, scale): def decode(self, z, scale): if isinstance(scale[0], torch.Tensor): scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] - z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( - 1, self.z_dim, 1, 1, 1) + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) else: scale = scale.to(dtype=z.dtype, device=z.device) z = z / scale[1] + scale[0] @@ -566,12 +595,11 @@ def decode(self, z, scale): def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs) + return custom_forward - x = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.conv2), z, use_reentrant=False) - out = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.decoder), x, use_reentrant=True) + x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.conv2), z, use_reentrant=False) + out = torch.utils.checkpoint.checkpoint(create_custom_forward(self.decoder), x, use_reentrant=True) return out def clear_cache(self): @@ -585,6 +613,7 @@ def clear_cache(self): # ── WanVideoVAE (top-level wrapper) ── + def _replace_conv_with_act_grad_only(model): """Replace Conv layers with act-grad-only versions (post-construction). @@ -592,12 +621,20 @@ def _replace_conv_with_act_grad_only(model): """ for name, module in model.named_modules(): if isinstance(module, CausalConv3d) and not isinstance(module, CausalConv3dActGradOnly): - parent_name, child_name = name.rsplit('.', 1) if '.' in name else ('', name) - parent = model if parent_name == '' else dict(model.named_modules())[parent_name] + parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model if parent_name == "" else dict(model.named_modules())[parent_name] new_module = CausalConv3dActGradOnly.__new__(CausalConv3dActGradOnly) - nn.Conv3d.__init__(new_module, module.in_channels, module.out_channels, - module.kernel_size, module.stride, (0, 0, 0), - module.dilation, module.groups, module.bias is not None) + nn.Conv3d.__init__( + new_module, + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + (0, 0, 0), + module.dilation, + module.groups, + module.bias is not None, + ) new_module._padding = module._padding new_module.padding = module.padding new_module.time_kernel_size = module.time_kernel_size @@ -609,12 +646,18 @@ def _replace_conv_with_act_grad_only(model): setattr(parent, child_name, new_module) elif isinstance(module, nn.Conv2d) and not isinstance(module, Conv2dActGradOnly): - parent_name, child_name = name.rsplit('.', 1) if '.' in name else ('', name) - parent = model if parent_name == '' else dict(model.named_modules())[parent_name] + parent_name, child_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model if parent_name == "" else dict(model.named_modules())[parent_name] new_module = Conv2dActGradOnly( - module.in_channels, module.out_channels, module.kernel_size, - module.stride, module.padding, module.dilation, - module.groups, module.bias is not None) + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + module.padding, + module.dilation, + module.groups, + module.bias is not None, + ) new_module.weight = module.weight new_module.weight.requires_grad = False if module.bias is not None: @@ -679,17 +722,44 @@ class WanVideoVAE(nn.Module): - Optional act-grad-only conv optimization """ - def __init__(self, z_dim=16, use_nested_grad_checkpoint=True, - use_act_grad_only_conv=True): + def __init__(self, z_dim=16, use_nested_grad_checkpoint=True, use_act_grad_only_conv=True): super().__init__() mean = [ - -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, - 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, ] std = [ - 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, - 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, ] self.mean = torch.tensor(mean) self.std = torch.tensor(std) @@ -705,16 +775,20 @@ def __init__(self, z_dim=16, use_nested_grad_checkpoint=True, scale_factor_temporal=4, scale_factor_spatial=8, latents_mean=mean, - latents_std=[1.0 / s for s in std], + latents_std=std, z_dim=z_dim, scaling_factor=1.0, ) # Build model with standard conv, then optionally replace with act-grad-only - self.model = VideoVAE_( - z_dim=z_dim, - use_nested_grad_checkpoint=use_nested_grad_checkpoint, - ).eval().requires_grad_(False) + self.model = ( + VideoVAE_( + z_dim=z_dim, + use_nested_grad_checkpoint=use_nested_grad_checkpoint, + ) + .eval() + .requires_grad_(False) + ) if use_act_grad_only_conv: _replace_conv_with_act_grad_only(self.model) @@ -742,8 +816,7 @@ def single_decode(self, hidden_state, device): video = self.model.decode(hidden_state, self.scale) return video.clamp_(-1, 1) - def encode(self, videos, device=None, tiled=False, - tile_size=(34, 34), tile_stride=(18, 16)): + def encode(self, videos, device=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): """Encode videos to latent space. Two calling conventions are supported (so this class can drop into @@ -768,10 +841,7 @@ def encode(self, videos, device=None, tiled=False, latents = self._encode_batched(videos, target_device, tiled, tile_size, tile_stride) return _EncoderOutput(latents) if device is None: - raise ValueError( - "WanVideoVAE.encode: ``device`` is required for the " - "batched/list calling convention." - ) + raise ValueError("WanVideoVAE.encode: ``device`` is required for the batched/list calling convention.") return self._encode_batched(videos, device, tiled, tile_size, tile_stride) def _encode_batched(self, videos, device, tiled, tile_size, tile_stride): @@ -782,18 +852,15 @@ def _encode_batched(self, videos, device, tiled, tile_size, tile_stride): for video in videos: video = video.unsqueeze(0) if tiled: - tile_size_px = (tile_size[0] * self.upsampling_factor, - tile_size[1] * self.upsampling_factor) - tile_stride_px = (tile_stride[0] * self.upsampling_factor, - tile_stride[1] * self.upsampling_factor) + tile_size_px = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride_px = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) hidden_state = self.tiled_encode(video, device, tile_size_px, tile_stride_px) else: hidden_state = self.single_encode(video, device) hidden_states.append(hidden_state.squeeze(0)) return torch.stack(hidden_states) - def decode(self, hidden_states, device=None, tiled=True, sp_group=None, - tile_size=(34, 34), tile_stride=(18, 16)): + def decode(self, hidden_states, device=None, tiled=True, sp_group=None, tile_size=(34, 34), tile_stride=(18, 16)): """Decode latents to video. ``device`` defaults to the latent's own device (so callers @@ -806,11 +873,9 @@ def decode(self, hidden_states, device=None, tiled=True, sp_group=None, device = hidden_states.device if tiled: if sp_group is not None: - video = self.tiled_parallel_decode( - hidden_states, device, tile_size, tile_stride, sp_group) + video = self.tiled_parallel_decode(hidden_states, device, tile_size, tile_stride, sp_group) else: - video = self.tiled_decode( - hidden_states, device, tile_size, tile_stride) + video = self.tiled_decode(hidden_states, device, tile_size, tile_stride) else: video = self.single_decode(hidden_states, device) return video @@ -822,8 +887,7 @@ def build_1d_mask(self, length, left_bound, right_bound, border_width): if not left_bound: x[:border_width] = (torch.arange(border_width) + 1) / border_width if not right_bound: - x[-border_width:] = torch.flip( - (torch.arange(border_width) + 1) / border_width, dims=(0,)) + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) return x def build_mask(self, data, is_bound, border_width): @@ -853,12 +917,16 @@ def tiled_decode(self, hidden_states, device, tile_size, tile_stride): tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) out_T = T * 4 - 3 - weight = torch.zeros((B, 1, out_T, H * self.upsampling_factor, - W * self.upsampling_factor), - dtype=hidden_states.dtype, device=device) - values = torch.zeros((B, 3, out_T, H * self.upsampling_factor, - W * self.upsampling_factor), - dtype=hidden_states.dtype, device=device) + weight = torch.zeros( + (B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + values = torch.zeros( + (B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) for h, h_, w, w_ in tasks: batch = hidden_states[:, :, :, h:h_, w:w_].to(device) @@ -866,22 +934,23 @@ def tiled_decode(self, hidden_states, device, tile_size, tile_stride): mask = self.build_mask( batch, is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), - border_width=((size_h - stride_h) * self.upsampling_factor, - (size_w - stride_w) * self.upsampling_factor) + border_width=( + (size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor, + ), ).to(dtype=hidden_states.dtype, device=device) th = h * self.upsampling_factor tw = w * self.upsampling_factor - value_slice = values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] - weight_slice = weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] + value_slice = values[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] + weight_slice = weight[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] value_slice += batch * mask weight_slice += mask.expand_as(weight_slice) values = values / weight return values.clamp_(-1, 1) - def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, - sp_group): + def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, sp_group): """Tiled decode with SP parallelism — each rank decodes a subset of tiles. Args: @@ -895,12 +964,16 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, tasks = self._make_tile_tasks(H, W, size_h, size_w, stride_h, stride_w) out_T = T * 4 - 3 - weight = torch.zeros((B, 1, out_T, H * self.upsampling_factor, - W * self.upsampling_factor), - dtype=hidden_states.dtype, device=device) - values = torch.zeros((B, 3, out_T, H * self.upsampling_factor, - W * self.upsampling_factor), - dtype=hidden_states.dtype, device=device) + weight = torch.zeros( + (B, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) + values = torch.zeros( + (B, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=device, + ) world_size = dist.get_world_size(sp_group) sp_rank = dist.get_rank(sp_group) @@ -912,8 +985,7 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, batch = hidden_states[:, :, :, h:h_, w:w_].to(device) batch = self.model.decode(batch, self.scale).to(device) # Pad to uniform tile size for all_gather - padding = (0, tile_img_w - batch.shape[-1], - 0, tile_img_h - batch.shape[-2]) + padding = (0, tile_img_w - batch.shape[-1], 0, tile_img_h - batch.shape[-2]) if padding != (0, 0, 0, 0): batch = F.pad(batch, padding) all_decoded.append(batch) @@ -922,17 +994,20 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, local_stack = torch.stack(all_decoded, dim=0) # (num_local_tasks, B, C, T, H, W) gathered = [torch.empty_like(local_stack) for _ in range(world_size)] # Pad to max local tasks across ranks - max_tasks = max((task_end - task_start) - for s, e in [_get_task_range(len(tasks), world_size, r) - for r in range(world_size)] - for task_start, task_end in [(s, e)]) + max_tasks = max( + (task_end - task_start) + for s, e in [_get_task_range(len(tasks), world_size, r) for r in range(world_size)] + for task_start, task_end in [(s, e)] + ) if local_stack.shape[0] < max_tasks: pad_n = max_tasks - local_stack.shape[0] - local_stack = torch.cat([ - local_stack, - torch.zeros(pad_n, *local_stack.shape[1:], - dtype=local_stack.dtype, device=local_stack.device) - ], dim=0) + local_stack = torch.cat( + [ + local_stack, + torch.zeros(pad_n, *local_stack.shape[1:], dtype=local_stack.dtype, device=local_stack.device), + ], + dim=0, + ) gathered = [torch.empty_like(local_stack) for _ in range(world_size)] dist.all_gather(gathered, local_stack.contiguous(), group=sp_group) @@ -951,16 +1026,16 @@ def tiled_parallel_decode(self, hidden_states, device, tile_size, tile_stride, mask = self.build_mask( decoded_tile, is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), - border_width=((size_h - stride_h) * self.upsampling_factor, - (size_w - stride_w) * self.upsampling_factor) + border_width=( + (size_h - stride_h) * self.upsampling_factor, + (size_w - stride_w) * self.upsampling_factor, + ), ).to(dtype=hidden_states.dtype, device=device) th = h * self.upsampling_factor tw = w * self.upsampling_factor - value_slice = values[:, :, :, th:th + decoded_tile.shape[3], - tw:tw + decoded_tile.shape[4]] - weight_slice = weight[:, :, :, th:th + decoded_tile.shape[3], - tw:tw + decoded_tile.shape[4]] + value_slice = values[:, :, :, th : th + decoded_tile.shape[3], tw : tw + decoded_tile.shape[4]] + weight_slice = weight[:, :, :, th : th + decoded_tile.shape[3], tw : tw + decoded_tile.shape[4]] value_slice += decoded_tile * mask weight_slice += mask.expand_as(weight_slice) @@ -975,12 +1050,16 @@ def tiled_encode(self, video, device, tile_size, tile_stride): out_T = (T + 3) // 4 data_device = "cpu" - weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, - W // self.upsampling_factor), - dtype=video.dtype, device=data_device) - values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, - W // self.upsampling_factor), - dtype=video.dtype, device=data_device) + weight = torch.zeros( + (1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=data_device, + ) + values = torch.zeros( + (1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=data_device, + ) for h, h_, w, w_ in tasks: batch = video[:, :, :, h:h_, w:w_].to(device) @@ -988,14 +1067,16 @@ def tiled_encode(self, video, device, tile_size, tile_stride): mask = self.build_mask( batch, is_bound=(h == 0, h_ >= H, w == 0, w_ >= W), - border_width=((size_h - stride_h) // self.upsampling_factor, - (size_w - stride_w) // self.upsampling_factor) + border_width=( + (size_h - stride_h) // self.upsampling_factor, + (size_w - stride_w) // self.upsampling_factor, + ), ).to(dtype=video.dtype, device=data_device) th = h // self.upsampling_factor tw = w // self.upsampling_factor - values[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += batch * mask - weight[:, :, :, th:th + batch.shape[3], tw:tw + batch.shape[4]] += mask + values[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] += batch * mask + weight[:, :, :, th : th + batch.shape[3], tw : tw + batch.shape[4]] += mask return values / weight @@ -1016,6 +1097,7 @@ def load_from_diffusers(cls, pretrained_path, **kwargs): safetensors_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") if os.path.exists(safetensors_path): from safetensors.torch import load_file + hf_sd = load_file(safetensors_path) else: bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") @@ -1029,6 +1111,7 @@ def load_from_diffusers(cls, pretrained_path, **kwargs): # ── State dict converter ── + def convert_diffusers_state_dict(hf_sd: dict) -> OrderedDict: """Convert HuggingFace AutoencoderKLWan state dict to WanVideoVAE format. @@ -1068,19 +1151,17 @@ def _convert_single_key(hf_key: str, resblock_map: dict) -> str: # quant_conv / post_quant_conv if hf_key.startswith("quant_conv."): - return "model.conv1." + hf_key[len("quant_conv."):] + return "model.conv1." + hf_key[len("quant_conv.") :] if hf_key.startswith("post_quant_conv."): - return "model.conv2." + hf_key[len("post_quant_conv."):] + return "model.conv2." + hf_key[len("post_quant_conv.") :] # Encoder if hf_key.startswith("encoder."): - return "model.encoder." + _convert_encoder_key( - hf_key[len("encoder."):], resblock_map) + return "model.encoder." + _convert_encoder_key(hf_key[len("encoder.") :], resblock_map) # Decoder if hf_key.startswith("decoder."): - return "model.decoder." + _convert_decoder_key( - hf_key[len("decoder."):], resblock_map) + return "model.decoder." + _convert_decoder_key(hf_key[len("decoder.") :], resblock_map) raise ValueError(f"Unknown HF key: {hf_key}") @@ -1088,35 +1169,35 @@ def _convert_single_key(hf_key: str, resblock_map: dict) -> str: def _convert_encoder_key(key: str, rb_map: dict) -> str: # conv_in → conv1 if key.startswith("conv_in."): - return "conv1." + key[len("conv_in."):] + return "conv1." + key[len("conv_in.") :] # norm_out → head.0 if key.startswith("norm_out."): - return "head.0." + key[len("norm_out."):] + return "head.0." + key[len("norm_out.") :] # conv_out → head.2 if key.startswith("conv_out."): - return "head.2." + key[len("conv_out."):] + return "head.2." + key[len("conv_out.") :] # down_blocks.{i}.{suffix} if key.startswith("down_blocks."): - rest = key[len("down_blocks."):] + rest = key[len("down_blocks.") :] dot = rest.index(".") block_idx = rest[:dot] - suffix = rest[dot + 1:] + suffix = rest[dot + 1 :] return "downsamples." + block_idx + "." + _convert_resblock_suffix(suffix, rb_map) # mid_block: HF [attn, res0, res1] → ROLL [res0(middle.0), attn(middle.1), res1(middle.2)] if key.startswith("mid_block."): - rest = key[len("mid_block."):] + rest = key[len("mid_block.") :] if rest.startswith("resnets.0."): - suffix = rest[len("resnets.0."):] + suffix = rest[len("resnets.0.") :] return "middle.0." + _convert_resblock_suffix(suffix, rb_map) if rest.startswith("attentions.0."): - suffix = rest[len("attentions.0."):] + suffix = rest[len("attentions.0.") :] return "middle.1." + suffix if rest.startswith("resnets.1."): - suffix = rest[len("resnets.1."):] + suffix = rest[len("resnets.1.") :] return "middle.2." + _convert_resblock_suffix(suffix, rb_map) raise ValueError(f"Unknown encoder key: {key}") @@ -1125,33 +1206,33 @@ def _convert_encoder_key(key: str, rb_map: dict) -> str: def _convert_decoder_key(key: str, rb_map: dict) -> str: # conv_in → conv1 if key.startswith("conv_in."): - return "conv1." + key[len("conv_in."):] + return "conv1." + key[len("conv_in.") :] # norm_out → head.0 if key.startswith("norm_out."): - return "head.0." + key[len("norm_out."):] + return "head.0." + key[len("norm_out.") :] # conv_out → head.2 if key.startswith("conv_out."): - return "head.2." + key[len("conv_out."):] + return "head.2." + key[len("conv_out.") :] # mid_block: same reorder as encoder if key.startswith("mid_block."): - rest = key[len("mid_block."):] + rest = key[len("mid_block.") :] if rest.startswith("resnets.0."): - suffix = rest[len("resnets.0."):] + suffix = rest[len("resnets.0.") :] return "middle.0." + _convert_resblock_suffix(suffix, rb_map) if rest.startswith("attentions.0."): - suffix = rest[len("attentions.0."):] + suffix = rest[len("attentions.0.") :] return "middle.1." + suffix if rest.startswith("resnets.1."): - suffix = rest[len("resnets.1."):] + suffix = rest[len("resnets.1.") :] return "middle.2." + _convert_resblock_suffix(suffix, rb_map) # up_blocks.{block_i}.resnets.{res_j}.{suffix} → upsamples.{flat_idx}.{suffix} # up_blocks.{block_i}.upsamplers.0.{suffix} → upsamples.{flat_idx}.{suffix} if key.startswith("up_blocks."): - return _convert_upblock_key(key[len("up_blocks."):], rb_map) + return _convert_upblock_key(key[len("up_blocks.") :], rb_map) raise ValueError(f"Unknown decoder key: {key}") @@ -1176,22 +1257,22 @@ def _convert_upblock_key(key: str, rb_map: dict) -> str: """ dot = key.index(".") block_idx = int(key[:dot]) - rest = key[dot + 1:] + rest = key[dot + 1 :] # Base offset: each block contributes 4 items (3 resnets + 1 upsampler) # except the last block which has 3 items (3 resnets, no upsampler) base = block_idx * 4 # works for blocks 0,1,2; block 3 won't have upsampler if rest.startswith("resnets."): - rest2 = rest[len("resnets."):] + rest2 = rest[len("resnets.") :] dot2 = rest2.index(".") res_idx = int(rest2[:dot2]) - suffix = rest2[dot2 + 1:] + suffix = rest2[dot2 + 1 :] flat_idx = base + res_idx return f"upsamples.{flat_idx}." + _convert_resblock_suffix(suffix, rb_map) if rest.startswith("upsamplers.0."): - suffix = rest[len("upsamplers.0."):] + suffix = rest[len("upsamplers.0.") :] flat_idx = base + 3 # upsampler comes after 3 resnets return f"upsamples.{flat_idx}." + suffix diff --git a/unirl/train/optim.py b/unirl/train/optim.py index 26243013a..5438cf653 100644 --- a/unirl/train/optim.py +++ b/unirl/train/optim.py @@ -166,7 +166,7 @@ def lr_lambda(step: int) -> float: return max(0.0, 1.0 - (step - warmup_steps) / (total_steps - warmup_steps)) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) - + if scheduler_type == "linear_warmup": constant = torch.optim.lr_scheduler.LinearLR( optimizer, From ac36b309fe25db5dc354ab68e8a4ec34e6800e3b Mon Sep 17 00:00:00 2001 From: haonan3 Date: Wed, 29 Jul 2026 18:11:37 +0800 Subject: [PATCH 10/24] refactor(recipes/refl): rebase onto sample-native core and drop the roles DSL - Rebase onto current main (#214): RolloutReq/RolloutInputs are gone; the actor role now consumes Texts/Images primitives plus per-sample metadata records straight from the data-source Sample. - Drop recipes/common/ (role-list orchestration): REFLTrainer subclasses BaseTrainer directly and wires actor + reward with placement()+remote_hydra, mirroring RewardBackpropTrainer (the SD3 image-ReFL driver). ReflActorRole mirrors ReFLPolicy's family-agnostic contract (pipeline_target + model_config + from_config; FSDPBackend composed in initialize()). - Re-root both configs from the roles: list to the repo-wide flat schema (actor:/reward:/data_source:/sampling:/logging:), the same shape as examples/diffusion/refl_sd3.yaml. - KL correctness across DP shards: diffuse_with_grad now returns per-sample [B] KL (concat field) instead of a per-shard scalar shared field, so DP_SCATTER merge/re-shard round-trips each shard's own KL. Previously the driver collapsed all shards to one value (loss/logging skew; gradient flow was unaffected because dKL/dkl is the constant kl_weight). - I2V condition assembly moves out of the role into Wan21/Wan22ReflPipeline.build_refl_conditions (mirrors each mainline pipeline's generate); negative prompts ride sampler_kwargs. - Seed scheme now matches ReFLPolicy (base + 1000*rollout_id + dp_rank); previously every rollout redrew the same init noise. - reward/service.py: restore the List typing import lost in the merge. --- recipes/common/__init__.py | 1 - recipes/common/trainer.py | 421 ------------------ .../configs/wan21_t2v_videoalign_refl.yaml | 204 ++++----- recipes/refl/configs/wan22_i2v_face_refl.yaml | 186 ++++---- recipes/refl/models/wan21.py | 60 ++- recipes/refl/models/wan22.py | 58 ++- recipes/refl/roles.py | 338 +++++++------- recipes/refl/trainer.py | 238 +++++++--- unirl/models/types/diffusion.py | 7 +- unirl/reward/service.py | 2 +- 10 files changed, 649 insertions(+), 866 deletions(-) delete mode 100644 recipes/common/__init__.py delete mode 100644 recipes/common/trainer.py diff --git a/recipes/common/__init__.py b/recipes/common/__init__.py deleted file mode 100644 index d80502c7d..000000000 --- a/recipes/common/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared helpers for recipe-level implementations.""" diff --git a/recipes/common/trainer.py b/recipes/common/trainer.py deleted file mode 100644 index 1a020bb69..000000000 --- a/recipes/common/trainer.py +++ /dev/null @@ -1,421 +0,0 @@ -"""Trainer: config-driven Remote role orchestration.""" - -from __future__ import annotations - -import json -import logging -import os -from dataclasses import dataclass -from typing import Any, Dict, Iterable, List, Optional - -from hydra.utils import get_method -from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict - -from unirl.distributed.group.remote import Remote -from unirl.reward.service import RewardService -from unirl.trainer.base import BaseTrainer, build_sampling_dict - -logger = logging.getLogger(__name__) - -_RESERVED_ROLE_NAMES = { - "cfg", - "pool", - "roles", - "role_specs", - "role_device_ids", - "role_slot_ids", - "data_source", - "sampling_params", - "wandb_logger", -} - - -@dataclass(frozen=True) -class RoleSpec: - """Driver-side parsed role declaration.""" - - name: str - target: str - placement: DictConfig - raw_cfg: DictConfig - index: int - - -class Trainer(BaseTrainer): - """Base trainer that creates recipe roles as UniRL Remote handles. - - ``Trainer`` owns the generic role lifecycle: - parse ``cfg.roles`` → topo-sort placement dependencies → create Remote - handles with explicit ``device_ids`` / ``slot_id`` → initialize roles. - Recipe trainers only implement ``build_req`` and ``train_step``. - """ - - def __init__(self, *, cfg: DictConfig, logging_cfg: Optional[DictConfig] = None) -> None: - self.cfg = cfg - self.role_specs: List[RoleSpec] = self.parse_role_specs(cfg) - self._role_specs_by_name: Dict[str, RoleSpec] = {spec.name: spec for spec in self.role_specs} - self._sorted_role_specs: List[RoleSpec] = self.topological_roles() - self._simulated_slot_ids: Dict[str, int] = self._simulate_role_slots(self._sorted_role_specs) - self._ensure_workers_per_device(cfg, self._simulated_slot_ids) - - super().__init__(cfg=cfg, logging_cfg=logging_cfg if logging_cfg is not None else cfg.get("logging")) - - self.batch_size = int(cfg.get("batch_size", 1)) - self.data_source = self.instantiate_data_source(cfg) - self.sampling_params = build_sampling_dict(cfg.sampling) if cfg.get("sampling") is not None else {} - self.roles: Dict[str, Any] = {} - self.role_device_ids: Dict[str, List[int]] = {} - self.role_slot_ids: Dict[str, int] = {} - self._next_colocate_slot = 1 - self._roles_initialized = False - - self.setup_roles() - self.initialize_roles() - self.validate_config() - - # ------------------------------------------------------------------ - # Config parsing and placement planning. - # ------------------------------------------------------------------ - - def parse_role_specs(self, cfg: DictConfig) -> List[RoleSpec]: - roles_cfg = cfg.get("roles") - if roles_cfg is None: - raise ValueError("Trainer requires cfg.roles as a list of role declarations.") - if not isinstance(roles_cfg, (list, tuple, ListConfig)): - raise TypeError(f"cfg.roles must be a list, got {type(roles_cfg).__name__}.") - - specs: List[RoleSpec] = [] - seen: set[str] = set() - for idx, role_cfg in enumerate(roles_cfg): - if not OmegaConf.is_config(role_cfg): - role_cfg = OmegaConf.create(role_cfg) - name = str(role_cfg.get("name") or "").strip() - if not name: - raise ValueError(f"roles[{idx}] is missing required field `name`.") - if name in seen: - raise ValueError(f"Duplicate role name {name!r} in cfg.roles.") - if name in _RESERVED_ROLE_NAMES or hasattr(self.__class__, name): - raise ValueError(f"Role name {name!r} is reserved; choose a different role name.") - target = str(role_cfg.get("_target_") or "").strip() - if not target: - raise ValueError(f"roles[{idx}] ({name}) is missing required field `_target_`.") - placement = role_cfg.get("placement") or OmegaConf.create({}) - if not OmegaConf.is_config(placement): - placement = OmegaConf.create(placement) - if placement.get("share_with") and placement.get("colocate_with"): - raise ValueError(f"Role {name!r}: placement cannot set both share_with and colocate_with.") - specs.append(RoleSpec(name=name, target=target, placement=placement, raw_cfg=role_cfg, index=idx)) - seen.add(name) - return specs - - def topological_roles(self) -> List[RoleSpec]: - specs = self.role_specs - by_name = {spec.name: spec for spec in specs} - visiting: set[str] = set() - visited: set[str] = set() - ordered: List[RoleSpec] = [] - - def parent_of(spec: RoleSpec) -> Optional[str]: - parent = spec.placement.get("share_with") or spec.placement.get("colocate_with") - return str(parent) if parent else None - - def visit(spec: RoleSpec) -> None: - if spec.name in visited: - return - if spec.name in visiting: - raise ValueError(f"Cycle detected in role placement dependencies at role {spec.name!r}.") - visiting.add(spec.name) - parent = parent_of(spec) - if parent: - if parent not in by_name: - raise ValueError(f"Role {spec.name!r} placement references unknown role {parent!r}.") - visit(by_name[parent]) - visiting.remove(spec.name) - visited.add(spec.name) - ordered.append(spec) - - for spec in specs: - visit(spec) - return ordered - - def _simulate_role_slots(self, specs: Iterable[RoleSpec]) -> Dict[str, int]: - slots: Dict[str, int] = {} - next_colocate_slot = 1 - for spec in specs: - if spec.placement.get("share_with"): - slots[spec.name] = slots[str(spec.placement.share_with)] - elif spec.placement.get("colocate_with"): - slots[spec.name] = next_colocate_slot - next_colocate_slot += 1 - else: - slots[spec.name] = 0 - return slots - - def _ensure_workers_per_device(self, cfg: DictConfig, slots: Dict[str, int]) -> None: - required = max(slots.values(), default=0) + 1 - current = int(cfg.get("workers_per_device", 1)) - if required <= current: - return - transport_kind = str(cfg.get("transport_kind", "colocate_store")) - if transport_kind in ("colocate_store", "colocate"): - raise ValueError( - "placement.colocate_with requires multiple worker slots per GPU, but " - f"transport_kind={transport_kind!r} only supports workers_per_device=1. " - "Set transport_kind='gpu_store' (or another multi-slot-capable transport) " - f"and workers_per_device>={required}." - ) - with open_dict(cfg): - cfg.workers_per_device = required - - # ------------------------------------------------------------------ - # Role creation. - # ------------------------------------------------------------------ - - def setup_roles(self) -> None: - self.roles = {} - self.role_device_ids = {} - self.role_slot_ids = {} - self._next_colocate_slot = 1 - - for spec in self._sorted_role_specs: - device_ids, slot_id = self.resolve_role_placement(spec) - handle = self.create_remote_role(spec, device_ids=device_ids, slot_id=slot_id) - self.roles[spec.name] = handle - self.role_device_ids[spec.name] = device_ids - self.role_slot_ids[spec.name] = slot_id - setattr(self, spec.name, handle) - - def resolve_role_placement(self, spec: RoleSpec) -> tuple[List[int], int]: - p = spec.placement or OmegaConf.create({}) - if p.get("share_with"): - parent = str(p.share_with) - parent_ids = self.role_device_ids[parent] - n = int(p.get("n_devices") or len(parent_ids)) - return self.step_subset(parent_ids, n), self.role_slot_ids[parent] - if p.get("colocate_with"): - parent = str(p.colocate_with) - parent_ids = self.role_device_ids[parent] - n = int(p.get("n_devices") or len(parent_ids)) - return self.step_subset(parent_ids, n), self.allocate_colocate_slot(parent) - - if p.get("device_ids") is not None: - return [int(d) for d in list(p.device_ids)], 0 - n_devices = int(p.get("n_devices") or spec.raw_cfg.get("n_devices") or self.cfg.num_devices) - if n_devices <= 0: - raise ValueError(f"Role {spec.name!r}: placement.n_devices must be positive, got {n_devices}.") - return self.pool.allocate(n_devices), 0 - - @staticmethod - def step_subset(device_ids: List[int], n_devices: int) -> List[int]: - if n_devices <= 0: - raise ValueError(f"n_devices must be positive, got {n_devices}.") - if n_devices > len(device_ids): - raise ValueError(f"Cannot take {n_devices} devices from parent device slab {device_ids}.") - if n_devices == len(device_ids): - return list(device_ids) - if n_devices == 1: - return [device_ids[0]] - last = len(device_ids) - 1 - return [device_ids[round(i * last / (n_devices - 1))] for i in range(n_devices)] - - def allocate_colocate_slot(self, parent: str) -> int: - del parent # The current implementation allocates globally unique colocate slots. - slot = max(self._next_colocate_slot, max(self.role_slot_ids.values(), default=0) + 1) - self._next_colocate_slot = slot + 1 - return slot - - def resolve_role_cls(self, target: str) -> type: - role_cls = get_method(target) - if not isinstance(role_cls, type) or not issubclass(role_cls, Remote): - raise TypeError(f"Role target {target!r} must resolve to a Remote subclass, got {role_cls!r}.") - return role_cls - - def prepare_role_cfg(self, spec: RoleSpec) -> DictConfig: - container = OmegaConf.to_container(spec.raw_cfg, resolve=True) - if not isinstance(container, dict): - raise TypeError(f"Role {spec.name!r} config must resolve to a mapping.") - for key in ("name", "_target_", "placement"): - container.pop(key, None) - return OmegaConf.create(container) - - def create_remote_role(self, spec: RoleSpec, *, device_ids: List[int], slot_id: int): - role_cls = self.resolve_role_cls(spec.target) - if issubclass(role_cls, RewardService): - # RewardService takes a materialized ``backend`` (not a ``cfg`` blob): - # pass the role's own fields as plain-dict kwargs so the worker's - # ``_resolve_init_kwargs`` walker instantiates the nested ``_target_`` - # backend in its own CUDA context. - init_kwargs = OmegaConf.to_container(self.prepare_role_cfg(spec), resolve=True) - else: - init_kwargs = {"cfg": self.prepare_role_cfg(spec)} - return self.pool.create_remote( - role_cls, - device_ids=device_ids, - slot_id=slot_id, - role_name=spec.name, - init_kwargs=init_kwargs, - ) - - def initialize_roles(self) -> None: - if self._roles_initialized: - return - for spec in self._sorted_role_specs: - self.roles[spec.name].initialize() - self._roles_initialized = True - - # ------------------------------------------------------------------ - # Data, validation, training loop. - # ------------------------------------------------------------------ - - def instantiate_data_source(self, cfg: DictConfig): - if cfg.get("data_source") is None: - return None - from hydra.utils import instantiate - - return instantiate(cfg.data_source) - - def validate_config(self) -> None: - if self.data_source is None: - raise ValueError("Trainer requires cfg.data_source.") - if not self.sampling_params: - raise ValueError("Trainer requires cfg.sampling.") - - def build_req(self, inputs: Any, rollout_id: int) -> Any: - raise NotImplementedError - - def train_step(self, req: Any, *, training_progress: float = 0.0, rollout_id: int = 0) -> Dict[str, Any]: - raise NotImplementedError - - def train( - self, - *, - num_rollouts: Optional[int] = None, - save_interval: Optional[int] = None, - save_dir: Optional[str] = None, - load_dir: Optional[str] = None, - save_mode: Optional[str] = None, - ) -> None: - num_rollouts = int(num_rollouts if num_rollouts is not None else self.cfg.get("num_rollouts", 100)) - save_interval = int(save_interval if save_interval is not None else self.cfg.get("save_interval", 0)) - save_dir = save_dir if save_dir is not None else self.cfg.get("save_dir") - load_dir = load_dir if load_dir is not None else self.cfg.get("load_dir") - save_mode = save_mode if save_mode is not None else self.cfg.get("save_mode", "auto") - - start_rollout = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) - data_source = self.data_source - if data_source is None: - raise ValueError("Trainer requires cfg.data_source.") - for _ in range(start_rollout): - data_source.get_samples(self.batch_size) - self._init_wandb(num_rollouts=num_rollouts) - try: - for rollout_id in range(start_rollout, num_rollouts): - training_progress = rollout_id / max(1, num_rollouts - 1) - inputs = data_source.get_samples(self.batch_size) - req = self.build_req(inputs, rollout_id) - metrics = self.train_step(req, training_progress=training_progress, rollout_id=rollout_id) - self.log_metrics(metrics, rollout_id=rollout_id, num_rollouts=num_rollouts) - self.maybe_save_checkpoint( - rollout_id, - num_rollouts, - save_interval=save_interval, - save_dir=save_dir, - save_mode=save_mode, - ) - finally: - self._finish_wandb() - - def log_metrics(self, metrics: Dict[str, Any], *, rollout_id: int, num_rollouts: int) -> None: - trainer_name = self.__class__.__name__.removesuffix("Trainer") or self.__class__.__name__ - metric_text = " ".join(f"{key}={self._format_metric_for_log(value)}" for key, value in metrics.items()) - if metric_text: - logger.info("%s rollout %d/%d %s", trainer_name, rollout_id + 1, num_rollouts, metric_text) - else: - logger.info("%s rollout %d/%d", trainer_name, rollout_id + 1, num_rollouts) - - wb = self.wandb_logger - if wb is not None and wb.initialized: - wb.log_step( - step=rollout_id + 1, - metrics={k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))}, - prefix=str((self.logging_cfg or {}).get("metric_prefix", "train/")), - ) - - @staticmethod - def _format_metric_for_log(value: Any) -> str: - if isinstance(value, bool): - return str(value) - if isinstance(value, int): - return str(value) - if isinstance(value, float): - abs_value = abs(value) - if value != 0.0 and (abs_value < 1e-3 or abs_value >= 1e4): - return f"{value:.6e}" - return f"{value:.4f}" - return repr(value) - - # ------------------------------------------------------------------ - # Role-aware checkpointing. - # ------------------------------------------------------------------ - - def checkpoint_roles(self, path: str, *, step: int, mode: str) -> None: - for name, role in self.roles.items(): - if hasattr(role, "save_checkpoint"): - role_path = os.path.join(path, name) - role.save_checkpoint(role_path, step=step, mode=mode) - - def load_checkpoint_roles(self, path: str) -> int: - starts: List[int] = [] - for name, role in self.roles.items(): - if hasattr(role, "load_checkpoint"): - role_path = os.path.join(path, name) - if os.path.exists(role_path): - result = role.load_checkpoint(role_path) - if isinstance(result, list): - result = result[0] if result else 0 - starts.append(int(result or 0)) - return max(starts, default=0) - - def maybe_save_checkpoint( - self, - rollout_id: int, - num_rollouts: int, - *, - save_interval: int, - save_dir: Optional[str], - save_mode: str = "auto", - ) -> None: - if save_interval <= 0: - return - step = rollout_id + 1 - if step % save_interval != 0 and step < num_rollouts: - return - base_dir = os.path.abspath(save_dir) if save_dir else os.path.join(os.getcwd(), "checkpoints") - path = os.path.join(base_dir, f"checkpoint-{step}") - os.makedirs(path, exist_ok=True) - logger.info("Saving role checkpoint at rollout %d/%d -> %s", step, num_rollouts, path) - self.checkpoint_roles(path, step=step, mode=save_mode) - with open(os.path.join(path, "trainer_state.json"), "w") as f: - json.dump({"wandb_run_id": self.wandb_logger.run_id, "optimizer_step": self.wandb_logger.optimizer_step}, f) - - def maybe_load_checkpoint(self, load_dir: Optional[str], *, num_rollouts: Optional[int] = None) -> int: - if not load_dir: - return 0 - load_dir = os.path.abspath(load_dir) - logger.info("Loading role checkpoint from %s", load_dir) - start = self.load_checkpoint_roles(load_dir) - state_path = os.path.join(load_dir, "trainer_state.json") - if os.path.exists(state_path): - with open(state_path) as f: - self._resume_state = json.load(f) - logger.info("Checkpoint restored; resuming at rollout %d", start) - if num_rollouts is not None and start >= num_rollouts: - logger.warning( - "Checkpoint step %d >= num_rollouts %d — nothing left to train.", - start, - num_rollouts, - ) - return start - - -__all__ = ["Trainer", "RoleSpec"] diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 7ba583585..2ba89ecba 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -1,121 +1,106 @@ # @package _global_ -# REFL WAN 2.1 T2V — VideoAlign (Qwen2-VL VQ/MQ/TA) reward — role config. +# REFL WAN 2.1 T2V — VideoAlign (Qwen2-VL VQ/MQ/TA) reward. +# +# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a +# ReflActorRole (FSDP WAN 2.1 + grad BPTT sampling + optimizer) and the frozen +# differentiable VideoAlign reward, colocated on the actor's worker slab so +# decoded video never leaves the GPU. GradContext backprops reward → VAE → +# DiT LoRA across the sibling RPC boundary. num_devices: 8 batch_size: 8 num_rollouts: 1000 save_interval: 100 save_dir: ${oc.env:OUTPUT_DIR,outputs/wan21_t2v_videoalign_refl} +save_mode: adapter +max_grad_norm: 1.0 -roles: - - name: actor - _target_: recipes.refl.roles.ReflActorRole - placement: - n_devices: ${num_devices} - - model: - _target_: unirl.models.wan21.bundle.WAN21Bundle.from_config - config: - _target_: unirl.models.wan21.config.WAN21PipelineConfig - pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} - model_precision: bf16 - shift: 5.0 - max_sequence_length: 512 - - pipeline: - _target_: recipes.refl.models.wan21.Wan21ReflPipeline - shift: 5.0 - autocast_precision: bf16 - trajectory_precision: bf16 - logprob_precision: fp32 - max_sequence_length: 512 - strategy: - # eta=0.0 below reduces FlowSDE to deterministic ODE — REFL wants a - # deterministic transition on the differentiable path. - _target_: unirl.sde.kernels.FlowSDEStrategy - - backend: - _target_: unirl.train.backend.fsdp.FSDPBackend - block_class_names: ["WanTransformerBlock"] - trainable_attr: transformer - fsdp_cfg: - _target_: unirl.train.configs.FSDPConfig - param_dtype: bf16 - master_dtype: fp32 - cpu_offload: false - mixed_precision: true - fsdp_mode: full - reshard_after_forward: true - # BPTT keeps the full 24→final grad window alive; activation - # checkpointing is mandatory for any non-trivial - # num_frames × num_inference_steps product. - activation_checkpointing: true - use_torch_compile: false - optimizer_cfg: - _target_: unirl.train.backend.base.OptimizerConfig - learning_rate: 5.0e-6 - adam_beta1: 0.9 - adam_beta2: 0.999 - adam_epsilon: 1.0e-8 - weight_decay: 0.0 - scheduler_cfg: - _target_: unirl.train.backend.base.LrSchedulerConfig - type: linear_warmup - warmup_steps: 10 - total_steps: ${num_rollouts} - lora_cfg: - _target_: unirl.train.configs.LoraConfig - rank: 64 - alpha: 128 - dropout: 0.0 - bias: none - task_type: FEATURE_EXTRACTION - # Peft path takes suffix matches → these six suffixes catch both - # attn1 (self-attention) and attn2 (cross-attention) LoRA targets - # inside every WanTransformerBlock. - target_modules: - - to_q - - to_k - - to_v - - to_out.0 - - ffn.net.0.proj - - ffn.net.2 - - algorithm: ${algorithm} - - - name: reward - _target_: unirl.reward.service.RewardService - placement: - # Co-locate reward with the actor to avoid shipping decoded video - # across processes. NOT an autograd requirement — GradContext - # backprops rewards across workers via RPC — purely an efficiency choice. - share_with: actor - - backend: - _target_: recipes.refl.rewards.videoalign.VideoAlignRewardScorer - base_device: cuda - config: - _target_: recipes.refl.rewards.videoalign.VideoAlignSpec - reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} - device: cuda - batch_size: 1 - resize_height: 336 - resize_width: 588 - micro_batch_size: 1 - reward_num_frames: 36 # 81 → 36 uniform frames before scoring - use_norm: true - w_vq: 1.0 - w_mq: 1.0 - w_ta: 1.0 - differentiable: true # REFL requires grad through the reward forward - -algorithm: +actor: + _target_: recipes.refl.roles.ReflActorRole + # Family selector — swap pipeline_target + model_config for another family, + # no code changes (mirrors ReFLPolicy's pipeline_target contract). + pipeline_target: recipes.refl.models.wan21.Wan21ReflPipeline + block_class_names: ["WanTransformerBlock"] + # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.25 reward_baseline: 0.0 reward_scale: 1.0 kl_weight: 0.0 - max_grad_norm: 1.0 - sampling_params: ${sampling} + strategy: + # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + model_config: + _target_: unirl.models.wan21.config.WAN21PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + shift: 5.0 + max_sequence_length: 512 + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + # BPTT keeps the full mid→final grad window alive; activation + # checkpointing is mandatory for any non-trivial + # num_frames × num_inference_steps product. + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 5.0e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 64 + alpha: 128 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + # Peft path takes suffix matches → these six suffixes catch both + # attn1 (self-attention) and attn2 (cross-attention) LoRA targets + # inside every WanTransformerBlock. + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: recipes.refl.rewards.videoalign.VideoAlignRewardScorer + base_device: cuda + config: + _target_: recipes.refl.rewards.videoalign.VideoAlignSpec + reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} + device: cuda + batch_size: 1 + resize_height: 336 + resize_width: 588 + micro_batch_size: 1 + reward_num_frames: 36 # 81 → 36 uniform frames before scoring + use_norm: true + w_vq: 1.0 + w_mq: 1.0 + w_ta: 1.0 + differentiable: true # REFL requires grad through the reward forward data_source: _target_: unirl.data.data_source.MultimodalRLDataSource @@ -149,14 +134,13 @@ sampling: # the final step traverses the reward backward pass (DRaFT-1 / ReFL). mid_timestep: 24 final_timestep: 24 - # KL branch off — outer algorithm.kl_weight also 0. Single-pass forward - # per step, no LoRA-disabled reference recompute. + # KL branch off — actor.kl_weight is also 0. Single-pass forward per + # step, no LoRA-disabled reference recompute. kl_weight: 0.0 logging: report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} - project_name: ${oc.env:WANDB_PROJECT,unirl-refl-t2v} + project_name: ${oc.env:WANDB_PROJECT,unirl-refl} run_name: ${oc.env:WANDB_RUN_NAME,wan21_t2v_videoalign_refl} tags: ["wan21", "t2v", "refl", "bptt", "videoalign"] log_media: false - metric_prefix: refl_recipe/ diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index 8ca767858..72d420086 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -1,110 +1,99 @@ # @package _global_ -# REFL WAN 2.2 I2V — Remote role config. +# REFL WAN 2.2 I2V — Face-identity reward. +# +# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a +# ReflActorRole (FSDP WAN 2.2 dual-DiT + grad BPTT sampling + optimizer, +# LoRA restricted to the low_noise DiT) and the frozen differentiable Face +# reward, colocated on the actor's worker slab. The I2V first frame rides the +# data source's (image, condition) MediaRef; the face reference video path +# rides per-sample metadata (``ref_video_path``) into the reward's records. num_devices: 8 batch_size: 8 num_rollouts: 1000 save_interval: 100 save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_i2v_face_refl} +save_mode: adapter +max_grad_norm: 1.0 -roles: - - name: actor - _target_: recipes.refl.roles.ReflActorRole - placement: - n_devices: ${num_devices} - - model: - _target_: unirl.models.wan22.bundle.WAN22Bundle.from_config - config: - _target_: unirl.models.wan22.config.WAN22PipelineConfig - pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} - model_precision: bf16 - shift: 5.0 - max_sequence_length: 512 - boundary_ratio: 0.9 - num_train_timesteps: 1000 - - pipeline: - _target_: recipes.refl.models.wan22.Wan22ReflPipeline - shift: 5.0 - autocast_precision: bf16 - trajectory_precision: bf16 - logprob_precision: fp32 - max_sequence_length: 512 - strategy: - # eta=0.0 below reduces FlowSDE to deterministic ODE — REFL wants a - # deterministic transition on the differentiable path. - _target_: unirl.sde.kernels.FlowSDEStrategy - - backend: - _target_: unirl.train.backend.fsdp.FSDPBackend - block_class_names: ["WanTransformerBlock"] - trainable_attr: transformer - fsdp_cfg: - _target_: unirl.train.configs.FSDPConfig - param_dtype: bf16 - master_dtype: fp32 - cpu_offload: false - mixed_precision: true - fsdp_mode: full - reshard_after_forward: true - activation_checkpointing: true - use_torch_compile: false - optimizer_cfg: - _target_: unirl.train.backend.base.OptimizerConfig - learning_rate: 2.5e-6 - adam_beta1: 0.9 - adam_beta2: 0.999 - adam_epsilon: 1.0e-8 - weight_decay: 0.0 - scheduler_cfg: - _target_: unirl.train.backend.base.LrSchedulerConfig - type: linear_warmup - warmup_steps: 10 - total_steps: ${num_rollouts} - lora_cfg: - _target_: unirl.train.configs.LoraConfig - rank: 32 - alpha: 32 - dropout: 0.0 - bias: none - task_type: FEATURE_EXTRACTION - module_prefix: low_noise - target_modules: - - to_q - - to_k - - to_v - - to_out.0 - - ffn.net.0.proj - - ffn.net.2 - - algorithm: ${algorithm} - - - name: reward - _target_: unirl.reward.service.RewardService - placement: - share_with: actor - - backend: - _target_: recipes.refl.rewards.face.FaceRewardScorer - base_device: cuda - config: - _target_: recipes.refl.rewards.face.FaceRewardSpec - model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} - device: cuda - batch_size: 1 - image_size: 112 - ref_max_frames: 81 - ref_max_pixels: 230400 - differentiable: true - -algorithm: +actor: + _target_: recipes.refl.roles.ReflActorRole + pipeline_target: recipes.refl.models.wan22.Wan22ReflPipeline + block_class_names: ["WanTransformerBlock"] + # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.1 reward_baseline: 0.54 reward_scale: 0.16 kl_weight: 1.0 - max_grad_norm: 1.0 - sampling_params: ${sampling} + strategy: + # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a + # deterministic transition on the differentiable path. + _target_: unirl.sde.kernels.FlowSDEStrategy + model_config: + _target_: unirl.models.wan22.config.WAN22PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + shift: 5.0 + max_sequence_length: 512 + boundary_ratio: 0.9 + num_train_timesteps: 1000 + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: fp32 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 2.5e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: linear_warmup + warmup_steps: 10 + total_steps: ${num_rollouts} + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 32 + alpha: 32 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + # Inject LoRA only into the low-noise DiT (the final-timestep expert the + # BPTT window trains); the high-noise DiT stays frozen. + module_prefix: low_noise + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ffn.net.0.proj + - ffn.net.2 + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: recipes.refl.rewards.face.FaceRewardScorer + base_device: cuda + config: + _target_: recipes.refl.rewards.face.FaceRewardSpec + model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} + device: cuda + batch_size: 1 + image_size: 112 + ref_max_frames: 81 + ref_max_pixels: 230400 + differentiable: true data_source: _target_: unirl.data.data_source.MultimodalRLDataSource @@ -130,14 +119,17 @@ sampling: seed: 42 init_same_noise: false sampler_kwargs: + # Truncated BPTT window: grad on from mid_timestep..final_timestep — the + # last 4 of 8 steps, i.e. the low-noise DiT's sigma range. mid_timestep: 4 final_timestep: 7 + # Per-step KL against the LoRA-disabled reference (weighted again by + # actor.kl_weight at the loss site). kl_weight: 1.0 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,wan22_i2v_face_refl} - tags: ["wan22", "i2v", "refl", "face", "recipe"] + tags: ["wan22", "i2v", "refl", "bptt", "face"] log_media: false - metric_prefix: refl_recipe/ diff --git a/recipes/refl/models/wan21.py b/recipes/refl/models/wan21.py index d5fdc6834..43d061ed6 100644 --- a/recipes/refl/models/wan21.py +++ b/recipes/refl/models/wan21.py @@ -19,6 +19,7 @@ from __future__ import annotations +import dataclasses from contextlib import nullcontext from typing import Any, Dict, Optional, Tuple @@ -26,10 +27,13 @@ from unirl.models.types.diffusion import DiffuseWithGradResult from unirl.models.wan21.bundle import WAN21Bundle +from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.diffusion import WAN21DiffusionStage, WAN21DiffusionStep +from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage from unirl.models.wan21.pipeline import WAN21Pipeline from unirl.train.lora import adapters_disabled +from unirl.types.primitives import Images, Texts from unirl.types.sampling import DiffusionSamplingParams # Matches the mainline module-level constant in unirl/models/wan21/diffusion.py. @@ -161,7 +165,7 @@ def diffuse_with_grad( """Differentiable WAN 2.1 T2V sampling for REFL-style BPTT training. Returns :class:`DiffuseWithGradResult` with the live-grad - ``z_final`` + scalar ``kl_loss``. + ``z_final`` + per-sample ``kl_loss`` ``[B]``. BPTT knobs (read from ``params.sampler_kwargs``): @@ -172,9 +176,10 @@ def diffuse_with_grad( - ``final_timestep`` (int, default ``num_inference_steps - 1``): early stop. The loop breaks once ``i >= final_timestep``. - ``kl_weight`` (float, default 0.0): when non-zero, per-step KL - ``mean((pred - ref_pred)**2 / (2 * sigma**2))`` is accumulated - and returned in ``kl_loss``. The trainer multiplies it by its own - ``kl_weight`` at the loss site. + ``(pred - ref_pred)**2 / (2 * sigma**2)`` is reduced per sample + and accumulated into the ``[B]`` ``kl_loss``. Per-sample (not a + scalar) so DP_SCATTER merge/re-shard round-trips each shard's own + KL. The actor multiplies it by its ``kl_weight`` at the loss site. ``initial_latents`` follows the same contract as :meth:`diffuse`: when provided, used verbatim and the internal RNG path is @@ -235,7 +240,7 @@ def diffuse_with_grad( sigma_max = float(schedule[1].item()) if int(schedule.shape[0]) > 1 else 0.99 transformer = self.model.transformer - kl_total = torch.zeros((), device=device, dtype=torch.float32) + kl_total = torch.zeros(batch_size, device=device, dtype=torch.float32) kl_steps = 0 guidance_scale = float(params.guidance_scale) @@ -303,7 +308,7 @@ def diffuse_with_grad( branch="cond", ) sigma_f32 = sigma.to(dtype=torch.float32) - kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).mean() + kl_step = ((kl_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 @@ -354,5 +359,48 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: logprob_precision=old.logprob_dtype, ) + def build_refl_conditions( + self, + texts: Texts, + *, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> WAN21Conditions: + """Full REFL conditioning: text + CFG negative + optional I2V image. + + Mirrors the condition assembly of :meth:`WAN21Pipeline.generate` + (text via the public ``build_conditions``, image slots attached with + the diffusion geometry from ``params``), plus the REFL convention + that an explicit negative prompt rides + ``params.sampler_kwargs['negative_prompt']``. + """ + 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 + ) + conds = self.build_conditions(texts, negatives=negatives, guidance_scale=guidance) + + if images is not None: + if images.pixels is None or int(images.pixels.shape[0]) != len(texts.texts): + raise ValueError( + f"Wan21ReflPipeline.build_refl_conditions: image count " + f"{None if images.pixels is None else int(images.pixels.shape[0])} " + f"!= text count {len(texts.texts)}" + ) + image_latent = WAN21ImageLatentEncodeStage( + self.bundle, + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ).encode(images) + image_embed = ( + WAN21CLIPVisionEncodeStage(self.bundle).encode(images) + if getattr(self.bundle, "uses_clip_vision", False) + else None + ) + conds = dataclasses.replace(conds, image_latent=image_latent, image_embed=image_embed) + return conds + __all__ = ["Wan21ReflDiffusionStep", "Wan21ReflDiffusionStage", "Wan21ReflPipeline"] diff --git a/recipes/refl/models/wan22.py b/recipes/refl/models/wan22.py index 1c671a8f8..60ceadb3e 100644 --- a/recipes/refl/models/wan22.py +++ b/recipes/refl/models/wan22.py @@ -2,17 +2,21 @@ from __future__ import annotations +import dataclasses from contextlib import nullcontext from typing import Any, Dict, Optional, Tuple import torch from unirl.models.types.diffusion import DiffuseWithGradResult +from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions +from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage from unirl.models.wan22.bundle import WAN22Bundle from unirl.models.wan22.diffusion import WAN22DiffusionStage, WAN22DiffusionStep from unirl.models.wan22.pipeline import WAN22Pipeline from unirl.train.lora import adapters_disabled +from unirl.types.primitives import Images, Texts from unirl.types.sampling import DiffusionSamplingParams # Matches the mainline module-level constant in unirl/models/wan22/diffusion.py. @@ -146,7 +150,7 @@ def diffuse_with_grad( """Differentiable WAN 2.2 sampling for REFL-style BPTT training. Returns :class:`DiffuseWithGradResult` with the live-grad - ``z_final`` + scalar ``kl_loss``. + ``z_final`` + per-sample ``kl_loss`` ``[B]``. """ if conditions.text is None or conditions.text.embeds is None: @@ -215,7 +219,7 @@ def diffuse_with_grad( ) dual = self.model.transformer - kl_total = torch.zeros((), device=device, dtype=torch.float32) + kl_total = torch.zeros(batch_size, device=device, dtype=torch.float32) kl_steps = 0 for i in range(T): @@ -283,7 +287,7 @@ def diffuse_with_grad( use_high_noise=use_high_noise, ) sigma_f32 = sigma.to(dtype=torch.float32) - kl_step = ((kl_pred.float() - ref_pred.float()) ** 2 / (2.0 * sigma_f32**2)).mean() + kl_step = ((kl_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 @@ -329,5 +333,53 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: logprob_precision=old.logprob_dtype, ) + def build_refl_conditions( + self, + texts: Texts, + *, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + ) -> WAN21Conditions: + """Full REFL conditioning: text + CFG negative + optional I2V image. + + Mirrors the condition assembly of :meth:`WAN22Pipeline.generate` — + the negative branch is encoded against the **effective** guidance + ``max(guidance_scale, guidance_scale_2)`` (WAN22 routes CFG by sigma + across the low/high-noise DiTs), and image slots are attached with + the diffusion geometry from ``params``. An explicit negative prompt + rides ``params.sampler_kwargs['negative_prompt']``. + """ + primary_g = float(params.guidance_scale) + low_g = float(params.guidance_scale_2) if params.guidance_scale_2 is not None else primary_g + effective_guidance = max(primary_g, low_g) + negative_prompt = (params.sampler_kwargs or {}).get("negative_prompt") + negatives = ( + Texts(texts=[str(negative_prompt)] * len(texts.texts)) + if negative_prompt and effective_guidance > 1.0 + else None + ) + conds = self.build_conditions(texts, negatives=negatives, guidance_scale=effective_guidance) + + if images is not None: + if images.pixels is None or int(images.pixels.shape[0]) != len(texts.texts): + raise ValueError( + f"Wan22ReflPipeline.build_refl_conditions: image count " + f"{None if images.pixels is None else int(images.pixels.shape[0])} " + f"!= text count {len(texts.texts)}" + ) + image_latent = WAN21ImageLatentEncodeStage( + self.bundle, + num_frames=int(params.num_frames), + height=int(params.height), + width=int(params.width), + ).encode(images) + image_embed = ( + WAN21CLIPVisionEncodeStage(self.bundle).encode(images) + if getattr(self.bundle, "uses_clip_vision", False) + else None + ) + conds = dataclasses.replace(conds, image_latent=image_latent, image_embed=image_embed) + return conds + __all__ = ["Wan22ReflDiffusionStep", "Wan22ReflDiffusionStage", "Wan22ReflPipeline"] diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index ee48355de..97cdbb742 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -1,30 +1,57 @@ -"""Remote roles for the WAN22 REFL recipe.""" +"""ReflActorRole — family-agnostic REFL/BPTT actor Remote for the refl recipe. + +Mirrors :class:`unirl.train.refl.policy.ReFLPolicy` (the SD3 image-ReFL actor): +a **config-chosen** ``Pipeline`` (``pipeline_target`` + ``model_config``, no +per-family imports), FSDP-wrapped in place via ``FSDPBackend``, driven by three +driver RPCs per step under the distributed ``enable_grad()`` context:: + + gen = actor.generate_samples(texts=…, images=…, params=…) # grad BPTT sample + VAE decode + rew = reward.score_differentiable(gen.decoded, prompts, records) + actor.forward_backward_loss(rewards=rew, kl_loss=gen.kl_loss) # loss seed → local backward + actor.step(max_grad_norm=…) # ctx exit routed grads → optimizer + +The pipeline named by ``pipeline_target`` must expose the recipe contract +(see ``recipes.refl.models``): ``build_refl_conditions(texts, images=…, +params=…)`` plus a ``diffusion`` stage with ``diffuse_with_grad`` and a +``vae_decode`` stage with ``decode_with_grad``. +""" from __future__ import annotations import dataclasses -from collections.abc import Mapping +import logging from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional, Tuple import torch -from hydra.utils import instantiate -from omegaconf import OmegaConf +import torch.distributed as dist +from hydra.utils import get_class -from unirl.distributed.group.dispatch import Dispatch, distributed +from unirl.distributed.group.dispatch import Dispatch, Execute, distributed from unirl.distributed.group.remote import Remote -from unirl.distributed.tensor.batch import Batch, concat_field, shared_field +from unirl.distributed.tensor.batch import Batch, concat_field from unirl.sde.runtime import get_sigma_schedule +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.rollout_req import RolloutReq +from unirl.types.sampling import DiffusionSamplingParams + +logger = logging.getLogger(__name__) @dataclass class REFLGenerated(Batch): - """Generated BPTT payload carrying live-grad decoded pixels and KL loss.""" + """Generated BPTT payload: live-grad decoded pixels + per-sample KL. + + Both fields are batch-aligned concat fields (``decoded`` ``[B, …]``, + ``kl_loss`` ``[B]``) so DP_SCATTER merge/re-shard keeps each DP shard's + own values — the KL fed back into ``forward_backward_loss`` is the KL + that shard actually accumulated. + """ decoded: torch.Tensor = concat_field(default_factory=lambda: torch.empty(0)) - kl_loss: torch.Tensor = shared_field(default_factory=lambda: torch.zeros(1)) + kl_loss: torch.Tensor = concat_field(default_factory=lambda: torch.empty(0)) @dataclass @@ -37,184 +64,144 @@ class REFLLossMetrics(Batch): reward_mean: List[float] = concat_field(default_factory=list) -def _maybe_instantiate(value: Any) -> Any: - if OmegaConf.is_config(value) and value.get("_target_") is not None: - return instantiate(value) - return value - - -@dataclass(frozen=True) -class RoleStepResult: - """Generic result of one role-local optimizer step.""" - - metrics: Mapping[str, object] - grad_norm: float - lr: float - - class ReflActorRole(Remote): - """Actor role: bundle + pipeline + backend + REFL BPTT logic.""" + """Family-agnostic REFL actor: config-chosen pipeline + FSDP + BPTT loss.""" - bundle: Any - pipeline: Any - backend: Any - algo_cfg: Any - sampling_params: Any - - def __init__(self, cfg: Any) -> None: + 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, ...] = ("WanTransformerBlock",), + reward_weight: float = 1.0, + reward_baseline: float = 0.0, + reward_scale: float = 1.0, + kl_weight: float = 0.0, + ) -> None: super().__init__() - self.cfg = cfg + 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.reward_weight = float(reward_weight) + self.reward_baseline = float(reward_baseline) + self.reward_scale = float(reward_scale) + self.kl_weight = float(kl_weight) def initialize(self) -> None: - self.bundle = instantiate(self.cfg.get("model")) - self.pipeline = instantiate(self.cfg.get("pipeline"), bundle=self.bundle) - self.backend = instantiate( - self.cfg.get("backend"), - bundle=self.bundle, - device=torch.device(self.device), - rank=int(self.rank_info.rank), - ) - self.algo_cfg = self.cfg.algorithm - self.sampling_params = _maybe_instantiate(self.algo_cfg.get("sampling_params")) - - @distributed - def step(self) -> RoleStepResult: - """Clip gradients and run one backend optimizer step.""" - if not hasattr(self, "backend") or not hasattr(self.backend, "optimizer_step"): - raise RuntimeError(f"{type(self).__name__}.step requires a backend with optimizer_step(...).") - grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(self.algo_cfg.get("max_grad_norm", 1.0)))) - lr = 0.0 + torch.cuda.set_device(self.device) + # Default PG over the actor role's workers (env:// from Remote.setup's + # dist_env); FSDP2 fully_shard wraps over it. Same order as ReFLPolicy. + 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: - sched = getattr(self.backend, "scheduler", None) - if sched is not None: - last = sched.get_last_lr() - lr = float(last[0]) if last else 0.0 + self._model_config.device = self.device # runtime device injection except Exception: - lr = 0.0 - return RoleStepResult( - metrics={"grad_norm": grad_norm, "lr": lr}, - grad_norm=grad_norm, - lr=lr, - ) - - @distributed(dispatch_mode=Dispatch.BROADCAST) - def save_checkpoint(self, path: str, step: Optional[int] = None, mode: str = "auto") -> None: - """Save backend checkpoint when the role backend supports it.""" - if hasattr(self, "backend") and hasattr(self.backend, "save"): - self.backend.save(path, step=step, mode=mode) - - @distributed(dispatch_mode=Dispatch.BROADCAST) - def load_checkpoint(self, path: str) -> int: - """Load backend checkpoint when the role backend supports it.""" - if hasattr(self, "backend") and hasattr(self.backend, "load"): - return int(self.backend.load(path) or 0) - return 0 - - @distributed - def generate_samples(self, req: RolloutReq) -> REFLGenerated: - """Run live-grad diffusion sampling and VAE decode.""" - stage = self.pipeline.diffusion - decode_stage = self.pipeline.vae_decode - if not hasattr(stage, "diffuse_with_grad"): - raise RuntimeError("ReflActorRole: pipeline.diffusion lacks diffuse_with_grad(...).") - if not hasattr(decode_stage, "decode_with_grad"): - raise RuntimeError("ReflActorRole: pipeline.vae_decode lacks decode_with_grad(...).") - - texts = req.primitives.get("text") if req.primitives else None - if not isinstance(texts, Texts): - raise TypeError( - f"ReflActorRole.generate_samples: req.primitives['text'] must be Texts, " - f"got {type(texts).__name__ if texts is not None else 'None'}" - ) - negatives_raw = req.primitives.get("negative_text") if req.primitives else None - negatives = negatives_raw if isinstance(negatives_raw, Texts) else None - if negatives is not None and len(negatives.texts) != len(texts.texts): - raise ValueError( - f"ReflActorRole.generate_samples: negative_text length {len(negatives.texts)} " - f"!= text length {len(texts.texts)}" - ) + pass - params = self.sampling_params - primary_g = float(getattr(params, "guidance_scale", 1.0)) - secondary_g = getattr(params, "guidance_scale_2", None) - effective_guidance = max(primary_g, float(secondary_g)) if secondary_g is not None else primary_g - conditions = self.pipeline.build_conditions(texts, negatives=negatives, guidance_scale=effective_guidance) + pipeline_cls = get_class(self._pipeline_target) + self.pipeline = pipeline_cls.from_config(self._model_config, strategy=self._strategy) - images_prim = req.primitives.get("image") if req.primitives else None - if images_prim is not None: - if not isinstance(images_prim, Images): + for stage_attr, method in (("diffusion", "diffuse_with_grad"), ("vae_decode", "decode_with_grad")): + stage = getattr(self.pipeline, stage_attr, None) + if not hasattr(stage, method): raise TypeError( - f"ReflActorRole.generate_samples: req.primitives['image'] must be Images, " - f"got {type(images_prim).__name__}" - ) - if int(images_prim.pixels.shape[0]) != len(texts.texts): - raise ValueError( - f"ReflActorRole.generate_samples: image count {images_prim.pixels.shape[0]} " - f"!= text count {len(texts.texts)}" + f"ReflActorRole: pipeline {self._pipeline_target} .{stage_attr} lacks {method}(...); " + f"use a recipes.refl.models pipeline (or implement the REFL contract)." ) - from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage - from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage - - image_latent_cond = WAN21ImageLatentEncodeStage( - self.pipeline.bundle, - num_frames=int(params.num_frames), - height=int(params.height), - width=int(params.width), - ).encode(images_prim) - image_embed_cond = ( - WAN21CLIPVisionEncodeStage(self.pipeline.bundle).encode(images_prim) - if getattr(self.pipeline.bundle, "uses_clip_vision", False) - else None + if not hasattr(self.pipeline, "build_refl_conditions"): + raise TypeError( + f"ReflActorRole: pipeline {self._pipeline_target} lacks build_refl_conditions(...); " + f"use a recipes.refl.models pipeline." ) - if image_latent_cond is not None or image_embed_cond is not None: - conditions = dataclasses.replace( - conditions, - image_latent=image_latent_cond, - image_embed=image_embed_cond, - ) - device = getattr(getattr(self.pipeline, "bundle", None), "device", None) - schedule = get_sigma_schedule( - int(params.num_inference_steps), - shift=float(getattr(self.pipeline, "shift", 5.0)), - device=device, + # FSDP-wrap pipeline.bundle.transformer in place + LoRA + optimizer. + # The pipeline's stages reference the same bundle, so sampling runs + # through 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( + "ReflActorRole initialized: pipeline=%s reward_weight=%.3f kl_weight=%.3f", + self._pipeline_target, + self.reward_weight, + self.kl_weight, ) - train_model = getattr(self.backend, "model", None) - if train_model is not None and hasattr(train_model, "train"): - train_model.train() + # ------------------------------------------------------------------ + # Grad chain (run under the driver's enable_grad() context) + # ------------------------------------------------------------------ + + @distributed(dispatch_mode=Dispatch.DP_SCATTER) + def generate_samples( + self, + *, + texts: Texts, + images: Optional[Images] = None, + params: DiffusionSamplingParams, + rollout_id: int = 0, + ) -> REFLGenerated: + """Grad-enabled BPTT sampling + in-graph VAE decode. + + ``texts`` / ``images`` are the data-source conditioning primitives + (``images`` is the I2V first frame; ``None`` for T2V). Negative + prompts ride ``params.sampler_kwargs['negative_prompt']`` and are + expanded by the pipeline's ``build_refl_conditions``. + """ + self.backend.model.train() self.backend.zero_grad() - if bool(getattr(params, "init_same_noise", False)) and not getattr(params, "noise_group_ids", None): - params = dataclasses.replace(params, noise_group_ids=list(req.group_ids)) + # Decorrelate init noise across DP shards and rollouts (same scheme as + # ReFLPolicy): the config seed is the base, not the per-step value. + dp_rank = int(self.rank_info.dp_rank) if self.rank_info is not None else 0 + base_seed = int(params.seed) if params.seed is not None else 42 + params = dataclasses.replace(params, seed=base_seed + 1000 * int(rollout_id) + dp_rank) - result = stage.diffuse_with_grad( - conditions, - schedule=schedule, - params=params, + conditions = self.pipeline.build_refl_conditions(texts, images=images, params=params) + schedule = get_sigma_schedule( + int(params.num_inference_steps), + shift=float(getattr(self.pipeline, "shift", 5.0)), + device=self.pipeline.bundle.device, ) - kl_loss = result.kl_loss - pixels = decode_stage.decode_with_grad(result.z_final) - return REFLGenerated(decoded=pixels, kl_loss=kl_loss.unsqueeze(0) if kl_loss.ndim == 0 else kl_loss) + result = self.pipeline.diffusion.diffuse_with_grad(conditions, schedule=schedule, params=params) + pixels = self.pipeline.vae_decode.decode_with_grad(result.z_final) + return REFLGenerated(decoded=pixels, kl_loss=result.kl_loss) - @distributed + @distributed(dispatch_mode=Dispatch.DP_SCATTER) def forward_backward_loss( self, *, rewards: torch.Tensor, kl_loss: Optional[torch.Tensor] = None, ) -> REFLLossMetrics: - """Assemble REFL loss and run backward on the actor graph.""" - algo = self.algo_cfg - reward_weight = float(algo.get("reward_weight", 1.0)) - reward_baseline = float(algo.get("reward_baseline", 0.0)) - reward_scale = float(algo.get("reward_scale", 1.0)) - kl_weight = float(algo.get("kl_weight", 0.0)) + """Assemble the REFL loss and backward on this shard's actor graph. + ``rewards`` and ``kl_loss`` are RPC inputs marked as grad leaves by the + framework; on ``enable_grad()`` exit their grads route back through the + reward / generate calls into the sampling graph in ONE backward pass. + """ reward = rewards.to(dtype=torch.bfloat16) - reward_loss = (-(reward - reward_baseline) / reward_scale * reward_weight).mean() - if kl_loss is not None and kl_weight != 0.0: - kl_term = kl_weight * kl_loss.squeeze() + reward_loss = (-(reward - self.reward_baseline) / self.reward_scale * self.reward_weight).mean() + if kl_loss is not None and self.kl_weight != 0.0: + kl_term = self.kl_weight * kl_loss.float().mean() else: kl_term = torch.zeros((), device=reward.device, dtype=reward_loss.dtype) loss = reward_loss + kl_term @@ -223,8 +210,35 @@ def forward_backward_loss( loss=[float(loss.detach().item())], reward_loss=[float(reward_loss.detach().item())], kl_loss=[float(kl_term.detach().item())], - reward_mean=[float(reward.detach().mean().item())], + reward_mean=[float(reward.detach().float().mean().item())], ) + # ------------------------------------------------------------------ + # Optimizer / checkpoint (delegate to the composed FSDPBackend) + # ------------------------------------------------------------------ + + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.ALL) + def step(self, *, max_grad_norm: float) -> Dict[str, float]: + """Clip + one optimizer step; returns grad_norm and current lr.""" + grad_norm = float(self.backend.optimizer_step(max_grad_norm=float(max_grad_norm))) + lr = 0.0 + sched = getattr(self.backend, "scheduler", None) + if sched is not None: + last = sched.get_last_lr() + lr = float(last[0]) if last else 0.0 + return {"grad_norm": grad_norm, "lr": lr} + + @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__ = ["ReflActorRole", "RoleStepResult"] +__all__ = ["ReflActorRole", "REFLGenerated", "REFLLossMetrics"] diff --git a/recipes/refl/trainer.py b/recipes/refl/trainer.py index 50b26fe2e..eb82a8fb9 100644 --- a/recipes/refl/trainer.py +++ b/recipes/refl/trainer.py @@ -1,82 +1,192 @@ -"""REFLTrainer — recipe-local trainer for WAN22 REFL/BPTT.""" +"""REFLTrainer — recipe driver for WAN REFL/BPTT (video reward backprop). + +The video sibling of :class:`unirl.trainer.refl.RewardBackpropTrainer`: two +roles, always — a :class:`recipes.refl.roles.ReflActorRole` (FSDP WAN + +grad BPTT sampling + optimizer) and a frozen differentiable video reward +(:class:`unirl.reward.service.RewardService`), colocated on the same worker +slab so decoded video never leaves the GPU. Each step runs, under the +distributed ``enable_grad()`` context:: + + gen = actor.generate_samples(texts, images, params) # grad through BPTT window + VAE + rew = reward.score_differentiable(gen.decoded, prompts, records) + actor.forward_backward_loss(rewards=rew, kl_loss=gen.kl_loss) + # ctx exit → grads route reward → decode → DiT LoRA params + actor.step(max_grad_norm) + +No advantages / replay / ratio / rollout-engine / weight-sync — REFL is not a +PG-RL loop. Success signal: the reward curve rises. +""" from __future__ import annotations +import logging import time -from collections.abc import Mapping -from typing import Any, Dict +from typing import Any, Dict, List, Optional import numpy as np +from hydra.utils import instantiate +from omegaconf import DictConfig -from recipes.common.trainer import Trainer +from unirl.distributed.group.placement import placement from unirl.distributed.tensor.grad_context import enable_grad -from unirl.types.prompts import RolloutInputs -from unirl.types.rollout_req import RolloutReq -from unirl.types.sampling import total_samples_per_prompt - - -class REFLTrainer(Trainer): - """REFL / BPTT recipe trainer: role-driven 3-RPC train step.""" - - def build_req(self, inputs: RolloutInputs, rollout_id: int) -> RolloutReq: - """Build one RolloutReq from data-source samples.""" - inputs = inputs.expand(total_samples_per_prompt(self.sampling_params)) - primitives: Dict[str, Any] = dict(inputs.primitives) - - diff_params = self.sampling_params.get("diffusion") - guidance_scale = float(getattr(diff_params, "guidance_scale", 1.0)) if diff_params is not None else 1.0 - sampler_kwargs = getattr(diff_params, "sampler_kwargs", {}) if diff_params is not None else {} - negative_prompt = sampler_kwargs.get("negative_prompt") if isinstance(sampler_kwargs, Mapping) else None - if negative_prompt is None and diff_params is not None: - negative_prompt = getattr(diff_params, "negative_prompt", None) - if negative_prompt is not None and guidance_scale > 1.0 and "negative_text" not in primitives: - texts = primitives.get("text") - if not hasattr(texts, "texts"): - raise TypeError( - "REFLTrainer.build_req: sampling negative_prompt requires " - "req.primitives['text'] to expose a `texts` field." - ) - text_list = getattr(texts, "texts") - text_cls: Any = type(texts) - primitives["negative_text"] = text_cls(texts=[str(negative_prompt)] * len(text_list)) - - return RolloutReq( - sample_ids=list(inputs.sample_ids), - group_ids=list(inputs.group_ids), - primitives=primitives, - request_conditions={}, - task_config={}, - sampling_params=dict(self.sampling_params), - metadata=list(inputs.metadata) if inputs.metadata else [], - init_noise_group_ids=[], - init_noise_latent_shape=None, +from unirl.trainer.base import BaseTrainer, build_sampling_dict +from unirl.types.primitives import Images, Texts +from unirl.types.sample import Sample +from unirl.types.sampling import DiffusionSamplingParams, total_samples_per_prompt +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``.""" + prompts = inputs.parts[0].primitives.get("text") if inputs.parts else None + if not isinstance(prompts, Texts): + raise TypeError(f"REFL data-source root requires Texts, got {type(prompts).__name__}.") + return prompts + + +def _image_inputs(inputs: Sample) -> Optional[Images]: + """Read the optional I2V first-frame conditioning (a chained input Part).""" + for part in inputs.parts[1:]: + image = part.primitives.get("image") + if isinstance(image, Images): + return image + return None + + +def _records(inputs: Sample) -> Optional[List[Optional[Dict[str, Any]]]]: + """Per-sample data-source metadata (e.g. ``ref_video_path`` for Face reward).""" + metadata = inputs.parts[0].metadata if inputs.parts else None + return list(metadata) if metadata else None + + +class REFLTrainer(BaseTrainer): + """REFL/BPTT recipe driver: actor + differentiable reward, 3-RPC step.""" + + def __init__(self, *, cfg: DictConfig) -> None: + super().__init__(cfg=cfg, logging_cfg=cfg.get("logging")) + self.batch_size = int(cfg.batch_size) + self.max_grad_norm = float(cfg.get("max_grad_norm", 1.0)) + self.data_source = instantiate(cfg.data_source) + + sampling = build_sampling_dict(cfg.sampling) + params = sampling.get("diffusion") + if not isinstance(params, DiffusionSamplingParams): + raise TypeError("REFLTrainer requires cfg.sampling to build DiffusionSamplingParams.") + if total_samples_per_prompt(sampling) != 1: + raise ValueError( + "REFLTrainer trains one sample per prompt (no advantage groups); set sampling.samples_per_prompt=1." + ) + self.sampling_params = params + + # Reward shares the actor's workers: decoded video stays on-GPU for + # scoring. (Cross-slab reward placement — DiffusionTrainer's + # reward_fraction — is deliberately not supported here; add it only + # when a recipe actually cannot colocate.) + with placement(self.pool, fraction=1.0, shared_workers=True): + self.actor = remote_hydra(cfg.actor) + self.reward = remote_hydra(cfg.reward) + self.actor.initialize() + # BaseTrainer.maybe_save/load_checkpoint operate on ``self.backend``. + self.backend = self.actor + + adp, rdp = self.actor.dp_size, self.reward.dp_size + if self.batch_size % adp or self.batch_size % rdp: + raise ValueError(f"batch_size={self.batch_size} must be divisible by actor dp={adp} and reward dp={rdp}") + logger.info( + "REFLTrainer ready: actor dp=%d reward dp=%d batch=%d max_grad_norm=%.2f", + adp, + rdp, + self.batch_size, + self.max_grad_norm, ) - def train_step(self, req: RolloutReq, *, training_progress: float = 0.0, rollout_id: int = 0) -> Dict[str, Any]: - """One REFL step: actor generate → reward score → actor backward → actor step.""" + def train_step(self, inputs: Sample, *, rollout_id: int) -> Dict[str, float]: + """One enable_grad() generate → score → backward, then optimizer step.""" t0 = time.perf_counter() - prompts = list(req.primitives["text"].texts) - records = list(req.metadata) if req.metadata else None + texts = _text_inputs(inputs) + images = _image_inputs(inputs) + records = _records(inputs) with enable_grad(): - gen = self.actor.generate_samples(req) - rewards = self.reward.score_differentiable(gen.decoded, prompts, records) - loss_metrics = self.actor.forward_backward_loss( - rewards=rewards, - kl_loss=gen.kl_loss, + gen = self.actor.generate_samples( + texts=texts, + images=images, + params=self.sampling_params, + rollout_id=rollout_id, ) - step_result = self.actor.step() - - metrics: Dict[str, Any] = { - "loss": np.mean(loss_metrics.loss), - "reward_loss": np.mean(loss_metrics.reward_loss), - "kl_loss": np.mean(loss_metrics.kl_loss), - "reward_mean": np.mean(loss_metrics.reward_mean), - "grad_norm": np.mean(step_result.metrics.get("grad_norm")) if step_result.metrics else 0.0, - "lr": np.mean(step_result.metrics.get("lr")) if step_result.metrics else 0.0, + rewards = self.reward.score_differentiable(gen.decoded, list(texts.texts), records) + loss_metrics = self.actor.forward_backward_loss(rewards=rewards, kl_loss=gen.kl_loss) + step_result = self.actor.step(max_grad_norm=self.max_grad_norm) + if isinstance(step_result, list): # BROADCAST → one result per worker + step_result = step_result[0] + + return { + "loss": float(np.mean(loss_metrics.loss)), + "reward_loss": float(np.mean(loss_metrics.reward_loss)), + "kl_loss": float(np.mean(loss_metrics.kl_loss)), + "reward_mean": float(np.mean(loss_metrics.reward_mean)), + "grad_norm": float(step_result.get("grad_norm", 0.0)), + "lr": float(step_result.get("lr", 0.0)), "step_time_s": time.perf_counter() - t0, - "training_progress": float(training_progress), } - return metrics + + def train( + self, + *, + num_rollouts: Optional[int] = None, + save_interval: Optional[int] = None, + save_dir: Optional[str] = None, + load_dir: Optional[str] = None, + save_mode: Optional[str] = None, + ) -> None: + cfg = self.cfg + num_rollouts = int(num_rollouts if num_rollouts is not None else cfg.get("num_rollouts", 100)) + save_interval = int(save_interval if save_interval is not None else cfg.get("save_interval", 0)) + save_dir = save_dir if save_dir is not None else cfg.get("save_dir") + load_dir = load_dir if load_dir is not None else cfg.get("load_dir") + save_mode = str(save_mode if save_mode is not None else cfg.get("save_mode", "adapter")) + + start = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) + for _ in range(start): # fast-forward the data stream on resume + self.data_source.get_samples(self.batch_size) + self._init_wandb(num_rollouts=num_rollouts) + try: + for rollout_id in range(start, num_rollouts): + inputs = self.data_source.get_samples(self.batch_size) + metrics = self.train_step(inputs, rollout_id=rollout_id) + logger.info( + "rollout %d/%d reward=%.4f loss=%.4f kl=%.4f grad_norm=%.4f %.1fs", + rollout_id + 1, + num_rollouts, + metrics["reward_mean"], + metrics["loss"], + metrics["kl_loss"], + metrics["grad_norm"], + metrics["step_time_s"], + ) + self.wandb_logger.log_step( + rollout_id + 1, + { + "rollout/mean_reward": metrics["reward_mean"], + "train/loss": metrics["loss"], + "train/reward_loss": metrics["reward_loss"], + "train/kl_loss": metrics["kl_loss"], + "train/grad_norm": metrics["grad_norm"], + "train/lr": metrics["lr"], + "perf/step_time_s": metrics["step_time_s"], + }, + prefix="", + ) + 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__ = ["REFLTrainer"] diff --git a/unirl/models/types/diffusion.py b/unirl/models/types/diffusion.py index e632ecc0d..be69f493f 100644 --- a/unirl/models/types/diffusion.py +++ b/unirl/models/types/diffusion.py @@ -170,7 +170,12 @@ def diffuse_with_grad( @dataclass class DiffuseWithGradResult: - """Output of :meth:`DiffusionStage.diffuse_with_grad`.""" + """Output of :meth:`DiffusionStage.diffuse_with_grad`. + + ``kl_loss`` is per-sample ``[B]`` (zeros when the KL branch is off) so + DP-scattered consumers round-trip each shard's own KL, never a + cross-shard aggregate. + """ z_final: torch.Tensor kl_loss: torch.Tensor diff --git a/unirl/reward/service.py b/unirl/reward/service.py index 291562a34..7893293de 100644 --- a/unirl/reward/service.py +++ b/unirl/reward/service.py @@ -10,7 +10,7 @@ from __future__ import annotations import logging -from typing import Dict, Optional +from typing import Dict, List, Optional import torch From ec561c7a70d712ff132a3bad210401c8c7f2a192 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 15:26:33 +0800 Subject: [PATCH 11/24] test(recipes/refl): verify KL DP-batching semantics across topologies Regression verification for the P1 review finding: the original per-shard scalar shared_field KL collapsed to rank 0 on DP collect and only lined up on the batch_size == actor_dp == 8 topology. kl_loss is now a per-sample [B] concat column; scripts/verify_refl_kl_batching.py pins chunk/cat round-trips, rewards/KL payload lockstep, per-shard backward grad shapes, unequal actor/reward dp re-chunking, and numeric equivalence with the legacy scalar mean at the pytree wire layer (CPU, no Ray). Shipped as a standalone runnable script rather than a tests/ tree per the #99/#267 no-unenforced-test-suite policy. --- scripts/verify_refl_kl_batching.py | 166 +++++++++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100755 scripts/verify_refl_kl_batching.py diff --git a/scripts/verify_refl_kl_batching.py b/scripts/verify_refl_kl_batching.py new file mode 100755 index 000000000..69f62d764 --- /dev/null +++ b/scripts/verify_refl_kl_batching.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Verify KL DP-batching semantics for the refl recipe (CPU, no Ray, no GPU). + +Regression verification for the P1 review finding on PR #210: the original +``REFLGenerated.kl_loss`` was a per-shard *scalar* ``shared_field`` — DP +collect kept only rank 0's KL and re-broadcast it to every actor rank, which +happened to "work" only on the verified ``batch_size == actor_dp == 8`` +topology (logs duplicated rank 0; other B/dp splits risked a hard shape +mismatch between the routed KL grad and each rank's saved scalar). + +The fix makes ``kl_loss`` a batch-aligned per-sample ``[B]`` concat field. +This script pins the invariants at the exact wire layer DP dispatch uses +(``pytree_chunk`` / ``pytree_cat`` / ``infer_batch_size``) across the +topologies called out in review: B == dp, B > dp, non-power-of-two, dp == 1, +and unequal actor/reward dp. + +Standalone by design (repo policy after #99/#267 is no unenforced test tree): +run it directly whenever the refl recipe's KL/reward wire types change:: + + python scripts/verify_refl_kl_batching.py +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from recipes.refl.roles import REFLGenerated, REFLLossMetrics # noqa: E402 +from unirl.distributed.tensor.pytree import infer_batch_size, pytree_cat, pytree_chunk # noqa: E402 + +# (batch_size, dp) — B == dp (the only previously-verified shape), B > dp, +# non-power-of-two, and the degenerate dp == 1. +TOPOLOGIES = [(8, 8), (8, 2), (6, 3), (12, 4), (4, 1)] + + +def _generated(batch: int) -> REFLGenerated: + decoded = torch.arange(batch * 6, dtype=torch.float32).reshape(batch, 2, 3) + kl = torch.arange(batch, dtype=torch.float32) + 1.0 # distinct per sample + return REFLGenerated(decoded=decoded, kl_loss=kl) + + +def check_refl_generated_chunk_cat_roundtrip(batch: int, dp: int) -> None: + """Each DP shard sees exactly its own KL rows; merge restores the batch.""" + gen = _generated(batch) + assert infer_batch_size((gen,), {}) == batch + + shards = pytree_chunk(gen, dp, batch) + assert len(shards) == dp + per = batch // dp + for rank, shard in enumerate(shards): + expect = gen.kl_loss[rank * per : (rank + 1) * per] + assert torch.equal(shard.kl_loss, expect), f"rank {rank} got foreign KL rows" + assert torch.equal(shard.decoded, gen.decoded[rank * per : (rank + 1) * per]) + + merged = pytree_cat(shards) + assert torch.equal(merged.kl_loss, gen.kl_loss) + assert torch.equal(merged.decoded, gen.decoded) + + +def check_forward_backward_payload_alignment(batch: int, dp: int) -> None: + """rewards and kl_loss chunk in lockstep — the forward_backward_loss wire. + + The old scalar-shared KL made this payload rewards=[B] + kl=[1]; chunking + by the inferred batch could not keep the two aligned off the B == dp + topology. Per-sample KL makes both first-class batch columns. + """ + kwargs = { + "rewards": torch.randn(batch), + "kl_loss": torch.arange(batch, dtype=torch.float32), + } + assert infer_batch_size((), kwargs) == batch + shards = pytree_chunk(kwargs, dp, batch) + per = batch // dp + for rank, shard in enumerate(shards): + assert shard["rewards"].shape == (per,) + assert shard["kl_loss"].shape == (per,) + assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) + + +def check_per_shard_backward_grad_shape(batch: int, dp: int) -> None: + """The KL grad each rank produces matches its saved generate output rows. + + Mirrors ReflActorRole.forward_backward_loss on one shard: the KL input is + a grad leaf of shape [B/dp]; after backward its .grad must be the same + shape, because GradContext routes it as out_grads onto the SAME rank's + saved kl tensor from generate_samples. With the old scalar KL this pairing + was [broadcast scalar] vs [rank-local scalar] and only lined up by luck. + """ + per = batch // dp + for _rank in range(dp): + rewards = torch.randn(per, requires_grad=True) + kl = torch.rand(per, requires_grad=True) + reward_loss = (-(rewards.to(torch.bfloat16) - 0.5) / 0.25 * 1.0).mean() + loss = reward_loss + 1.0 * kl.float().mean() + loss.backward() + assert kl.grad is not None and kl.grad.shape == kl.shape + assert rewards.grad is not None and rewards.grad.shape == rewards.shape + + +def check_unequal_actor_reward_dp_stays_aligned() -> None: + """B-length columns survive reward-dp merge → actor-dp re-chunk. + + The recipe colocates actor and reward (equal dp by construction), but the + wire contract must not depend on that: scoring merged at reward dp=2 and + re-scattered at actor dp=4 must hand every actor rank the reward/KL rows + of its own samples. + """ + batch, rdp, adp = 8, 2, 4 + rewards = torch.arange(batch, dtype=torch.float32) + reward_shards = pytree_chunk({"r": rewards}, rdp, batch) + merged = pytree_cat(reward_shards)["r"] + assert torch.equal(merged, rewards) + + kwargs = {"rewards": merged, "kl_loss": torch.arange(batch, dtype=torch.float32) * 10.0} + actor_shards = pytree_chunk(kwargs, adp, batch) + per = batch // adp + for rank, shard in enumerate(actor_shards): + assert torch.equal(shard["rewards"], rewards[rank * per : (rank + 1) * per]) + assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) + + +def check_per_sample_kl_equals_legacy_scalar_mean() -> None: + """Per-sample reduction then batch-mean == the legacy global scalar mean.""" + torch.manual_seed(0) + kl_pred = torch.randn(3, 4, 2, 5, 5) + ref = torch.randn(3, 4, 2, 5, 5) + sigma = torch.tensor(0.7) + per_sample = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).flatten(1).mean(dim=1) + legacy = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).mean() + assert per_sample.shape == (3,) + assert torch.allclose(per_sample.mean(), legacy, atol=1e-6) + + +def check_loss_metrics_concat_keeps_every_shard() -> None: + """REFLLossMetrics concat lists must surface every shard's scalars.""" + shards = [ + REFLLossMetrics(loss=[0.1], reward_loss=[0.2], kl_loss=[0.3], reward_mean=[0.4]), + REFLLossMetrics(loss=[1.1], reward_loss=[1.2], kl_loss=[1.3], reward_mean=[1.4]), + ] + merged = pytree_cat(shards) + assert merged.loss == [0.1, 1.1] + assert merged.kl_loss == [0.3, 1.3] + + +def main() -> int: + for batch, dp in TOPOLOGIES: + check_refl_generated_chunk_cat_roundtrip(batch, dp) + check_forward_backward_payload_alignment(batch, dp) + check_per_shard_backward_grad_shape(batch, dp) + print(f"[ok] topology B={batch} dp={dp}: chunk/cat, payload lockstep, backward shapes") + check_unequal_actor_reward_dp_stays_aligned() + print("[ok] unequal actor/reward dp (rdp=2 → adp=4) stays row-aligned") + check_per_sample_kl_equals_legacy_scalar_mean() + print("[ok] per-sample KL batch-mean equals legacy scalar mean") + check_loss_metrics_concat_keeps_every_shard() + print("[ok] REFLLossMetrics keeps every shard's scalars") + print("verify-refl-kl-batching: ALL CHECKS PASSED") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From e0497399159e2c5f51f67b709c08a6ef2aa3eff1 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 16:46:54 +0800 Subject: [PATCH 12/24] =?UTF-8?q?refactor(recipes/refl):=20close=20review?= =?UTF-8?q?=20blockers=20=E2=80=94=20deps=20contract,=20Hub=20VAE=20load,?= =?UTF-8?q?=20recipe-local=20BPTT=20contract?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review response (second pass): - VideoAlign now runs on the UniRL core stack. Colocated reward+actor share one Python process, so a recipe requirements file can only ADD packages, never re-pin the core stack; the old file pinned transformers==4.45.2 / peft==0.10.0 / flash-attn==2.5.8, which would downgrade the core env (core: transformers>=5.6, peft>=0.14). Enable the wrapper's documented 5.x shim (explicit mm_token_type_ids black-list pop — never signature filtering, which PEFT breaks) and fall back to SDPA when flash-attn is absent. requirements.txt is additive-only (torchvision). - WanVideoVAE keeps loading Hub repo ids: load_from_diffusers reads local files only, so WAN21Bundle resolves non-local paths through the HF cache (snapshot_download) first — mainline repo-id configs (the wan21_t2v.yaml default) keep working. The decode-implementation swap itself stays in core deliberately: memory/numeric VAE optimizations are model assets shared by GRPO and ReFL, not algorithm property. - BPTT contract is recipe-local: concrete stages inherit the DiffusionStage Protocol explicitly, so a protocol-level stub becomes a real None-returning method on every stage and defeats hasattr capability checks. DiffuseWithGradResult moves to recipes/refl/models/types.py; promote later as a separate opt-in protocol (the DifferentiableReward idiom) once a second out-of-recipe BPTT consumer exists. - Single KL knob: actor.kl_weight owns on/off + weight; the actor injects it into sampler_kwargs for the stage, and a stale sampling.sampler_kwargs.kl_weight now raises instead of being silently overridden. Sampling keeps only sampling-shape knobs (mid/final window). - Hygiene: validate 0 <= mid_timestep <= final_timestep < T at diffuse time; launch scripts get a shebang + set -euo pipefail; the Face reference-embedding cache becomes a bounded LRU (64). --- .../configs/wan21_t2v_videoalign_refl.yaml | 5 +- recipes/refl/configs/wan22_i2v_face_refl.yaml | 7 +- recipes/refl/models/types.py | 40 ++++++++++ recipes/refl/models/wan21.py | 8 +- recipes/refl/models/wan22.py | 8 +- recipes/refl/rewards/face/scorer.py | 11 ++- .../refl/rewards/videoalign/requirements.txt | 20 +++-- recipes/refl/rewards/videoalign/wrapper.py | 75 ++++++------------- recipes/refl/roles.py | 17 ++++- recipes/refl/scripts/start_wan21_t2v.sh | 3 +- recipes/refl/scripts/start_wan22_i2v.sh | 3 +- unirl/models/types/diffusion.py | 31 +------- unirl/models/wan21/bundle.py | 13 +++- 13 files changed, 132 insertions(+), 109 deletions(-) create mode 100644 recipes/refl/models/types.py diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 2ba89ecba..3aa8cc7b9 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -134,9 +134,8 @@ sampling: # the final step traverses the reward backward pass (DRaFT-1 / ReFL). mid_timestep: 24 final_timestep: 24 - # KL branch off — actor.kl_weight is also 0. Single-pass forward per - # step, no LoRA-disabled reference recompute. - kl_weight: 0.0 + # KL on/off + weight live on actor.kl_weight (single knob; 0 here → + # single-pass forward per step, no LoRA-disabled reference recompute). logging: report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index 72d420086..b4f39736e 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -120,12 +120,11 @@ sampling: init_same_noise: false sampler_kwargs: # Truncated BPTT window: grad on from mid_timestep..final_timestep — the - # last 4 of 8 steps, i.e. the low-noise DiT's sigma range. + # last 4 of 8 steps, i.e. the low-noise DiT's sigma range. The per-step + # KL against the LoRA-disabled reference is switched/weighted by the + # single actor.kl_weight knob. mid_timestep: 4 final_timestep: 7 - # Per-step KL against the LoRA-disabled reference (weighted again by - # actor.kl_weight at the loss site). - kl_weight: 1.0 logging: report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}} diff --git a/recipes/refl/models/types.py b/recipes/refl/models/types.py new file mode 100644 index 000000000..1ba537877 --- /dev/null +++ b/recipes/refl/models/types.py @@ -0,0 +1,40 @@ +"""Recipe-local BPTT stage contract for the refl recipe. + +``diffuse_with_grad`` is deliberately NOT part of the core +:class:`~unirl.models.types.diffusion.DiffusionStage` protocol: concrete +stages inherit that Protocol *explicitly*, so a protocol-level stub would +become a real ``None``-returning method on every diffusion stage in the +repo and make ``hasattr``-based capability checks meaningless. While REFL +is the only BPTT consumer, the contract lives here; if a second consumer +appears outside ``recipes/refl``, promote it to core as a separate opt-in +``@runtime_checkable`` protocol (the ``DifferentiableReward`` / +``LatentShapeProvider`` idiom), not as a method on ``DiffusionStage``. + +Implementors (``Wan21ReflDiffusionStage`` / ``Wan22ReflDiffusionStage``) +provide:: + + diffuse_with_grad(conditions, *, schedule, params, initial_latents=None) + -> DiffuseWithGradResult +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class DiffuseWithGradResult: + """Output of a recipe stage's ``diffuse_with_grad``. + + ``kl_loss`` is per-sample ``[B]`` (zeros when the KL branch is off) so + DP-scattered consumers round-trip each shard's own KL, never a + cross-shard aggregate. + """ + + z_final: torch.Tensor + kl_loss: torch.Tensor + + +__all__ = ["DiffuseWithGradResult"] diff --git a/recipes/refl/models/wan21.py b/recipes/refl/models/wan21.py index 43d061ed6..f0d77aaec 100644 --- a/recipes/refl/models/wan21.py +++ b/recipes/refl/models/wan21.py @@ -25,7 +25,7 @@ import torch -from unirl.models.types.diffusion import DiffuseWithGradResult +from recipes.refl.models.types import DiffuseWithGradResult from unirl.models.wan21.bundle import WAN21Bundle from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions @@ -231,6 +231,12 @@ def diffuse_with_grad( 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"Wan21ReflDiffusionStage.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) diff --git a/recipes/refl/models/wan22.py b/recipes/refl/models/wan22.py index 60ceadb3e..b2a6638bb 100644 --- a/recipes/refl/models/wan22.py +++ b/recipes/refl/models/wan22.py @@ -8,7 +8,7 @@ import torch -from unirl.models.types.diffusion import DiffuseWithGradResult +from recipes.refl.models.types import DiffuseWithGradResult from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage @@ -199,6 +199,12 @@ def diffuse_with_grad( 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"Wan22ReflDiffusionStage.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) diff --git a/recipes/refl/rewards/face/scorer.py b/recipes/refl/rewards/face/scorer.py index b9b04c672..c1e28a285 100644 --- a/recipes/refl/rewards/face/scorer.py +++ b/recipes/refl/rewards/face/scorer.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections import OrderedDict from dataclasses import dataclass from typing import List, Optional @@ -17,6 +18,9 @@ logger = logging.getLogger(__name__) +# Reference-embedding LRU bound (see FaceRewardScorer._ref_cache). +_REF_CACHE_MAX = 64 + # --------------------------------------------------------------------------- # Reference-video loader (imageio + resize-to-cover + center_crop) @@ -106,7 +110,9 @@ def _load_model(self) -> None: self._ref_max_frames = int(self.model_kwargs.get("ref_max_frames", 81)) self._ref_max_pixels = int(self.model_kwargs.get("ref_max_pixels", 480 * 480)) self._differentiable = bool(self.model_kwargs.get("differentiable", True)) - self._ref_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + # Bounded LRU: reference embeddings are small, but a many-identity + # dataset must not grow GPU residency without bound. + self._ref_cache: OrderedDict[str, tuple[torch.Tensor, torch.Tensor]] = OrderedDict() def _compute_model_rewards(self, request: RewardRequest) -> List[float]: raise NotImplementedError("FaceRewardScorer is REFL-only; use compute_rewards_differentiable().") @@ -182,6 +188,7 @@ def _extract_face_embeddings( def _get_ref_embeddings(self, ref_video_path: str) -> tuple[torch.Tensor, torch.Tensor]: cached = self._ref_cache.get(ref_video_path) if cached is not None: + self._ref_cache.move_to_end(ref_video_path) return cached with torch.no_grad(): video = _load_ref_video_frames( @@ -193,6 +200,8 @@ def _get_ref_embeddings(self, ref_video_path: str) -> tuple[torch.Tensor, torch. # Detach (defence-in-depth) before caching. cached = (ref_emb.detach(), ref_mask.detach()) self._ref_cache[ref_video_path] = cached + while len(self._ref_cache) > _REF_CACHE_MAX: + self._ref_cache.popitem(last=False) return cached # ------------------------------------------------------------------ diff --git a/recipes/refl/rewards/videoalign/requirements.txt b/recipes/refl/rewards/videoalign/requirements.txt index 458920c6b..83b2e094b 100644 --- a/recipes/refl/rewards/videoalign/requirements.txt +++ b/recipes/refl/rewards/videoalign/requirements.txt @@ -1,11 +1,9 @@ -accelerate==0.34.0 -decord==0.6.0 -einops -flash-attn==2.5.8 -huggingface_hub -opencv-python-headless -peft==0.10.0 -safetensors -setuptools<70 -transformers==4.45.2 -trl==0.8.6 +# VideoAlign reward extras — ADDITIVE ONLY on top of the UniRL core stack. +# +# The reward runs in the SAME Python process as the actor (colocated Remote +# siblings), so a requirements file cannot "isolate" versions — any pin that +# overlaps the core stack would downgrade the core environment itself. Core +# already provides transformers / peft / safetensors / huggingface-hub / +# einops (see pyproject.toml); never re-pin those here. flash-attn is +# optional: the model factory falls back to SDPA when the wheel is absent. +torchvision # frame preprocessing transforms; matches the engine extra's torch build diff --git a/recipes/refl/rewards/videoalign/wrapper.py b/recipes/refl/rewards/videoalign/wrapper.py index fe7bc2f05..d80f2fa60 100644 --- a/recipes/refl/rewards/videoalign/wrapper.py +++ b/recipes/refl/rewards/videoalign/wrapper.py @@ -25,16 +25,15 @@ from __future__ import annotations -# NOTE: ``inspect`` was previously used to filter ``batch`` against the -# reward backbone's ``forward`` signature in order to drop the -# ``mm_token_type_ids`` kwarg that transformers>=4.58/5.x injects via -# ``Qwen2VLProcessor``. We have aligned the runtime back to transformers -# 4.54 (matching the mmrl baseline), where that kwarg is not produced, so -# the filter is no longer necessary — and under PEFT it can spuriously -# strip ``pixel_values_videos`` / ``video_grid_thw`` when ``base.forward`` -# turns out to be ``LoraModel.forward(*args, **kwargs)``. The import is -# kept commented for future re-enablement once we move past 4.58. -# import inspect +# NOTE: this wrapper runs on the UniRL core stack (transformers>=5.6 — the +# reward shares one Python process with the actor, so there is no separate +# environment to pin). Newer ``Qwen2VLProcessor`` versions inject +# ``mm_token_type_ids``, which the 4.45-era reward backbone does not accept; +# ``compute_scores`` drops it by explicit black-list. Signature-based +# filtering (``inspect``) is deliberately NOT used: under PEFT it can resolve +# to ``LoraModel.forward(*args, **kwargs)`` and silently strip +# ``pixel_values_videos`` / ``video_grid_thw``. +import importlib.util import json import logging import os @@ -141,7 +140,9 @@ def __init__( load_from_pretrained=checkpoint_dir, load_from_pretrained_step=-1, gradient_checkpointing=False, - disable_flash_attn2=False, + # flash-attn 2 is not part of the UniRL core stack; fall back to + # SDPA automatically when the wheel is absent. + disable_flash_attn2=importlib.util.find_spec("flash_attn") is None, bf16=(dtype == torch.bfloat16), fp16=(dtype == torch.float16), output_dir="", @@ -313,50 +314,16 @@ def forward_scores( for start in range(0, len(video_tensors), self.micro_batch_size): end = start + self.micro_batch_size batch = self.prepare_batch_from_frames(video_tensors[start:end], prompts[start:end]) - # Aligned with mmrl's wrapper: pass the processor batch - # straight through to the model. transformers 4.54 (current - # pinned version) does not inject ``mm_token_type_ids``, so - # no filtering is required. + # transformers>=4.58/5.x ``Qwen2VLProcessor`` injects + # ``mm_token_type_ids`` (text/image/video ids for the rewritten + # ``get_rope_index``); ``Qwen2VLRewardModelBT.forward`` was + # authored against 4.45 and does not accept it. Drop it by + # explicit black-list — NEVER ``inspect.signature`` filtering, + # which under PEFT can resolve to ``LoraModel.forward(*args, + # **kwargs)`` and silently strip ``pixel_values_videos`` / + # ``video_grid_thw``, leaving the reward blind to the video. + batch.pop("mm_token_type_ids", None) logits = self.model(**batch, return_dict=True)["logits"] # (B, 3) - - # ---------------------------------------------------------- - # transformers>=4.58/5.x compatibility shim (DISABLED). - # - # Newer ``Qwen2VLProcessor`` versions inject ``mm_token_type_ids`` - # (text=0 / image=1 / video=2) into the batch dict for the - # rewritten ``get_rope_index``. ``Qwen2VLRewardModelBT`` was - # authored against transformers 4.45 and its ``forward`` does - # NOT list this kwarg, so ``self.model(**batch)`` would blow - # up with ``TypeError: ... got an unexpected keyword argument - # 'mm_token_type_ids'`` once upgraded. - # - # The previous implementation used ``inspect.signature`` on - # the unwrapped base forward to filter the batch, but PEFT - # wraps the model as ``PeftModel -> LoraModel -> Qwen2VLRewardModelBT`` - # and depending on the PEFT version ``get_base_model()`` may - # return ``LoraModel`` whose forward signature is - # ``(*args, **kwargs)`` — the resulting filter would silently - # drop ``pixel_values_videos`` / ``video_grid_thw``, leaving - # the reward model "blind" to the video. - # - # When we re-upgrade past 4.58, the safer replacement is an - # explicit black-list (NOT signature inspection): - # - # _DROP = {"mm_token_type_ids"} - # filtered_batch = {k: v for k, v in batch.items() if k not in _DROP} - # logits = self.model(**filtered_batch, return_dict=True)["logits"] - # - # Original code preserved below for reference: - # - # base = ( - # self.model.get_base_model() - # if hasattr(self.model, "get_base_model") - # else self.model - # ) - # allowed = set(inspect.signature(base.forward).parameters) - # filtered_batch = {k: v for k, v in batch.items() if k in allowed} - # logits = self.model(**filtered_batch, return_dict=True)["logits"] - # ---------------------------------------------------------- vq, mq, ta = logits[:, 0], logits[:, 1], logits[:, 2] if use_norm: vq, mq, ta = self._norm(vq, mq, ta) diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index 97cdbb742..b02ae6039 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -169,11 +169,26 @@ def generate_samples( self.backend.model.train() self.backend.zero_grad() + # Single KL knob: the actor owns kl_weight (loss-side), the sampling + # config owns only sampling-shape knobs (mid/final window). A stale + # sampler_kwargs.kl_weight is an error, not a silent override. + sampler_kwargs = dict(params.sampler_kwargs or {}) + if "kl_weight" in sampler_kwargs: + raise ValueError( + "ReflActorRole: sampling.sampler_kwargs.kl_weight is no longer read; " + "configure actor.kl_weight instead (single source of truth)." + ) + sampler_kwargs["kl_weight"] = self.kl_weight + # Decorrelate init noise across DP shards and rollouts (same scheme as # ReFLPolicy): the config seed is the base, not the per-step value. dp_rank = int(self.rank_info.dp_rank) if self.rank_info is not None else 0 base_seed = int(params.seed) if params.seed is not None else 42 - params = dataclasses.replace(params, seed=base_seed + 1000 * int(rollout_id) + dp_rank) + params = dataclasses.replace( + params, + seed=base_seed + 1000 * int(rollout_id) + dp_rank, + sampler_kwargs=sampler_kwargs, + ) conditions = self.pipeline.build_refl_conditions(texts, images=images, params=params) schedule = get_sigma_schedule( diff --git a/recipes/refl/scripts/start_wan21_t2v.sh b/recipes/refl/scripts/start_wan21_t2v.sh index 96046457e..84fffb826 100755 --- a/recipes/refl/scripts/start_wan21_t2v.sh +++ b/recipes/refl/scripts/start_wan21_t2v.sh @@ -1,4 +1,5 @@ -set -u +#!/usr/bin/env bash +set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "${SCRIPT_DIR}/../../.." diff --git a/recipes/refl/scripts/start_wan22_i2v.sh b/recipes/refl/scripts/start_wan22_i2v.sh index 9e581b180..00912d188 100755 --- a/recipes/refl/scripts/start_wan22_i2v.sh +++ b/recipes/refl/scripts/start_wan22_i2v.sh @@ -1,4 +1,5 @@ -set -u +#!/usr/bin/env bash +set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "${SCRIPT_DIR}/../../.." diff --git a/unirl/models/types/diffusion.py b/unirl/models/types/diffusion.py index be69f493f..114978141 100644 --- a/unirl/models/types/diffusion.py +++ b/unirl/models/types/diffusion.py @@ -14,7 +14,6 @@ from __future__ import annotations -from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional, Protocol, Tuple, TypeVar, runtime_checkable import torch @@ -152,33 +151,5 @@ def predict_noise_at_step( """ ... - # ------------------------------------------------------------------ - # BPTT path (REFL): generate-and-train in a single forward. - # ------------------------------------------------------------------ - def diffuse_with_grad( - self, - conditions: C, - *, - schedule: torch.Tensor, - params: object, - initial_latents: Optional[torch.Tensor] = None, - ) -> "DiffuseWithGradResult": - """Differentiable sampling: ``C → (z_final, kl_loss)``.""" - ... - - -@dataclass -class DiffuseWithGradResult: - """Output of :meth:`DiffusionStage.diffuse_with_grad`. - - ``kl_loss`` is per-sample ``[B]`` (zeros when the KL branch is off) so - DP-scattered consumers round-trip each shard's own KL, never a - cross-shard aggregate. - """ - - z_final: torch.Tensor - kl_loss: torch.Tensor - - -__all__ = ["DiffusionStage", "DiffusionStep", "DiffuseWithGradResult", "ReplayResult"] +__all__ = ["DiffusionStage", "DiffusionStep", "ReplayResult"] diff --git a/unirl/models/wan21/bundle.py b/unirl/models/wan21/bundle.py index d7520cd39..a5c157228 100644 --- a/unirl/models/wan21/bundle.py +++ b/unirl/models/wan21/bundle.py @@ -137,9 +137,20 @@ def from_config(cls, config: WAN21PipelineConfig) -> "WAN21Bundle": if config.load_vae: from .wan_video_vae import WanVideoVAE + # ``load_from_diffusers`` reads local files only. Hub repo ids + # (e.g. ``Wan-AI/Wan2.1-T2V-1.3B-Diffusers``, the mainline WAN + # config default) resolve through the HF cache first, preserving + # the loading semantics of the previous + # ``AutoencoderKLWan.from_pretrained`` path. + vae_src = vae_path + if not os.path.isdir(vae_src): + from huggingface_hub import snapshot_download + + vae_src = snapshot_download(repo_id=vae_src, allow_patterns=["vae/*"]) + vae = ( WanVideoVAE.load_from_diffusers( - vae_path, + vae_src, use_nested_grad_checkpoint=True, use_act_grad_only_conv=True, ) From d96800dcef2ea35e44dec7c5e1ee417f5d0acf1a Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 17:11:17 +0800 Subject: [PATCH 13/24] =?UTF-8?q?fix(recipes/refl):=20store=20cfg=20on=20R?= =?UTF-8?q?EFLTrainer=20=E2=80=94=20train()=20reads=20run=20defaults=20fro?= =?UTF-8?q?m=20it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BaseTrainer does not retain the cfg it receives; the first GPU smoke hit AttributeError at train(). Caught by the 8xH20 smoke, invisible to compose-check (which never enters train()). --- recipes/refl/trainer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/recipes/refl/trainer.py b/recipes/refl/trainer.py index eb82a8fb9..4dac2fd16 100644 --- a/recipes/refl/trainer.py +++ b/recipes/refl/trainer.py @@ -66,6 +66,7 @@ class REFLTrainer(BaseTrainer): def __init__(self, *, cfg: DictConfig) -> None: super().__init__(cfg=cfg, logging_cfg=cfg.get("logging")) + self.cfg = cfg # train() reads run-length/checkpoint defaults from it self.batch_size = int(cfg.batch_size) self.max_grad_norm = float(cfg.get("max_grad_norm", 1.0)) self.data_source = instantiate(cfg.data_source) From 40b3f4c91d75db66a410e2b68c16d2c8ac8e2c60 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 18:14:26 +0800 Subject: [PATCH 14/24] =?UTF-8?q?fix(recipes/refl):=20drop=20fp32=20LoRA?= =?UTF-8?q?=20master=20dtype=20=E2=80=94=20trips=20FSDP2=20uniform-dtype?= =?UTF-8?q?=20assert?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fp32 LoRA params over the bf16 base fail torch 2.7 FSDP2's uniform-original-dtype assertion at first forward (documented in examples/diffusion/refl_sd3.yaml). Omit master_dtype so LoRA stays bf16, matching the SD3 refl recipe's portable choice. Caught by the 8xH20 fleet smoke; the contributor's environment tolerated the fp32 mix. --- recipes/refl/configs/wan21_t2v_videoalign_refl.yaml | 4 +++- recipes/refl/configs/wan22_i2v_face_refl.yaml | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 3aa8cc7b9..75e1753ba 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -42,7 +42,9 @@ actor: fsdp_cfg: _target_: unirl.train.configs.FSDPConfig param_dtype: bf16 - master_dtype: fp32 + # master_dtype deliberately omitted → LoRA stays bf16 (uniform dtype). + # fp32 LoRA over a bf16 base trips torch FSDP2's uniform-original-dtype + # assert (same note as examples/diffusion/refl_sd3.yaml). cpu_offload: false mixed_precision: true fsdp_mode: full diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index b4f39736e..32e924450 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -43,7 +43,9 @@ actor: fsdp_cfg: _target_: unirl.train.configs.FSDPConfig param_dtype: bf16 - master_dtype: fp32 + # master_dtype deliberately omitted → LoRA stays bf16 (uniform dtype). + # fp32 LoRA over a bf16 base trips torch FSDP2's uniform-original-dtype + # assert (same note as examples/diffusion/refl_sd3.yaml). cpu_offload: false mixed_precision: true fsdp_mode: full From 811be0f97354a237fcd14c1d1961b534419e11ea Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 20:22:29 +0800 Subject: [PATCH 15/24] fix(recipes/refl): make VideoAlign correct on the declared transformers 5.6 stack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validated on 8xH20 in an isolated transformers==5.6.2 + peft==0.20.0 venv (real VideoReward checkpoint load + differentiable forward/backward: rewards finite, grad_abs_mean=2.2e-3): - vision features: prefer pooler_output — on 5.6 it holds the merged features that fill media placeholders; last_hidden_state is the PRE-merger states (4x tokens, vision dim) and silently corrupts the reward if scattered. Raw-Tensor branch kept for the 4.56 fleet image (TODO drop when the image moves to the locked stack); all other fallbacks removed per the pin-one-version policy. - checkpoint loader: apply the old→new Qwen2VL layout remap in the LoRA-split branch too (module paths are embedded in LoRA keys); guard against double-prefixing already-new-layout keys. - factory: stop forwarding use_cache through from_pretrained — 5.x passes unknown kwargs to the model ctor (TypeError); set it on the config after load, as the author's own disabled-shim note prescribed. - slim stale multi-era compat comments down to the load-bearing facts. --- .../rewards/videoalign/model/checkpoint.py | 59 ++++++++++--------- .../refl/rewards/videoalign/model/factory.py | 30 +--------- .../rewards/videoalign/model/reward_model.py | 23 +++----- recipes/refl/rewards/videoalign/wrapper.py | 24 +++----- 4 files changed, 51 insertions(+), 85 deletions(-) diff --git a/recipes/refl/rewards/videoalign/model/checkpoint.py b/recipes/refl/rewards/videoalign/model/checkpoint.py index 0f9609860..bfc2c797a 100644 --- a/recipes/refl/rewards/videoalign/model/checkpoint.py +++ b/recipes/refl/rewards/videoalign/model/checkpoint.py @@ -111,34 +111,37 @@ def load_model_from_checkpoint( lora_ckpt = os.path.join(checkpoint_path, "adapter_model.safetensors") non_lora_ckpt = os.path.join(checkpoint_path, "non_lora_state_dict.pth") + # Upstream checkpoints use the old Qwen2VL layout (LM at + # base_model.model.model.*, vision at base_model.model.visual.*); + # transformers>=5 nests them under model.language_model / model.visual. + # Detected from the TARGET model's keys; applied to every loaded dict — + # LoRA keys embed module paths too. + target_keys = model.state_dict().keys() + needs_remap = any( + k.startswith("base_model.model.model.language_model.") or k.startswith("base_model.model.model.visual.") + for k in target_keys + ) + + def _remap_qwen_layout(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + if not needs_remap: + return state_dict + new_state_dict: Dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("base_model.model.model.language_model.") or key.startswith( + "base_model.model.model.visual." + ): + new_state_dict[key] = value # already the new layout + elif key.startswith("base_model.model.model"): + new_state_dict["base_model.model.model.language_model" + key[len("base_model.model.model") :]] = value + elif key.startswith("base_model.model.visual"): + new_state_dict["base_model.model.model.visual" + key[len("base_model.model.visual") :]] = value + else: + new_state_dict[key] = value + return new_state_dict + if os.path.exists(full_ckpt): model_state_dict = torch.load(full_ckpt, map_location="cpu", weights_only=True) - # The upstream checkpoints were written against the old Qwen2VL - # submodule layout (LM under ``base_model.model.model.*``, vision - # under ``base_model.model.visual.*``). transformers>=5 moved - # language layers under - # ``base_model.model.model.language_model.*`` and the visual tower - # under ``base_model.model.model.visual.*``. Detect by inspecting - # the target model's own state_dict keys and remap when needed. - target_keys = model.state_dict().keys() - needs_remap = any( - k.startswith("base_model.model.model.language_model.") or k.startswith("base_model.model.model.visual.") - for k in target_keys - ) - if needs_remap: - new_state_dict: Dict[str, torch.Tensor] = {} - for key, value in model_state_dict.items(): - if key.startswith("base_model.model.model"): - new_key = "base_model.model.model.language_model" + key[len("base_model.model.model") :] - new_state_dict[new_key] = value - elif key.startswith("base_model.model.visual"): - new_key = "base_model.model.model.visual" + key[len("base_model.model.visual") :] - new_state_dict[new_key] = value - else: - new_state_dict[key] = value - model.load_state_dict(new_state_dict) - else: - model.load_state_dict(model_state_dict) + model.load_state_dict(_remap_qwen_layout(model_state_dict)) else: # LoRA branch — merge LoRA adapter + non-LoRA tensors. if not os.path.exists(lora_ckpt) or not os.path.exists(non_lora_ckpt): @@ -147,8 +150,8 @@ def load_model_from_checkpoint( f"({lora_ckpt!r} + {non_lora_ckpt!r}) was found under {checkpoint_path!r}." ) - lora_state_dict = safetensors.torch.load_file(lora_ckpt) - non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu") + lora_state_dict = _remap_qwen_layout(safetensors.torch.load_file(lora_ckpt)) + non_lora_state_dict = _remap_qwen_layout(torch.load(non_lora_ckpt, map_location="cpu")) lora_state_dict = _insert_adapter_name_into_state_dict( lora_state_dict, diff --git a/recipes/refl/rewards/videoalign/model/factory.py b/recipes/refl/rewards/videoalign/model/factory.py index d65fab275..1afce4df0 100644 --- a/recipes/refl/rewards/videoalign/model/factory.py +++ b/recipes/refl/rewards/videoalign/model/factory.py @@ -98,14 +98,6 @@ def create_model_and_processor( # Build the reward model. Quantisation is intentionally not supported # here — the reward path expects full-precision (or bf16/fp16) weights. - # - # Aligned with mmrl's ``create_model_and_processor``: forward - # ``revision`` and ``use_cache`` directly into ``from_pretrained``. - # ``use_cache`` is a standard ``PretrainedConfig`` field, so HF's - # ``from_pretrained`` absorbs it into the config rather than passing - # it down to ``Qwen2VLRewardModelBT.__init__`` — no ``TypeError``. - # ``revision`` is a no-op for local-path checkpoints but matches the - # mmrl call site verbatim. model = Qwen2VLRewardModelBT.from_pretrained( model_config.model_name_or_path, output_dim=model_config.output_dim, @@ -115,26 +107,10 @@ def create_model_and_processor( attn_implementation=("flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa"), cache_dir=cache_dir, revision=getattr(model_config, "model_revision", "main"), - use_cache=True if training_args.gradient_checkpointing else False, ) - - # ------------------------------------------------------------------ - # transformers>=4.58/5.x compatibility shim (DISABLED). - # - # Older comment / behaviour preserved here for reference: under some - # future transformers versions, ``from_pretrained`` may stop treating - # ``use_cache`` as a known ``PretrainedConfig`` field and pass it - # through to ``Qwen2VLRewardModelBT.__init__``, which only accepts - # ``output_dim`` / ``reward_token`` / ``special_token_ids`` and would - # raise ``TypeError: unexpected keyword argument 'use_cache'``. If - # that happens, drop ``use_cache`` / ``revision`` from the - # ``from_pretrained`` call above and re-enable the post-hoc setter - # below: - # - # model.config.use_cache = bool( - # True if training_args.gradient_checkpointing else False - # ) - # ------------------------------------------------------------------ + # transformers 5.x forwards unknown from_pretrained kwargs to the model + # ctor (4.x absorbed config fields like use_cache) — set it on the config. + model.config.use_cache = bool(training_args.gradient_checkpointing) if model_config.use_special_tokens: model.resize_token_embeddings(len(processor.tokenizer)) diff --git a/recipes/refl/rewards/videoalign/model/reward_model.py b/recipes/refl/rewards/videoalign/model/reward_model.py index d0e4765f0..a04ef7107 100644 --- a/recipes/refl/rewards/videoalign/model/reward_model.py +++ b/recipes/refl/rewards/videoalign/model/reward_model.py @@ -160,27 +160,22 @@ def forward( # an LM head — only the regression head over the final hidden # states. # - # Backwards-compat shim for transformers >= 4.52: the Qwen2-VL - # vision tower's forward used to return a raw ``Tensor`` (the - # flattened visual hidden states) but now returns a - # ``BaseModelOutputWithPooling`` dataclass. We accept both shapes - # by extracting ``.last_hidden_state`` when present. + # transformers 5.6 returns BaseModelOutputWithPooling: pooler_output + # = merged features (fills the placeholders); last_hidden_state = + # PRE-merger states (4x tokens, vision dim) — never usable here. + # transformers 4.56 (fleet image; TODO drop once the image moves to + # the locked 5.6 stack) returns the merged features as a raw Tensor. def _as_tensor(visual_out): if isinstance(visual_out, torch.Tensor): return visual_out - # transformers >= 4.52 returns a ModelOutput-like dataclass; - # the visual hidden states live on ``last_hidden_state``. - t = getattr(visual_out, "last_hidden_state", None) + t = getattr(visual_out, "pooler_output", None) if t is not None: return t - # Last resort: some forks return a tuple where the first - # element is the hidden-states tensor. - if isinstance(visual_out, (tuple, list)) and len(visual_out) > 0: - return visual_out[0] raise TypeError( "Qwen2-VL vision tower returned an unsupported type: " - f"{type(visual_out).__name__}. Expected Tensor or " - "BaseModelOutputWithPooling (transformers>=4.52)." + f"{type(visual_out).__name__} (no pooler_output). This code " + "targets the locked transformers stack — align the environment " + "instead of widening this shim." ) if inputs_embeds is None: diff --git a/recipes/refl/rewards/videoalign/wrapper.py b/recipes/refl/rewards/videoalign/wrapper.py index d80f2fa60..697a234a7 100644 --- a/recipes/refl/rewards/videoalign/wrapper.py +++ b/recipes/refl/rewards/videoalign/wrapper.py @@ -25,14 +25,12 @@ from __future__ import annotations -# NOTE: this wrapper runs on the UniRL core stack (transformers>=5.6 — the -# reward shares one Python process with the actor, so there is no separate -# environment to pin). Newer ``Qwen2VLProcessor`` versions inject -# ``mm_token_type_ids``, which the 4.45-era reward backbone does not accept; -# ``compute_scores`` drops it by explicit black-list. Signature-based -# filtering (``inspect``) is deliberately NOT used: under PEFT it can resolve -# to ``LoraModel.forward(*args, **kwargs)`` and silently strip -# ``pixel_values_videos`` / ``video_grid_thw``. +# Runs on the shared core stack — reward and actor share one process, so +# there is no separate env to pin. transformers>=4.58 processors inject +# ``mm_token_type_ids``, which the reward backbone does not accept; +# ``compute_scores`` pops it explicitly. Never filter by forward-signature: +# under PEFT that resolves to ``LoraModel.forward(*args, **kwargs)`` and +# silently drops the video tensors. import importlib.util import json import logging @@ -314,14 +312,8 @@ def forward_scores( for start in range(0, len(video_tensors), self.micro_batch_size): end = start + self.micro_batch_size batch = self.prepare_batch_from_frames(video_tensors[start:end], prompts[start:end]) - # transformers>=4.58/5.x ``Qwen2VLProcessor`` injects - # ``mm_token_type_ids`` (text/image/video ids for the rewritten - # ``get_rope_index``); ``Qwen2VLRewardModelBT.forward`` was - # authored against 4.45 and does not accept it. Drop it by - # explicit black-list — NEVER ``inspect.signature`` filtering, - # which under PEFT can resolve to ``LoraModel.forward(*args, - # **kwargs)`` and silently strip ``pixel_values_videos`` / - # ``video_grid_thw``, leaving the reward blind to the video. + # 5.x processors inject mm_token_type_ids; the backbone doesn't + # accept it (see module NOTE — pop, never signature-filter). batch.pop("mm_token_type_ids", None) logits = self.model(**batch, return_dict=True)["logits"] # (B, 3) vq, mq, ta = logits[:, 0], logits[:, 1], logits[:, 2] From 9087a6716f53e9a87ed36d5ff067fd2586790504 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 21:00:15 +0800 Subject: [PATCH 16/24] =?UTF-8?q?refactor(recipes/refl):=20single-stack=20?= =?UTF-8?q?policy=20=E2=80=94=20VideoAlign=20targets=20locked=20transforme?= =?UTF-8?q?rs=205.6=20only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review direction: no version-compat branches; a wrong environment fails loudly and the user aligns the env, not the code. - _as_tensor: pooler_output only (drop the 4.56 raw-Tensor concession) - visual property: single 5.6 location (self.model.visual) - checkpoint remap: always applied (target layout is 5.6 by policy), idempotent per key - processor: rely on 5.6 fast-by-default; drop the force-swap + use_fast - attention: hardcode sdpa — flash-attn is not part of the locked stack - pyproject peft floor 0.14 → 0.20: older peft imports transformers cache symbols removed in 5.x and fails at import against the 5.6 pin Gate on 8xH20 (isolated transformers==5.6.2 + peft==0.20.0 venv, real VideoReward ckpt): load + differentiable forward/backward PASS, rewards=[-3.5938, -3.375], grad_abs_mean=3.5e-3. --- pyproject.toml | 5 +- .../rewards/videoalign/model/checkpoint.py | 14 +-- .../rewards/videoalign/model/reward_model.py | 86 ++++++------------- .../refl/rewards/videoalign/requirements.txt | 12 +-- recipes/refl/rewards/videoalign/wrapper.py | 33 +++---- 5 files changed, 51 insertions(+), 99 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9eded6e3c..6428210fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,10 @@ dependencies = [ # in transformer_boogu.py / attention_processor.py / rope.py. Pure-python, # engine-agnostic. "einops>=0.7", - "peft>=0.14.0", # exclude_modules support used by LoRA/NFT adapter injection + # 0.20 floor: older peft imports transformers cache symbols (HybridCache) + # removed in 5.x and fails at import against the transformers pin above. + # exclude_modules support (>=0.14) used by LoRA/NFT adapter injection. + "peft>=0.20", "safetensors>=0.4", "Pillow>=10", "requests>=2.31", diff --git a/recipes/refl/rewards/videoalign/model/checkpoint.py b/recipes/refl/rewards/videoalign/model/checkpoint.py index bfc2c797a..b592177ad 100644 --- a/recipes/refl/rewards/videoalign/model/checkpoint.py +++ b/recipes/refl/rewards/videoalign/model/checkpoint.py @@ -113,18 +113,10 @@ def load_model_from_checkpoint( # Upstream checkpoints use the old Qwen2VL layout (LM at # base_model.model.model.*, vision at base_model.model.visual.*); - # transformers>=5 nests them under model.language_model / model.visual. - # Detected from the TARGET model's keys; applied to every loaded dict — - # LoRA keys embed module paths too. - target_keys = model.state_dict().keys() - needs_remap = any( - k.startswith("base_model.model.model.language_model.") or k.startswith("base_model.model.model.visual.") - for k in target_keys - ) - + # transformers 5.6 nests them under model.language_model / model.visual. + # Applied to every loaded dict — LoRA keys embed module paths too — and + # idempotent for already-new-layout keys. def _remap_qwen_layout(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - if not needs_remap: - return state_dict new_state_dict: Dict[str, torch.Tensor] = {} for key, value in state_dict.items(): if key.startswith("base_model.model.model.language_model.") or key.startswith( diff --git a/recipes/refl/rewards/videoalign/model/reward_model.py b/recipes/refl/rewards/videoalign/model/reward_model.py index a04ef7107..095103851 100644 --- a/recipes/refl/rewards/videoalign/model/reward_model.py +++ b/recipes/refl/rewards/videoalign/model/reward_model.py @@ -25,15 +25,9 @@ def _cfg_get(config: Any, name: str) -> Any: - """Read a config attribute that may live on the top-level config or - on the nested ``text_config`` sub-config. - - transformers>=4.52 refactored ``Qwen2VLConfig`` so that fields like - ``hidden_size`` / ``image_token_id`` / ``video_token_id`` / ``pad_token_id`` - moved under ``config.text_config``. Older transformers exposed them at - the top level. This helper transparently supports both layouts and - raises a friendly error when the field is genuinely absent. - """ + """Read a ``Qwen2VLConfig`` field from wherever 5.6 keeps it: media + token ids live on the top-level config, LM fields like ``hidden_size`` + under the nested ``text_config``.""" if hasattr(config, name): val = getattr(config, name) if val is not None: @@ -88,44 +82,21 @@ def __init__( if self.special_token_ids is not None: self.reward_token = "special" - # ------------------------------------------------------------------ - # Backwards-compat shim: vision tower location - # ------------------------------------------------------------------ - # transformers <= 4.51 exposed the Qwen2-VL vision tower as - # ``self.visual`` directly on ``Qwen2VLForConditionalGeneration``. - # transformers >= 4.52 refactored it under ``self.model.visual`` - # (the LM backbone now owns the vision encoder). The forward code - # below was written against the old layout; rather than fork it, - # we expose a thin property that resolves to whichever location - # the currently-installed transformers uses. - # - # Implementation notes: - # - Read from ``self._modules`` (the underlying ``OrderedDict``) - # instead of ``getattr``/``hasattr`` to avoid recursing back into - # ``nn.Module.__getattr__`` which would re-trigger this property - # and yield a misleading "no attribute 'visual'" error. - # - We deliberately do *not* define a setter: ``super().__init__`` - # registers any ``self.visual = module`` assignment via - # ``nn.Module.__setattr__`` (which writes into ``_modules``), and - # the property here just reads that slot back out — so old-style - # checkpoints still load and the ``state_dict`` layout is - # completely unchanged. + # The forward code below addresses ``self.visual``; transformers 5.6 + # owns the vision tower at ``self.model.visual``. Read via ``_modules`` + # to avoid recursing into ``nn.Module.__getattr__``, which would + # re-trigger this property with a misleading "no attribute" error. @property def visual(self): # type: ignore[override] - own = self._modules.get("visual", None) - if own is not None: - return own # transformers <= 4.51 layout inner = self._modules.get("model", None) - if inner is not None: - inner_visual = getattr(inner, "visual", None) - if inner_visual is not None: - return inner_visual # transformers >= 4.52 layout - raise AttributeError( - "Qwen2VLRewardModelBT: vision tower not found on either " - "self.visual (transformers<=4.51) or self.model.visual " - "(transformers>=4.52). The installed transformers version " - "may be incompatible with this checkpoint." - ) + visual = getattr(inner, "visual", None) if inner is not None else None + if visual is None: + raise AttributeError( + "Qwen2VLRewardModelBT: vision tower not found at " + "self.model.visual — this code targets the locked " + "transformers 5.6 stack; align the environment." + ) + return visual # ------------------------------------------------------------------ # Forward @@ -160,23 +131,20 @@ def forward( # an LM head — only the regression head over the final hidden # states. # - # transformers 5.6 returns BaseModelOutputWithPooling: pooler_output - # = merged features (fills the placeholders); last_hidden_state = - # PRE-merger states (4x tokens, vision dim) — never usable here. - # transformers 4.56 (fleet image; TODO drop once the image moves to - # the locked 5.6 stack) returns the merged features as a raw Tensor. + # transformers 5.6: the vision tower returns BaseModelOutputWithPooling + # whose pooler_output holds the merged features that fill the media + # placeholders (what get_video_features reads). last_hidden_state is + # the PRE-merger states (4x tokens, vision dim) — never usable here. def _as_tensor(visual_out): - if isinstance(visual_out, torch.Tensor): - return visual_out t = getattr(visual_out, "pooler_output", None) - if t is not None: - return t - raise TypeError( - "Qwen2-VL vision tower returned an unsupported type: " - f"{type(visual_out).__name__} (no pooler_output). This code " - "targets the locked transformers stack — align the environment " - "instead of widening this shim." - ) + if t is None: + raise TypeError( + "Qwen2-VL vision tower returned " + f"{type(visual_out).__name__} with no pooler_output. This " + "code targets the locked transformers 5.6 stack — align " + "the environment instead of widening this path." + ) + return t if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) diff --git a/recipes/refl/rewards/videoalign/requirements.txt b/recipes/refl/rewards/videoalign/requirements.txt index 83b2e094b..bcd1c63c0 100644 --- a/recipes/refl/rewards/videoalign/requirements.txt +++ b/recipes/refl/rewards/videoalign/requirements.txt @@ -1,9 +1,9 @@ # VideoAlign reward extras — ADDITIVE ONLY on top of the UniRL core stack. # -# The reward runs in the SAME Python process as the actor (colocated Remote -# siblings), so a requirements file cannot "isolate" versions — any pin that -# overlaps the core stack would downgrade the core environment itself. Core -# already provides transformers / peft / safetensors / huggingface-hub / -# einops (see pyproject.toml); never re-pin those here. flash-attn is -# optional: the model factory falls back to SDPA when the wheel is absent. +# The reward shares one Python process with the actor (colocated Remote +# siblings), so pins here cannot "isolate" versions — an overlapping pin +# would downgrade the core environment itself. Core already provides +# transformers / peft / safetensors / huggingface-hub / einops (see +# pyproject.toml); never re-pin those here. Attention runs SDPA — flash-attn +# is not part of the locked stack. torchvision # frame preprocessing transforms; matches the engine extra's torch build diff --git a/recipes/refl/rewards/videoalign/wrapper.py b/recipes/refl/rewards/videoalign/wrapper.py index 697a234a7..d4e2c9732 100644 --- a/recipes/refl/rewards/videoalign/wrapper.py +++ b/recipes/refl/rewards/videoalign/wrapper.py @@ -17,21 +17,19 @@ Gradient flow ------------- The Qwen2-VL vision encoder is fully differentiable w.r.t. its pixel input -when the *fast* image processor (``Qwen2VLImageProcessorFast``) is used — -the slow PIL-based variant routes through ``numpy`` and silently cuts the -graph. We force-swap to the fast processor on construction; this is the -single line that makes REFL gradients work end-to-end. +through the *fast* (tensor-native) image processor, which transformers 5.6 +loads by default — the slow PIL-based variant routes through ``numpy`` and +silently cuts the graph. """ from __future__ import annotations # Runs on the shared core stack — reward and actor share one process, so -# there is no separate env to pin. transformers>=4.58 processors inject +# there is no separate env to pin. The 5.6 processor injects # ``mm_token_type_ids``, which the reward backbone does not accept; # ``compute_scores`` pops it explicitly. Never filter by forward-signature: # under PEFT that resolves to ``LoraModel.forward(*args, **kwargs)`` and # silently drops the video tensors. -import importlib.util import json import logging import os @@ -41,7 +39,6 @@ import torch import torchvision.transforms.functional as TF from torchvision.transforms import InterpolationMode -from transformers import AutoImageProcessor from .model import ( ModelConfig, @@ -87,8 +84,7 @@ class VideoRewardWrapper: :func:`create_model_and_processor`. 3. Load weights from ``checkpoint-K`` via :func:`load_model_from_checkpoint`. - 4. Replace the slow image processor with the autograd-friendly fast one. - 5. Move to the requested device + dtype, ``eval()`` + freeze + 4. Move to the requested device + dtype, ``eval()`` + freeze parameters (RL gradients flow *through* the activations, not into the reward weights). """ @@ -138,14 +134,17 @@ def __init__( load_from_pretrained=checkpoint_dir, load_from_pretrained_step=-1, gradient_checkpointing=False, - # flash-attn 2 is not part of the UniRL core stack; fall back to - # SDPA automatically when the wheel is absent. - disable_flash_attn2=importlib.util.find_spec("flash_attn") is None, + # sdpa — deterministic on the locked stack; flash-attn 2 is not + # part of it. + disable_flash_attn2=True, bf16=(dtype == torch.bfloat16), fp16=(dtype == torch.float16), output_dir="", ) + # transformers 5.6 loads the fast (tensor-native) image processor by + # default; gradient flow requires it — the slow variant round-trips + # through PIL and severs autograd. model, processor, _ = create_model_and_processor( model_config=model_config, peft_lora_config=peft_lora_config, @@ -153,16 +152,6 @@ def __init__( ) model, _ = load_model_from_checkpoint(model, checkpoint_dir, -1) - # CRITICAL: force the *fast* image processor. The slow variant - # converts inputs to numpy/PIL and (a) rejects float pixels outside - # [0, 1] via ``_rescale_for_pil_conversion``, and (b) silently - # severs autograd even when you push the values past the range - # check. The fast variant operates on torch tensors end-to-end via - # torchvision.transforms.v2.functional.resize, so gradients flow - # from pixels through the processor into the vision encoder. - fast_ip = AutoImageProcessor.from_pretrained(model_config.model_name_or_path, use_fast=True) - processor.image_processor = fast_ip - model.to(self.device) model.eval() model.requires_grad_(False) From b253be2573b3c8478691997fe01855edea1d769f Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 21:38:29 +0800 Subject: [PATCH 17/24] refactor(recipes/refl): model_adaptor naming, drop launch scripts, add README MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review decisions: - recipes/refl/models/ → recipes/refl/model_adaptor/ — the directory holds algorithm-side adaptations wrapping core models to the recipe's BPTT contract, not model definitions; the old name collided with unirl/models semantics and blurred the graduation boundary. - Launch .sh wrappers deleted: zero logic beyond env-var placeholders; the launch surface is one documented command. recipes/refl/README.md now carries the launch examples, layout map, environment policy, and the verification table. Static at this head: compose x2 rc=0; 2270 _target_ paths resolve; KL topology checks pass; full pre-commit green. --- recipes/refl/README.md | 57 +++++++++++++++++++ .../configs/wan21_t2v_videoalign_refl.yaml | 2 +- recipes/refl/configs/wan22_i2v_face_refl.yaml | 2 +- .../{models => model_adaptor}/__init__.py | 0 .../refl/{models => model_adaptor}/types.py | 0 .../refl/{models => model_adaptor}/wan21.py | 4 +- .../refl/{models => model_adaptor}/wan22.py | 2 +- recipes/refl/roles.py | 6 +- recipes/refl/scripts/start_wan21_t2v.sh | 33 ----------- recipes/refl/scripts/start_wan22_i2v.sh | 21 ------- 10 files changed, 65 insertions(+), 62 deletions(-) create mode 100644 recipes/refl/README.md rename recipes/refl/{models => model_adaptor}/__init__.py (100%) rename recipes/refl/{models => model_adaptor}/types.py (100%) rename recipes/refl/{models => model_adaptor}/wan21.py (99%) rename recipes/refl/{models => model_adaptor}/wan22.py (99%) delete mode 100755 recipes/refl/scripts/start_wan21_t2v.sh delete mode 100755 recipes/refl/scripts/start_wan22_i2v.sh diff --git a/recipes/refl/README.md b/recipes/refl/README.md new file mode 100644 index 000000000..9f5564689 --- /dev/null +++ b/recipes/refl/README.md @@ -0,0 +1,57 @@ +# recipes/refl — WAN ReFL/BPTT (differentiable reward backprop) + +The first self-contained UniRL recipe: direct reward backprop (ReFL / DRaFT) +for WAN video models. Two colocated roles — a `ReflActorRole` (FSDP WAN + +grad BPTT sampling + optimizer) and a frozen differentiable reward +(`RewardService`) — run a 3-RPC step under the distributed `enable_grad()` +context: generate → score → backward, then optimizer step. No advantages, +no replay, no rollout engine, no weight sync. + +## Launch + +One command per config; a Ray cluster must be up (`ray start --head`). + +```bash +# WAN 2.1 T2V + VideoAlign (Qwen2-VL VQ/MQ/TA) reward +export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers \ + VIDEOALIGN_MODEL_PATH=/path/to/VideoReward \ + DATA_PATH=/path/to/prompts.txt +RAY_ADDRESS=auto python -m recipes.refl.run --config-name=wan21_t2v_videoalign_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 recipes/refl/rewards/face/requirements.txt +export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers \ + FACE_MODEL_PATH=/path/to/antelodev2 \ + DATA_PATH=/path/to/i2v_prompts.jsonl +RAY_ADDRESS=auto python -m recipes.refl.run --config-name=wan22_i2v_face_refl num_devices=8 +``` + +## Layout + +| Path | What | +|---|---| +| `trainer.py` | `REFLTrainer(BaseTrainer)` — driver: wiring + the 3-RPC train step | +| `roles.py` | `ReflActorRole(Remote)` — family-agnostic actor (`pipeline_target` + `model_config`) | +| `model_adaptor/` | Per-model adaptations of core pipelines to the BPTT contract (`types.py` defines it): `wan21.py`, `wan22.py` | +| `rewards/` | Recipe-local differentiable rewards (VideoAlign, Face), each with an additive-only `requirements.txt` | +| `configs/` | Flat Hydra configs (repo-wide schema) | + +## Environment + +Targets the **locked core stack only** (`transformers>=5.6,<5.7`, +`peft>=0.20` — see `pyproject.toml`). There are deliberately no +version-compat branches: a wrong environment fails loudly; align the +environment, not the code. Reward and actor share one Python process, so +recipe `requirements.txt` files may only ADD packages, never re-pin the +core stack. + +## Verification + +| Config | Hardware | Head | Status | +| --- | --- | --- | --- | +| `wan21_t2v_videoalign_refl` (835 rollouts) | 8xH20 | pre-adjustment (`e3c6b940` lineage) | contributor long run — reward curve in PR #210 | +| `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` | +| `wan21_t2v_videoalign_refl` (150-rollout trend run, 5.6 stack) | 8xH20 | `9087a671` | running — curve to be attached | +| `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml index 75e1753ba..fd1e7b4ba 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -19,7 +19,7 @@ actor: _target_: recipes.refl.roles.ReflActorRole # Family selector — swap pipeline_target + model_config for another family, # no code changes (mirrors ReFLPolicy's pipeline_target contract). - pipeline_target: recipes.refl.models.wan21.Wan21ReflPipeline + pipeline_target: recipes.refl.model_adaptor.wan21.Wan21ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.25 diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/recipes/refl/configs/wan22_i2v_face_refl.yaml index 32e924450..0a06ca286 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/recipes/refl/configs/wan22_i2v_face_refl.yaml @@ -18,7 +18,7 @@ max_grad_norm: 1.0 actor: _target_: recipes.refl.roles.ReflActorRole - pipeline_target: recipes.refl.models.wan22.Wan22ReflPipeline + pipeline_target: recipes.refl.model_adaptor.wan22.Wan22ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.1 diff --git a/recipes/refl/models/__init__.py b/recipes/refl/model_adaptor/__init__.py similarity index 100% rename from recipes/refl/models/__init__.py rename to recipes/refl/model_adaptor/__init__.py diff --git a/recipes/refl/models/types.py b/recipes/refl/model_adaptor/types.py similarity index 100% rename from recipes/refl/models/types.py rename to recipes/refl/model_adaptor/types.py diff --git a/recipes/refl/models/wan21.py b/recipes/refl/model_adaptor/wan21.py similarity index 99% rename from recipes/refl/models/wan21.py rename to recipes/refl/model_adaptor/wan21.py index f0d77aaec..e4e4cf498 100644 --- a/recipes/refl/models/wan21.py +++ b/recipes/refl/model_adaptor/wan21.py @@ -1,6 +1,6 @@ """Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. -Mirrors ``recipes.refl.models.wan22`` but targets the WAN 2.1 T2V +Mirrors ``recipes.refl.model_adaptor.wan22`` but targets the WAN 2.1 T2V single-DiT stack. The REFL-specific pieces live here: - :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch @@ -25,7 +25,7 @@ import torch -from recipes.refl.models.types import DiffuseWithGradResult +from recipes.refl.model_adaptor.types import DiffuseWithGradResult from unirl.models.wan21.bundle import WAN21Bundle from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions diff --git a/recipes/refl/models/wan22.py b/recipes/refl/model_adaptor/wan22.py similarity index 99% rename from recipes/refl/models/wan22.py rename to recipes/refl/model_adaptor/wan22.py index b2a6638bb..215ee7b61 100644 --- a/recipes/refl/models/wan22.py +++ b/recipes/refl/model_adaptor/wan22.py @@ -8,7 +8,7 @@ import torch -from recipes.refl.models.types import DiffuseWithGradResult +from recipes.refl.model_adaptor.types import DiffuseWithGradResult from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage diff --git a/recipes/refl/roles.py b/recipes/refl/roles.py index b02ae6039..f4e4edd73 100644 --- a/recipes/refl/roles.py +++ b/recipes/refl/roles.py @@ -11,7 +11,7 @@ actor.step(max_grad_norm=…) # ctx exit routed grads → optimizer The pipeline named by ``pipeline_target`` must expose the recipe contract -(see ``recipes.refl.models``): ``build_refl_conditions(texts, images=…, +(see ``recipes.refl.model_adaptor``): ``build_refl_conditions(texts, images=…, params=…)`` plus a ``diffusion`` stage with ``diffuse_with_grad`` and a ``vae_decode`` stage with ``decode_with_grad``. """ @@ -117,12 +117,12 @@ def initialize(self) -> None: if not hasattr(stage, method): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} .{stage_attr} lacks {method}(...); " - f"use a recipes.refl.models pipeline (or implement the REFL contract)." + f"use a recipes.refl.model_adaptor pipeline (or implement the REFL contract)." ) if not hasattr(self.pipeline, "build_refl_conditions"): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} lacks build_refl_conditions(...); " - f"use a recipes.refl.models pipeline." + f"use a recipes.refl.model_adaptor pipeline." ) # FSDP-wrap pipeline.bundle.transformer in place + LoRA + optimizer. diff --git a/recipes/refl/scripts/start_wan21_t2v.sh b/recipes/refl/scripts/start_wan21_t2v.sh deleted file mode 100755 index 84fffb826..000000000 --- a/recipes/refl/scripts/start_wan21_t2v.sh +++ /dev/null @@ -1,33 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "${SCRIPT_DIR}/../../.." - -# === Data === -# Newline-delimited prompts .txt (or {"prompt": ...} .jsonl). VideoAlign is -# text-conditioned; no reference video / first-frame needed. -export DATA_PATH=${DATA_PATH:-/path/to/wan21_prompts.txt} -export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} - -# === Models === -# WAN 2.1 T2V 1.3B base checkpoint -export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/path/to/Wan2.1-T2V-1.3B-Diffusers} -# VideoAlign Qwen2-VL reward checkpoint -export VIDEOALIGN_MODEL_PATH=${VIDEOALIGN_MODEL_PATH:-/path/to/VideoReward} - -# === Output / Logging === -export OUTPUT_DIR=${OUTPUT_DIR:-./outputs/wan21_t2v_videoalign_refl} -export REPORT_TO_WANDB=${REPORT_TO_WANDB:-true} -export WANDB_PROJECT=${WANDB_PROJECT:-unirl-refl} -export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan21_t2v_videoalign_refl_recipe_opt} - -mkdir -p "${OUTPUT_DIR}" logs - -LOG_FILE="logs/wan21_t2v_videoalign_refl_$(date +%Y%m%d_%H%M%S).log" -echo "=== launching wan21 t2v videoalign refl, log → ${LOG_FILE} ===" - -RAY_ADDRESS=auto python -u -m recipes.refl.run \ - --config-name=wan21_t2v_videoalign_refl \ - num_devices=8 \ - 2>&1 | tee "${LOG_FILE}" diff --git a/recipes/refl/scripts/start_wan22_i2v.sh b/recipes/refl/scripts/start_wan22_i2v.sh deleted file mode 100755 index 00912d188..000000000 --- a/recipes/refl/scripts/start_wan22_i2v.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "${SCRIPT_DIR}/../../.." - -export PRETRAINED_MODEL=${PRETRAINED_MODEL:-/path/to/Wan2.2-I2V-A14B-Diffusers} -export DATA_PATH=${DATA_PATH:-/path/to/wan22_i2v_face_refl_prompts.jsonl} -export EVAL_DATA_PATH=${EVAL_DATA_PATH:-${DATA_PATH}} -export FACE_MODEL_PATH=${FACE_MODEL_PATH:-/path/to/antelodev2} -export OUTPUT_DIR=${OUTPUT_DIR:-outputs/wan22_i2v_face_refl_recipe_opt} - -export REPORT_TO_WANDB=${REPORT_TO_WANDB:-true} -export WANDB_PROJECT=${WANDB_PROJECT:-unirl-refl} -export WANDB_RUN_NAME=${WANDB_RUN_NAME:-wan22_i2v_face_refl_recipe_opt} - -mkdir -p logs - -RAY_ADDRESS=auto python -m recipes.refl.run \ - num_devices=8 \ - 2>&1 | tee ./logs/wan22_i2v_refl.log From 74f58bfd2fb2cbbeff4e299ae3ad30897f5e86f5 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 21:47:15 +0800 Subject: [PATCH 18/24] =?UTF-8?q?refactor(experimental):=20rename=20the=20?= =?UTF-8?q?extension=20layer=20recipes/=20=E2=86=92=20experimental/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Core is not just components — the official training flows live in core (train_*.py + the trainers). A layer named 'recipes' reads as the official, recommended usage and pulls that attention to the non-official tier; 'recipe' also already means the flat YAML in this repo's vocabulary (examples/ is the former recipes/ tree). experimental/ names what this tier actually is: the incubation side of a two-way flow — packages start here, mainstream ones get absorbed and solidified into core, ill-fitting core paths move down or out. Internal/private packages live here uncommitted. All dotpaths, the target checker, and docs updated; upstream provenance references (mmrl/recipes/...) intentionally untouched. --- {recipes => experimental}/__init__.py | 0 {recipes => experimental}/refl/README.md | 8 ++++---- {recipes => experimental}/refl/__init__.py | 0 .../refl/configs/wan21_t2v_videoalign_refl.yaml | 8 ++++---- .../refl/configs/wan22_i2v_face_refl.yaml | 8 ++++---- {recipes => experimental}/refl/model_adaptor/__init__.py | 0 {recipes => experimental}/refl/model_adaptor/types.py | 2 +- {recipes => experimental}/refl/model_adaptor/wan21.py | 4 ++-- {recipes => experimental}/refl/model_adaptor/wan22.py | 2 +- {recipes => experimental}/refl/rewards/__init__.py | 0 {recipes => experimental}/refl/rewards/face/__init__.py | 0 {recipes => experimental}/refl/rewards/face/face_tools.py | 0 .../refl/rewards/face/requirements.txt | 0 {recipes => experimental}/refl/rewards/face/scorer.py | 0 .../refl/rewards/videoalign/__init__.py | 2 +- .../refl/rewards/videoalign/model/__init__.py | 0 .../refl/rewards/videoalign/model/checkpoint.py | 0 .../refl/rewards/videoalign/model/configs.py | 2 +- .../refl/rewards/videoalign/model/factory.py | 0 .../refl/rewards/videoalign/model/prompt_template.py | 0 .../refl/rewards/videoalign/model/reward_model.py | 2 +- .../refl/rewards/videoalign/requirements.txt | 0 .../refl/rewards/videoalign/scorer.py | 2 +- .../refl/rewards/videoalign/wrapper.py | 2 +- {recipes => experimental}/refl/roles.py | 6 +++--- {recipes => experimental}/refl/run.py | 4 ++-- {recipes => experimental}/refl/trainer.py | 2 +- scripts/check_recipe_targets.py | 4 ++-- scripts/verify_refl_kl_batching.py | 2 +- 29 files changed, 30 insertions(+), 30 deletions(-) rename {recipes => experimental}/__init__.py (100%) rename {recipes => experimental}/refl/README.md (88%) rename {recipes => experimental}/refl/__init__.py (100%) rename {recipes => experimental}/refl/configs/wan21_t2v_videoalign_refl.yaml (94%) rename {recipes => experimental}/refl/configs/wan22_i2v_face_refl.yaml (94%) rename {recipes => experimental}/refl/model_adaptor/__init__.py (100%) rename {recipes => experimental}/refl/model_adaptor/types.py (94%) rename {recipes => experimental}/refl/model_adaptor/wan21.py (99%) rename {recipes => experimental}/refl/model_adaptor/wan22.py (99%) rename {recipes => experimental}/refl/rewards/__init__.py (100%) rename {recipes => experimental}/refl/rewards/face/__init__.py (100%) rename {recipes => experimental}/refl/rewards/face/face_tools.py (100%) rename {recipes => experimental}/refl/rewards/face/requirements.txt (100%) rename {recipes => experimental}/refl/rewards/face/scorer.py (100%) rename {recipes => experimental}/refl/rewards/videoalign/__init__.py (88%) rename {recipes => experimental}/refl/rewards/videoalign/model/__init__.py (100%) rename {recipes => experimental}/refl/rewards/videoalign/model/checkpoint.py (100%) rename {recipes => experimental}/refl/rewards/videoalign/model/configs.py (98%) rename {recipes => experimental}/refl/rewards/videoalign/model/factory.py (100%) rename {recipes => experimental}/refl/rewards/videoalign/model/prompt_template.py (100%) rename {recipes => experimental}/refl/rewards/videoalign/model/reward_model.py (99%) rename {recipes => experimental}/refl/rewards/videoalign/requirements.txt (100%) rename {recipes => experimental}/refl/rewards/videoalign/scorer.py (99%) rename {recipes => experimental}/refl/rewards/videoalign/wrapper.py (99%) rename {recipes => experimental}/refl/roles.py (97%) rename {recipes => experimental}/refl/run.py (74%) rename {recipes => experimental}/refl/trainer.py (99%) diff --git a/recipes/__init__.py b/experimental/__init__.py similarity index 100% rename from recipes/__init__.py rename to experimental/__init__.py diff --git a/recipes/refl/README.md b/experimental/refl/README.md similarity index 88% rename from recipes/refl/README.md rename to experimental/refl/README.md index 9f5564689..5ab80bc3a 100644 --- a/recipes/refl/README.md +++ b/experimental/refl/README.md @@ -1,4 +1,4 @@ -# recipes/refl — WAN ReFL/BPTT (differentiable reward backprop) +# experimental/refl — WAN ReFL/BPTT (differentiable reward backprop) The first self-contained UniRL recipe: direct reward backprop (ReFL / DRaFT) for WAN video models. Two colocated roles — a `ReflActorRole` (FSDP WAN + @@ -16,15 +16,15 @@ One command per config; a Ray cluster must be up (`ray start --head`). export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers \ VIDEOALIGN_MODEL_PATH=/path/to/VideoReward \ DATA_PATH=/path/to/prompts.txt -RAY_ADDRESS=auto python -m recipes.refl.run --config-name=wan21_t2v_videoalign_refl num_devices=8 +RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan21_t2v_videoalign_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 recipes/refl/rewards/face/requirements.txt +pip install -r experimental/refl/rewards/face/requirements.txt export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers \ FACE_MODEL_PATH=/path/to/antelodev2 \ DATA_PATH=/path/to/i2v_prompts.jsonl -RAY_ADDRESS=auto python -m recipes.refl.run --config-name=wan22_i2v_face_refl num_devices=8 +RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan22_i2v_face_refl num_devices=8 ``` ## Layout diff --git a/recipes/refl/__init__.py b/experimental/refl/__init__.py similarity index 100% rename from recipes/refl/__init__.py rename to experimental/refl/__init__.py diff --git a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml b/experimental/refl/configs/wan21_t2v_videoalign_refl.yaml similarity index 94% rename from recipes/refl/configs/wan21_t2v_videoalign_refl.yaml rename to experimental/refl/configs/wan21_t2v_videoalign_refl.yaml index fd1e7b4ba..c425adcc3 100644 --- a/recipes/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/experimental/refl/configs/wan21_t2v_videoalign_refl.yaml @@ -16,10 +16,10 @@ save_mode: adapter max_grad_norm: 1.0 actor: - _target_: recipes.refl.roles.ReflActorRole + _target_: experimental.refl.roles.ReflActorRole # Family selector — swap pipeline_target + model_config for another family, # no code changes (mirrors ReFLPolicy's pipeline_target contract). - pipeline_target: recipes.refl.model_adaptor.wan21.Wan21ReflPipeline + pipeline_target: experimental.refl.model_adaptor.wan21.Wan21ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.25 @@ -87,10 +87,10 @@ actor: reward: _target_: unirl.reward.service.RewardService backend: - _target_: recipes.refl.rewards.videoalign.VideoAlignRewardScorer + _target_: experimental.refl.rewards.videoalign.VideoAlignRewardScorer base_device: cuda config: - _target_: recipes.refl.rewards.videoalign.VideoAlignSpec + _target_: experimental.refl.rewards.videoalign.VideoAlignSpec reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} device: cuda batch_size: 1 diff --git a/recipes/refl/configs/wan22_i2v_face_refl.yaml b/experimental/refl/configs/wan22_i2v_face_refl.yaml similarity index 94% rename from recipes/refl/configs/wan22_i2v_face_refl.yaml rename to experimental/refl/configs/wan22_i2v_face_refl.yaml index 0a06ca286..6fce57fcf 100644 --- a/recipes/refl/configs/wan22_i2v_face_refl.yaml +++ b/experimental/refl/configs/wan22_i2v_face_refl.yaml @@ -17,8 +17,8 @@ save_mode: adapter max_grad_norm: 1.0 actor: - _target_: recipes.refl.roles.ReflActorRole - pipeline_target: recipes.refl.model_adaptor.wan22.Wan22ReflPipeline + _target_: experimental.refl.roles.ReflActorRole + pipeline_target: experimental.refl.model_adaptor.wan22.Wan22ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.1 @@ -85,10 +85,10 @@ actor: reward: _target_: unirl.reward.service.RewardService backend: - _target_: recipes.refl.rewards.face.FaceRewardScorer + _target_: experimental.refl.rewards.face.FaceRewardScorer base_device: cuda config: - _target_: recipes.refl.rewards.face.FaceRewardSpec + _target_: experimental.refl.rewards.face.FaceRewardSpec model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} device: cuda batch_size: 1 diff --git a/recipes/refl/model_adaptor/__init__.py b/experimental/refl/model_adaptor/__init__.py similarity index 100% rename from recipes/refl/model_adaptor/__init__.py rename to experimental/refl/model_adaptor/__init__.py diff --git a/recipes/refl/model_adaptor/types.py b/experimental/refl/model_adaptor/types.py similarity index 94% rename from recipes/refl/model_adaptor/types.py rename to experimental/refl/model_adaptor/types.py index 1ba537877..096c8f2cf 100644 --- a/recipes/refl/model_adaptor/types.py +++ b/experimental/refl/model_adaptor/types.py @@ -6,7 +6,7 @@ become a real ``None``-returning method on every diffusion stage in the repo and make ``hasattr``-based capability checks meaningless. While REFL is the only BPTT consumer, the contract lives here; if a second consumer -appears outside ``recipes/refl``, promote it to core as a separate opt-in +appears outside ``experimental/refl``, promote it to core as a separate opt-in ``@runtime_checkable`` protocol (the ``DifferentiableReward`` / ``LatentShapeProvider`` idiom), not as a method on ``DiffusionStage``. diff --git a/recipes/refl/model_adaptor/wan21.py b/experimental/refl/model_adaptor/wan21.py similarity index 99% rename from recipes/refl/model_adaptor/wan21.py rename to experimental/refl/model_adaptor/wan21.py index e4e4cf498..7e76b0a73 100644 --- a/recipes/refl/model_adaptor/wan21.py +++ b/experimental/refl/model_adaptor/wan21.py @@ -1,6 +1,6 @@ """Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. -Mirrors ``recipes.refl.model_adaptor.wan22`` but targets the WAN 2.1 T2V +Mirrors ``experimental.refl.model_adaptor.wan22`` but targets the WAN 2.1 T2V single-DiT stack. The REFL-specific pieces live here: - :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch @@ -25,7 +25,7 @@ import torch -from recipes.refl.model_adaptor.types import DiffuseWithGradResult +from experimental.refl.model_adaptor.types import DiffuseWithGradResult from unirl.models.wan21.bundle import WAN21Bundle from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions diff --git a/recipes/refl/model_adaptor/wan22.py b/experimental/refl/model_adaptor/wan22.py similarity index 99% rename from recipes/refl/model_adaptor/wan22.py rename to experimental/refl/model_adaptor/wan22.py index 215ee7b61..ad486eb7b 100644 --- a/recipes/refl/model_adaptor/wan22.py +++ b/experimental/refl/model_adaptor/wan22.py @@ -8,7 +8,7 @@ import torch -from recipes.refl.model_adaptor.types import DiffuseWithGradResult +from experimental.refl.model_adaptor.types import DiffuseWithGradResult from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage diff --git a/recipes/refl/rewards/__init__.py b/experimental/refl/rewards/__init__.py similarity index 100% rename from recipes/refl/rewards/__init__.py rename to experimental/refl/rewards/__init__.py diff --git a/recipes/refl/rewards/face/__init__.py b/experimental/refl/rewards/face/__init__.py similarity index 100% rename from recipes/refl/rewards/face/__init__.py rename to experimental/refl/rewards/face/__init__.py diff --git a/recipes/refl/rewards/face/face_tools.py b/experimental/refl/rewards/face/face_tools.py similarity index 100% rename from recipes/refl/rewards/face/face_tools.py rename to experimental/refl/rewards/face/face_tools.py diff --git a/recipes/refl/rewards/face/requirements.txt b/experimental/refl/rewards/face/requirements.txt similarity index 100% rename from recipes/refl/rewards/face/requirements.txt rename to experimental/refl/rewards/face/requirements.txt diff --git a/recipes/refl/rewards/face/scorer.py b/experimental/refl/rewards/face/scorer.py similarity index 100% rename from recipes/refl/rewards/face/scorer.py rename to experimental/refl/rewards/face/scorer.py diff --git a/recipes/refl/rewards/videoalign/__init__.py b/experimental/refl/rewards/videoalign/__init__.py similarity index 88% rename from recipes/refl/rewards/videoalign/__init__.py rename to experimental/refl/rewards/videoalign/__init__.py index 90ee502ad..468b776a8 100644 --- a/recipes/refl/rewards/videoalign/__init__.py +++ b/experimental/refl/rewards/videoalign/__init__.py @@ -3,7 +3,7 @@ Self-contained port of the VideoAlign reward family for the REFL WAN recipe. The reward model code (``Qwen2VLRewardModelBT``), prompt templates, config dataclasses and checkpoint-loading helpers all live under -:mod:`recipes.refl.rewards.videoalign.model`, matching the recipe-local +:mod:`experimental.refl.rewards.videoalign.model`, matching the recipe-local layout used by the WAN22 face reward. Public API diff --git a/recipes/refl/rewards/videoalign/model/__init__.py b/experimental/refl/rewards/videoalign/model/__init__.py similarity index 100% rename from recipes/refl/rewards/videoalign/model/__init__.py rename to experimental/refl/rewards/videoalign/model/__init__.py diff --git a/recipes/refl/rewards/videoalign/model/checkpoint.py b/experimental/refl/rewards/videoalign/model/checkpoint.py similarity index 100% rename from recipes/refl/rewards/videoalign/model/checkpoint.py rename to experimental/refl/rewards/videoalign/model/checkpoint.py diff --git a/recipes/refl/rewards/videoalign/model/configs.py b/experimental/refl/rewards/videoalign/model/configs.py similarity index 98% rename from recipes/refl/rewards/videoalign/model/configs.py rename to experimental/refl/rewards/videoalign/model/configs.py index e9845c437..f52285c81 100644 --- a/recipes/refl/rewards/videoalign/model/configs.py +++ b/experimental/refl/rewards/videoalign/model/configs.py @@ -36,7 +36,7 @@ class TrainingConfig: Only ``bf16`` / ``fp16`` / ``gradient_checkpointing`` / ``disable_flash_attn2`` are actually consumed by the inference path - (see :func:`recipes.refl.rewards.videoalign.model.factory.create_model_and_processor`). + (see :func:`experimental.refl.rewards.videoalign.model.factory.create_model_and_processor`). The other fields are kept so ``TrainingConfig(**model_config_json["training_args"])`` still works when someone wants to introspect the saved config. """ diff --git a/recipes/refl/rewards/videoalign/model/factory.py b/experimental/refl/rewards/videoalign/model/factory.py similarity index 100% rename from recipes/refl/rewards/videoalign/model/factory.py rename to experimental/refl/rewards/videoalign/model/factory.py diff --git a/recipes/refl/rewards/videoalign/model/prompt_template.py b/experimental/refl/rewards/videoalign/model/prompt_template.py similarity index 100% rename from recipes/refl/rewards/videoalign/model/prompt_template.py rename to experimental/refl/rewards/videoalign/model/prompt_template.py diff --git a/recipes/refl/rewards/videoalign/model/reward_model.py b/experimental/refl/rewards/videoalign/model/reward_model.py similarity index 99% rename from recipes/refl/rewards/videoalign/model/reward_model.py rename to experimental/refl/rewards/videoalign/model/reward_model.py index 095103851..5159e6479 100644 --- a/recipes/refl/rewards/videoalign/model/reward_model.py +++ b/experimental/refl/rewards/videoalign/model/reward_model.py @@ -12,7 +12,7 @@ The forward signature is preserved so the same checkpoint state_dict loads cleanly (no key remapping needed beyond the standard transformers>=5 ``base_model.model.model.language_model.*`` shim handled by -:func:`recipes.refl.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). +:func:`experimental.refl.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). """ from __future__ import annotations diff --git a/recipes/refl/rewards/videoalign/requirements.txt b/experimental/refl/rewards/videoalign/requirements.txt similarity index 100% rename from recipes/refl/rewards/videoalign/requirements.txt rename to experimental/refl/rewards/videoalign/requirements.txt diff --git a/recipes/refl/rewards/videoalign/scorer.py b/experimental/refl/rewards/videoalign/scorer.py similarity index 99% rename from recipes/refl/rewards/videoalign/scorer.py rename to experimental/refl/rewards/videoalign/scorer.py index 5c01adf30..ee55cdeaa 100644 --- a/recipes/refl/rewards/videoalign/scorer.py +++ b/experimental/refl/rewards/videoalign/scorer.py @@ -21,7 +21,7 @@ This scorer no longer requires the sibling ``mmrl`` repo on disk. The Qwen2-VL reward backbone, prompt template, checkpoint loader and inference wrapper all live under -:mod:`recipes.refl.rewards.videoalign.model` / :mod:`...wrapper`. The +:mod:`experimental.refl.rewards.videoalign.model` / :mod:`...wrapper`. The ``mmrl_repo_root`` Spec field has been removed; ``MMRL_REPO_ROOT`` env var is now irrelevant. """ diff --git a/recipes/refl/rewards/videoalign/wrapper.py b/experimental/refl/rewards/videoalign/wrapper.py similarity index 99% rename from recipes/refl/rewards/videoalign/wrapper.py rename to experimental/refl/rewards/videoalign/wrapper.py index d4e2c9732..8a7e0d61a 100644 --- a/recipes/refl/rewards/videoalign/wrapper.py +++ b/experimental/refl/rewards/videoalign/wrapper.py @@ -2,7 +2,7 @@ Self-contained re-implementation — does NOT touch ``sys.path`` . The Qwen2-VL reward backbone, its prompt template and the checkpoint loader all live under -:mod:`recipes.refl.rewards.videoalign.model`. +:mod:`experimental.refl.rewards.videoalign.model`. Public API ---------- diff --git a/recipes/refl/roles.py b/experimental/refl/roles.py similarity index 97% rename from recipes/refl/roles.py rename to experimental/refl/roles.py index f4e4edd73..eb60ccc0e 100644 --- a/recipes/refl/roles.py +++ b/experimental/refl/roles.py @@ -11,7 +11,7 @@ actor.step(max_grad_norm=…) # ctx exit routed grads → optimizer The pipeline named by ``pipeline_target`` must expose the recipe contract -(see ``recipes.refl.model_adaptor``): ``build_refl_conditions(texts, images=…, +(see ``experimental.refl.model_adaptor``): ``build_refl_conditions(texts, images=…, params=…)`` plus a ``diffusion`` stage with ``diffuse_with_grad`` and a ``vae_decode`` stage with ``decode_with_grad``. """ @@ -117,12 +117,12 @@ def initialize(self) -> None: if not hasattr(stage, method): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} .{stage_attr} lacks {method}(...); " - f"use a recipes.refl.model_adaptor pipeline (or implement the REFL contract)." + f"use a experimental.refl.model_adaptor pipeline (or implement the REFL contract)." ) if not hasattr(self.pipeline, "build_refl_conditions"): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} lacks build_refl_conditions(...); " - f"use a recipes.refl.model_adaptor pipeline." + f"use a experimental.refl.model_adaptor pipeline." ) # FSDP-wrap pipeline.bundle.transformer in place + LoRA + optimizer. diff --git a/recipes/refl/run.py b/experimental/refl/run.py similarity index 74% rename from recipes/refl/run.py rename to experimental/refl/run.py index de7f6129f..efafd930e 100755 --- a/recipes/refl/run.py +++ b/experimental/refl/run.py @@ -1,12 +1,12 @@ #!/usr/bin/env python -"""Hydra entry point for the recipes.refl REFL recipe.""" +"""Hydra entry point for the experimental.refl REFL recipe.""" from __future__ import annotations import hydra from omegaconf import DictConfig -from recipes.refl.trainer import REFLTrainer +from experimental.refl.trainer import REFLTrainer @hydra.main(version_base=None, config_path="configs", config_name="wan22_i2v_face_refl") diff --git a/recipes/refl/trainer.py b/experimental/refl/trainer.py similarity index 99% rename from recipes/refl/trainer.py rename to experimental/refl/trainer.py index 4dac2fd16..15bba2ea7 100644 --- a/recipes/refl/trainer.py +++ b/experimental/refl/trainer.py @@ -1,7 +1,7 @@ """REFLTrainer — recipe driver for WAN REFL/BPTT (video reward backprop). The video sibling of :class:`unirl.trainer.refl.RewardBackpropTrainer`: two -roles, always — a :class:`recipes.refl.roles.ReflActorRole` (FSDP WAN + +roles, always — a :class:`experimental.refl.roles.ReflActorRole` (FSDP WAN + grad BPTT sampling + optimizer) and a frozen differentiable video reward (:class:`unirl.reward.service.RewardService`), colocated on the same worker slab so decoded video never leaves the GPU. Each step runs, under the diff --git a/scripts/check_recipe_targets.py b/scripts/check_recipe_targets.py index 176ba4097..eeb1a7d69 100755 --- a/scripts/check_recipe_targets.py +++ b/scripts/check_recipe_targets.py @@ -23,11 +23,11 @@ ROOT = Path(__file__).resolve().parents[1] # YAML trees that hold recipes / stage configs with ``_target_`` entries. -SCAN_DIRS = ["examples", "recipes", "CPPO", "DRPO", "FlowDPPO", "unirl"] +SCAN_DIRS = ["examples", "experimental", "CPPO", "DRPO", "FlowDPPO", "unirl"] # Vendored / sub-project trees kept byte-pristine (mirror .pre-commit-config exclude). SKIP_PARTS = {".git", "vendor"} -_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?((?:unirl|recipes)\.[A-Za-z0-9_.]+)['"]?\s*$""") +_TARGET_RE = re.compile(r"""^\s*_target_:\s*['"]?((?:unirl|experimental)\.[A-Za-z0-9_.]+)['"]?\s*$""") @lru_cache(maxsize=None) diff --git a/scripts/verify_refl_kl_batching.py b/scripts/verify_refl_kl_batching.py index 69f62d764..cd76e979b 100755 --- a/scripts/verify_refl_kl_batching.py +++ b/scripts/verify_refl_kl_batching.py @@ -29,7 +29,7 @@ import torch # noqa: E402 -from recipes.refl.roles import REFLGenerated, REFLLossMetrics # noqa: E402 +from experimental.refl.roles import REFLGenerated, REFLLossMetrics # noqa: E402 from unirl.distributed.tensor.pytree import infer_batch_size, pytree_cat, pytree_chunk # noqa: E402 # (batch_size, dp) — B == dp (the only previously-verified shape), B > dp, From de2e24ee196d0873a0d849cf14b35c9001c7b3d8 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 22:07:31 +0800 Subject: [PATCH 19/24] =?UTF-8?q?refactor(experimental):=20mirror=20core?= =?UTF-8?q?=20names=20for=20same-kind=20content=20=E2=80=94=20examples/,?= =?UTF-8?q?=20reward/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Naming rule for the experimental tier (two-way flow): directories holding the SAME kind of content as a core location take the core name, so promotion/demotion is a structural no-op — refl/configs/ → refl/examples/ (mirrors the top-level examples/; graduates into examples/refl/) and refl/rewards/ → refl/reward/ (mirrors unirl/reward/; graduates into unirl/reward/local/). Different-kind content keeps a distinct name on purpose (model_adaptor/ wraps core models to the BPTT contract — it is not a models/ and graduates by merging, not by moving). --- experimental/refl/README.md | 6 +++--- .../{configs => examples}/wan21_t2v_videoalign_refl.yaml | 4 ++-- .../refl/{configs => examples}/wan22_i2v_face_refl.yaml | 4 ++-- experimental/refl/{rewards => reward}/__init__.py | 0 experimental/refl/{rewards => reward}/face/__init__.py | 0 experimental/refl/{rewards => reward}/face/face_tools.py | 0 experimental/refl/{rewards => reward}/face/requirements.txt | 0 experimental/refl/{rewards => reward}/face/scorer.py | 0 .../refl/{rewards => reward}/videoalign/__init__.py | 2 +- .../refl/{rewards => reward}/videoalign/model/__init__.py | 0 .../refl/{rewards => reward}/videoalign/model/checkpoint.py | 0 .../refl/{rewards => reward}/videoalign/model/configs.py | 2 +- .../refl/{rewards => reward}/videoalign/model/factory.py | 0 .../{rewards => reward}/videoalign/model/prompt_template.py | 0 .../{rewards => reward}/videoalign/model/reward_model.py | 2 +- .../refl/{rewards => reward}/videoalign/requirements.txt | 0 experimental/refl/{rewards => reward}/videoalign/scorer.py | 2 +- experimental/refl/{rewards => reward}/videoalign/wrapper.py | 2 +- experimental/refl/run.py | 2 +- 19 files changed, 13 insertions(+), 13 deletions(-) rename experimental/refl/{configs => examples}/wan21_t2v_videoalign_refl.yaml (97%) rename experimental/refl/{configs => examples}/wan22_i2v_face_refl.yaml (97%) rename experimental/refl/{rewards => reward}/__init__.py (100%) rename experimental/refl/{rewards => reward}/face/__init__.py (100%) rename experimental/refl/{rewards => reward}/face/face_tools.py (100%) rename experimental/refl/{rewards => reward}/face/requirements.txt (100%) rename experimental/refl/{rewards => reward}/face/scorer.py (100%) rename experimental/refl/{rewards => reward}/videoalign/__init__.py (88%) rename experimental/refl/{rewards => reward}/videoalign/model/__init__.py (100%) rename experimental/refl/{rewards => reward}/videoalign/model/checkpoint.py (100%) rename experimental/refl/{rewards => reward}/videoalign/model/configs.py (98%) rename experimental/refl/{rewards => reward}/videoalign/model/factory.py (100%) rename experimental/refl/{rewards => reward}/videoalign/model/prompt_template.py (100%) rename experimental/refl/{rewards => reward}/videoalign/model/reward_model.py (99%) rename experimental/refl/{rewards => reward}/videoalign/requirements.txt (100%) rename experimental/refl/{rewards => reward}/videoalign/scorer.py (99%) rename experimental/refl/{rewards => reward}/videoalign/wrapper.py (99%) diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 5ab80bc3a..1c9eabd82 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -20,7 +20,7 @@ RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan21_t2v_videoal # 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/rewards/face/requirements.txt +pip install -r experimental/refl/reward/face/requirements.txt export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers \ FACE_MODEL_PATH=/path/to/antelodev2 \ DATA_PATH=/path/to/i2v_prompts.jsonl @@ -34,8 +34,8 @@ 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`) | | `model_adaptor/` | Per-model adaptations of core pipelines to the BPTT contract (`types.py` defines it): `wan21.py`, `wan22.py` | -| `rewards/` | Recipe-local differentiable rewards (VideoAlign, Face), each with an additive-only `requirements.txt` | -| `configs/` | Flat Hydra configs (repo-wide schema) | +| `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) | ## Environment diff --git a/experimental/refl/configs/wan21_t2v_videoalign_refl.yaml b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml similarity index 97% rename from experimental/refl/configs/wan21_t2v_videoalign_refl.yaml rename to experimental/refl/examples/wan21_t2v_videoalign_refl.yaml index c425adcc3..126c8cbbd 100644 --- a/experimental/refl/configs/wan21_t2v_videoalign_refl.yaml +++ b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml @@ -87,10 +87,10 @@ actor: reward: _target_: unirl.reward.service.RewardService backend: - _target_: experimental.refl.rewards.videoalign.VideoAlignRewardScorer + _target_: experimental.refl.reward.videoalign.VideoAlignRewardScorer base_device: cuda config: - _target_: experimental.refl.rewards.videoalign.VideoAlignSpec + _target_: experimental.refl.reward.videoalign.VideoAlignSpec reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH} device: cuda batch_size: 1 diff --git a/experimental/refl/configs/wan22_i2v_face_refl.yaml b/experimental/refl/examples/wan22_i2v_face_refl.yaml similarity index 97% rename from experimental/refl/configs/wan22_i2v_face_refl.yaml rename to experimental/refl/examples/wan22_i2v_face_refl.yaml index 6fce57fcf..70628d0f3 100644 --- a/experimental/refl/configs/wan22_i2v_face_refl.yaml +++ b/experimental/refl/examples/wan22_i2v_face_refl.yaml @@ -85,10 +85,10 @@ actor: reward: _target_: unirl.reward.service.RewardService backend: - _target_: experimental.refl.rewards.face.FaceRewardScorer + _target_: experimental.refl.reward.face.FaceRewardScorer base_device: cuda config: - _target_: experimental.refl.rewards.face.FaceRewardSpec + _target_: experimental.refl.reward.face.FaceRewardSpec model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt} device: cuda batch_size: 1 diff --git a/experimental/refl/rewards/__init__.py b/experimental/refl/reward/__init__.py similarity index 100% rename from experimental/refl/rewards/__init__.py rename to experimental/refl/reward/__init__.py diff --git a/experimental/refl/rewards/face/__init__.py b/experimental/refl/reward/face/__init__.py similarity index 100% rename from experimental/refl/rewards/face/__init__.py rename to experimental/refl/reward/face/__init__.py diff --git a/experimental/refl/rewards/face/face_tools.py b/experimental/refl/reward/face/face_tools.py similarity index 100% rename from experimental/refl/rewards/face/face_tools.py rename to experimental/refl/reward/face/face_tools.py diff --git a/experimental/refl/rewards/face/requirements.txt b/experimental/refl/reward/face/requirements.txt similarity index 100% rename from experimental/refl/rewards/face/requirements.txt rename to experimental/refl/reward/face/requirements.txt diff --git a/experimental/refl/rewards/face/scorer.py b/experimental/refl/reward/face/scorer.py similarity index 100% rename from experimental/refl/rewards/face/scorer.py rename to experimental/refl/reward/face/scorer.py diff --git a/experimental/refl/rewards/videoalign/__init__.py b/experimental/refl/reward/videoalign/__init__.py similarity index 88% rename from experimental/refl/rewards/videoalign/__init__.py rename to experimental/refl/reward/videoalign/__init__.py index 468b776a8..8b1b998fc 100644 --- a/experimental/refl/rewards/videoalign/__init__.py +++ b/experimental/refl/reward/videoalign/__init__.py @@ -3,7 +3,7 @@ Self-contained port of the VideoAlign reward family for the REFL WAN recipe. The reward model code (``Qwen2VLRewardModelBT``), prompt templates, config dataclasses and checkpoint-loading helpers all live under -:mod:`experimental.refl.rewards.videoalign.model`, matching the recipe-local +:mod:`experimental.refl.reward.videoalign.model`, matching the recipe-local layout used by the WAN22 face reward. Public API diff --git a/experimental/refl/rewards/videoalign/model/__init__.py b/experimental/refl/reward/videoalign/model/__init__.py similarity index 100% rename from experimental/refl/rewards/videoalign/model/__init__.py rename to experimental/refl/reward/videoalign/model/__init__.py diff --git a/experimental/refl/rewards/videoalign/model/checkpoint.py b/experimental/refl/reward/videoalign/model/checkpoint.py similarity index 100% rename from experimental/refl/rewards/videoalign/model/checkpoint.py rename to experimental/refl/reward/videoalign/model/checkpoint.py diff --git a/experimental/refl/rewards/videoalign/model/configs.py b/experimental/refl/reward/videoalign/model/configs.py similarity index 98% rename from experimental/refl/rewards/videoalign/model/configs.py rename to experimental/refl/reward/videoalign/model/configs.py index f52285c81..f2b32003b 100644 --- a/experimental/refl/rewards/videoalign/model/configs.py +++ b/experimental/refl/reward/videoalign/model/configs.py @@ -36,7 +36,7 @@ class TrainingConfig: Only ``bf16`` / ``fp16`` / ``gradient_checkpointing`` / ``disable_flash_attn2`` are actually consumed by the inference path - (see :func:`experimental.refl.rewards.videoalign.model.factory.create_model_and_processor`). + (see :func:`experimental.refl.reward.videoalign.model.factory.create_model_and_processor`). The other fields are kept so ``TrainingConfig(**model_config_json["training_args"])`` still works when someone wants to introspect the saved config. """ diff --git a/experimental/refl/rewards/videoalign/model/factory.py b/experimental/refl/reward/videoalign/model/factory.py similarity index 100% rename from experimental/refl/rewards/videoalign/model/factory.py rename to experimental/refl/reward/videoalign/model/factory.py diff --git a/experimental/refl/rewards/videoalign/model/prompt_template.py b/experimental/refl/reward/videoalign/model/prompt_template.py similarity index 100% rename from experimental/refl/rewards/videoalign/model/prompt_template.py rename to experimental/refl/reward/videoalign/model/prompt_template.py diff --git a/experimental/refl/rewards/videoalign/model/reward_model.py b/experimental/refl/reward/videoalign/model/reward_model.py similarity index 99% rename from experimental/refl/rewards/videoalign/model/reward_model.py rename to experimental/refl/reward/videoalign/model/reward_model.py index 5159e6479..b5179a097 100644 --- a/experimental/refl/rewards/videoalign/model/reward_model.py +++ b/experimental/refl/reward/videoalign/model/reward_model.py @@ -12,7 +12,7 @@ The forward signature is preserved so the same checkpoint state_dict loads cleanly (no key remapping needed beyond the standard transformers>=5 ``base_model.model.model.language_model.*`` shim handled by -:func:`experimental.refl.rewards.videoalign.model.checkpoint.load_model_from_checkpoint`). +:func:`experimental.refl.reward.videoalign.model.checkpoint.load_model_from_checkpoint`). """ from __future__ import annotations diff --git a/experimental/refl/rewards/videoalign/requirements.txt b/experimental/refl/reward/videoalign/requirements.txt similarity index 100% rename from experimental/refl/rewards/videoalign/requirements.txt rename to experimental/refl/reward/videoalign/requirements.txt diff --git a/experimental/refl/rewards/videoalign/scorer.py b/experimental/refl/reward/videoalign/scorer.py similarity index 99% rename from experimental/refl/rewards/videoalign/scorer.py rename to experimental/refl/reward/videoalign/scorer.py index ee55cdeaa..f19fcc94e 100644 --- a/experimental/refl/rewards/videoalign/scorer.py +++ b/experimental/refl/reward/videoalign/scorer.py @@ -21,7 +21,7 @@ This scorer no longer requires the sibling ``mmrl`` repo on disk. The Qwen2-VL reward backbone, prompt template, checkpoint loader and inference wrapper all live under -:mod:`experimental.refl.rewards.videoalign.model` / :mod:`...wrapper`. The +:mod:`experimental.refl.reward.videoalign.model` / :mod:`...wrapper`. The ``mmrl_repo_root`` Spec field has been removed; ``MMRL_REPO_ROOT`` env var is now irrelevant. """ diff --git a/experimental/refl/rewards/videoalign/wrapper.py b/experimental/refl/reward/videoalign/wrapper.py similarity index 99% rename from experimental/refl/rewards/videoalign/wrapper.py rename to experimental/refl/reward/videoalign/wrapper.py index 8a7e0d61a..e5ebc2397 100644 --- a/experimental/refl/rewards/videoalign/wrapper.py +++ b/experimental/refl/reward/videoalign/wrapper.py @@ -2,7 +2,7 @@ Self-contained re-implementation — does NOT touch ``sys.path`` . The Qwen2-VL reward backbone, its prompt template and the checkpoint loader all live under -:mod:`experimental.refl.rewards.videoalign.model`. +:mod:`experimental.refl.reward.videoalign.model`. Public API ---------- diff --git a/experimental/refl/run.py b/experimental/refl/run.py index efafd930e..2e28a1328 100755 --- a/experimental/refl/run.py +++ b/experimental/refl/run.py @@ -9,7 +9,7 @@ from experimental.refl.trainer import REFLTrainer -@hydra.main(version_base=None, config_path="configs", config_name="wan22_i2v_face_refl") +@hydra.main(version_base=None, config_path="examples", config_name="wan22_i2v_face_refl") def main(cfg: DictConfig) -> None: trainer = REFLTrainer(cfg=cfg) trainer.train() From c1653356782a6e501c9859c3b3de804eef463894 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Thu, 30 Jul 2026 22:21:46 +0800 Subject: [PATCH 20/24] =?UTF-8?q?refactor(experimental):=20name=20by=20gra?= =?UTF-8?q?duation=20destination=20=E2=80=94=20model=5Fadaptor/=20?= =?UTF-8?q?=E2=86=92=20models/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The mirror rule settles on one criterion: a directory is named after the core home its content graduates into. wan21.py / wan22.py define pipeline and stage classes — the same kind of artifact unirl/models// holds — and graduate by merging into those model packages, so the directory mirrors unirl/models/ (plural, matching core exactly; reward/ is singular because unirl/reward/ is). The earlier model_adaptor name predates the mirror rule; the fork-confusion concern it addressed is now handled by the graduation notes in the layout table and module docstrings. --- experimental/refl/README.md | 2 +- experimental/refl/examples/wan21_t2v_videoalign_refl.yaml | 2 +- experimental/refl/examples/wan22_i2v_face_refl.yaml | 2 +- experimental/refl/{model_adaptor => models}/__init__.py | 0 experimental/refl/{model_adaptor => models}/types.py | 0 experimental/refl/{model_adaptor => models}/wan21.py | 4 ++-- experimental/refl/{model_adaptor => models}/wan22.py | 2 +- experimental/refl/roles.py | 6 +++--- 8 files changed, 9 insertions(+), 9 deletions(-) rename experimental/refl/{model_adaptor => models}/__init__.py (100%) rename experimental/refl/{model_adaptor => models}/types.py (100%) rename experimental/refl/{model_adaptor => models}/wan21.py (99%) rename experimental/refl/{model_adaptor => models}/wan22.py (99%) diff --git a/experimental/refl/README.md b/experimental/refl/README.md index 1c9eabd82..764c093ec 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -33,7 +33,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`) | -| `model_adaptor/` | Per-model adaptations of core pipelines to the BPTT contract (`types.py` defines it): `wan21.py`, `wan22.py` | +| `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) | | `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) | diff --git a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml index 126c8cbbd..88de60c50 100644 --- a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml +++ b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml @@ -19,7 +19,7 @@ actor: _target_: experimental.refl.roles.ReflActorRole # Family selector — swap pipeline_target + model_config for another family, # no code changes (mirrors ReFLPolicy's pipeline_target contract). - pipeline_target: experimental.refl.model_adaptor.wan21.Wan21ReflPipeline + pipeline_target: experimental.refl.models.wan21.Wan21ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.25 diff --git a/experimental/refl/examples/wan22_i2v_face_refl.yaml b/experimental/refl/examples/wan22_i2v_face_refl.yaml index 70628d0f3..b08a424f2 100644 --- a/experimental/refl/examples/wan22_i2v_face_refl.yaml +++ b/experimental/refl/examples/wan22_i2v_face_refl.yaml @@ -18,7 +18,7 @@ max_grad_norm: 1.0 actor: _target_: experimental.refl.roles.ReflActorRole - pipeline_target: experimental.refl.model_adaptor.wan22.Wan22ReflPipeline + pipeline_target: experimental.refl.models.wan22.Wan22ReflPipeline block_class_names: ["WanTransformerBlock"] # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL. reward_weight: 0.1 diff --git a/experimental/refl/model_adaptor/__init__.py b/experimental/refl/models/__init__.py similarity index 100% rename from experimental/refl/model_adaptor/__init__.py rename to experimental/refl/models/__init__.py diff --git a/experimental/refl/model_adaptor/types.py b/experimental/refl/models/types.py similarity index 100% rename from experimental/refl/model_adaptor/types.py rename to experimental/refl/models/types.py diff --git a/experimental/refl/model_adaptor/wan21.py b/experimental/refl/models/wan21.py similarity index 99% rename from experimental/refl/model_adaptor/wan21.py rename to experimental/refl/models/wan21.py index 7e76b0a73..508289c67 100644 --- a/experimental/refl/model_adaptor/wan21.py +++ b/experimental/refl/models/wan21.py @@ -1,6 +1,6 @@ """Recipe-local WAN 2.1 T2V step + stage + pipeline for REFL BPTT. -Mirrors ``experimental.refl.model_adaptor.wan22`` but targets the WAN 2.1 T2V +Mirrors ``experimental.refl.models.wan22`` but targets the WAN 2.1 T2V single-DiT stack. The REFL-specific pieces live here: - :class:`Wan21ReflDiffusionStep` — strict recipe-local single-branch @@ -25,7 +25,7 @@ import torch -from experimental.refl.model_adaptor.types import DiffuseWithGradResult +from experimental.refl.models.types import DiffuseWithGradResult from unirl.models.wan21.bundle import WAN21Bundle from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions diff --git a/experimental/refl/model_adaptor/wan22.py b/experimental/refl/models/wan22.py similarity index 99% rename from experimental/refl/model_adaptor/wan22.py rename to experimental/refl/models/wan22.py index ad486eb7b..74d865d1a 100644 --- a/experimental/refl/model_adaptor/wan22.py +++ b/experimental/refl/models/wan22.py @@ -8,7 +8,7 @@ import torch -from experimental.refl.model_adaptor.types import DiffuseWithGradResult +from experimental.refl.models.types import DiffuseWithGradResult from unirl.models.wan21.clip_vision_encode import WAN21CLIPVisionEncodeStage from unirl.models.wan21.conditions import WAN21Conditions from unirl.models.wan21.image_encode import WAN21ImageLatentEncodeStage diff --git a/experimental/refl/roles.py b/experimental/refl/roles.py index eb60ccc0e..7031eafbf 100644 --- a/experimental/refl/roles.py +++ b/experimental/refl/roles.py @@ -11,7 +11,7 @@ actor.step(max_grad_norm=…) # ctx exit routed grads → optimizer The pipeline named by ``pipeline_target`` must expose the recipe contract -(see ``experimental.refl.model_adaptor``): ``build_refl_conditions(texts, images=…, +(see ``experimental.refl.models``): ``build_refl_conditions(texts, images=…, params=…)`` plus a ``diffusion`` stage with ``diffuse_with_grad`` and a ``vae_decode`` stage with ``decode_with_grad``. """ @@ -117,12 +117,12 @@ def initialize(self) -> None: if not hasattr(stage, method): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} .{stage_attr} lacks {method}(...); " - f"use a experimental.refl.model_adaptor pipeline (or implement the REFL contract)." + f"use a experimental.refl.models pipeline (or implement the REFL contract)." ) if not hasattr(self.pipeline, "build_refl_conditions"): raise TypeError( f"ReflActorRole: pipeline {self._pipeline_target} lacks build_refl_conditions(...); " - f"use a experimental.refl.model_adaptor pipeline." + f"use a experimental.refl.models pipeline." ) # FSDP-wrap pipeline.bundle.transformer in place + LoRA + optimizer. From f3f63c004d5d93ba41121ffbb9a8c4909d526bfc Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 06:53:24 +0800 Subject: [PATCH 21/24] docs(experimental/refl): record the 150-rollout trend-run verdict; sync main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mechanics PASS on the declared 5.6 stack (9.5h, 228s/step, clean exit, no reward collapse, grads healthy). Reward is FLAT at this horizon on substitute assets (pickscore prompts + VideoReward-2B): segment means 1.91/1.95/1.91/2.07/1.80/1.95, first25→last25 +2.4%, OLS slope ~0. Learning-effect verdict deferred to longer horizons / original assets; lr-sensitivity diagnostic tracked in PR #210. --- 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 764c093ec..c7ddacd1d 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -53,5 +53,5 @@ core stack. | `wan21_t2v_videoalign_refl` (835 rollouts) | 8xH20 | pre-adjustment (`e3c6b940` lineage) | contributor long run — reward curve in PR #210 | | `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` | -| `wan21_t2v_videoalign_refl` (150-rollout trend run, 5.6 stack) | 8xH20 | `9087a671` | running — curve to be attached | +| `wan21_t2v_videoalign_refl` (150-rollout trend, 5.6 stack, substitute assets: pickscore prompts + VideoReward-2B) | 8xH20 | `9087a671` | mechanics PASS (9.5h, 228s/step, clean exit, no collapse, grads healthy); reward FLAT at this horizon — segment means 1.91/1.95/1.91/2.07/1.80/1.95, first25→last25 +2.4%. Learning-effect verdict deferred to longer horizons / original assets; lr-sensitivity diagnostic tracked in PR #210 | | `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | From 379ce6aae899c3fb0282762a80bb6f32d97edb17 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 10:14:33 +0800 Subject: [PATCH 22/24] docs(experimental/refl): drop the trend-run row from the verification table The 150-rollout run was a maintainer-side sanity diagnostic on substitute assets; its record lives in the PR #210 discussion, not in the package's verification table, which lists only runs with product-level standing. --- experimental/refl/README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/experimental/refl/README.md b/experimental/refl/README.md index c7ddacd1d..9c38aaeb9 100644 --- a/experimental/refl/README.md +++ b/experimental/refl/README.md @@ -53,5 +53,4 @@ core stack. | `wan21_t2v_videoalign_refl` (835 rollouts) | 8xH20 | pre-adjustment (`e3c6b940` lineage) | contributor long run — reward curve in PR #210 | | `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` | -| `wan21_t2v_videoalign_refl` (150-rollout trend, 5.6 stack, substitute assets: pickscore prompts + VideoReward-2B) | 8xH20 | `9087a671` | mechanics PASS (9.5h, 228s/step, clean exit, no collapse, grads healthy); reward FLAT at this horizon — segment means 1.91/1.95/1.91/2.07/1.80/1.95, first25→last25 +2.4%. Learning-effect verdict deferred to longer horizons / original assets; lr-sensitivity diagnostic tracked in PR #210 | | `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) | From fe2788c6711d7ed184d985513478d77ef01bd633 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 10:42:00 +0800 Subject: [PATCH 23/24] fix(experimental/refl): restore the contributor's fixed-noise seed semantics Checking pass on the flat trend run found this adjustment-era divergence: the role was rewriting params.seed per rollout/rank (decorrelated noise), while the contributor's verified 835-rollout curve trains DRaFT on a fixed initial noise (seed used verbatim, eta=0 ODE). Ship his regime; varying noise is a semantics change that needs its own evidence. Drops the now consumer-less rollout_id plumbing from generate_samples/train_step. --- experimental/refl/roles.py | 15 +++++---------- experimental/refl/trainer.py | 5 ++--- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/experimental/refl/roles.py b/experimental/refl/roles.py index 7031eafbf..2265695d6 100644 --- a/experimental/refl/roles.py +++ b/experimental/refl/roles.py @@ -157,7 +157,6 @@ def generate_samples( texts: Texts, images: Optional[Images] = None, params: DiffusionSamplingParams, - rollout_id: int = 0, ) -> REFLGenerated: """Grad-enabled BPTT sampling + in-graph VAE decode. @@ -180,15 +179,11 @@ def generate_samples( ) sampler_kwargs["kl_weight"] = self.kl_weight - # Decorrelate init noise across DP shards and rollouts (same scheme as - # ReFLPolicy): the config seed is the base, not the per-step value. - dp_rank = int(self.rank_info.dp_rank) if self.rank_info is not None else 0 - base_seed = int(params.seed) if params.seed is not None else 42 - params = dataclasses.replace( - params, - seed=base_seed + 1000 * int(rollout_id) + dp_rank, - sampler_kwargs=sampler_kwargs, - ) + # Fixed init noise across rollouts and ranks — the contributor's + # verified DRaFT regime (params.seed used verbatim by the stage). + # Varying noise per rollout is a training-semantics change; do not + # introduce it here without its own evidence. + params = dataclasses.replace(params, sampler_kwargs=sampler_kwargs) conditions = self.pipeline.build_refl_conditions(texts, images=images, params=params) schedule = get_sigma_schedule( diff --git a/experimental/refl/trainer.py b/experimental/refl/trainer.py index 15bba2ea7..3217fb3d8 100644 --- a/experimental/refl/trainer.py +++ b/experimental/refl/trainer.py @@ -103,7 +103,7 @@ def __init__(self, *, cfg: DictConfig) -> None: self.max_grad_norm, ) - def train_step(self, inputs: Sample, *, rollout_id: int) -> Dict[str, float]: + def train_step(self, inputs: Sample) -> Dict[str, float]: """One enable_grad() generate → score → backward, then optimizer step.""" t0 = time.perf_counter() texts = _text_inputs(inputs) @@ -114,7 +114,6 @@ def train_step(self, inputs: Sample, *, rollout_id: int) -> Dict[str, float]: texts=texts, images=images, params=self.sampling_params, - rollout_id=rollout_id, ) rewards = self.reward.score_differentiable(gen.decoded, list(texts.texts), records) loss_metrics = self.actor.forward_backward_loss(rewards=rewards, kl_loss=gen.kl_loss) @@ -155,7 +154,7 @@ def train( try: for rollout_id in range(start, num_rollouts): inputs = self.data_source.get_samples(self.batch_size) - metrics = self.train_step(inputs, rollout_id=rollout_id) + metrics = self.train_step(inputs) logger.info( "rollout %d/%d reward=%.4f loss=%.4f kl=%.4f grad_norm=%.4f %.1fs", rollout_id + 1, From f8914721e43c89dcf51b48cf7774f729141b1a78 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 10:52:44 +0800 Subject: [PATCH 24/24] =?UTF-8?q?fix(experimental/refl):=20restore=20maste?= =?UTF-8?q?r=5Fdtype=20fp32=20=E2=80=94=20bf16=20master=20freezes=20lora?= =?UTF-8?q?=5FA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Checking follow-up on the flat trend run. The contributor's configs carried master_dtype: fp32; the adjustment dropped it to dodge FSDP2's uniform-original-dtype assert on the fleet image's torch 2.7.1. That assert is an environment-misalignment symptom, not a code constraint: the pinned torch family (2.11+, via the sglang extra) checks dtype uniformity over trainable params only, and the FSDP backend's master_dtype path implements exactly this bf16-base + fp32-LoRA-master regime. Checkpoint forensics on the 150-rollout run show the bf16 master froze lora_A entirely (delta ~0.0001% over 100 steps; AdamW step 5e-6 < bf16 ULP at |A|~0.1) while the contributor's fp32 runs train the full (A, B). Align the environment, not the config. --- experimental/refl/examples/wan21_t2v_videoalign_refl.yaml | 8 +++++--- experimental/refl/examples/wan22_i2v_face_refl.yaml | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml index 88de60c50..e229d3435 100644 --- a/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml +++ b/experimental/refl/examples/wan21_t2v_videoalign_refl.yaml @@ -42,9 +42,11 @@ 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-original-dtype - # assert (same note as examples/diffusion/refl_sd3.yaml). + # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below + # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A. + # 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 diff --git a/experimental/refl/examples/wan22_i2v_face_refl.yaml b/experimental/refl/examples/wan22_i2v_face_refl.yaml index b08a424f2..2693ae5ae 100644 --- a/experimental/refl/examples/wan22_i2v_face_refl.yaml +++ b/experimental/refl/examples/wan22_i2v_face_refl.yaml @@ -43,9 +43,11 @@ 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-original-dtype - # assert (same note as examples/diffusion/refl_sd3.yaml). + # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below + # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A. + # 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