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
…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 rescales the LoRA to alpha*(t) x gain per denoise step, where alpha*(t) is the measured systematic fraction of the drift-induced error at each solver timestep. 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, 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; the gate profile resolves by nearest-t lookup. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: wenqingw <wenqingw@nvidia.com>
Greptile SummaryThis PR adds a Clean Forcing drift corrector for Omnidreams single-view rollouts: a frozen-base LoRA (r16, ~7.3M params on self-attention q/k/v/o projections) trained via counterfactual teacher + DAgger + drift-contraction and gated per denoise step at
Confidence Score: 4/5Safe to merge with the caveat that the production checkpoint-loading path uses The core mechanism is well-designed: Files Needing Attention:
|
| Filename | Overview |
|---|---|
| integrations/omnidreams/omnidreams/_drift_corrector.py | Production deploy hook: wraps self-attn linears with LoRA, loads checkpoint, patches predict_flow with an alpha*(t) gate. Two issues: weights_only=False in torch.load allows arbitrary code execution from a user-supplied path; and gated_pf silently converts a missing timestep arg into a None-dereference AttributeError. |
| integrations/omnidreams/omnidreams/runner.py | Adds drift_corrector and drift_corrector_gain fields to OmnidreamsRunnerConfig and calls apply_drift_corrector in run(). The corrector is deployed unconditionally before the save-embeddings early return, wasting ~88 MB of load and permanently patching the transformer even when the AR loop is never entered. |
| integrations/omnidreams/drift_correction/_train_attn.py | Functional, grad-friendly self-attention toggle for training. The _INJECT.pop(0) pattern runs O(n) per block during contraction training steps, but otherwise the context-manager design is clean and idempotent. |
| integrations/omnidreams/drift_correction/_lora.py | Training-side LoRA helper: apply_lora, save_lora, load_lora. Uses weights_only=False in load_lora (training-only). Initialization sets scale=1.0 vs deploy-side 0.0; intentional difference since deploy always sets scale before inference. |
| integrations/omnidreams/drift_correction/_host.py | Omnidreams host adapter for rollout capture, KV-cache replay, counterfactual probes, clip I/O. load_clip uses weights_only=False. The os.environ.setdefault at module import time is a side-effect that runs on any import. |
| integrations/omnidreams/drift_correction/train_v1.py | v1 counterfactual-teacher LoRA trainer. Includes resume, exclusion logging, and startup replay-equivalence check. Well-structured. |
| integrations/omnidreams/drift_correction/train_v2.py | v2 DAgger + drift-contraction trainer. Accesses private scheduler internals (_full_timesteps, _full_sigmas). Val-peak snapshot logic is correct. |
| integrations/omnidreams/drift_correction/gate_faithful.py | Step-0 systematicity gate: computes alpha*(t) per timestep and depth, writes JSON summary, prints GO/NO-GO. Clean and read-only. |
| integrations/omnidreams/tests/test_drift_corrector.py | Six CPU-only unit tests covering identity at zero scale/B, linearity of delta, target matching, scale propagation, and gate profile lookup. Good deploy-contract coverage. |
Sequence Diagram
sequenceDiagram
participant Runner as OmnidreamsRunner.run()
participant DC as apply_drift_corrector()
participant Net as CosmosTransformer.network
participant PF as transformer.predict_flow (gated_pf)
participant Orig as original predict_flow
Runner->>DC: checkpoint, gain
DC->>Net: _apply_lora(network) wrap Linear to _LoRALinear
DC->>Net: load checkpoint tensors into A/B params
DC->>PF: "monkey-patch transformer.predict_flow = gated_pf"
DC-->>Runner: log line
Note over Runner: AR rollout loop or save-embeddings early return
Runner->>PF: predict_flow(noisy_latent, timestep, cache, input)
PF->>PF: "nearest-t lookup alpha*(t)"
PF->>Net: "_set_scale(network, alpha * gain)"
PF->>Orig: "orig_pf(*args, **kwargs)"
Orig-->>PF: flow
PF-->>Runner: corrected flow
Reviews (1): Last reviewed commit: "omnidreams: add Clean Forcing drift corr..." | Re-trigger Greptile
![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!
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)
corrgate025alpha*(t)gate prevents merging the LoRA into one weight set, and the hook runs the rank-16 delta path in fp32 with a per-step rescale. Added parameters: 7.34M (0.36% of the 2.06B base, ~56 MB of weights); peak activation VRAM unchanged. A bf16 delta path / compiled gate would reclaim most of this headroom.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