Add export and Dynamo capture fallback

[aryan5v]

Stack

This PR is intentionally stacked on fx-capture-hooks / PR #19. It should be reviewed and merged after that foundation branch.

What changed

• Default repeated-module capture to auto: symbolic FX, non-strict torch.export, then Dynamo.
• Retain bounded in-memory example args/kwargs and transient attention context per observed shape only until finalization; clear all live references before JSON serialization.
• Preserve nested tensor input metadata and record capture mode, attempts, mode breakdown, and sanitized failures.
• Temporarily expose FastVideo ModuleHookManager originals, compiler-disabled attention forwards, and observed forward context only inside capture.
• Restore the exact wrapped forward, parameter references, and layerwise-offload cache after every capture attempt, including failures.
• Reject unknown module hooks and keep reports metadata-only.

Root cause

torch.export cannot inspect FastVideo's functools.partial offload wrapper, and Dynamo intentionally skips compiler-disabled offload/attention entry points. The first bypass also called the offload pre-hook, which prefetches the next layer and polluted the next normal generation. Capture now loads only the current module state and restores the complete prior state observationally.

Validation

• FastVideo optimization integration: 31 passed.
• Ruff on changed capture/test files: passed.
• MotionKernel CPU suite before each push: 474 passed, 10 deselected.
• Wan GB200 job 920: clean generation, dominant transformer captured through export; 18 regions total (17 symbolic, 1 export), no capture errors, 3.8560-second clean median, 30,734.01 MiB peak CUDA memory.
• LTX GB200 job 923: clean generation; 59 regions total (35 symbolic, 24 export), no capture errors, 4.4961-second clean median, 67,804.14 MiB peak CUDA memory. The dominant LTX transformer remains a reported graph break.

No weights, prompts, activations, tensor values, source, or credentials are serialized.

Greptile Summary

This PR adds Dynamo as a third capture backend, implements a symbolic→export→dynamo fallback chain in auto mode, and introduces three context managers that temporarily bypass FastVideo's offload wrapper, re-establish the observed forward context, and expose compiler-disabled attention forwards — all with complete state restoration in finally blocks regardless of trace outcome.

• Fallback chain (_regions_for): each shape variant is tried against all modes in order; per-variant tensor references (example_args, example_kwargs, observed_context) are cleared in a guaranteed finally, with a second outer finally in finalize() covering variants that were never reached due to an early exception.
• Offload bypass (_capture_ready_module): saves parameter data and gpu_named_parameters, calls wait_and_replace_params() eagerly, exposes original_forward, then restores the exact wrapped state; rejects unknown ModuleHookManager hooks to prevent silent semantic changes.
• New output fields: capture_mode, capture_attempts, capture_failures, and capture_mode_breakdown give downstream consumers full visibility into how each region was captured.

Confidence Score: 4/5

Safe to merge after the schema version question is resolved; all new context managers restore state correctly and every tensor reference is guaranteed to be cleared before finalization returns.

The fallback logic, state-restoration guarantees, and test coverage are thorough. The one open question is whether forward_batch=None in the reproduced forward context can cause a different control-flow branch to be captured for production attention modules — acceptable for metadata profiling but worth a conscious decision.

Files Needing Attention: fastvideo/optimization/fx_capture.py — specifically _capture_forward_context (forward_batch omission) and _trace (dynamo return-type handling).

Important Files Changed

Filename	Overview
fastvideo/optimization/fx_capture.py	Core capture logic overhauled: adds three new context managers (_capture_ready_module, _capture_forward_context, _traceable_module_forwards), a per-variant try/finally guarantee for tensor cleanup, and a cascading symbolic→export→dynamo fallback. forward_batch is dropped from the reproduced context and the dynamo return-type extraction is version-sensitive.
fastvideo/envs.py	Default tracer changed from 'symbolic' to 'auto' in both the TYPE_CHECKING block and the runtime lambda — straightforward and consistent.
fastvideo/optimization/profiler.py	Exception message in _finish_capture truncated to type name only, preventing source snippets from leaking into the error output — safe, minimal change.
fastvideo/tests/optimization/test_fx_capture.py	New tests cover the full fallback chain, offload-wrapper bypass, context restoration, each explicit tracer mode, invalid-tracer cleanup, and nested-dataclass input metadata; all assertions are well-targeted and consistent with the implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[_regions_for: iterate variants] --> B[_mode_order returns attempt list]
    B -->|auto| C[symbolic, export, dynamo]
    B -->|explicit mode| D[single mode]
    B -->|invalid| E[ValueError - finally clears variant refs]
    C --> F[_trace: symbolic]
    D --> F
    F -->|success| K[extract_graph check operations]
    F -->|failure| G[record failure - next mode]
    G --> H[_trace: export via _capture_ready_module + _traceable_module_forwards + _capture_forward_context]
    H -->|success| K
    H -->|failure| I[_trace: dynamo]
    I -->|success| K
    I -->|all failed| J[capture_mode=None - continue]
    K -->|empty ops| J
    K -->|ops present| L[capture_mode set - build region dict]
    L --> M[finally: clear example_args/kwargs/context]
    J --> M
    E --> N[finalize finally: clear all remaining variants]
    M --> O[append region to results]

Loading

_{Reviews (2): Last reviewed commit: "fix: reject undocumented FX tracer alias..." | Re-trigger Greptile}

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 9a884454-71b3-4a36-a806-12e321310ad1

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}