SolitaryThinker
21e59268d8
Sweep all of v2/ through pre-commit (was previously only run on changed files). Fixes surfaced across untouched modules: - ruff: isinstance-tuple -> X | Y (UP038), try/except/pass -> contextlib.suppress (SIM105), negated-return (SIM103), %-format -> f-string (UP031). - mypy: add annotations for no-untyped-call + var-annotated across recipes, training methods, torch/toy backends, and serving. - yapf reflow of the SF-Wan KV-cache call sites (semantics unchanged). yapf/ruff/codespell/mypy all pass; v2 tests 237 passed / 5 pre-existing (bucket_c torch-absent on CPU venv).
235 lines
12 KiB
Python
235 lines
12 KiB
Python
"""WorkflowRLMethod — end-to-end RL over a cross-model workflow.
|
|
|
|
Train BOTH stages of the T2I→I2V workflow (``flux-t2i`` and ``wan-i2v`` — two separate models) from one
|
|
final-video reward. The rollout spans two instances: roll out the T2I diffusion loop with SDE capture →
|
|
decode an image → condition and roll out the I2V diffusion loop with SDE capture → decode a video → score
|
|
it. One group-relative advantage from the final video drives a FlowGRPO PPO update on each stage's
|
|
transformer, with two independent WeightSyncPlans on two different instances.
|
|
|
|
The interesting part is credit assignment: the reward is at the end (the video) yet credits an earlier
|
|
model (the T2I generator), so the T2I stage learns to produce images that lead to better videos. This
|
|
shows "rollout == serve + capture" holds for a multi-instance workflow with no new primitive — the method
|
|
drives the same shared loops the engine serves and threads the artifact between them as the serve-time
|
|
Workflow does.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from v2._enums import ConsistencyLevel, ExecutionProfile
|
|
from v2.loop.sampler import flow_sde_ml_velocity, flow_sde_step_with_logprob
|
|
from v2.recipes.common import cached_text_encode
|
|
from v2.platform import FLOW_SDE_STEP
|
|
from v2.request import DiffusionParams, TaskType, make_request
|
|
from v2.training.rollout import rollout_loop
|
|
from v2.training.weight_sync import WeightRole, WeightSyncPlan
|
|
from v2.training.methods.base import TrainingMethod, new_instance
|
|
|
|
|
|
def _flowgrpo_ppo_step(dit: Any,
|
|
ref_dit: Any,
|
|
behavior: Any,
|
|
emb: Any,
|
|
advantage: float,
|
|
*,
|
|
lr: float,
|
|
ppo_clip: float,
|
|
beta: float,
|
|
noise_scale: float,
|
|
sde_step: Any = None) -> tuple[float, float, float, float]:
|
|
"""One FlowGRPO PPO pass over an SDE trajectory: per-step ratio + KL, then move the policy toward
|
|
the max-likelihood velocity of each realized sample (advantage-weighted). Returns
|
|
(ppo_loss, kl, ratio, grad_norm). Shared by both stages. ``sde_step`` is the platform-resolved SDE
|
|
kernel (pins the recompute to the rollout's kernel, C2); defaults to the numpy reference."""
|
|
sde_step = sde_step or flow_sde_step_with_logprob
|
|
ppo_losses, kls, ratios, gnorms = [], [], [], []
|
|
for rec in behavior or []:
|
|
if "sde_logprob" not in rec:
|
|
continue
|
|
prev, sample = rec["prev"], rec["sample"]
|
|
st, sn = float(rec["sigma_t"]), float(rec["sigma_next"])
|
|
v_cur = dit(prev, emb, st)
|
|
_, logp_cur, mean_cur, eff_std = sde_step(prev, v_cur, st, sn, prev_sample=sample, noise_scale=noise_scale)
|
|
v_ref = ref_dit(prev, emb, st)
|
|
_, _, mean_ref, _ = sde_step(prev, v_ref, st, sn, prev_sample=sample, noise_scale=noise_scale)
|
|
ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0)))
|
|
clipped = float(np.clip(ratio, 1.0 - ppo_clip, 1.0 + ppo_clip))
|
|
ppo_losses.append(max(-advantage * ratio, -advantage * clipped))
|
|
kls.append(float(np.mean((mean_cur - mean_ref)**2) / (2.0 * max(eff_std**2, 1e-8))))
|
|
ratios.append(ratio)
|
|
v_target = flow_sde_ml_velocity(prev, sample, st, sn, noise_scale=noise_scale)
|
|
_, gn = dit.mse_grad_step(prev, emb, st, v_target, lr * float(np.clip(advantage, -1.0, 1.0)))
|
|
gnorms.append(gn)
|
|
|
|
def m(xs: Any) -> float:
|
|
return float(np.mean(xs)) if xs else 0.0
|
|
|
|
return m(ppo_losses), m(kls), m(ratios), m(gnorms)
|
|
|
|
|
|
class WorkflowRLMethod(TrainingMethod):
|
|
name = "workflow_rl"
|
|
consistency = ConsistencyLevel.C2
|
|
|
|
def __init__(self,
|
|
t2i_student,
|
|
i2v_student,
|
|
t2i_reference,
|
|
i2v_reference,
|
|
*,
|
|
num_samples_per_prompt: int = 4,
|
|
rollout_steps: int = 4,
|
|
t2i_lr: float = 0.02,
|
|
i2v_lr: float = 0.02,
|
|
beta: float = 0.01,
|
|
adv_clip_max: float = 5.0,
|
|
ppo_clip: float = 1e-2,
|
|
sde_noise_scale: float = 0.7,
|
|
**kw):
|
|
super().__init__(t2i_student, lr=t2i_lr, **kw) # base "student" = the T2I stage
|
|
self.t2i, self.i2v = t2i_student, i2v_student
|
|
self.t2i_ref, self.i2v_ref = t2i_reference, i2v_reference
|
|
self.K = num_samples_per_prompt
|
|
self.rollout_steps = rollout_steps
|
|
self.t2i_lr, self.i2v_lr = t2i_lr, i2v_lr
|
|
self.beta, self.adv_clip_max, self.ppo_clip = beta, adv_clip_max, ppo_clip
|
|
self.sde_noise_scale = sde_noise_scale
|
|
# one WeightSyncPlan per stage's generator — on two DIFFERENT instances (independent versions)
|
|
self.t2i_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer", ))
|
|
self.i2v_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer", ))
|
|
self.t2i_ref.component("transformer").copy_from(self.t2i.component("transformer"))
|
|
self.i2v_ref.component("transformer").copy_from(self.i2v.component("transformer"))
|
|
|
|
def manages_optimization(self) -> bool:
|
|
return True
|
|
|
|
def get_grad_clip_targets(self, iteration: int = 0) -> dict:
|
|
return {
|
|
"t2i.transformer": self.t2i.component("transformer"),
|
|
"i2v.transformer": self.i2v.component("transformer")
|
|
}
|
|
|
|
def _sde_params(self, **extra: Any) -> DiffusionParams:
|
|
return DiffusionParams(num_steps=self.rollout_steps,
|
|
guidance_scale=1.0,
|
|
sde_rollout=True,
|
|
sde_noise_scale=self.sde_noise_scale,
|
|
**extra)
|
|
|
|
def _rollout_t2i(self, prompt: Any, seed: Any) -> tuple[Any, Any, Any]:
|
|
emb = cached_text_encode(self.t2i, prompt)
|
|
neg = cached_text_encode(self.t2i, "")
|
|
req = make_request(TaskType.T2I,
|
|
self.t2i.card.model_id,
|
|
prompt,
|
|
diffusion=self._sde_params(num_frames=1, seed=seed))
|
|
res = rollout_loop(self.t2i,
|
|
"t2i_denoise",
|
|
req,
|
|
slots={
|
|
"text_embeds": emb,
|
|
"neg_text_embeds": neg
|
|
},
|
|
profile=ExecutionProfile.ROLLOUT)
|
|
image = self.t2i.component("vae").decode(res.outputs["latents"])
|
|
return res, emb, image
|
|
|
|
def _rollout_i2v(self, prompt: Any, image: Any, seed: Any) -> tuple[Any, Any, Any]:
|
|
emb = cached_text_encode(self.i2v, prompt)
|
|
neg = cached_text_encode(self.i2v, "")
|
|
cond = self.i2v.component("vae").encode(np.asarray(image, dtype="float32")) # image conditioning
|
|
emb = (np.asarray(emb, dtype="float32") + float(np.tanh(np.mean(cond)))).astype("float32")
|
|
req = make_request(TaskType.I2V,
|
|
self.i2v.card.model_id,
|
|
prompt,
|
|
diffusion=self._sde_params(num_frames=81, seed=seed))
|
|
res = rollout_loop(self.i2v,
|
|
"i2v_denoise",
|
|
req,
|
|
slots={
|
|
"text_embeds": emb,
|
|
"neg_text_embeds": neg
|
|
},
|
|
profile=ExecutionProfile.ROLLOUT)
|
|
video = self.i2v.component("vae").decode(res.outputs["latents"])
|
|
return res, emb, video
|
|
|
|
@staticmethod
|
|
def _reward(video: Any) -> float:
|
|
return float(np.tanh(-np.std(np.asarray(video, dtype="float64")))) # lower-spread video = better
|
|
|
|
def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
|
self.iteration = iteration
|
|
prompts = batch["prompts"]
|
|
seeds = batch.get("seeds", list(range(len(prompts))))
|
|
t2i_loss, i2v_loss, t2i_gn, i2v_gn, ratios, all_reward = [], [], [], [], [], []
|
|
|
|
for pi, prompt in enumerate(prompts):
|
|
samples = []
|
|
for k in range(self.K):
|
|
seed = seeds[pi] * 1000 + k
|
|
res_t, emb_t, image = self._rollout_t2i(prompt, seed) # stage 1 rollout (SDE)
|
|
res_i, emb_i, video = self._rollout_i2v(prompt, image, seed) # stage 2 rollout (SDE)
|
|
samples.append((res_t, emb_t, res_i, emb_i, self._reward(video)))
|
|
|
|
rewards = np.array([s[4] for s in samples], dtype="float64")
|
|
adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), -self.adv_clip_max,
|
|
self.adv_clip_max) # one final-video advantage
|
|
all_reward.extend(rewards.tolist())
|
|
|
|
for k, (res_t, emb_t, res_i, emb_i, _r) in enumerate(samples):
|
|
a = float(adv[k])
|
|
# the SAME final advantage credits BOTH stages (end-to-end across the model boundary)
|
|
pt, _kt, rt, gt = _flowgrpo_ppo_step(self.t2i.component("transformer"),
|
|
self.t2i_ref.component("transformer"),
|
|
res_t.behavior,
|
|
emb_t,
|
|
a,
|
|
lr=self.t2i_lr,
|
|
ppo_clip=self.ppo_clip,
|
|
beta=self.beta,
|
|
noise_scale=self.sde_noise_scale,
|
|
sde_step=self.t2i.platform.kernels.get(FLOW_SDE_STEP))
|
|
pi_, _ki, ri, gi = _flowgrpo_ppo_step(self.i2v.component("transformer"),
|
|
self.i2v_ref.component("transformer"),
|
|
res_i.behavior,
|
|
emb_i,
|
|
a,
|
|
lr=self.i2v_lr,
|
|
ppo_clip=self.ppo_clip,
|
|
beta=self.beta,
|
|
noise_scale=self.sde_noise_scale,
|
|
sde_step=self.i2v.platform.kernels.get(FLOW_SDE_STEP))
|
|
t2i_loss.append(pt)
|
|
i2v_loss.append(pi_)
|
|
t2i_gn.append(gt)
|
|
i2v_gn.append(gi)
|
|
ratios.extend([rt, ri])
|
|
|
|
t2i_ver = self.t2i_sync.apply(self.t2i.component("transformer"), self.t2i.component("transformer"), self.t2i)
|
|
i2v_ver = self.i2v_sync.apply(self.i2v.component("transformer"), self.i2v.component("transformer"), self.i2v)
|
|
|
|
def m(xs: Any) -> float:
|
|
return float(np.mean(xs)) if xs else 0.0
|
|
|
|
metrics = {
|
|
"reward_mean": m(all_reward),
|
|
"grad_norm/t2i": m(t2i_gn),
|
|
"grad_norm/i2v": m(i2v_gn),
|
|
"ppo_ratio_mean": m(ratios) if ratios else 1.0,
|
|
"t2i_weights_version": t2i_ver,
|
|
"i2v_weights_version": i2v_ver,
|
|
}
|
|
return {"t2i_pg_loss": m(t2i_loss), "i2v_pg_loss": m(i2v_loss)}, metrics
|
|
|
|
def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
|
return self.managed_train_step(batch, iteration)
|
|
|
|
|
|
def build_workflow_rl(t2i_card, i2v_card, **kw) -> WorkflowRLMethod:
|
|
"""Four roles: trainable T2I + I2V students and their frozen references (two instances each side)."""
|
|
return WorkflowRLMethod(new_instance(t2i_card), new_instance(i2v_card), new_instance(t2i_card),
|
|
new_instance(i2v_card), **kw)
|