[feat] route all diffusion loops + RL recompute through the kernel table
Finish the kernel seam across the board so the platform's KernelTable is the universal solver-dispatch path, not just wan21. - Loops: ltx2, wan_causal, adapters, adaptive now resolve the flow-match solver via model.platform.kernels.get(FLOW_MATCH_STEP) instead of importing the numpy sampler directly (wan21 already did). On CPU this is bit-identical (the cpu kernel IS the old function); a GPU/accel backend now overrides every loop. - RL: the FlowGRPO log-prob recompute in unified_rl / joint_multi_rl / workflow_rl dispatches FLOW_SDE_STEP through the platform, pinned to the SAME kernel the rollout used (C2 kernel-pinning — otherwise the PPO ratio biases on a real GPU where rollout and recompute kernels could differ). - accel backend: add a kind-generic AccelComponent wrapper and override the vae component too (text_encoder left unregistered to keep the device->cpu fallback demonstrated), closing the "only dit is overridden" gap. - tests: vae override assertion; a second-loop (wan-causal chunk rollout) parity oracle proving accel == cpu bit-identical beyond wan21. README scope updated. 184 tests pass.
This commit is contained in:
SolitaryThinker
committed by
Peiyuan Zhang
Peiyuan Zhang
parent
2f044c02dd
commit
634f0828a2
+12
-11
@@ -17,14 +17,15 @@ scheduler, caches, parity, training, and workflow code are unchanged.
|
||||
|
||||
Honest scope of what is *wired* today (the substrate supports more than it demonstrates): a pure-
|
||||
python `accel` stand-in backend proves the dispatch is genuinely device-generic (cross-device
|
||||
resolution + arch fallback + the parity oracle, bit-identical to numpy) without a GPU. The wan21
|
||||
denoise loop routes its solver ops (`flow_match_step`/`flow_sde_step`) through the kernel table, and
|
||||
the `dit` component has a backend override; the other model loops still call the numpy samplers
|
||||
directly and only the `dit` kind is overridden — each is a one-line change to adopt. The full
|
||||
torch/CUDA path and a cudagraph capture-safety (`workspace_bytes`) contract are declared/deferred,
|
||||
not implemented.
|
||||
resolution + arch fallback + the parity oracle, bit-identical to numpy) without a GPU. All diffusion
|
||||
loops (wan21, ltx2, wan-causal, adapters, adaptive) and the FlowGRPO RL log-prob recompute route
|
||||
their solver ops (`flow_match_step`/`flow_sde_step`) through the kernel table — the recompute is
|
||||
pinned to the rollout's kernel for C2 correctness. The `dit` and `vae` components have `accel`
|
||||
overrides (`text_encoder` demonstrates the device→cpu fallback). Not yet routed: the AR/vocoder op
|
||||
families (omni/qwen-omni) have no KERNELS entries, the full torch/CUDA path is declared-unavailable,
|
||||
and a cudagraph capture-safety (`workspace_bytes`) contract is deferred.
|
||||
|
||||
**183 tests, 30 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only.
|
||||
**184 tests, 30 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only.
|
||||
|
||||
## Scope
|
||||
|
||||
@@ -100,7 +101,7 @@ Stress tests (the design held — see **designv4 §9**):
|
||||
|
||||
```bash
|
||||
cd /Users/willlin/src/FastVideo-mini
|
||||
python3 -m pytest v2/tests/ -q # 127 tests
|
||||
python3 -m pytest v2/tests/ -q # 184 tests
|
||||
python3 v2/run_tests.py # same suite, ZERO deps (no pytest needed)
|
||||
python3 -m v2.examples # the worked examples
|
||||
```
|
||||
@@ -115,9 +116,9 @@ Each `ComponentSpec` records the real `load_id` (e.g.
|
||||
`vae.decode`; `text_encoder.encode`), and the device kernels for the solver/primitive ops. They
|
||||
flip to `available=True` when torch+CUDA are present, and `Platform.detect()` resolves them; the
|
||||
numpy `ComponentSpec.factory` stays as the CPU terminal fallback rung;
|
||||
3. for the loops that don't yet route through the kernel table, swap the direct numpy
|
||||
flow-match/SDE sampler calls for `model.platform.kernels.get(...)` (the wan21 loop shows the
|
||||
pattern — a one-line change per call site).
|
||||
3. the diffusion loops + RL recompute already dispatch their solver ops via
|
||||
`model.platform.kernels.get(...)`, so the registered `cuda` kernels are picked up automatically;
|
||||
only the not-yet-routed op families (AR sampling, vocoder synth) need the same one-line swap.
|
||||
The loops, policies, scheduler, caches, parity gates, training methods, and workflows are unchanged.
|
||||
|
||||
## Serving & fleet (our own version — Dynamo is an option, not a dependency)
|
||||
|
||||
@@ -11,7 +11,7 @@ from __future__ import annotations
|
||||
import numpy as np
|
||||
|
||||
from ...loop.contracts import Done
|
||||
from ...loop.sampler import flow_match_euler_step
|
||||
from ...platform import FLOW_MATCH_STEP
|
||||
from ..wan21.loop import WanDenoiseLoop
|
||||
|
||||
|
||||
@@ -39,7 +39,8 @@ class AdapterDenoiseLoop(WanDenoiseLoop):
|
||||
for aid in adapters: # apply each active adapter's delta
|
||||
v = v + model.component(aid).delta(x, control)
|
||||
res.output["noise_pred"] = v
|
||||
res.output["latents"] = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next).astype("float32")
|
||||
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
|
||||
res.output["latents"] = fm(prec.cast(x), v, sigma_t, sigma_next).astype("float32")
|
||||
return res
|
||||
|
||||
plan.run = run
|
||||
|
||||
@@ -17,7 +17,7 @@ from __future__ import annotations
|
||||
import numpy as np
|
||||
|
||||
from ...loop.contracts import Done, StepResult
|
||||
from ...loop.sampler import flow_match_euler_step
|
||||
from ...platform import FLOW_MATCH_STEP
|
||||
from ..wan21.loop import WanDenoiseLoop
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ class CacheDiTDenoiseLoop(WanDenoiseLoop):
|
||||
v = np.asarray(last_v, dtype="float32")
|
||||
|
||||
def run(model, override=None):
|
||||
x_next = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next)
|
||||
x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(prec.cast(x), v, sigma_t, sigma_next)
|
||||
return StepResult(output={"noise_pred": v, "latents": x_next.astype("float32"), "skipped": True})
|
||||
|
||||
plan.run = run
|
||||
|
||||
@@ -28,7 +28,7 @@ from ...loop.contracts import (
|
||||
StepResult,
|
||||
WorkPlan,
|
||||
)
|
||||
from ...loop.sampler import flow_match_euler_step
|
||||
from ...platform import FLOW_MATCH_STEP
|
||||
from ..backend import LATENT_CHANNELS
|
||||
|
||||
BASE_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
|
||||
@@ -99,6 +99,7 @@ class LTX2DenoiseLoop:
|
||||
au = st.latents.get("audio")
|
||||
|
||||
def run(model, override=None):
|
||||
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
|
||||
if override is not None and "noise_pred" in override:
|
||||
velocity = np.asarray(override["noise_pred"], dtype="float32")
|
||||
else:
|
||||
@@ -110,14 +111,14 @@ class LTX2DenoiseLoop:
|
||||
velocity = (pos
|
||||
+ (v_guidance - 1.0) * (pos - neg)
|
||||
+ stg_scale * (pos - ptb))
|
||||
x_next = flow_match_euler_step(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next)
|
||||
x_next = fm(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next)
|
||||
out = {"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")}
|
||||
if want_audio and au is not None: # joint audio denoise, conditioned on video
|
||||
dit = model.component("transformer")
|
||||
a_pos = dit(au, pe, sigma_t, context=x) # context=x ⇒ audio synced to the video
|
||||
a_neg = dit(au, ne, sigma_t, context=x)
|
||||
a_vel = a_pos + (a_guidance - 1.0) * (a_pos - a_neg) # audio's OWN guidance scale
|
||||
out["audio_latents"] = flow_match_euler_step(au, a_vel, sigma_t, sigma_next).astype("float32")
|
||||
out["audio_latents"] = fm(au, a_vel, sigma_t, sigma_next).astype("float32")
|
||||
return StepResult(output=out)
|
||||
|
||||
res = ResourceRequest(
|
||||
|
||||
@@ -24,7 +24,7 @@ from ...loop.contracts import (
|
||||
StepResult,
|
||||
WorkPlan,
|
||||
)
|
||||
from ...loop.sampler import flow_match_euler_step
|
||||
from ...platform import FLOW_MATCH_STEP
|
||||
from ...request.streams import StreamChunk
|
||||
from ..backend import LATENT_CHANNELS
|
||||
|
||||
@@ -80,13 +80,14 @@ class ChunkRolloutLoop:
|
||||
cfg, precision = self.cfg, self.precision
|
||||
|
||||
def run(model, override=None):
|
||||
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
|
||||
if override is not None and "noise_pred" in override:
|
||||
velocity = np.asarray(override["noise_pred"], dtype="float32")
|
||||
else:
|
||||
dit = model.component("transformer")
|
||||
preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, context=context) for b in branches}
|
||||
velocity = cfg.combine(preds, scale, sctx, cfg_state)
|
||||
x_next = flow_match_euler_step(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next)
|
||||
x_next = fm(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next)
|
||||
return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"),
|
||||
"latents": x_next.astype("float32")})
|
||||
|
||||
|
||||
@@ -58,6 +58,31 @@ def _build_accel_dit(spec: Any, instance: Any, platform: Any) -> AccelDiT:
|
||||
register_component("dit", _build_accel_dit, device="accel", source="accel(stand-in):AccelDiT")
|
||||
|
||||
|
||||
class AccelComponent:
|
||||
"""Generic device-tagged wrapper for non-callable component kinds (vae, audio_vae, …). Delegates
|
||||
the full call surface (``decode``/``encode``/…) to the inner toy, so accel ≡ cpu numerically."""
|
||||
|
||||
def __init__(self, inner: Any):
|
||||
self._inner = inner
|
||||
self.device = "accel"
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self._inner, name)
|
||||
|
||||
|
||||
def _build_accel_component(spec: Any, instance: Any, platform: Any) -> AccelComponent:
|
||||
if getattr(spec, "factory", None) is None:
|
||||
raise RuntimeError(f"accel: no factory to wrap for kind={spec.kind!r}")
|
||||
return AccelComponent(spec.factory(instance))
|
||||
|
||||
|
||||
# Override more than one component kind on accel (proving the seam is kind-generic, not dit-special).
|
||||
# text_encoder is deliberately left UNregistered so the device→cpu fallback path stays demonstrated.
|
||||
for _kind in ("vae", "audio_vae"):
|
||||
register_component(_kind, _build_accel_component, device="accel",
|
||||
source="accel(stand-in):AccelComponent")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Kernels: independent accel implementations that match the numpy reference. #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@@ -19,6 +19,7 @@ from v2.loop.driver import LoopRunner
|
||||
from v2.models.backend import ToyDiT, ToyTextEncoder, ToyVAE
|
||||
from v2.models.common import text_encode_node_fn
|
||||
from v2.models.wan21 import build_wan21_card
|
||||
from v2.models.wan_causal import build_wan_causal_card
|
||||
from v2.parity import bit_identical # the typed C1 instrument (ladder.py)
|
||||
from v2.platform import (
|
||||
FLOW_MATCH_STEP,
|
||||
@@ -27,7 +28,7 @@ from v2.platform import (
|
||||
component_matrix,
|
||||
kernel_matrix,
|
||||
)
|
||||
from v2.platform.backends.accel import AccelDiT
|
||||
from v2.platform.backends.accel import AccelComponent, AccelDiT
|
||||
from v2.request import DiffusionParams, TaskType, make_request
|
||||
from v2.request.streams import Stream
|
||||
from v2.runtime import Engine
|
||||
@@ -41,8 +42,8 @@ def _instance(platform):
|
||||
return load_card(card, cache_manager=CacheManager.from_card(card), platform=platform)
|
||||
|
||||
|
||||
def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False):
|
||||
"""Drive the wan21 denoise loop step-by-step on ``inst``'s platform; return the final latent."""
|
||||
def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False, loop_id="diffusion_denoise"):
|
||||
"""Drive a loop step-by-step on ``inst``'s platform; return the final latent."""
|
||||
req = make_request(TaskType.T2V, inst.card.model_id, prompt,
|
||||
diffusion=DiffusionParams(num_steps=steps, seed=seed, sde_rollout=sde))
|
||||
slots: dict = {}
|
||||
@@ -51,7 +52,7 @@ def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False):
|
||||
ctx = RuntimeLoopContext(inst, observers=_HUB.observers, interceptors=_HUB.interceptors,
|
||||
slots=slots, stream=Stream(req.request_id), cancel_scope=None,
|
||||
profile=profile, metrics={}, request_id=req.request_id)
|
||||
runner = LoopRunner(inst.loop("diffusion_denoise"), ctx, req, inst)
|
||||
runner = LoopRunner(inst.loop(loop_id), ctx, req, inst)
|
||||
while not runner.done:
|
||||
runner.step()
|
||||
return np.asarray(runner.result.outputs["latents"])
|
||||
@@ -134,9 +135,10 @@ def test_accel_component_override_and_device_fallback():
|
||||
inst = _instance(Platform.accel("sm90"))
|
||||
dit = inst.component("transformer") # registered for accel → AccelDiT wrapper
|
||||
assert isinstance(dit, AccelDiT) and dit.device == "accel"
|
||||
# text_encoder / vae are NOT registered for accel → fall back over accel→cpu to the factory toy
|
||||
vae = inst.component("vae") # vae is ALSO overridden on accel (kind-generic)
|
||||
assert isinstance(vae, AccelComponent) and vae.device == "accel"
|
||||
# text_encoder is NOT registered for accel → falls back over accel→cpu to the numpy factory toy
|
||||
assert isinstance(inst.component("text_encoder"), ToyTextEncoder)
|
||||
assert isinstance(inst.component("vae"), ToyVAE)
|
||||
|
||||
|
||||
def test_cuda_component_falls_back_to_factory_when_unavailable():
|
||||
@@ -170,6 +172,17 @@ def test_accel_denoise_still_produces_valid_output():
|
||||
assert out.shape[0] == 4 and np.all(np.isfinite(out))
|
||||
|
||||
|
||||
def test_parity_oracle_accel_equals_cpu_second_loop():
|
||||
"""A SECOND loop family (wan-causal chunk rollout) also routes its solver through the kernel
|
||||
table — so the seam is finished across loops, not just wan21 denoise. accel ≡ cpu, bit-identical."""
|
||||
def drive(platform):
|
||||
card = build_wan_causal_card()
|
||||
inst = load_card(card, cache_manager=CacheManager.from_card(card), platform=platform)
|
||||
loop_id = next(iter(inst.card.loops)) # the chunk-rollout loop
|
||||
return _drive_denoise(inst, seed=5, loop_id=loop_id)
|
||||
assert bit_identical(drive(Platform.cpu()), drive(Platform.accel("sm90")))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Enumerable matrix — including unavailable backends, without importing torch #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@@ -28,8 +28,8 @@ import numpy as np
|
||||
|
||||
from ..._enums import ConsistencyLevel, ExecutionProfile
|
||||
from ...loop.sampler import flow_sde_ml_velocity as _ml_velocity
|
||||
from ...loop.sampler import flow_sde_step_with_logprob
|
||||
from ...models.common import cached_text_encode
|
||||
from ...platform import FLOW_SDE_STEP
|
||||
from ...request import DiffusionParams, TaskType, make_request
|
||||
from ..rollout import rollout_loop
|
||||
from ..weight_sync import WeightRole, WeightSyncPlan
|
||||
@@ -116,6 +116,8 @@ class JointMultiExpertRL(TrainingMethod):
|
||||
|
||||
def _dit_ppo(self, behavior, emb, advantage):
|
||||
ref_dit = self.reference.component(self.generator_id)
|
||||
# Pin the log-prob recompute to the SAME SDE kernel the rollout used (C2 kernel-pinning).
|
||||
sde_step = self.student.platform.kernels.get(FLOW_SDE_STEP)
|
||||
ppo_losses, kls, ratios, gnorms = [], [], [], []
|
||||
for rec in behavior or []:
|
||||
if "sde_logprob" not in rec:
|
||||
@@ -123,10 +125,10 @@ class JointMultiExpertRL(TrainingMethod):
|
||||
prev, sample = rec["prev"], rec["sample"]
|
||||
st, sn = float(rec["sigma_t"]), float(rec["sigma_next"])
|
||||
v_cur = self.dit(prev, emb, st)
|
||||
_, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob(
|
||||
_, logp_cur, mean_cur, eff_std = sde_step(
|
||||
prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale)
|
||||
v_ref = ref_dit(prev, emb, st)
|
||||
_, _, mean_ref, _ = flow_sde_step_with_logprob(
|
||||
_, _, mean_ref, _ = sde_step(
|
||||
prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_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 - self.ppo_clip, 1.0 + self.ppo_clip))
|
||||
|
||||
@@ -32,8 +32,8 @@ import numpy as np
|
||||
|
||||
from ..._enums import ConsistencyLevel, ExecutionProfile
|
||||
from ...loop.sampler import flow_sde_ml_velocity as _ml_velocity
|
||||
from ...loop.sampler import flow_sde_step_with_logprob
|
||||
from ...models.common import cached_text_encode
|
||||
from ...platform import FLOW_SDE_STEP
|
||||
from ...request import DiffusionParams, TaskType, make_request
|
||||
from ..rollout import rollout_loop
|
||||
from ..weight_sync import WeightRole, WeightSyncPlan
|
||||
@@ -112,6 +112,9 @@ class UnifiedRLMethod(TrainingMethod):
|
||||
"""FlowGRPO PPO over one rollout's SDE trajectory: per-step ratio + KL, then a toy DiT step.
|
||||
Returns (mean_ppo_loss, mean_kl, mean_ratio, grad_norm)."""
|
||||
ref_dit = self.reference.component("transformer")
|
||||
# Pin the log-prob recompute to the SAME SDE kernel the rollout used (C2 kernel-pinning):
|
||||
# on a GPU box the rollout's logp and this ratio must come from one kernel or the PPO ratio biases.
|
||||
sde_step = self.student.platform.kernels.get(FLOW_SDE_STEP)
|
||||
ppo_losses, kls, ratios, gnorms = [], [], [], []
|
||||
for rec in behavior or []:
|
||||
if "sde_logprob" not in rec:
|
||||
@@ -119,10 +122,10 @@ class UnifiedRLMethod(TrainingMethod):
|
||||
prev, sample = rec["prev"], rec["sample"]
|
||||
st, sn = float(rec["sigma_t"]), float(rec["sigma_next"])
|
||||
v_cur = self.dit(prev, emb, st)
|
||||
_, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob(
|
||||
_, logp_cur, mean_cur, eff_std = sde_step(
|
||||
prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale)
|
||||
v_ref = ref_dit(prev, emb, st)
|
||||
_, _, mean_ref, _ = flow_sde_step_with_logprob(
|
||||
_, _, mean_ref, _ = sde_step(
|
||||
prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_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 - self.ppo_clip, 1.0 + self.ppo_clip))
|
||||
|
||||
@@ -21,16 +21,20 @@ import numpy as np
|
||||
from ..._enums import ConsistencyLevel, ExecutionProfile
|
||||
from ...loop.sampler import flow_sde_ml_velocity, flow_sde_step_with_logprob
|
||||
from ...models.common import cached_text_encode
|
||||
from ...platform import FLOW_SDE_STEP
|
||||
from ...request import DiffusionParams, TaskType, make_request
|
||||
from ..rollout import rollout_loop
|
||||
from ..weight_sync import WeightRole, WeightSyncPlan
|
||||
from .base import TrainingMethod, new_instance
|
||||
|
||||
|
||||
def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale):
|
||||
def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale,
|
||||
sde_step=None):
|
||||
"""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."""
|
||||
(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:
|
||||
@@ -38,10 +42,10 @@ def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip,
|
||||
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 = flow_sde_step_with_logprob(
|
||||
_, 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, _ = flow_sde_step_with_logprob(
|
||||
_, _, 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))
|
||||
@@ -141,11 +145,13 @@ class WorkflowRLMethod(TrainingMethod):
|
||||
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)
|
||||
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)
|
||||
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])
|
||||
|
||||
Reference in New Issue
Block a user