nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/basic/basic_fasth3.py
+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] = []