SolitaryThinker
9489c6c1dd
The remaining design_v3 probes (all except 8b speculative decoding). All fit with no new runtime primitive (148 tests pass). #6 LTX-2 joint audio+video (§9.11) — LTX-2 declared an audio_vae required_for t2vs but never used it; now a single 2-stage denoise carries a synchronized audio latent (conditioned on video), applies per-modality CFG (guidance_per_modality), and decodes via video VAE + audio VAE → video + audio. Gated on requesting audio, so the T2V path is byte-identical (existing tests untouched). ToyAudioVAE; build_ltx2_av_program. 5 tests. #4 Live weight-sync under in-flight serving (§9.14) — WeightSyncController makes the freeze → drain → transfer → bump version + invalidate → resume lifecycle explicit. Tests: a mid-flight swap corrupts (the hazard); draining first leaves the in-flight request bit-identical to baseline while a post-sync request reflects new weights; transformer-only sync, so the frozen text-encoder cache survives. The RL flywheel's hardest correctness. 3 tests. #5 Reward-model-as-a-served-card (§9.15) — a reward model is a card (scorer + a score loop emitting REWARD_BATCH units); ServedRewardScorer drop-in-replaces the numpy scorer so any RL method becomes RLHF/RLAIF with no method change. ToyRewardModel; models/reward/. 4 tests. #7 Content-adaptive control flow (§9.12) — CacheDiTDenoiseLoop (isolated WanDenoiseLoop subclass) reuses the cached velocity when predictions barely change (cache-dit skip) and early-exits on convergence — variable step count; interleave parity holds across ragged loops. models/adaptive/. 4 tests. #8a Nested workflows (§9.13) — a workflow stage can invoke another workflow (engine.run routes ids); requires/validate recurse; cycles caught at registration + a run-time guard (engine._wf_running). build_t2i_i2v_extend_workflow. 5 tests. designv4 §9.11–§9.15 + falsifier/layout/closing updates (148 tests, 25 files). Also removed a pre-existing unused import in ltx2/loop.py.
65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
"""RewardLoop — a served reward model scoring samples as ``REWARD_BATCH`` WorkUnits (design_v3 §10).
|
|
|
|
The reward in RL is computed by a *model* (PickScore/CLIP/a VLM verifier), not a heuristic. This loop
|
|
makes that a first-class part of the runtime: it reads a batch of K rollout media from a slot and emits
|
|
``REWARD_BATCH`` work units (chunked) that score them through the engine — so reward compute is admitted,
|
|
priced, and interleavable exactly like a denoise step or an AR token. A GRPO group's K samples score in
|
|
one (or a few) batched units — the homogeneous-batch win the scheduler likes.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from ..._enums import WorkUnitKind
|
|
from ...loop.contracts import (
|
|
Done,
|
|
LoopResult,
|
|
LoopState,
|
|
ResourceRequest,
|
|
ShapeSignature,
|
|
StepResult,
|
|
WorkPlan,
|
|
)
|
|
|
|
|
|
class RewardLoop:
|
|
def __init__(self, *, loop_id, scorer_id="scorer", cost, media_slot="media", batch=4):
|
|
self.loop_id = loop_id
|
|
self.scorer_id = scorer_id
|
|
self.cost = cost
|
|
self.media_slot = media_slot
|
|
self.batch = batch
|
|
|
|
def init(self, req, model, ctx) -> LoopState:
|
|
media = list(ctx.slots.get(self.media_slot) or [])
|
|
st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id,
|
|
request_id=req.request_id, profile=ctx.profile)
|
|
st.scratch["batches"] = [media[i:i + self.batch] for i in range(0, len(media), self.batch)] or [[]]
|
|
st.scratch["scores"] = []
|
|
return st
|
|
|
|
def next(self, st: LoopState):
|
|
i = st.step_idx
|
|
batches = st.scratch["batches"]
|
|
if i >= len(batches):
|
|
return Done()
|
|
chunk, sid = batches[i], self.scorer_id
|
|
|
|
def run(model, override=None):
|
|
scorer = model.component(sid)
|
|
return StepResult(output={"scores": [float(scorer.score(m)) for m in chunk], "i": i})
|
|
|
|
return WorkPlan(
|
|
loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.REWARD_BATCH,
|
|
shape_sig=ShapeSignature(WorkUnitKind.REWARD_BATCH, dims=(len(chunk),)),
|
|
resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)),
|
|
resident_bytes=64 * max(1, len(chunk)), peak_activation_bytes=64),
|
|
payload={"reward_batch": i}, run=run, label=f"reward.batch{i}")
|
|
|
|
def advance(self, st: LoopState, result: StepResult) -> LoopState:
|
|
st.scratch["scores"].extend(result.output["scores"])
|
|
st.step_idx += 1
|
|
return st
|
|
|
|
def finalize(self, st: LoopState) -> LoopResult:
|
|
return LoopResult(outputs={"rewards": list(st.scratch["scores"])},
|
|
metrics={"reward_batches": float(st.step_idx)})
|