[feat]: run FastH3 across two DGX Sparks with Ray sequence parallel (#1803)

Co-authored-by: Kyle <shh075@ucsd.edu>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Satyam Srivastava <srivastavasatyam53@gmail.com>
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
2026-09-01 01:37:23 -07:00
committed by GitHub
co-authored by Kyle Claude Opus 5 Satyam Srivastava Aryan Kumar
parent cf6a00b9be
commit aaef49bfc6
37 changed files with 2645 additions and 321 deletions
+24
View File
@@ -48,6 +48,14 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
# License review completes. A local snapshot can be passed here instead.
parser.add_argument("--prompt", required=True)
parser.add_argument("--output", default="outputs/fasth3")
parser.add_argument("--lazy-module-load",
action=argparse.BooleanOptionalAction,
default=None,
help="load each heavy component on first use and free it after the last stage that "
"needs it, so peak memory is the largest overlapping set instead of the sum of every "
"component. Omit for auto (on for unified-memory devices such as GB10; off on discrete "
"GPUs). Costs a reload per generation; pass --no-lazy-module-load to keep every "
"component resident")
parser.add_argument("--profile",
choices=("all", "strict"),
default="all",
@@ -69,6 +77,13 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
default=True,
help="run one excluded request before timing")
parser.add_argument("--num-gpus", type=int, default=4)
parser.add_argument(
"--execution-backend",
choices=("mp", "ray"),
default=None,
help="mp for one node; ray for a Ray cluster (two DGX Sparks). "
"Default: ray when RAY_ADDRESS is set, otherwise mp",
)
parser.add_argument("--vsa-sparsity",
type=float,
default=0.9,
@@ -231,6 +246,12 @@ def validate_profile_dependencies(args: argparse.Namespace) -> None:
"`cd fastvideo-kernel && ./build.sh`), or pass --vsa-kernel triton.")
def _execution_backend(args: argparse.Namespace) -> str:
if args.execution_backend is not None:
return args.execution_backend
return "ray" if os.environ.get("RAY_ADDRESS") else "mp"
def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
use_vsa = _uses_vsa(args)
experimental: dict[str, object] = {
@@ -261,6 +282,7 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
),
engine=EngineConfig(
num_gpus=args.num_gpus,
execution_backend=_execution_backend(args),
use_fsdp_inference=args.num_gpus > 1 and not args.replicated_dit,
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
offload=OffloadConfig(
@@ -269,6 +291,7 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
text_encoder=True,
vae=True,
pin_cpu_memory=args.pin_cpu_memory,
lazy_module_load=args.lazy_module_load,
),
compile=CompileConfig(
enabled=args.torch_compile,
@@ -331,6 +354,7 @@ def run(args: argparse.Namespace) -> list[float]:
f"Denoising contract override: {args.steps} sigma points = {args.steps - 1} DiT forwards")
print("Profile environment: " + " ".join(f"{key}={value if value is not None else '<unset>'}"
for key, value in environment.items()))
print(f"Execution backend: {_execution_backend(args)}")
generator = VideoGenerator.from_config(build_generator_config(args))
measured_wall_times: list[float] = []
@@ -0,0 +1,59 @@
# FastH3 on two DGX Sparks (one GPU each) over QSFP RoCE.
# Bring up the Ray cluster first: docs/getting_started/installation/spark_pair.md
#
# request.sampling below is an example, not a required recipe. Change height,
# width, num_frames, num_inference_steps, seed, and prompt. Legal H3 frame
# counts are 17n+5, max 345 (15 s).
#
# source examples/inference/optimizations/spark_pair_env.sh
# export RAY_ADDRESS=<qsfp-head>:6379
# export FASTVIDEO_HOST_IP=<qsfp-head>
# FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \
# FASTVIDEO_VAE_PARALLEL_DECODE=1 FASTVIDEO_STAGE_LOGGING=1 \
# fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yaml
generator:
model_path: FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree
engine:
num_gpus: 2
execution_backend: ray
use_fsdp_inference: false
parallelism:
tp_size: 1
sp_size: 2
offload:
dit: false
dit_layerwise: false
text_encoder: true
vae: true
pin_cpu_memory: true
lazy_module_load: true
compile:
enabled: false
vae_enabled: true
pipeline:
experimental:
attention_backend: VIDEO_SPARSE_ATTN_H3
VSA_sparsity: 0.9
VSA_tile_size: 64
inference_torch_compile: true
vae_parallel_decode: true
vae_parallel_decode_strategy: gather
h3_sequential_load: true
request:
prompt: >-
A wide cinematic shot of an alpine meadow at sunrise, pale pink mountain
peaks above a blue valley filled with thin morning mist.
negative_prompt: ""
sampling:
seed: 2026
height: 768
width: 1344
num_frames: 124
fps: 24
num_inference_steps: 5
guidance_scale: 1.0
batch_cfg: false
output:
output_path: outputs/fasth3_spark_pair/
save_video: true
return_frames: false
@@ -4,6 +4,7 @@
from __future__ import annotations
import argparse
import os
from pathlib import Path
from fastvideo import VideoGenerator
@@ -38,7 +39,19 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--num-gpus", type=int, default=4)
parser.add_argument(
"--execution-backend",
choices=("mp", "ray"),
default=None,
help="mp for one node; ray for a Ray cluster (two DGX Sparks). "
"Default: ray when RAY_ADDRESS is set, otherwise mp",
)
parser.add_argument("--torch-compile", action="store_true", help="torch.compile the DiT transformer path")
parser.add_argument("--compile-vae",
action=argparse.BooleanOptionalAction,
default=True,
help="compile the video VAE decoder independently of the DiT (on by default; "
"the Spark lazy-load path needs this registered before first materialize)")
parser.add_argument("--compile-mode",
default=None,
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
@@ -48,6 +61,14 @@ def parse_args() -> argparse.Namespace:
"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("--lazy-module-load",
action=argparse.BooleanOptionalAction,
default=None,
help="load each heavy component on first use and free it after the last stage that "
"needs it, so peak memory is the largest overlapping set instead of the sum of every "
"component. Omit for auto (on for unified-memory devices such as GB10; off on discrete "
"GPUs). Costs a reload per generation; pass --no-lazy-module-load to keep every "
"component resident")
parser.add_argument("--repeats",
type=int,
default=1,
@@ -67,12 +88,14 @@ def main() -> None:
if args.inference_torch_compile:
experimental["inference_torch_compile"] = True
execution_backend = args.execution_backend or ("ray" if os.environ.get("RAY_ADDRESS") else "mp")
generator = VideoGenerator.from_config(
GeneratorConfig(
model_path=args.model_path,
pipeline=PipelineSelection(experimental=experimental),
engine=EngineConfig(
num_gpus=args.num_gpus,
execution_backend=execution_backend,
use_fsdp_inference=args.num_gpus > 1,
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
offload=OffloadConfig(
@@ -81,10 +104,12 @@ def main() -> None:
text_encoder=True,
vae=True,
pin_cpu_memory=False,
lazy_module_load=args.lazy_module_load,
),
compile=CompileConfig(
enabled=args.torch_compile,
mode=args.compile_mode,
vae_enabled=args.compile_vae,
),
),
))