SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
166 lines
8.2 KiB
Python
166 lines
8.2 KiB
Python
"""DiffusionNFT — likelihood-free RL (repo: train/methods/rl/diffusion_nft).
|
||
|
||
Likelihood-free (no log-probs to match), so consistency is the C2 behavioral rung: seeded final
|
||
sample plus prediction-space identity (``old_deviate`` / ref-MSE), declared on the recipe. Sampling
|
||
is from the old (decay-blended behavior) policy, NOT the student — a caveat the WeightSyncPlan must
|
||
carry. K samples of one prompt form a homogeneous group; group-relative advantages drive the update;
|
||
the shared prompt encodes once via the feature cache.
|
||
|
||
NFT objective (repo diffusion_nft.py):
|
||
positive = β·forward + (1−β)·old ; implicit_neg = (1+β)·old − β·forward
|
||
x0⁺ = xt − t·positive ; x0⁻ = xt − t·implicit_neg
|
||
policy_loss = r·‖x0⁺−x0‖²/wf⁺/β + (1−r)·‖x0⁻−x0‖²/wf⁻/β (r = clamp(adv/clip/2+0.5, 0, 1))
|
||
loss = policy_loss + kl_β·‖forward − ref‖² ; old_deviate = ‖forward − old‖²
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from v2._enums import ConsistencyLevel
|
||
from v2.recipes.common import cached_text_encode
|
||
from v2.request import DiffusionParams, TaskType, make_request
|
||
from v2.training.rewards import build_multi_reward_scorer
|
||
from v2.training.rollout import Trajectory, TrajectoryBuffer
|
||
from v2.training.weight_sync import WeightRole, WeightSyncPlan
|
||
from v2.training.methods.base import TrainingMethod, new_instance
|
||
|
||
|
||
class DiffusionNFTMethod(TrainingMethod):
|
||
name = "diffusion_nft"
|
||
consistency = ConsistencyLevel.C2 # likelihood-free behavioral identity
|
||
|
||
def __init__(self,
|
||
student_instance,
|
||
old_instance,
|
||
reference_instance,
|
||
*,
|
||
reward_fn: dict | None = None,
|
||
num_video_per_prompt: int = 4,
|
||
rollout_steps: int = 3,
|
||
beta: float = 0.1,
|
||
kl_beta: float = 1e-4,
|
||
adv_clip_max: float = 5.0,
|
||
num_inner_timesteps: int = 2,
|
||
lr: float = 0.05,
|
||
decay_type: int = 1,
|
||
**kw):
|
||
super().__init__(student_instance, lr=lr, **kw)
|
||
self.old = old_instance # behavior policy (sampled from!)
|
||
self.reference = reference_instance # frozen KL anchor
|
||
self.K = num_video_per_prompt
|
||
self.rollout_steps = rollout_steps
|
||
self.beta = beta
|
||
self.kl_beta = kl_beta
|
||
self.adv_clip_max = adv_clip_max
|
||
self.num_inner_timesteps = num_inner_timesteps
|
||
self.decay_type = decay_type
|
||
self.scorer = build_multi_reward_scorer(reward_fn or {"pickscore": 1.0, "clipscore": 1.0})
|
||
self.old_sync = WeightSyncPlan(role=WeightRole.OLD_POLICY)
|
||
# init old + reference from student (NFT inits old from student)
|
||
self.old.component("transformer").copy_from(self.student_dit)
|
||
self.reference.component("transformer").copy_from(self.student_dit)
|
||
|
||
def manages_optimization(self) -> bool: # RL owns its sample→score→train cadence
|
||
return True
|
||
|
||
def get_grad_clip_targets(self, iteration: int = 0) -> dict:
|
||
return {"student": self.student_dit} # NFT clips internally (matches repo)
|
||
|
||
@staticmethod
|
||
def _decay(step: int, decay_type: int) -> float:
|
||
if decay_type == 1:
|
||
return min(step * 0.001, 0.5)
|
||
if decay_type == 2:
|
||
return min(max(0, step - 75) * 0.0075, 0.999)
|
||
return 0.0
|
||
|
||
def _sample_group(self, prompt: str, base_seed: int) -> TrajectoryBuffer:
|
||
"""K rollouts of one prompt FROM THE OLD POLICY via the shared loop (homogeneous group)."""
|
||
buf = TrajectoryBuffer()
|
||
for k in range(self.K):
|
||
seed = base_seed * 1000 + k
|
||
req = make_request(TaskType.T2V,
|
||
self.old.card.model_id,
|
||
prompt,
|
||
diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed))
|
||
res = self._rollout(req, instance=self.old) # behavior policy = OLD
|
||
buf.add(
|
||
Trajectory(request_id=req.request_id,
|
||
prompt=prompt,
|
||
seed=seed,
|
||
latents=np.asarray(res.outputs["latents"], dtype="float32"),
|
||
behavior=res.behavior,
|
||
weights_version=self.old.weights_version))
|
||
return buf
|
||
|
||
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))))
|
||
policy_losses, kl_losses, old_devs, gnorms, all_adv, all_reward = [], [], [], [], [], []
|
||
|
||
for pi, prompt in enumerate(prompts):
|
||
group = self._sample_group(prompt, seeds[pi])
|
||
media = [t.latents for t in group]
|
||
scored = self.scorer.score(media, [prompt] * len(group))
|
||
rewards = scored["avg"]
|
||
adv = (rewards - rewards.mean()) / (rewards.std() + 1e-4) # group-relative
|
||
adv = np.clip(adv, -self.adv_clip_max, self.adv_clip_max)
|
||
all_reward.extend(rewards.tolist())
|
||
all_adv.extend(adv.tolist())
|
||
|
||
emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group
|
||
old_dit, ref_dit = self.old.component("transformer"), self.reference.component("transformer")
|
||
for k, traj in enumerate(group):
|
||
x0 = traj.latents
|
||
r = float(np.clip(adv[k] / self.adv_clip_max / 2.0 + 0.5, 0.0, 1.0))
|
||
rng = np.random.default_rng(traj.seed ^ 0x4F7)
|
||
for _ in range(self.num_inner_timesteps):
|
||
t = float(rng.uniform(0.05, 0.95))
|
||
noise = rng.standard_normal(x0.shape).astype("float32")
|
||
xt = ((1.0 - t) * x0 + t * noise).astype("float32")
|
||
fwd = self.student_dit(xt, emb, t)
|
||
old_p = old_dit(xt, emb, t)
|
||
ref_p = ref_dit(xt, emb, t)
|
||
positive = self.beta * fwd + (1 - self.beta) * old_p
|
||
neg = (1.0 + self.beta) * old_p - self.beta * fwd
|
||
x0p = xt - t * positive
|
||
x0n = xt - t * neg
|
||
wfp = max(float(np.mean(np.abs(x0p - x0))), 1e-5)
|
||
wfn = max(float(np.mean(np.abs(x0n - x0))), 1e-5)
|
||
pos_loss = float(np.mean((x0p - x0)**2) / wfp)
|
||
neg_loss = float(np.mean((x0n - x0)**2) / wfn)
|
||
policy_loss = (r * pos_loss / self.beta + (1 - r) * neg_loss / self.beta) * self.adv_clip_max
|
||
kl = float(np.mean((fwd - ref_p)**2))
|
||
policy_losses.append(policy_loss)
|
||
kl_losses.append(kl)
|
||
old_devs.append(float(np.mean((fwd - old_p)**2)))
|
||
# update student toward x0 along the (reward-weighted) positive direction
|
||
v_target = ((xt - x0) / max(t, 1e-3)).astype("float32")
|
||
_, gn = self.student_dit.mse_grad_step(xt, emb, t, v_target, self.lr * max(r, 0.05))
|
||
gnorms.append(gn)
|
||
|
||
# update OLD ← decay-blend student (the behavior-policy role); bump version + invalidate caches
|
||
self.old_sync.decay = self._decay(iteration, self.decay_type)
|
||
new_ver = self.old_sync.apply(self.student_dit, self.old.component("transformer"), self.old)
|
||
|
||
metrics = {
|
||
"policy_loss": float(np.mean(policy_losses)),
|
||
"kl_div_loss": float(np.mean(kl_losses)),
|
||
"old_deviate": float(np.mean(old_devs)),
|
||
"grad_norm/student": float(np.mean(gnorms)),
|
||
"reward_mean": float(np.mean(all_reward)),
|
||
"advantage_std": float(np.std(all_adv)),
|
||
"old_decay": self.old_sync.decay,
|
||
"old_weights_version": new_ver,
|
||
}
|
||
loss_map = {"policy_loss": metrics["policy_loss"], "kl_div_loss": metrics["kl_div_loss"]}
|
||
return loss_map, metrics
|
||
|
||
def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
||
return self.managed_train_step(batch, iteration)
|
||
|
||
|
||
def build_diffusion_nft(card, **kw) -> DiffusionNFTMethod:
|
||
return DiffusionNFTMethod(new_instance(card), new_instance(card), new_instance(card), **kw)
|