[perf] inference: regional fullgraph torch.compile of the DiT blocks (port of #1718 to the inference path)
Extend the training-side regional-compile port (internal 863e87342, GO
verdict vsa_gate/compile_ab/VERDICT.md) to inference so every H3 run can
compile the 52 transformer blocks. Opt-in and off by default:
- New FastVideoArgs.inference_torch_compile (CLI --inference-torch-compile),
env FASTVIDEO_INFERENCE_TORCH_COMPILE=1 folded in __post_init__ (the
attention_backend parse-once pattern), reachable through
PipelineSelection.experimental {"inference_torch_compile": true} exactly
like the VSA_sparsity / VSA_tile_size knobs.
- maybe_load_fsdp_model applies the compile right after the transformer
loads: per-_compile_conditions block, fullgraph=True + inductor
options.emulate_precision_casts injected, no user kwargs needed (the
compile-A/B verdict recipe). prepare_for_compile still runs first, so the
H3 fusion-inertness warning of #1735 is preserved.
- _regional_compile_unsupported_reason ported from the training loader:
VSA / VSA-H3 backends and the FASTVIDEO_DISABLE_ATTENTION_COMPILE=1
escape hatch degrade the transformer to eager with one warning; FLASH_ATTN
on flash-attn 3 is rejected with an actionable message. FA4/FA2/SDPA
compile through their existing custom-op boundaries (forward-only is
enough at inference; the training port's backward custom ops are not
needed under no_grad).
- FASTVIDEO_DISABLE_ATTENTION_COMPILE default flipped to 0 (attention is
traced by default), matching the training port and upstream #1718 —
eager runs are unaffected (torch.compiler.disable is inert outside
dynamo).
- enable_torch_compile + inference_torch_compile together skip the
pipeline-level DiT compile (the loader already owns those forwards).
- examples: --inference-torch-compile on basic_minimax_h3_t2v.py (dense
target) and basic_fasth3.py (exercises the VSA degrade path).
- CPU-safe contract tests for the guard and the kwargs injection.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
2f3d407406
commit
129fcaffab
@@ -85,6 +85,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--compile-mode",
|
||||
default=None,
|
||||
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
|
||||
parser.add_argument("--inference-torch-compile",
|
||||
action="store_true",
|
||||
help="regional fullgraph torch.compile of each DiT block after load. NOTE: this "
|
||||
"script always runs the VSA-H3 backend, which is not fullgraph-traceable — the "
|
||||
"loader logs one warning and keeps the transformer eager. The flag is exposed "
|
||||
"here to exercise exactly that guard")
|
||||
parser.add_argument("--repeats",
|
||||
type=int,
|
||||
default=1,
|
||||
@@ -119,6 +125,8 @@ def main() -> None:
|
||||
}
|
||||
if args.vsa_sparsity > 0.0:
|
||||
experimental["VSA_sparsity"] = args.vsa_sparsity
|
||||
if args.inference_torch_compile:
|
||||
experimental["inference_torch_compile"] = True
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
|
||||
@@ -15,6 +15,7 @@ from fastvideo.api import (
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--compile-mode",
|
||||
default=None,
|
||||
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
|
||||
parser.add_argument("--inference-torch-compile",
|
||||
action="store_true",
|
||||
help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
|
||||
"training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
|
||||
"First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
|
||||
"the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
|
||||
parser.add_argument("--repeats",
|
||||
type=int,
|
||||
default=1,
|
||||
@@ -54,9 +61,16 @@ def main() -> None:
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Boot-time run configuration folded into FastVideoArgs (the same
|
||||
# experimental-dict route basic_fasth3.py uses for the VSA knobs).
|
||||
experimental: dict[str, object] = {}
|
||||
if args.inference_torch_compile:
|
||||
experimental["inference_torch_compile"] = True
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
pipeline=PipelineSelection(experimental=experimental),
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.num_gpus > 1,
|
||||
|
||||
Reference in New Issue
Block a user