From 57662a9ea5fe91cd1bc79624aa0d2189c59f9bbd Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 20 Oct 2026 22:28:39 +0000 Subject: [PATCH] [fix] Sync PAN2 padding parity with latest upstream (#1961) --- docs-site/navigation.json | 4 + docs/inference/pan2.md | 28 ++++- examples/inference/basic/basic_pan2.py | 2 +- tests/local_tests/pan2/PORT_STATUS.md | 14 +-- tests/local_tests/pan2/README.md | 17 ++- tests/local_tests/pan2/reference.py | 3 +- .../pipelines/test_pan2_pipeline_parity.py | 61 +++++---- .../test_pan2_transformer_parity.py | 118 +++++++++++------- 8 files changed, 154 insertions(+), 93 deletions(-) diff --git a/docs-site/navigation.json b/docs-site/navigation.json index d9a18ed76..fe5f0338a 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -123,6 +123,10 @@ ], "display": "group" }, + { + "label": "PAN 2 text/image to video", + "href": "/inference/pan2" + }, { "label": "Examples", "href": "/inference/examples/examples_inference_index" diff --git a/docs/inference/pan2.md b/docs/inference/pan2.md index 9e8fc5ba7..849f690ad 100644 --- a/docs/inference/pan2.md +++ b/docs/inference/pan2.md @@ -1,6 +1,6 @@ -# PAN2 Text/Image to Video +# PAN 2 Text/Image to Video -PAN2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the +PAN 2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the Wan2.2 VAE. One pipeline serves text-to-video and image-to-video: pass `image_path` to condition on a first frame and leave it out for text only. @@ -21,7 +21,7 @@ from the `pan2_t2v` and `pan2_i2v` presets: | Steps | 50, uniform sigmas from 1 to 0 with flow shift 7 | | Guidance | 3.0, classifier-free guidance with an empty negative prompt | -PAN2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw +PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw prompts work, with weaker results. ## Usage @@ -37,6 +37,7 @@ generator = VideoGenerator.from_pretrained( "IFM/PAN2", num_gpus=1, dit_cpu_offload=False, + dit_layerwise_offload=False, text_encoder_cpu_offload=True, ) generator.generate_video( @@ -55,6 +56,27 @@ With `num_gpus > 1`, FastVideo runs Ulysses sequence parallelism: the video toke text tokens are replicated. The number of GPUs must divide the 32 attention heads. Sequence-parallel padding of the video tokens is trimmed before attention, so no attention mask is needed. +Padding is always excluded, without an opt-in setting. At 704x1248 with 313 frames, the 67,782 video tokens need +two padding tokens at SP=4. They are removed before joining the text tokens in every attention block. This provides +the same exclusion as vLLM-Omni's unconditional padding mask in +[revision `e0c4c731`](https://github.com/vllm-project/vllm-omni/commit/e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74). +The trimmed sequence can use either FlashAttention or Torch SDPA. + +```bash +PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py +``` + +## Memory and offloading + +The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards +video tokens; each rank still holds the DiT weights. Increasing `num_gpus` alone does not shard those weights. +For devices that cannot hold the DiT, see [CPU offloading](offloading.md) for `dit_cpu_offload` and +`dit_layerwise_offload`. Hardware recipes and timings reported by vLLM-Omni use its own offload implementation; +they have not been validated for FastVideo. + +The full `IFM/PAN2` checkpoint is not public yet. Local parity tests use the public random-weight +`wuqing157/tiny-pan2-modular-pipe` checkpoint; it verifies the integration and produces noise rather than useful video. + ## Differences from the vLLM-Omni integration - The optional prompt and video guardrails (`pan2-guardrail`) are not included. diff --git a/examples/inference/basic/basic_pan2.py b/examples/inference/basic/basic_pan2.py index 8dd05e0c9..10600ca96 100644 --- a/examples/inference/basic/basic_pan2.py +++ b/examples/inference/basic/basic_pan2.py @@ -22,12 +22,12 @@ def main(): num_gpus=int(os.environ.get("PAN2_NUM_GPUS", "1")), use_fsdp_inference=False, dit_cpu_offload=False, + dit_layerwise_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, ) - # Defaults come from the pan2 preset: 704x1248, 313 frames (13 s at 24 fps), 50 steps, guidance 3.0. generator.generate_video( PROMPT, image_path=IMAGE_PATH, diff --git a/tests/local_tests/pan2/PORT_STATUS.md b/tests/local_tests/pan2/PORT_STATUS.md index 4b03b1ae2..7c5b53a3b 100644 --- a/tests/local_tests/pan2/PORT_STATUS.md +++ b/tests/local_tests/pan2/PORT_STATUS.md @@ -3,7 +3,7 @@ ## Summary - model_family: pan2 - workload_types: T2V, I2V (one pipeline; I2V when the request carries a first frame) -- official_ref: vllm-project/vllm-omni#8602 at `59dd8c9a` (the diffusers integration is private) +- official_ref: vllm-project/vllm-omni#8602 at `e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74` (the diffusers integration is private) - official_ref_dir: none; rebuilt in `tests/local_tests/pan2/reference.py` - hf_weights_path: `IFM/PAN2` (not public); tests use `wuqing157/tiny-pan2-modular-pipe@3ff3130b` - local_weights_dir: HF cache, or `PAN2_TINY_MODEL_DIR` @@ -14,7 +14,7 @@ - phase: 11 (post-parity review and handoff) - status: in_progress (GPU and real-weight validation pending, see I001 and I002) - owner: orchestrator -- last_updated: 2026-10-15 +- last_updated: 2026-10-20 ## Component Matrix | Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | @@ -36,24 +36,24 @@ ## Parity Commands | Scope | Command | Last Result | Notes | |---|---|---|---| -| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 12 passed (CPU) | includes SP=2 vs one rank over gloo | +| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 13 passed (CPU) | SP=2/4 vs one rank over gloo, T2V/I2V, one/two pad tokens; max abs 7.7e-6 | | Text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | 5 passed (CPU) | exact | -| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | passed (CPU) | exact | -| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | passed (CPU) | exact | +| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | 9 passed (CPU) | exact | +| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | 9 passed (CPU) | exact | | Pipeline | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | 3 passed (CPU) | T2V and I2V latents exact | ## Open Questions | ID | Question | Owner | Needed By Phase | Status | Resolution | |---|---|---|---|---|---| | Q001 | Which reference, with the diffusers integration private? | orchestrator | 2 | resolved | vLLM-Omni #8602, rebuilt from diffusers/transformers blocks in `reference.py` | -| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: the package is not public and FastVideo has no guardrail stage. | +| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: FastVideo has no guardrail stage. | | Q003 | Port Cache-DiT `quality` and CFG-parallel? | orchestrator | 7 | resolved | Not ported: vLLM-Omni serving features. Ulysses SP is supported through `DistributedAttention`. | ## Issues And Blockers | ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | |---|---|---|---|---|---|---|---|---| | I001 | 9-10 | all | medium | No real-weight parity or SSIM regression | `IFM/PAN2` is not public | orchestrator | open | Add `fastvideo/tests/ssim/test_pan2_similarity.py` once weights are public. | -| I002 | 9 | transformer, text_encoder | medium | No GPU run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | All tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. | +| I002 | 9 | transformer, text_encoder | medium | No GPU parity run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | Local parity tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. | | I003 | 7 | pipeline | low | The shared denoiser casts latents to the model-output dtype in the scheduler step; PAN2 samples on fp32 latents | pipeline parity fails (8e-2) without the fix | orchestrator | resolved | `DenoisingStage.prepare_scheduler_model_output` hook; `PAN2DenoisingStage` returns `noise_pred.float()`. | ## Escape Hatches diff --git a/tests/local_tests/pan2/README.md b/tests/local_tests/pan2/README.md index b521313d4..2fac8e256 100644 --- a/tests/local_tests/pan2/README.md +++ b/tests/local_tests/pan2/README.md @@ -7,7 +7,7 @@ not part of the default CI suite. PAN2's diffusers integration (`PAN2Transformer3DModel`, `PAN2ModularPipeline`) is not public yet, and the full `IFM/PAN2` weights are not released. The reference is the vLLM-Omni port, -[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `59dd8c9a`), which reports +[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `e0c4c731`), which reports parity with the diffusers implementation. `reference.py` rebuilds it from diffusers and transformers building blocks: the DiT from `AdaLayerNormZero`, `AdaLayerNormContinuous`, `FeedForward` (gelu-approximate and swiglu), `Timesteps`, `TimestepEmbedding` and the diffusers RoPE helpers, using the checkpoint's module names so the checkpoint loads @@ -36,7 +36,7 @@ when no GPU is present. On a GPU-less host, `import fastvideo` also needs the `f |---|---|---| | DiT: tiny checkpoint, strict load + forward | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k tiny_checkpoint` | max abs 6.0e-7 | | DiT: random weights, batch 1 and 2 | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k random_weights` | max abs 0 and 1.9e-5 (outputs ~16) | -| DiT: Ulysses SP=2 vs one rank, padded tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sp2` | max abs 5.7e-6 (gloo on CPU; NCCL on 2 GPUs) | +| DiT: Ulysses SP=2/4 vs one rank, T2V/I2V with 15/18 tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sequence_parallel` | max abs 4.8e-6 / 7.7e-6 (gloo on CPU; NCCL on 2/4 GPUs) | | Qwen3.5 text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | all hidden states exact (fp32, bf16, padded batch) | | Wan2.2 VAE encode/decode + normalization | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | exact; normalization 2.4e-7 | | Flow-match scheduler, PAN2 sigmas + shift | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | exact | @@ -48,6 +48,19 @@ confirm the threshold catches real deviations, three injected changes were each `hidden_states[-2]`, letting the latents fall to bf16 between steps, and shifting the first-frame crop by one pixel. Each moved the latents by about 8e-2 relative L2. +Upstream revision `e0c4c731` always masks video padding, fixing Ulysses=4 drift when padding was left unmasked. +FastVideo already excludes this padding by passing the original video length to `DistributedAttention`, which trims +it before RoPE and joint attention, then pads the output for the reverse all-to-all. The distributed test covers one +pad token (15 tokens at SP=2/4), two pad tokens (18 at SP=4, matching 704x1248x313 modulo 4), and an unpadded +control (18 at SP=2), with both zero T2V and first-frame I2V conditioning. No masking toggle is required. +The pipeline math and checkpoint layout are unchanged between upstream revisions `59dd8c9a` and `e0c4c731`. + +On 2026-10-20, all 39 PAN2 local tests passed on CPU (Torch 2.12.0, Transformers 5.15.0, Diffusers 0.39.0). +The broader API/contract suite had 499 passes and two failures also reproduced on clean `main` at `cdda1b753`: +the schema inventory lacks MiniMax-H3's `audio_scheduler_shift`/`video_scheduler_shift`, and +`test_load_kandinsky5.py` reads `FASTVIDEO_ATTENTION_BACKEND` outside the environment registry. +Importing the PAN2 tests and running their fixtures preserves both an unset and a pre-existing attention backend. + ## Not covered - Real-weight parity and an SSIM regression: `IFM/PAN2` is not public. diff --git a/tests/local_tests/pan2/reference.py b/tests/local_tests/pan2/reference.py index 85c47b710..e723eead5 100644 --- a/tests/local_tests/pan2/reference.py +++ b/tests/local_tests/pan2/reference.py @@ -6,7 +6,8 @@ assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZ ``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``, ``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``, ``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was -validated against the diffusers implementation. Module names are the checkpoint's, so checkpoints load strictly. +validated against the diffusers implementation (upstream revision ``e0c4c731``). Module names are the checkpoint's, +so checkpoints load strictly. Its unconditional SP padding mask is equivalent to unpadded reference attention. """ from __future__ import annotations diff --git a/tests/local_tests/pipelines/test_pan2_pipeline_parity.py b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py index 917bffa62..9f366b69c 100644 --- a/tests/local_tests/pipelines/test_pan2_pipeline_parity.py +++ b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py @@ -12,16 +12,15 @@ Runs on CPU when no GPU is present. Set ``PAN2_TINY_MODEL_DIR`` to a local copy from __future__ import annotations -import os import socket +from contextlib import ExitStack import PIL.Image import pytest import torch -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir # noqa: E402 +from fastvideo import envs +from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir HEIGHT, WIDTH, NUM_FRAMES = 64, 96, 9 NUM_INFERENCE_STEPS = 3 @@ -51,37 +50,35 @@ def pan2_runtime(): from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines import build_pipeline - monkeypatch = pytest.MonkeyPatch() - if not torch.cuda.is_available(): - # On a GPU-less Linux host fastvideo resolves the local device to mps; run on CPU instead. Swapping the code - # object reaches every module that imported the function by name. - def _cpu_device() -> torch.device: - return torch.device("cpu") + with ExitStack() as stack: + stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + monkeypatch = stack.enter_context(pytest.MonkeyPatch.context()) + if not torch.cuda.is_available(): + # Reach modules that imported the device helper by name on a GPU-less Linux host. + def _cpu_device() -> torch.device: + return torch.device("cpu") - monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__) - for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), - ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): - monkeypatch.setenv(name, os.environ.get(name, value)) - maybe_init_distributed_environment_and_model_parallel(1, 1) + monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__) + for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), + ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): + stack.enter_context(envs.override_external(name, value)) + stack.callback(cleanup_dist_env_and_memory) + maybe_init_distributed_environment_and_model_parallel(1, 1) - pipeline_config = PAN2PipelineConfig(dit_precision="bf16", - vae_precision="fp32", - text_encoder_precisions=("fp32", )) - args = FastVideoArgs(model_path=str(model_dir), - num_gpus=1, - pipeline_config=pipeline_config, - dit_cpu_offload=False, - text_encoder_cpu_offload=False, - vae_cpu_offload=False, - use_fsdp_inference=False) - pipeline = build_pipeline(args) - device = parallel_state.get_local_torch_device() - try: + pipeline_config = PAN2PipelineConfig(dit_precision="bf16", + vae_precision="fp32", + text_encoder_precisions=("fp32", )) + args = FastVideoArgs(model_path=str(model_dir), + num_gpus=1, + pipeline_config=pipeline_config, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + use_fsdp_inference=False) + pipeline = build_pipeline(args) + device = parallel_state.get_local_torch_device() yield pipeline, args, model_dir, device - finally: - del pipeline - cleanup_dist_env_and_memory() - monkeypatch.undo() def _fastvideo_generate(pipeline, args, image, guidance_scale: float = GUIDANCE_SCALE): diff --git a/tests/local_tests/transformers/test_pan2_transformer_parity.py b/tests/local_tests/transformers/test_pan2_transformer_parity.py index 600971403..5385aaa28 100644 --- a/tests/local_tests/transformers/test_pan2_transformer_parity.py +++ b/tests/local_tests/transformers/test_pan2_transformer_parity.py @@ -17,6 +17,7 @@ import os import socket import subprocess import sys +from contextlib import ExitStack from dataclasses import fields from pathlib import Path @@ -25,16 +26,14 @@ import torch import torch.nn as nn from torch.testing import assert_close -os.environ.setdefault("DISABLE_SP", "1") -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - -from fastvideo.forward_context import set_forward_context # noqa: E402 -from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict # noqa: E402 -from fastvideo.models.loader.utils import get_param_names_mapping # noqa: E402 -from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir # noqa: E402 +from fastvideo import envs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict +from fastvideo.models.loader.utils import get_param_names_mapping +from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir REPO_ROOT = Path(__file__).resolve().parents[3] -SP_WORLD_SIZE = 2 +SP_WORLD_SIZES = (2, 4) # -------------------------------------------------------------------------------------------------------------------- @@ -69,17 +68,25 @@ def _small_kwargs() -> dict: @pytest.fixture -def single_rank_dist(): +def single_rank_dist(monkeypatch): from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + import fastvideo.distributed.parallel_state as parallel_state - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", str(_free_port())) - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", "0") - maybe_init_distributed_environment_and_model_parallel(1, 1) - yield - cleanup_dist_env_and_memory() + if not torch.cuda.is_available(): + def _cpu_device() -> torch.device: + return torch.device("cpu") + + monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__) + with ExitStack() as stack: + stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), + ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): + stack.enter_context(envs.override_external(name, value)) + try: + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + finally: + cleanup_dist_env_and_memory() def _build_native(arch_kwargs: dict): @@ -218,17 +225,16 @@ def test_pan2_random_weights_forward_parity(single_rank_dist, batch_size: int): # -------------------------------------------------------------------------------------------------------------------- -# Sequence parallel: SP=2 must match one rank (video tokens are padded to a multiple of the SP size) +# Sequence parallel: padded video tokens must never enter joint video/text attention. # -------------------------------------------------------------------------------------------------------------------- -def _run_sp_worker(mode: str, output_path: Path) -> None: +def _run_sp_worker(sp_size: int, output_path: Path) -> None: from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel from fastvideo.models.dits.pan2 import PAN2Transformer3DModel import torch.distributed as dist - sp_size = 1 if mode == "single" else SP_WORLD_SIZE if torch.cuda.is_available(): device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") torch.cuda.set_device(device) @@ -254,58 +260,76 @@ def _run_sp_worker(mode: str, output_path: Path) -> None: native = PAN2Transformer3DModel(config=config, hf_config=kwargs).float().eval() _load_into_native(native, reference.state_dict().items()) native = native.to(device) - # 5x2x6 latents patchify to 15 video tokens, which SP=2 pads to 16. - inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"]) - inputs["hidden_states"] = torch.randn(2, kwargs["in_channels"], 5, 2, 6, generator=torch.Generator().manual_seed(5)) - inputs = {name: value.to(device) for name, value in inputs.items()} - with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): - output = native(**inputs) + outputs = {} + # 15 tokens need one pad at SP=2/4; 18 need two at SP=4, as at 704x1248x313. + for latent_frames in (5, 6): + for image_conditioned in (False, True): + inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"]) + video = torch.randn(2, kwargs["in_channels"], latent_frames, 2, 6, + generator=torch.Generator().manual_seed(5)) + video[:, kwargs["out_channels"]:] = 0 + if image_conditioned: + video[:, kwargs["out_channels"]:-1, :1] = torch.randn( + 2, kwargs["out_channels"], 1, 2, 6, generator=torch.Generator().manual_seed(6)) + video[:, -1:, :1] = 1 + inputs["hidden_states"] = video + inputs = {name: value.to(device) for name, value in inputs.items()} + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + output = native(**inputs) + outputs[f"{latent_frames}_{image_conditioned}"] = output.detach().cpu() if int(os.environ.get("RANK", "0")) == 0: - torch.save({"output": output.detach().cpu()}, output_path) + torch.save(outputs, output_path) dist.barrier() finally: cleanup_dist_env_and_memory() -def _run_torchrun(mode: str, nproc_per_node: int, output_path: Path) -> None: +def _run_torchrun(nproc_per_node: int, output_path: Path) -> None: command = [ sys.executable, "-m", "torch.distributed.run", "--nnodes", "1", "--nproc_per_node", str(nproc_per_node), "--master_port", str(_free_port()), - str(Path(__file__).resolve()), "--pan2-sp-worker", "--mode", mode, "--output", + str(Path(__file__).resolve()), "--pan2-sp-worker", "--sp-size", str(nproc_per_node), "--output", str(output_path) ] environment = os.environ.copy() environment["PYTHONPATH"] = os.pathsep.join(filter(None, [str(REPO_ROOT), environment.get("PYTHONPATH")])) - environment["DISABLE_SP"] = "0" environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" - process = subprocess.run(command, capture_output=True, text=True, env=environment) + process = subprocess.run(command, capture_output=True, text=True, env=environment, timeout=180) if process.returncode != 0: - raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + raise RuntimeError(f"SP={nproc_per_node} worker failed with code {process.returncode}\n" f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}") -def test_pan2_sp2_matches_single_rank(tmp_path: Path): - """Ulysses SP=2 on an odd video token count (padded to the SP size) must match one rank. Runs on CPU with gloo - when no GPU is present.""" - if torch.cuda.is_available() and torch.cuda.device_count() < SP_WORLD_SIZE: - pytest.skip(f"PAN2 SP parity requires {SP_WORLD_SIZE} CUDA devices.") - single_path, sp_path = tmp_path / "single.pt", tmp_path / "sp2.pt" - _run_torchrun("single", 1, single_path) - _run_torchrun("sp", SP_WORLD_SIZE, sp_path) - single = torch.load(single_path, map_location="cpu", weights_only=True)["output"] - sp = torch.load(sp_path, map_location="cpu", weights_only=True)["output"] - diff = (sp - single).abs() - print(f"sp2-vs-single diff_max={diff.max().item():.3e}") - assert_close(sp, single, atol=1e-5, rtol=1e-5) +@pytest.fixture(scope="module") +def single_rank_sp_outputs(tmp_path_factory): + path = tmp_path_factory.mktemp("pan2-single") / "single.pt" + _run_torchrun(1, path) + return torch.load(path, map_location="cpu", weights_only=True) + + +@pytest.mark.parametrize("sp_size", SP_WORLD_SIZES) +def test_pan2_sequence_parallel_matches_single_rank(tmp_path: Path, single_rank_sp_outputs, sp_size: int): + """T2V and I2V joint attention excludes padding at Ulysses SP=2/4. Uses gloo on CPU or NCCL on 2/4 GPUs.""" + if torch.cuda.is_available() and torch.cuda.device_count() < sp_size: + pytest.skip(f"PAN2 SP parity requires {sp_size} CUDA devices.") + sp_path = tmp_path / f"sp{sp_size}.pt" + _run_torchrun(sp_size, sp_path) + sp_outputs = torch.load(sp_path, map_location="cpu", weights_only=True) + assert sp_outputs.keys() == single_rank_sp_outputs.keys() + for case, single in single_rank_sp_outputs.items(): + sp = sp_outputs[case] + diff = (sp - single).abs() + print(f"sp{sp_size}-vs-single {case} diff_max={diff.max().item():.3e}") + assert_close(sp, single, atol=1e-5, rtol=1e-5) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--pan2-sp-worker", action="store_true") - parser.add_argument("--mode", choices=["single", "sp"], required=True) + parser.add_argument("--sp-size", type=int, choices=[1, *SP_WORLD_SIZES], required=True) parser.add_argument("--output", type=str, required=True) args = parser.parse_args() if not args.pan2_sp_worker: raise SystemExit("This module is intended to be run by pytest.") - _run_sp_worker(args.mode, Path(args.output)) + _run_sp_worker(args.sp_size, Path(args.output))