[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:
co-authored by
Kyle
Claude Opus 5
Satyam Srivastava
Aryan Kumar
parent
cf6a00b9be
commit
aaef49bfc6
@@ -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,
|
||||
),
|
||||
),
|
||||
))
|
||||
|
||||
Reference in New Issue
Block a user