hy_worldplay: Clean Forcing drift corrector (training recipe + content-keyed deploy) - #396
hy_worldplay: Clean Forcing drift corrector (training recipe + content-keyed deploy)#396wenqingw-nv wants to merge 1 commit into
Conversation
Greptile SummaryThis PR adds a Clean Forcing drift corrector for HY-WorldPlay autoregressive rollouts: a rank-16 LoRA trained on self-generated strafe-loop pairs (zero real video) that reduces long-rollout drift ~50% with no dynamics loss. It ships with a content-keyed deploy hook (static trajectories run untouched base weights; motion jobs apply the LoRA at
Confidence Score: 4/5Safe to merge for single-run use cases (the shipped production path); re-use of the same runner across multiple The pre-merged path introduced in this PR (now the default) calls Files Needing Attention: integrations/hy_worldplay/hy_worldplay/_drift_corrector.py — the pre-merged weight-swap path and its interaction with repeated Important Files Changed
Sequence DiagramsequenceDiagram
participant Runner
participant DriftCorrector as maybe_apply_drift_corrector
participant Scheduler as scheduler.sample (gated)
participant Transformer as transformer.predict_flow
participant Finalizer as finalize_kv_cache (gated)
Runner->>DriftCorrector: called inside run() [no re-entry guard]
DriftCorrector->>DriftCorrector: is_static_trajectory(pose)?
alt static trajectory
DriftCorrector-->>Runner: base (static trajectory) — no-op
else motion trajectory
DriftCorrector->>DriftCorrector: _target_linears(network)
DriftCorrector->>DriftCorrector: _premerge_weight_sets(linears, sd, gain) reads lin.weight.data — BUG on 2nd run()
DriftCorrector->>Scheduler: wrap scheduler.sample → gated_sample
DriftCorrector->>Finalizer: wrap finalize_kv_cache → gated_finalize
DriftCorrector-->>Runner: corrected (pre-merged N weight sets)
end
loop AR chunks
Runner->>Scheduler: scheduler.sample(initial_noise, predict_flow, rng)
loop solver steps
Scheduler->>Scheduler: "_swap(step_alphas[i]) — lin.weight.data = merged_w"
Scheduler->>Transformer: predict_flow(noisy, timestep)
Transformer-->>Scheduler: flow estimate
end
Scheduler-->>Runner: x0 chunk
Runner->>Finalizer: finalize_kv_cache(...)
Finalizer->>Finalizer: _swap(ctx_alpha)
Finalizer-->>Runner: done
end
Reviews (5): Last reviewed commit: "hy_worldplay: add Clean Forcing drift co..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| network = network._orig_mod | ||
| params = _apply_lora(network) | ||
|
|
||
| sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] |
There was a problem hiding this comment.
weights_only=False allows arbitrary code execution
torch.load(..., weights_only=False) uses Python's pickle deserializer, which can execute arbitrary code embedded in the file. checkpoint is a user-controlled path via HyWorldPlayWanI2VRunnerConfig.drift_corrector, so a replaced or tampered .pt file would run attacker code on load. The save_lora format stores only {"lora": {int: Tensor}} — all primitive Python types and tensors — so weights_only=True loads it without issue.
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | ||
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | ||
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | ||
| drift-induced error at each of the distilled solver's timesteps. The | ||
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step.""" |
There was a problem hiding this comment.
GATE_ALPHA dict — deploy and training scripts diverge silently
GATE_ALPHA is defined with identical values in both this file and drift_correction/_rollout.py. If a future gate rerun changes the alpha*(t) profile, only one copy is likely to be updated, causing the deployed per-step gain profile to silently mismatch the training configuration.
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | |
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | |
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | |
| drift-induced error at each of the distilled solver's timesteps. The | |
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step.""" | |
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | |
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | |
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | |
| drift-induced error at each of the distilled solver's timesteps. The | |
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step. | |
| NOTE: This dict is also defined in ``drift_correction/_rollout.py``. Both copies | |
| must be kept in sync; prefer sourcing from a single shared location.""" |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
58702aa to
d91feb2
Compare
|
Checkpoint availability: the trained v2 corrector LoRA (59 MB) is published as a release asset on the contributing fork for review/testing: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/clean-forcing-hy-v2 — happy to rehome it (release asset here / HF / internal artifact store) per maintainer preference. Note the deploy dial is inference-time: this single checkpoint serves all configurations (flat gain, alpha*(t) gate, gate x gain); the shipped setting is alpha*(t) x 0.5 with content-keyed selection. |
|
/ok to test d91feb2 |
d91feb2 to
ada922d
Compare
| if cfg.drift_corrector is not None: | ||
| from hy_worldplay._drift_corrector import maybe_apply_drift_corrector | ||
|
|
||
| mode = maybe_apply_drift_corrector( | ||
| self, cfg.drift_corrector, cfg.drift_corrector_gain | ||
| ) | ||
| logger.info(f"[{cfg.runner_name}] drift corrector: {mode}") |
There was a problem hiding this comment.
run() call, stale gain overrides the new one
maybe_apply_drift_corrector installs a gated_pf closure that captures gain from the call site. Each subsequent run() wraps the previous gated_pf in a new one: the outermost closure (latest call's gain) sets the scale, then returns via orig_pf, which is the previous gated_pf — which overrides the scale again with an older captured gain before calling the real predict_flow. If drift_corrector_gain is mutated between runs (e.g. runner.config.drift_corrector_gain = 0.3 after a first run at 0.7), the network executes every subsequent step at the first-call gain (0.7), not the current one, with no error raised.
Moving the application into __init__ (or adding a _drift_corrector_applied guard) would prevent the stacking entirely.
|
Domain scope note: the shipped corrector was trained on urban/outdoor content without people; person-heavy or stylized domains may need a pairs refresh (recipe included, ~1 GPU-day). We validated the refresh path: a 200-step continue-train with ~10% person/scene-diverse pairs improves the walking-person stress scene (drift −47%, best seam profile) while keeping — in fact improving — the original commanded-motion profile (bridge drift cut −58%). Checkpoint to follow with the hosting decision. |
|
Checkpoint update: the shipped LoRA is now v2c2 — a content-diversified refresh (200-step continue-train, 10% person/scene-diverse pairs) that improves BOTH axes: bridge drift cut −58% (was −48%) and the walking-person stress scene fixed (drift −47%, owner-verified side-by-sides). Release asset: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/clean-forcing-hy-v2c2 (md5 1de6023c...; loads via the identical deploy path, no code change). The earlier v2 asset remains available for reproducing the PR's tables. |
ada922d to
8afdaf1
Compare
…ntent-keyed deploy) Adds a trained drift corrector for the HY-WorldPlay WAN-5B integration: a frozen-base LoRA (r16 q/k/v/o, ~0.3% params) trained from the model's own strafe-loop rollouts with a counterfactual clean-history teacher -- zero real videos end-to-end. Deploy: HyWorldPlayWanI2VRunnerConfig.drift_corrector (checkpoint path) + drift_corrector_gain. Selection is content-keyed per job: static (locked-off camera) poses run untouched base weights; commanded-motion poses apply the LoRA at alpha*(t) x gain per denoise step, where alpha*(t) is the measured systematic fraction of the drift-induced error at each solver timestep. The hook pre-merges each distinct alpha*(t) x gain into its own cached copy of the target projection weights at load (3 sets, +8.1 GiB) and swaps the set per solver step, driven CPU-side from the denoising schedule -- the corrected forward issues the same kernels as base, measured +0.22% per chunk. DRIFT_CORRECTOR_UNFUSED=1 selects the runtime fp32 A/B-delta path instead. Measured on 24-chunk (~19 s) rollouts: -48% delta-drift vs base with dynamics ~= base and no added seam pulse; the base's late-horizon saturation runaway is eliminated. Training/eval scripts and the step-0 diagnostic live under integrations/hy_worldplay/drift_correction/ with a README covering the method and the full reproduction recipe. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Wenqing Wang <wenqingw@nvidia.com>
8afdaf1 to
a57b507
Compare
| linears = _target_linears(network) | ||
| assert len(sd) == 2 * len(linears), ( | ||
| f"corrector checkpoint has {len(sd)} LoRA tensors but the network " | ||
| f"exposes {2 * len(linears)}; rank or target mismatch." | ||
| ) | ||
| weight_sets, added_bytes = _premerge_weight_sets(linears, sd, gain) | ||
| current: list[float | None] = [None] | ||
|
|
||
| def _swap(alpha: float) -> None: | ||
| if alpha != current[0]: | ||
| for lin, w in zip(linears, weight_sets[alpha]): | ||
| lin.weight.data = w | ||
| current[0] = alpha | ||
|
|
||
| # Drive the gate CPU-side. The Euler scheduler makes exactly one | ||
| # ``predict_flow`` call per solver step in a Python loop, so each | ||
| # step's alpha resolves from the load-time schedule by call index — | ||
| # reading the timestep tensor back per step (the unfused path's | ||
| # ``float(timestep.max())``) would stall the CPU launch queue every | ||
| # solver step. | ||
| scheduler = runner.pipeline.diffusion_model.scheduler | ||
| n_steps = scheduler.config.num_inference_steps | ||
| step_alphas = [_nearest_alpha(t) for t in scheduler.timesteps.tolist()[:n_steps]] | ||
| ctx_alpha = _nearest_alpha( | ||
| float(runner.pipeline.diffusion_model.config.context_noise) | ||
| ) | ||
| orig_sample = scheduler.sample | ||
|
|
||
| def gated_sample(initial_noise, predict_flow, rng=None): | ||
| calls = [0] | ||
|
|
||
| def pf(noisy, timestep): | ||
| assert calls[0] < len(step_alphas), "predict_flow calls > solver steps" | ||
| _swap(step_alphas[calls[0]]) | ||
| calls[0] += 1 | ||
| return predict_flow(noisy, timestep) | ||
|
|
||
| return orig_sample(initial_noise=initial_noise, predict_flow=pf, rng=rng) | ||
|
|
||
| scheduler.sample = gated_sample | ||
| orig_finalize = transformer.finalize_kv_cache | ||
|
|
||
| def gated_finalize(*args, **kwargs): | ||
| _swap(ctx_alpha) | ||
| return orig_finalize(*args, **kwargs) | ||
|
|
||
| transformer.finalize_kv_cache = gated_finalize |
There was a problem hiding this comment.
run() call
maybe_apply_drift_corrector is invoked at the top of run() with no re-entry guard. On the first call, _premerge_weight_sets correctly merges W_base + gain * alpha * B @ A and stores those tensors. _swap then does lin.weight.data = w, so after the AR loop finishes each linear's weight.data points at one of those merged tensors (the one for ctx_alpha).
On the second run(), _target_linears returns the same nn.Linear objects, but now lin.weight.detach() reads the already-merged value. _premerge_weight_sets then computes:
W_merged_prev + gain * alpha * B @ A = W_base + gain * (ctx_alpha + alpha) * B @ A
The correction is applied twice — the network runs with 2× the intended delta on every step, silently, with no error raised. The wrapped scheduler.sample and finalize_kv_cache also stack a second closure layer on each re-invocation.
The same guard recommended in the previous thread for the unfused path (applying in __init__, or tracking _drift_corrector_applied as a runner attribute) would also prevent the double-correction here, since it would stop _premerge_weight_sets from reading the already-swapped weights.
1 participant
Long autoregressive HY-WorldPlay rollouts drift: each chunk conditions on self-generated history and small errors compound into saturation/contrast runaway. This PR adds a trained drift corrector (Clean Forcing), the zero-real-data recipe that produced it, and a content-keyed deploy hook.
drift_corrector=None(default) is byte-identical to current behavior.Contents
hy_worldplay/_drift_corrector.py+ two runner-config fields (drift_corrector,drift_corrector_gain) — content-keyed per job: static (locked-off camera) trajectories run untouched base weights (static scenes measure negative drift on this host — correction there is pure artifact cost); commanded-motion trajectories apply the LoRA atalpha*(t) × gainper denoise step.integrations/hy_worldplay/drift_correction/— full recipe, zero real videos: strafe-loop pair construction, step-0alpha*(t)systematicity gate (go/no-go + the deploy gain profile), v1 counterfactual-teacher training, v2 DAgger + drift-contraction, gain-sweep eval, drift/dynamics/progression/seam scoring, static-scene demo suite. README documents the ~2.5 GPU-day reproduction.tests/test_drift_corrector.py(ci_cpu, 6 tests) — LoRA wrapper is a strict identity until a checkpoint is loaded and gated on; content-keyed selection triggers on action labels.Docs
Measured (704×1280, 24-chunk ≈19 s rollouts, held-out trajectories, GB300)
alpha*(t) × 0.5Shipped config cuts the drift metric ~50% with no dynamics loss and no added chunk-boundary pulse; commanded trajectories keep advancing (owner-verified side-by-sides). The method/reproduction details:
drift_correction/README.md(quick start)Deploy notes
alpha*(t)gating cannot merge into one weight set → motion jobs run the LoRA unfused (~0.3% extra matmul); static jobs run the base at zero overhead.Inference overhead (GB300, same rollout protocol as Measured; 3 runs, steady-state mean ± std)
corrgate050(pre-merged)alpha*(t) × gaininto one cached copy of the target projection weights per distinct gate value at load (3 sets, +8.1 GiB) and swaps the set per solver step, driven CPU-side from the denoising schedule — the corrected forward issues the same kernels as base: no LoRA matmuls, no GPU timestep readback. Overhead is therefore GPU-independent (base runtime on any card).DRIFT_CORRECTOR_UNFUSED=1selects the previous runtime fp32 delta path (1461 ± 69 ms/chunk, +12%).Video Comparison
Left: FlashDreams-HY (base), Right: w/ Clean Forcing LoRA drift corrector (better details/color preservation across time):
sbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042mix.mp4
fsbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042.mp4
sbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042.mp4
VBench Ablation: