omnidreams: Clean Forcing drift corrector (training recipe + gated deploy) - #398
omnidreams: Clean Forcing drift corrector (training recipe + gated deploy)#398wenqingw-nv wants to merge 1 commit into
Conversation
Greptile SummaryThis PR ships the Clean Forcing drift corrector for Omnidreams: a frozen-base LoRA (r16 on self-attention q/k/v/o projections, ~7.3 M params) trained to close the velocity gap between drifted and clean-history predictions, deployed at
Confidence Score: 4/5The production deploy path is well-tested and Both training scripts contain an unbounded retry loop in Files Needing Attention: Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if cfg.drift_corrector is not None: | ||
| from omnidreams._drift_corrector import apply_drift_corrector | ||
|
|
||
| mode = apply_drift_corrector( | ||
| self, cfg.drift_corrector, cfg.drift_corrector_gain | ||
| ) | ||
| logger.info(f"[{cfg.runner_name}] drift corrector: {mode}") | ||
| if cfg.save_embeddings_path is not None: | ||
| self._run_save_embeddings(cfg.save_embeddings_path) | ||
| return |
There was a problem hiding this comment.
apply_drift_corrector runs unconditionally before the save_embeddings_path early-return check. When a user sets both drift_corrector and save_embeddings_path, the full corrector deployment happens (loading ~88 MB, wrapping every target linear, monkey-patching predict_flow), then _run_save_embeddings returns immediately without ever calling predict_flow. The patch persists on the transformer instance, leaving it in a permanently modified state even after that run exits — which matters if the runner or pipeline object is later inspected or reused. Moving the corrector setup after the mode-dispatch block (or gating it on embeddings_path is None and save_embeddings_path is None) would avoid the wasted work and the permanently-patched transformer.
| 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 from the checkpoint file
torch.load(..., weights_only=False) deserializes arbitrary Python objects via pickle, so any file pointed at by drift_corrector (a user-supplied Path in the runner config) can execute arbitrary code at load time. The checkpoint format used here is a plain {"lora": {int: Tensor}} dict — no custom classes or non-tensor objects — so weights_only=True would work after registering the built-in types with torch.serialization.add_safe_globals if needed. The same pattern appears in _host.py (load_clip) and _lora.py (load_lora), which are used in the training scripts; at minimum the production deploy path in _drift_corrector.py should use weights_only=True.
| def gated_pf(*args, **kwargs): | ||
| ts = kwargs.get("timestep", args[1] if len(args) > 1 else None) | ||
| t = float(ts.reshape(-1).max()) | ||
| alpha = min(GATE_ALPHA.items(), key=lambda kv: abs(kv[0] - t))[1] |
There was a problem hiding this comment.
None path in gated_pf gives an unhelpful AttributeError
If predict_flow is ever called without a timestep argument and with fewer than two positional args, ts resolves to None and ts.reshape(-1) raises AttributeError: 'NoneType' object has no attribute 'reshape' with no context about the actual problem. In all current call sites timestep is passed as a keyword argument, so the defensive fallback args[1] if len(args) > 1 else None works, but the None case should raise an explicit ValueError rather than deferring to an opaque AttributeError deep inside the corrector hook. A one-line guard (if ts is None: raise ValueError(...)) would make the failure immediately diagnosable.
| if _INJECT is not None: | ||
| ik, iv = _INJECT.pop(0) | ||
| span = ik.shape[-3] | ||
| pk = torch.cat([pk[..., :-span, :, :], ik], dim=-3) | ||
| pv = torch.cat([pv[..., :-span, :, :], iv], dim=-3) |
There was a problem hiding this comment.
_INJECT.pop(0) is O(n) per block, making injection O(n²) over all blocks
inject_kv sets _INJECT = list(kvs) and _functional_forward calls _INJECT.pop(0) once per transformer block. list.pop(0) shifts every remaining element left, so across all B blocks the total cost is O(B²). Using an iterator (e.g. iter(kvs)) and next() instead would give O(1) per block. This only affects training (the deploy path has no _INJECT active), but it runs in the inner training loop over every sample_losses call with a contraction term.
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!
…d deploy) Adds a trained drift corrector for the Omnidreams single-view integration: a frozen-base LoRA (r16 on the self-attn q/k/v/o projections, ~7.3M params) trained from the model's own tiled-HDMap loop rollouts with a counterfactual clean-history teacher in velocity space. Deploy: OmnidreamsRunnerConfig.drift_corrector (checkpoint path) + drift_corrector_gain (default 0.25). The hook pre-merges alpha*(t) x gain into one cached copy of the target projection weights per distinct gate value at load (2 sets, +1.8 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 ms/chunk added within noise. alpha*(t) is the measured systematic fraction of the drift-induced error at each solver timestep. DRIFT_CORRECTOR_UNFUSED=1 selects the runtime fp32 A/B-delta path instead. drift_corrector=None (default) is byte-identical to current behavior. Recipe under integrations/omnidreams/drift_correction/: pair construction (mixed-lap tiled-HDMap loops; loop-free fork variant as an ablation), the step-0 alpha*(t) systematicity gate, v1 counterfactual- teacher training, v2 DAgger + drift-contraction round with step-tagged and validation-peak snapshots, dial-grid eval, and a per-timestep reliability analysis utility. Measured on held-out 21.5 s rollouts: drift Delta-MUSIQ +0.99 vs base +2.44 (60% cut), late-window MUSIQ +5% over base, dynamics -17%. tests/test_drift_corrector.py (ci_cpu, 9 tests): the LoRA wrapper is a strict identity until a checkpoint is loaded and gated on; the wrap set matches the training-side module; the gate profile resolves by nearest-t lookup; the pre-merged weight sets match the unfused delta and leave the base weights untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: wenqingw <wenqingw@nvidia.com>
2e328b7 to
c6464f2
Compare
Driving the interactive demo for more than about a minute made the world fall apart: the road and lane markings dissolved into smeared vegetation and structures, unrecoverably. Cause is the KV cache layout. BlockKVCache is [sink | local window] and only the window rolls, but sink_size was 0, so nothing was pinned. The scene's real initial frame therefore left the cache after local_attn_size chunks, and from then on every chunk was conditioned purely on the model's own previous output -- errors became the new "truth" and compounded. sink_size was already a manifest field, but _select_config_name rejected any non-zero value with NotImplementedError and the value was never passed to the transformer. The underlying support was already complete: BlockKVCache documents "Sink tokens are never evicted" and _roll_local_window_left shifts only the region past sink_size. So this drops the guard and wires sink_size through _transformer_overrides as sink_size_t. A/B measured on an RTX 3090, same scene, same seed (1234), 512x288, auto-respawn disabled so a respawn could not silently reset the rollout, 300 chunks (~80 s of generated video) per arm: elapsed sink_size=0 sink_size=2 0 s clean clean 35 s oversaturated, stylised clean 75 s COLLAPSED (no road, no lanes) clean Mean saturation runs away without sinks (65.8 -> 106.3 -> 83.7) but stays flat with them (49.7 -> 49.0 -> 55.8); luminance std grows +65% versus +16%. Saturation runaway is an early indicator of the drift. Cost is 2.8% of steady-state FPS (29.1 -> 28.29, 274.9 -> 282.8 ms/chunk) and no measurable VRAM change -- the 16.26 GiB peak is set by the text encoder during init, not by the KV cache. This is a deploy-side mitigation, not a cure: it prevented collapse at 80 s but drift is a training-time problem, and a longer rollout may still degrade. The principled fix is the upstream Clean Forcing drift corrector (NVIDIA#398). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FuZXfbb4LV9tkFobCq6ZTg
1 participant
Long autoregressive Omnidreams rollouts drift: each chunk conditions on self-generated history and small errors compound into color/saturation runaway, texture mush, and content repetition. This PR adds a trained drift corrector (Clean Forcing), the recipe that produced it, and a gated deploy hook.
drift_corrector=None(default) is byte-identical to current behavior.Contents
omnidreams/_drift_corrector.py+ two runner-config fields (drift_corrector,drift_corrector_gain) — a frozen-base LoRA (r16 on the self-attention q/k/v/o projections, ~7.3M params) applied atalpha*(t) × gainper denoise step, wherealpha*(t)is the measured systematic fraction of the drift-induced error at each solver timestep (0.96 @ t=1000, 0.667 @ t=803).integrations/omnidreams/drift_correction/— the full recipe: mixed-lap tiled-HDMap loop pair construction (+ a gate-validated loop-free fork variant as an ablation), step-0alpha*(t)systematicity gate (go/no-go + the deploy gain profile), v1 counterfactual-teacher training in velocity space, v2 DAgger + drift-contraction round with step-tagged and validation-peak snapshots, dial-grid eval, per-timestep reliability analysis. README documents the ~1 GPU-day reproduction.tests/test_drift_corrector.py(ci_cpu, 6 tests) — the LoRA wrapper is a strict identity until a checkpoint is loaded and gated on; the wrap set matches the training-side module (same checkpoint order); the gate profile resolves by nearest-t lookup.Docs
Measured (704×1280, 81-chunk ≈21.5 s rollouts, held-out scenes, GB300)
alpha*(t) × 0.25Shipped config cuts the drift metric 60% with overall quality above base and visibly better tree/foliage detail and consistency across the rollout (human side-by-side review was held as the final acceptance bar throughout). The method/reproduction details:
drift_correction/README.md(quick start)Deploy notes
alpha*(t)gating cannot merge into one weight set → the LoRA runs unfused (~0.4% extra matmul);drift_corrector=Noneskips module surgery entirely.84beb8c014346833ce182310373b5d32) hosting: currently a fork release asset (https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/clean-forcing-omnidreams-v3vp); release asset here vs HF — maintainer preference welcome.Inference overhead (GB300, same rollout protocol as Measured; 3 runs, steady-state mean ± std)
corrgate025(pre-merged)alpha*(t) × gaininto one cached copy of the target projection weights per distinct gate value at load (2 sets, +1.8 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 (338.5 ± 5.9 ms/chunk, +45%).drift_corrector=None(default) runs the untouched base graph at zero overhead.Video Comparison
Left: Omnidreams (base), Right: w/ Clean Forcing (better color/detail consistency across time; 4× speed):
base_vs_cleanForcing1.mp4
base_vs_cleanForcing2.mp4