# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import argparse import os import socket import subprocess from pathlib import Path import pytest import torch import torch.distributed as dist from fastvideo.configs.models.dits.wanvideo import ( WanVideoArchConfig, WanVideoConfig, ) from fastvideo.distributed import ( cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel, ) from fastvideo.forward_context import set_forward_context from fastvideo.models.dits.wanvideo import WanTransformer3DModel from fastvideo.pipelines.pipeline_batch_info import ForwardBatch SP_WORLD_SIZE = 2 SEED = 2026 def _free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) def _seed_everything(seed: int) -> None: torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False def _build_tiny_wan_config() -> WanVideoConfig: arch_config = WanVideoArchConfig( num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4, text_dim=16, freq_dim=16, ffn_dim=64, num_layers=2, patch_size=(1, 2, 2), rope_max_seq_len=64, ) return WanVideoConfig(arch_config=arch_config) def _build_hf_config(config: WanVideoConfig) -> dict[str, object]: return { "num_attention_heads": config.num_attention_heads, "attention_head_dim": config.attention_head_dim, "in_channels": config.in_channels, "out_channels": config.out_channels, "text_dim": config.text_dim, "freq_dim": config.freq_dim, "ffn_dim": config.ffn_dim, "num_layers": config.num_layers, "patch_size": config.patch_size, "text_len": config.text_len, "rope_max_seq_len": config.rope_max_seq_len, } def _build_inputs( config: WanVideoConfig, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: generator = torch.Generator(device="cpu") generator.manual_seed(SEED + 1) hidden_states = torch.randn( 1, config.in_channels, 4, 8, 8, generator=generator, dtype=torch.float32, ).to(device) encoder_hidden_states = torch.randn( 1, 8, config.text_dim, generator=generator, dtype=torch.float32, ).to(device) timestep = torch.tensor([10], device=device, dtype=torch.long) return hidden_states, encoder_hidden_states, timestep def _collect_named_parameter_grads(model: torch.nn.Module, ) -> dict[str, torch.Tensor]: named_grads: dict[str, torch.Tensor] = {} for name, param in model.named_parameters(): grad = param.grad if grad is None: grad = torch.zeros_like(param) named_grads[name] = grad.detach().float().cpu() return named_grads def _initialize_model_parameters(model: torch.nn.Module) -> None: # ReplicatedLinear parameters are allocated with torch.empty and need an # explicit initialization in tests to avoid undefined values. torch.manual_seed(SEED + 3) torch.cuda.manual_seed_all(SEED + 3) with torch.no_grad(): for name, param in model.named_parameters(): if param.ndim <= 1: if name.endswith("weight") and "norm" in name: param.fill_(1.0) else: param.zero_() continue torch.nn.init.xavier_uniform_(param) def _assert_finite(name: str, tensor: torch.Tensor) -> None: if torch.isfinite(tensor).all(): return nan_count = int(torch.isnan(tensor).sum().item()) inf_count = int(torch.isinf(tensor).sum().item()) raise RuntimeError(f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})") def _run_worker(mode: str, output_path: Path) -> None: if mode not in {"single", "sp"}: raise ValueError(f"Unsupported mode: {mode}") local_rank = int(os.environ.get("LOCAL_RANK", "0")) rank = int(os.environ.get("RANK", "0")) sp_size = 1 if mode == "single" else SP_WORLD_SIZE device = torch.device(f"cuda:{local_rank}") torch.cuda.set_device(device) _seed_everything(SEED) try: maybe_init_distributed_environment_and_model_parallel(1, sp_size) config = _build_tiny_wan_config() hf_config = _build_hf_config(config) model = WanTransformer3DModel(config=config, hf_config=hf_config) model = model.to(device=device, dtype=torch.float32) _initialize_model_parameters(model) model.train() model.zero_grad(set_to_none=True) hidden_states, encoder_hidden_states, timestep = _build_inputs( config, device, ) forward_batch = ForwardBatch(data_type="dummy") with set_forward_context( current_timestep=0, attn_metadata=None, forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, timestep=timestep, ) _assert_finite("output", output) loss = output.float().square().mean() _assert_finite("loss", loss) loss.backward() for name, param in model.named_parameters(): if param.grad is not None: _assert_finite(f"grad[{name}]", param.grad) if mode == "sp": for param in model.parameters(): if param.grad is not None: dist.all_reduce(param.grad, op=dist.ReduceOp.AVG) grads = _collect_named_parameter_grads(model) if rank == 0: torch.save({"grads": grads}, output_path) dist.barrier() finally: cleanup_dist_env_and_memory() def _run_torchrun( script_path: Path, mode: str, nproc_per_node: int, output_path: Path, ) -> None: cmd = [ "torchrun", "--nnodes", "1", "--nproc_per_node", str(nproc_per_node), "--master_port", str(_free_port()), str(script_path), "--sp-grad-worker", "--mode", mode, "--output", str(output_path), ] env = os.environ.copy() env["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.returncode != 0: raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" f"STDOUT:\n{process.stdout}\n" f"STDERR:\n{process.stderr}") def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: if not torch.cuda.is_available(): pytest.skip("This test requires CUDA.") if torch.cuda.device_count() < SP_WORLD_SIZE: pytest.skip(f"This test requires at least {SP_WORLD_SIZE} CUDA devices.") script_path = Path(__file__).resolve() single_path = tmp_path / "single_rank_grads.pt" sp_path = tmp_path / f"sp{SP_WORLD_SIZE}_grads.pt" _run_torchrun( script_path=script_path, mode="single", nproc_per_node=1, output_path=single_path, ) _run_torchrun( script_path=script_path, mode="sp", nproc_per_node=SP_WORLD_SIZE, output_path=sp_path, ) single_grads: dict[str, torch.Tensor] = torch.load( single_path, map_location="cpu", )["grads"] sp_grads: dict[str, torch.Tensor] = torch.load( sp_path, map_location="cpu", )["grads"] single_names = set(single_grads) sp_names = set(sp_grads) single_only = sorted(single_names - sp_names) sp_only = sorted(sp_names - single_names) same_params: list[str] = [] different_params: list[str] = [] for name in sorted(single_names & sp_names): single_grad = single_grads[name] sp_grad = sp_grads[name] if torch.allclose(single_grad, sp_grad, rtol=1e-4, atol=1e-5): same_params.append(name) continue max_abs_diff = (single_grad - sp_grad).abs().max().item() different_params.append(f"{name} (max_abs_diff={max_abs_diff:.3e})") report = (f"single-only params ({len(single_only)}):\n" + "\n".join(single_only) + f"\n\nsp-only params ({len(sp_only)}):\n" + "\n".join(sp_only) + "\n\n" f"same params ({len(same_params)}):\n" + "\n".join(same_params) + f"\n\ndifferent params ({len(different_params)}):\n" + "\n".join(different_params)) print(report) assert not single_only and not sp_only and not different_params, report def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--sp-grad-worker", action="store_true") parser.add_argument("--mode", choices=["single", "sp"], default=None) parser.add_argument("--output", type=str, default=None) return parser.parse_args() # pytest -sv fastvideo/tests/distributed/test_sp_wan.py if __name__ == "__main__": args = _parse_args() if not args.sp_grad_worker: raise SystemExit("This module is intended to be run by pytest.") if args.mode is None or args.output is None: raise SystemExit("--mode and --output are required in worker mode.") _run_worker(mode=args.mode, output_path=Path(args.output))