diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..91bf522 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,53 @@ +name: Test + +on: + pull_request: + push: + branches: [main] + +permissions: + contents: read + +jobs: + minimum-dependencies: + runs-on: ubuntu-latest + timeout-minutes: 20 + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: Install minimum supported dependencies + run: | + python -m pip install \ + "numpy<2" \ + "torch==2.2.*" \ + "diffusers==0.30.3" \ + "huggingface-hub<0.26" \ + pytest \ + packaging + python -m pip install --no-deps -e . + - name: Test minimum dependency boundary + run: | + python -m pytest -vv --durations=20 \ + test/test_adapter_compatibility.py \ + test/test_public_vae_api.py \ + test/test_decoderadapter.py \ + test/test_encoderadapter.py + + latest-dependencies: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install latest dependencies + run: | + python -m pip install -e . + python -m pip install pytest + - name: Run test suite + run: python -m pytest -q diff --git a/README.md b/README.md index 94b2569..d249e1c 100644 --- a/README.md +++ b/README.md @@ -1,60 +1,227 @@ -# DistVAE: A patch parallelism distributed VAE implement for high resolution generation +# DistVAE -By providing a set of adapter interfaces, this project allows users to quickly convert vae-related implementations in the diffusers library into parallel versions on multiple gpu's, enabling non-intrusive parallelisation of the vae portion of an existing model, thus reducing the memory footprint of the image generation process, and avoiding vae-induced memory spikes. +DistVAE replaces supported diffusers VAE encoders and decoders with distributed adapters. The rest +of the diffusion pipeline stays unchanged. ## Installation -``` bash +```bash pip install distvae ``` -## Usage +Python 3.10 or newer, with `torch>=2.2` and `diffusers>=0.30.3`. Individual VAE families may require +a newer Diffusers release. -Refering to the file in `test/` directory. In general, you only need to use the corresponding adapter for the diffusers module to make it work on multiple gpu in parallel. +The pipeline quickstart also needs Transformers: -As an example, we can transform an initialised vae decoder into a parallel versions: +```bash +pip install "distvae[pipeline]" +``` +## Quickstart -``` python -from diffusers.models.autoencoders.vae import Decoder -from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter +Every rank builds the same pipeline, and DistVAE shards the VAE inside it. Save this as `decode.py`: + +```python +import os import torch -import random import torch.distributed as dist +from diffusers import DiffusionPipeline +from distvae import vae as vae_api + +dist.init_process_group(backend="nccl") +device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") +torch.cuda.set_device(device) + +# The group the VAE is split over. Every rank that enters the VAE call must be a +# member. If you create a subgroup, gate the pipeline call to those ranks too. +vae_group = dist.group.WORLD + +pipe = DiffusionPipeline.from_pretrained( + os.environ["MODEL_ID"], torch_dtype=torch.bfloat16 +).to(device) -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - -@torch.no_grad() -def main(): - # init - set_seed() - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - torch.device('cuda', rank) - - # input - hidden_state = torch.randn(1, 4, 128, 128, device=f"cuda:{rank}") - # create vae.decoder instance - decoder = Decoder( - in_channels=4, out_channels=3, - up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], - block_out_channels=(128, 256, 512, 512), layers_per_block=2, - norm_num_groups=32, act_fn="silu", - ).to(f"cuda:{rank}") - # transform vae.decoder to distvae.decoder - patch_decoder = DecoderAdapter(decoder).to(f"cuda:{rank}") - # forward - result = decoder(hidden_state) - patch_result = patch_decoder(hidden_state) - - print("result shape: ", patch_result.shape) - if rank == 0: - assert torch.allclose(result, patch_result, atol=1e-2), "two hidden states are not equal" - -if __name__ == "__main__": - main() +vae_api.parallelize_decoder(pipe.vae, vae_group) +vae_api.parallelize_encoder(pipe.vae, vae_group) + +image = pipe("A cat holding a sign that says hello world", height=1024, width=1024).images[0] +if dist.get_rank() == 0: + image.save("out.png") +``` + +Then launch it across your GPUs with any pipeline whose VAE DistVAE supports. For example, with a +recent Diffusers release: + +```bash +MODEL_ID=black-forest-labs/FLUX.2-dev torchrun --nproc_per_node=4 decode.py ``` + +Both calls raise if there is no adapter for the VAE, so an unsupported model fails at setup rather +than part way through a decode. + +## Supported VAEs + +Every family below supports both row sharding and tiling. Qwen-Image is listed with the video VAEs +because its Wan-derived autoencoder has a frame axis. + +| VAE | Frame axis | Tiles by | A tile is | +| ---------------- | ---------- | ----------------------- | -------------------------------------------------- | +| `AutoencoderKL` | no | overlap-derived strides | one decoder call | +| Flux.2 | no | overlap-derived strides | one decoder call | +| HunyuanVideo 1.5 | yes | overlap-derived strides | one decoder call | +| HunyuanVideo | yes | a stored stride | one decoder call per temporal chunk | +| LTX-2 | yes | a stored stride | one decoder call unless temporal tiling is enabled | +| Wan | yes | a stored stride | a call per frame, threading a causal cache | +| Qwen-Image | yes | a stored stride | a call per frame, threading a causal cache | + +Tile size affects the families differently. A smaller tile reduces tile-local activation memory when +one tile is one decoder call, but allocations outside the spatial tile can determine the measured +peak. Wan and Qwen-Image decode one frame at a time, so their peak memory is often set by temporal +state. + +`tile_overlap_plan` accepts exact output-pixel `(height, width)` values and maps them to each VAE's +stride settings. DistVAE owns the tiling loop for every family in the table. CogVideoX is excluded +because it tiles frames inside the spatial loop, so its spatial tiles are not independent. + +## Distributed decode strategies + +DistVAE provides two distributed decode strategies: + +- **Row sharding** gives each rank a band in every adapted layer. It exchanges convolution halos and + normalization statistics, preserves the unsharded result within numerical tolerance, and usually + reduces activation memory as ranks are added. +- **Whole-tile distribution** gives each rank complete windows. Ranks exchange tile-edge data and + gather decoded pieces for assembly. Peak activation memory usually follows the tile window, + including on one GPU, while overlap repeats work and tile-local normalization can change the + output. + +The figure compares the two distributed paths at two tile sizes. Each row reports peak activations, +decoded work, seams, load imbalance, and synchronization: + +![Row sharding and two whole-tile distributions for a 1024 by 1024 image on four GPUs, compared by peak activations, work, seams, load imbalance, and synchronization](docs/figure.png) + +[Choosing a decode path](docs/strategies.md) explains how VAE family, input shape, rank count, and +interconnect affect the choice. The [benchmark guide](bench/README.md) shows how to measure both +strategies against a vanilla unsharded Diffusers decode. + +## Usage + +The quickstart uses `distvae.vae`, which picks the adapter for a whole VAE. To shard a single +diffusers module instead, wrap it in its adapter: + +```python +import os + +import torch +import torch.distributed as dist +from diffusers.models.autoencoders.vae import Decoder +from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter + +dist.init_process_group(backend="nccl") +local_rank = int(os.environ["LOCAL_RANK"]) +device = torch.device(f"cuda:{local_rank}") +torch.cuda.set_device(device) +torch.manual_seed(42) # every rank must build the same weights and the same input + +decoder = Decoder( + in_channels=4, out_channels=3, + up_block_types=["UpDecoderBlock2D"] * 4, + block_out_channels=(128, 256, 512, 512), layers_per_block=2, + norm_num_groups=32, act_fn="silu", +).to(device) + +hidden_state = torch.randn(1, 4, 128, 128, device=device) +with torch.no_grad(): + expected = decoder(hidden_state) + +# The adapter takes ownership of decoder and replaces its distributed layers in +# place. Do not use decoder as an unmodified reference after this call. +patch_decoder = DecoderAdapter(decoder, dist.group.WORLD).to(device) +with torch.no_grad(): + assert torch.allclose(expected, patch_decoder(hidden_state), atol=1e-2) +``` + +There are more runnable examples in `test/`. + +### Tiling + +Diffusers decides whether to tile. DistVAE resizes the window and distributes the tiles across the +group: + +```python +from distvae import vae as vae_api + +vae_api.require_vae_support(pipe.vae, "tiling", "enable_tiling()") +pipe.vae.enable_tiling() + +# Optional: ask for an exact 192x192px window. Invalid shapes are refused rather +# than silently changed. +plan = vae_api.tile_shape_plan(pipe.vae, 192, 192) +if plan is None: + raise ValueError("this VAE cannot use a 192x192px tile shape") +vae_api.apply_tile_plan(pipe.vae, plan) + +# Optional: overlap neighbouring tiles by 32 output pixels vertically and 64 +# horizontally. This reads the window now set on the VAE, so apply it second. +step = vae_api.tile_overlap_plan( + pipe.vae, 32, 64, sample_shape=(1024, 1024) +) +if step is None: + raise ValueError("this VAE cannot use a 32x64px tile overlap") +vae_api.apply_tile_plan(pipe.vae, step) +replacement = vae_api.tiled_decode_for(pipe.vae) +if replacement is not None: + pipe.vae.tiled_decode = replacement + +# Decode the tiles across the group instead of one after another. +if not vae_api.supports_tile_parallel(pipe.vae): + raise ValueError("this VAE does not support distributed tiled decode") +dispatch, assemble = vae_api.sharing(vae_group) +tiled_decode = vae_api.tiled_decode_for(pipe.vae, dispatch, assemble) +if tiled_decode is None: + raise ValueError("no distributed tiled decode is available for this VAE") +pipe.vae.tiled_decode = tiled_decode +``` + +Window and overlap are separate controls in output pixels. The window sets the memory required for +one tile. The overlap reduces the stride and increases repeated work. + +Both planners return `None` when a request cannot be represented exactly. Apply `tile_shape_plan` +first because `tile_overlap_plan` reads the current tile shape. Requested overlap values are never +rounded. + +[Choosing a tile window](docs/tiling.md) explains rectangular windows, clipped edge tiles, and +overlap. + +### xDiT integration + +xDiT chooses the tile settings and calls the DistVAE planners. Supply `vae_tile_overlap_height` and +`vae_tile_overlap_width` together in output pixels. Use zero on an axis that is not tiled. +Installing new shape or overlap settings replaces the previous tiled decode callable. + +## Performance + +Latency and memory depend on the VAE family, input shape, rank count, device, and interconnect. The +benchmark chooses up to three rectangular plans and records their work, memory estimate, and load +imbalance before running them. See `bench/README.md` for the suite and its limits. + +## Development + +```bash +git clone https://github.com/xdit-project/DistVAE +cd DistVAE +pip install -e ".[dev]" +mdformat --extensions gfm --wrap 100 README.md bench/README.md docs/*.md +pytest +``` + +Tests marked `gloo` spawn several ranks over gloo and need no accelerator, so `pytest -m gloo` +exercises the distributed paths on a CPU-only machine. + +`docs/make_figure.py` regenerates `docs/figure.svg` and, when `cairosvg` is installed, +`docs/figure.png`. + +## License + +MIT. See `LICENSE.txt`. diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 0000000..52b9920 --- /dev/null +++ b/bench/README.md @@ -0,0 +1,212 @@ +# Benchmarking DistVAE + +`distvae_bench.py` measures a real diffusers VAE architecture without downloading a checkpoint. +Every cell rebuilds the architecture with seed 0 and creates its input with seed 1. The weights are +synthetic; layer shapes, memory use, collectives, and scheduling are real. + +Copy `bench/` to the target machine, install the DistVAE revision under test, and run the launcher +with `torchrun`. + +## Requirements + +- PyTorch with a working `torch.distributed` CUDA or ROCm build +- `diffusers` +- DistVAE installed from the revision being measured + +The report records package versions, the DistVAE revision, a source digest, and accelerator details +under `provenance.device`. Compare results only when this context is available. + +`HW_FAMILY` adds your own label alongside it, for naming a fleet or a node type: + +```bash +HW_FAMILY=mi355 torchrun --nproc_per_node=8 bench/distvae_bench.py ... +``` + +The label is null when the variable is absent; the measured device is recorded either way. + +## Reproducing a run elsewhere + +`--matrix` runs the family's canonical shapes rather than one: + +```bash +torchrun --nproc_per_node=4 bench/distvae_bench.py \ + --family wan --half decoder --matrix --out wan-decoder.json +``` + +| family | shapes (height x width x frames) | +| ------------------ | -------------------------------- | +| `flux2` | 1024x1024, 2048x2048 | +| `kl` | 1024x1024, 2048x2048 | +| `qwen_image` | 1024x1024x1, 2048x2048x1 | +| `wan` | 832x480x81, 1280x720x81 | +| `hunyuan_video` | 832x480x129, 1280x720x129 | +| `hunyuan_video_15` | 832x480x129, 1280x720x129 | +| `ltx2` | 1536x1024x121, 1920x1280x121 | + +Canonical shapes are versioned with their architectures in `harness/catalog.py`. `--shape` overrides +the matrix for a one-off run. Video families use their normal frame counts because their temporal +compression ratios differ. Qwen-Image uses one frame. + +LTX-2 uses larger spatial shapes because its 32× compression must still leave at least sixteen +latent units on a tile's narrow axis. Its 1920x1280 case also provides enough tiles for eight ranks. +Use `--shape` to test a different resolution. + +Large unsharded video cases may exceed device memory. The failure is recorded for that case and the +remaining cases continue. + +## Default suite + +This command runs the default decoder suite for one 2048×2048 input: + +```bash +torchrun --nproc_per_node=4 bench/distvae_bench.py \ + --family flux2 --half decoder --height 2048 --width 2048 \ + --out flux2-decoder-2048.json +``` + +The suite compares one vanilla Diffusers baseline with DistVAE's two distributed modes: + +1. baseline `unsharded`: every rank runs the complete untiled VAE half +1. DistVAE `row`: adapted layers process rank-local row bands, exchange halos, and reassemble the + output; stages without adapters remain replicated +1. DistVAE `tile-runs`: ranks run complete spatial tiles without row sharding inside them, then + assemble the output + +Five cases where the sample supports three plans, four where it supports two. + +The plans are named `coarse`, `balanced`, and `fine`. `coarse` has the largest selected window area, +and `fine` has the smallest. `balanced` minimizes the worst normalized window-area, decoded-area, +and rank-imbalance score among the remaining candidates. The names describe geometry; benchmark +results determine which is fastest or has the lowest peak memory on a device. Larger windows usually +create fewer seams and repeat less work. Smaller windows reduce tile-local activation area, though +temporal state and decoder allocations outside that area can still determine the measured peak. The +report's `beats_row_sharding` field shows whether a plan's window area is smaller than one +row-sharded rank's activation area. + +`--diagnostics` adds local tiling for each plan and row sharding inside the finest tile plan. +Applications do not normally use these combinations, and they add substantial runtime. Use them to +separate tile overhead from communication overhead. The local case has no collectives and shows the +minimum measured memory for that window. + +Tiling is decode-only. `--half encoder` runs the vanilla unsharded case and DistVAE's row-sharded +case. + +The planner considers grids with up to four tiles per rank and overlaps down to one quarter of the +window. DistVAE validates each rectangular window and absolute overlap. The planner removes +candidates that are worse in window area, decoded area, rank imbalance, and tile columns. It then +selects the coarsest plan, the finest plan, and a balanced plan between them. An untiled axis uses +zero overlap. The JSON records the objectives, candidate limit, and Pareto frontier size. + +Three constraints limit the search: + +- **Overlap is searched, not pinned.** A tile is a memory win over row sharding only when its window + area is under the `(height / ranks) * width` a rank already holds. Since window is pitch plus + overlap, pinning overlap at the VAE native value floors every window at that value and, on a + 1024x1024 sample at four ranks, made the whole suite memory-neutral by construction. +- **A blend is at least a quarter of its window.** Overlap decides whether a tile's tone drift from + its neighbours reads as a gradient or a band. At 128x1024 on FLUX.2, a 32px blend is clean and a + 16px blend bands. +- **A tile is at least sixteen latent on its narrow axis.** Below that a tile normalizes over + content too unrepresentative of the image, and no blend repairs it. + +Plan selection uses geometry only. Matching family, shape, and world size therefore produce the same +plans on different machines. + +Use `--shape` to request more input shapes explicitly: + +```bash +torchrun --nproc_per_node=4 bench/distvae_bench.py \ + --family wan --half decoder \ + --shape 720x1280x81 --shape 1080x1920x81 \ + --out wan-decoder.json +``` + +Each requested shape gets its own suite. Add only shapes needed for a specific comparison because +VAE runs are expensive. + +## Exact cases + +Repeat `--case` to bypass automatic selection. The untiled cases are the vanilla `unsharded` +baseline and DistVAE's `row` mode. A tiled case uses `MODE:WINDOW_HxW@OVERLAP_HxW`, where `MODE` is +`local`, `tile-runs`, or `row-tiled`. Window and overlap values are output pixels. + +```bash +torchrun --nproc_per_node=4 bench/distvae_bench.py \ + --family kl --half decoder --height 1024 --width 1536 \ + --case unsharded \ + --case row \ + --case 'local:480x736@64x32' \ + --case 'tile-runs:480x736@64x32' \ + --out kl-exact.json +``` + +Exact cases and `--shape` cannot be combined. Run a second command to change both the input and +execution mode. + +## Shape-cost mode and profiling + +`--tile-shape-costs` is decoder-only and separate from the ordinary suite. By default it measures +the three selected rectangular plans. Override them with latent-space windows: + +```bash +torchrun --nproc_per_node=4 bench/distvae_bench.py \ + --family kl --half decoder --height 2048 --width 2048 \ + --tile-shape-costs --tile-shape-windows 88x144,88x88,56x88 \ + --tile-shape-batch 4 --out kl-shape-costs.json +``` + +`--profile`, `--profile-trace`, and `--profile-memory` add one profiler call after timed +measurement. Artifacts go under `--profile-dir`; repeated names receive numeric suffixes. + +## Output and exit status + +`--out` writes schema 7 JSON. One exact case is an object; a suite is an array. Stdout contains +progress and compact human-readable summaries, not a recoverable copy of the JSON. Always supply +`--out` when collecting results from another machine. + +Every record includes versions, provenance, world size, dtype, execution mode, effective tile +settings, latency, peak accelerator memory, communication-operation counts, and agreement with an +unsharded reference when the reference-size limit permits one. Windows and overlaps are +`[height, width]`. + +### Communication counts + +After warmup, the harness runs one VAE invocation with logging enabled, then disables logging before +timed iterations. It wraps `all_reduce`, `all_gather`, `all_gather_into_tensor`, `broadcast`, +`isend`, `irecv`, `recv`, `send`, `barrier`, and `batch_isend_irecv`. + +Calls made by those wrappers from inside PyTorch's `distributed_c10d` module receive a `(batched)` +label. They contribute tensor bytes but do not increase `total_calls`. This includes the sends and +receives inside `batch_isend_irecv` and tensor collectives used internally by `all_gather_object`; +`all_gather_object` itself is not counted as an API call. Timing barriers run while logging is +disabled. + +The reported count is an operation-level comparison between benchmark cases, not a complete count of +every distributed action. Byte totals sum positional tensor buffers visible to the wrappers, +including serialized object buffers passed through an internal tensor collective. They do not +measure network traffic. + +The JSON stores rank 0's `total_calls`, the busiest rank's `total_calls_max`, and +`total_calls_by_rank`. Use `total_calls_max` when comparing cases with uneven per-rank work. + +The process exits nonzero for setup or execution errors and for enforced agreement failures. +Row-sharded numerical agreement is enforced. Numerical differences caused by tiling are measured and +reported but do not control the exit status. Structural failures still fail every mode. + +Run identical family, shape, world-size, dtype, and benchmark digests before comparing machines. + +## Limits + +Synthetic weights do not model activation distributions from a trained checkpoint. The harness does +not measure the diffusion pipeline, host memory, image quality, or visual seam quality. Peak memory +covers the selected VAE half. Use a real model run for end-to-end peak memory and quality decisions. + +## Glossary + +- **adapter:** DistVAE wrapper that gives a diffusers encoder or decoder distributed behavior +- **case:** one input shape and execution mode measured as a record +- **coverage:** decoded tile area divided by image area; overlap raises it above one +- **halo:** neighboring rows exchanged so a sharded convolution has its required context +- **overlap:** output pixels shared and blended between adjacent tiles +- **patchify:** split an activation into rank-local row bands +- **window:** output-pixel height and width decoded by one spatial tile diff --git a/bench/__init__.py b/bench/__init__.py new file mode 100644 index 0000000..d7fd9ba --- /dev/null +++ b/bench/__init__.py @@ -0,0 +1 @@ +"""Importable benchmark utilities for DistVAE.""" diff --git a/bench/distvae_bench.py b/bench/distvae_bench.py new file mode 100644 index 0000000..0634f99 --- /dev/null +++ b/bench/distvae_bench.py @@ -0,0 +1,10 @@ +"""Compatibility launcher for the importable DistVAE benchmark harness.""" + +if __package__: + from .harness.cli import main +else: + from harness.cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/harness/__init__.py b/bench/harness/__init__.py new file mode 100644 index 0000000..0d03cc4 --- /dev/null +++ b/bench/harness/__init__.py @@ -0,0 +1,5 @@ +"""Native DistVAE benchmark harness.""" + +from .cli import main + +__all__ = ["main"] diff --git a/bench/harness/cases.py b/bench/harness/cases.py new file mode 100644 index 0000000..921d5a0 --- /dev/null +++ b/bench/harness/cases.py @@ -0,0 +1,528 @@ +"""Bounded benchmark cases and deterministic rectangular tile-plan selection.""" + +import math + +from distvae import vae as vae_api +from distvae.vae.tile_parallel import shares +from distvae.vae.tiling import _latent_shape + + +# Named for tile count, which is a fact about the plan, rather than for an outcome, which is a +# claim about a device - see select_plans. +PROFILES = ("coarse", "balanced", "fine") +MODES = ("unsharded", "row", "local", "tile-runs", "row-tiled") + +# The smallest latent extent a tile may have on its narrower axis. Below roughly this, a tile +# normalizes over content too unrepresentative of the image and comes out at a different tone +# from its neighbours. The blend then ramps that difference across the overlap rather than +# stepping at the join, so it reads as banding and no seam metric detects it: the join is smooth, +# the tone is wrong. +# +# In latent units rather than pixels, deliberately, because that is what carries across families +# - a scale-16 VAE reaches the same bound at twice the pixel height a scale-8 one does. A +# fraction of the VAE's native window would NOT carry: FLUX.2's native tile is 128 latent and +# Wan's is 16, so one percentage would mean an eight-fold difference in strictness between them. +# A fraction of the sample would be wrong in a different way, making an identical tile legal at +# one canvas size and illegal at another when the tile's own statistics do not depend on the +# canvas it was cut from. +# +# 16 is where two unrelated families agree. Measured on FLUX.2 at 1024x1024 on four ranks, a +# 96px window is 12 latent and bands visibly while a 128px window is 16 and does not, which +# brackets the threshold at (12, 16]; and Wan's own native tile is exactly 16 latent, so raising +# this bound would reject a vendor default. The bracket has not been narrowed further - 13, 14 +# and 15 are untested - so treat 16 as the conservative end of a measurement, not a precise edge. +MIN_TILE_LATENT_EXTENT = 16 + + +def parse_pair(value, label): + """Parse an exact HEIGHTxWIDTH integer pair.""" + parts = value.lower().split("x") + if len(parts) != 2: + raise ValueError(f"{label} must be HEIGHTxWIDTH, not {value!r}") + try: + pair = tuple(int(part) for part in parts) + except ValueError: + raise ValueError(f"{label} must be HEIGHTxWIDTH, not {value!r}") from None + if any(axis < 0 for axis in pair): + raise ValueError(f"{label} axes must be non-negative") + return pair + + +def _cell(name, mode, height, width, frames, window=None, overlap=None, **facts): + sharding = "row" if mode in ("row", "row-tiled") else "unsharded" + distribution = "runs" if mode == "tile-runs" else None + return { + "name": name, + "mode": mode, + "sharding": sharding, + "window": window, + "overlap": overlap, + "tile_distribution": distribution, + "height": height, + "width": width, + "frames": frames, + **facts, + } + + +def parse_case(value, height, width, frames): + """Parse MODE or tiled MODE:HEIGHTxWIDTH@HEIGHTxWIDTH.""" + if value in ("unsharded", "row"): + return _cell(value, value, height, width, frames) + try: + mode, plan = value.split(":", 1) + window_text, overlap_text = plan.split("@", 1) + except ValueError: + raise ValueError( + "case must be unsharded, row, or " + "MODE:WINDOW_HEIGHTxWINDOW_WIDTH@OVERLAP_HEIGHTxOVERLAP_WIDTH" + ) from None + if mode not in ("local", "tile-runs", "row-tiled"): + raise ValueError(f"unknown tiled case mode {mode!r}") + window = parse_pair(window_text, "tile window") + overlap = parse_pair(overlap_text, "tile overlap") + if any(axis <= 0 for axis in window): + raise ValueError("tile window axes must be positive") + if any(overlap_axis >= window_axis for overlap_axis, window_axis in zip(overlap, window)): + raise ValueError("tile overlap must be smaller than its window") + label = f"{mode}-{window[0]}x{window[1]}-ov{overlap[0]}x{overlap[1]}" + return _cell(label, mode, height, width, frames, window, overlap) + + +def cells_from_args(args): + """Return exact requested cases; an empty list requests the default suite.""" + if args.case and args.shape: + raise ValueError("--shape cannot be combined with exact --case values") + return [ + parse_case(value, args.height, args.width, args.frames) + for value in (args.case or ()) + ] + + +def shapes_from_args(args): + """Return the requested sample shapes, most explicit request first. + + `--shape` beats `--matrix` beats the single `--height/--width/--frames`, so asking for one + shape by hand always overrides the family's matrix rather than being appended to it. + """ + if not args.shape: + if getattr(args, "matrix", False): + # Imported here rather than at module scope because catalog builds VAEs and so pulls + # in diffusers; nothing else in this module needs it, and the planner is exercised + # without a model. + from . import catalog + + return list(catalog.matrix_for(args.family)) + return [(args.height, args.width, args.frames)] + shapes = [] + for value in args.shape: + parts = value.lower().split("x") + if len(parts) not in (2, 3): + raise ValueError(f"--shape must be HxW or HxWxFRAMES, not {value!r}") + try: + height, width = (int(axis) for axis in parts[:2]) + frames = int(parts[2]) if len(parts) == 3 else args.frames + except ValueError: + raise ValueError( + f"--shape must be HxW or HxWxFRAMES, not {value!r}" + ) from None + if min(height, width, frames) <= 0: + raise ValueError("--shape axes and frames must be positive") + shapes.append((height, width, frames)) + return shapes + + +def _axis_window(length, overlap, count): + if count == 1: + return length + return math.ceil(length / count) + overlap + + +def _overlap_options(length, count, native): + """Overlap candidates for one axis, in output pixels, widest first. + + An inactive axis blends nothing, as before. On an active axis the pitch - the un-overlapped + share each tile advances by - is the only scale a blend means anything against, so the ladder + is a fraction of the pitch rather than one fixed pixel count. + + Using the native overlap for every candidate imposes a lower bound on each active window + axis because `window = pitch + overlap`. With a 256px native overlap, the smallest window for + a 1024x1024 sample on four ranks is 512x512. Its area equals the 262144 pixels assigned to one + row-sharded rank, so it cannot reduce memory relative to row sharding. Allowing smaller + overlaps makes 272x272 windows reachable on the same sample. + + The ladder stops at a quarter of the window, which is a measured bound and not a margin. + Tile size decides how far a tile's tone drifts from its neighbours'; overlap decides how far + that drift is ramped out, and so whether the eye reads a gradient or a band. On FLUX.2 at + 1024x1024 on four ranks, a 128px window blended 32px - a quarter - is clean, while the same + window blended 16px bands. The difference between those decodes is concentrated near tile + boundaries spaced 112px apart. Since window is pitch + overlap, a quarter of the window is a + third of the pitch. + + `native` stays in the set so the previous behaviour remains reachable and comparable, but it + is dropped where it would fall under that quarter. + """ + if count == 1: + return (0,) + pitch = math.ceil(length / count) + options = {native, pitch // 2, math.ceil(pitch / 3)} + return tuple(sorted( + (option for option in options if option > 0 and option * 3 >= pitch), + reverse=True, + )) + + +def topology_objectives(window, overlap, sample_shape, world_size): + """Compute clipped tile areas and deterministic scheduler imbalance.""" + axis_sizes = [] + for length, size, blend in zip(sample_shape, window, overlap): + stride = size - blend + axis_sizes.append( + [min(size, length - start) for start in range(0, length, stride)] + ) + weights = [ + height * width for height in axis_sizes[0] for width in axis_sizes[1] + ] + tile_count = len(weights) + owners = shares(weights, world_size) + loads = [ + sum(weight for weight, owner in zip(weights, owners) if owner == rank) + for rank in range(world_size) + ] + average = sum(loads) / world_size + window_area = window[0] * window[1] + # What a rank holds under plain row sharding, which is the baseline every tiled plan is + # really competing with - not the unsharded decode. Recording it makes "is this plan a + # memory win at all?" answerable from the report instead of by hand. + row_shard_area = math.ceil(sample_shape[0] / world_size) * sample_shape[1] + return { + "window_area": window_area, + "decoded_area": sum(weights), + "tile_count": tile_count, + "max_rank_area": max(loads), + "rank_imbalance": max(loads) / average - 1, + "tile_grid": tuple(len(sizes) for sizes in axis_sizes), + "row_shard_area": row_shard_area, + "beats_row_sharding": window_area < row_shard_area, + # How many tile columns the grid has, which is the one thing separating a plan from its + # transpose. Area, work and imbalance are all symmetric under transpose, so without this + # the model cannot tell a full-WIDTH strip from a full-HEIGHT one - and the hardware very + # much can. Measured on FLUX.2 at 1024x1024 on four ranks, at identical window area and + # tile count: 128x1024 costs 966 MB against 1024x128's 1126 MB under tile-runs, and + # 651 MB against 812 MB under local. A wide tile is a few long contiguous spans and a + # tall one is a row of short ones, so fewer columns is cheaper at the same area. + "tile_columns": len(axis_sizes[1]), + } + + +def _dominates(left, right): + keys = ("window_area", "decoded_area", "rank_imbalance", "tile_columns") + return all(left[key] <= right[key] for key in keys) and any( + left[key] < right[key] for key in keys + ) + + +def pareto_frontier(candidates): + """Return candidates not dominated on memory, work, imbalance, and tile columns.""" + return [ + candidate + for candidate in candidates + if not any( + other is not candidate + and _dominates(other["objectives"], candidate["objectives"]) + for other in candidates + ) + ] + + +def _balanced_key(candidate, frontier): + objectives = candidate["objectives"] + keys = ("window_area", "decoded_area", "rank_imbalance") + distances = [] + for key in keys: + values = [entry["objectives"][key] for entry in frontier] + low, high = min(values), max(values) + distances.append(0.0 if high == low else (objectives[key] - low) / (high - low)) + return max(distances), sum(distances), candidate["window"] + + +def select_plans(sample_shape, native_overlap, world_size, normalize): + """Bracket the tile axis with a coarse, a knee, and a fine plan. + + Coarse has the fewest tiles, fine has the most, and balanced is the Pareto knee between them. + Coarse minimizes blend boundaries; fine minimizes window area. Return two plans when no + distinct balanced plan exists. + """ + if world_size < 1: + raise ValueError("world size must be positive") + max_tiles = max(4, 4 * world_size) + min_tiles = max(2, world_size) + candidates = {} + for down in range(1, max_tiles + 1): + for across in range(1, max_tiles + 1): + requested_tiles = down * across + if not min_tiles <= requested_tiles <= max_tiles: + continue + # Overlap is a search dimension, not a constant. Reducing it shrinks the window + # without changing the grid - stride stays at the pitch either way - so it is the + # cheapest axis the planner has, and holding it fixed forfeited the whole region + # where tiling beats row sharding. See _overlap_options. + for down_overlap in _overlap_options( + sample_shape[0], down, native_overlap[0] + ): + for across_overlap in _overlap_options( + sample_shape[1], across, native_overlap[1] + ): + overlap = (down_overlap, across_overlap) + window = ( + _axis_window(sample_shape[0], overlap[0], down), + _axis_window(sample_shape[1], overlap[1], across), + ) + normalized = normalize(window, overlap) + if normalized is None: + continue + window, overlap = normalized + if any(blend >= size for blend, size in zip(overlap, window)): + continue + # Re-check the quarter-of-window bound against the window actually used. + # normalize() may enlarge the requested window. blend_for_window() increases + # the overlap when necessary, and this check validates the resulting window + # and overlap. A VAE may normalize the requested overlap in turn. + # An inactive axis blends nothing and is exempt. + if any( + 0 < blend * 4 < size for blend, size in zip(overlap, window) + ): + continue + objectives = topology_objectives( + window, overlap, sample_shape, world_size + ) + if not min_tiles <= objectives["tile_count"] <= max_tiles: + continue + candidates[(tuple(window), tuple(overlap))] = { + "window": tuple(window), + "overlap": tuple(overlap), + "objectives": objectives, + } + frontier = pareto_frontier(list(candidates.values())) + if len(frontier) < 2: + raise ValueError( + f"sample {sample_shape} produces only {len(frontier)} useful tile plans" + ) + # Select the largest- and smallest-window Pareto candidates without predicting which is + # faster. The plan space is primarily ordered by window size, equivalently tile count. Its + # limits are the minimum window that avoids banding and the maximum useful window. + # Performance within those limits depends on hardware and must be measured. + # + # So the profiles name geometry, not predicted outcome. An earlier pair named throughput and + # memory scored plans by decoded_area, least total work, which reliably chose the widest + # window because it overlaps its neighbours fewer times. On gfx1201 those configurations + # were slower and used more memory than row sharding: 5034 MB versus 3526 MB at 2048x2048 on + # four ranks. Geometry-based profile names remain accurate across devices. + coarse = max( + frontier, + key=lambda item: ( + item["objectives"]["window_area"], + -item["objectives"]["tile_columns"], + item["window"], + ), + ) + fine = min( + (item for item in frontier if item is not coarse), + key=lambda item: ( + item["objectives"]["window_area"], + item["objectives"]["decoded_area"], + item["objectives"]["rank_imbalance"], + item["objectives"]["tile_columns"], + item["window"], + ), + ) + # The fine candidate must have a smaller window area than the coarse candidate. A transposed + # window can have identical modeled area, work, and imbalance while using more measured + # memory; 1024x128 used 17% more than 128x1024 in the measured configuration. + if fine["objectives"]["window_area"] >= coarse["objectives"]["window_area"]: + fine = None + # Distinct by WINDOW, not by identity. Two frontier points can share a window and differ only + # in blend, and 832x128 blended 36px against the same window blended 34px is not two profiles + # worth two cases each. + taken = {plan["window"] for plan in (coarse, fine) if plan is not None} + remaining = [item for item in frontier if item["window"] not in taken] + balanced = ( + min(remaining, key=lambda item: _balanced_key(item, frontier)) + if remaining + else None + ) + selected = [ + (profile, plan) + for profile, plan in zip(PROFILES, (coarse, balanced, fine)) + if plan is not None + ] + return [ + { + **plan, + "profile": profile, + "selection": { + "pareto_optimal": True, + "frontier_size": len(frontier), + "candidate_limit": max_tiles, + }, + } + for profile, plan in selected + ] + + +def blend_for_window(overlap, window): + """Widen a blend that snapping to a legal window has left under a quarter of it. + + Some VAEs round requested windows upward without changing the requested overlap. Increase + each active overlap to at least one quarter of the normalized window before select_plans() + validates the result. + + This adjustment matters most for VAEs with coarse window increments. LTX-2 uses 256px + increments; without the adjustment, a 1088x1920 sample on eight ranks produces no candidates. + Increasing the overlap restores three plans at its native 16-latent tile size. + + An inactive axis blends nothing and stays at zero. + """ + return tuple( + blend if blend == 0 or blend * 4 >= size else -(-size // 4) + for blend, size in zip(overlap, window) + ) + + +def normalizer_for_vae(vae, sample_shape, world_size): + """Return a candidate normalizer backed by DistVAE's exact planners.""" + native = vae_api.tile_shape(vae) + if native is None: + raise ValueError(f"{type(vae).__name__} has no tile window") + + def normalize(window, overlap): + height_options = [ + value + for value in range(window[0], window[0] + native[0] + 1) + if vae_api.tile_shape_plan(vae, value, native[1]) is not None + ][:8] + width_options = [ + value + for value in range(window[1], window[1] + native[1] + 1) + if vae_api.tile_shape_plan(vae, native[0], value) is not None + ][:8] + for height in height_options: + for width in width_options: + shape_plan = vae_api.tile_shape_plan(vae, height, width) + if shape_plan is None: + continue + # A tile needs enough of its narrow latent axis both to shard across the ranks + # and to normalize over something representative; the second is the binding + # constraint at every world size we run. Without it the widened overlap search + # reaches genuinely small windows for the first time and the memory profile + # selects them - it picked 9 latent rows on FLUX.2 at 1024. + latent_shape = _latent_shape(vae, shape_plan) + extent = min(latent_shape) if latent_shape is not None else None + if extent is not None and extent < max(world_size, MIN_TILE_LATENT_EXTENT): + continue + blend = blend_for_window(overlap, (height, width)) + if any(size <= value for value, size in zip(blend, (height, width))): + continue + original = {} + missing = [] + for name, planned in shape_plan.items(): + if hasattr(vae, name): + original[name] = getattr(vae, name) + else: + missing.append(name) + setattr(vae, name, planned) + try: + overlap_plan = vae_api.tile_overlap_plan( + vae, *blend, sample_shape=sample_shape + ) + finally: + for name in missing: + delattr(vae, name) + for name, value in original.items(): + setattr(vae, name, value) + if overlap_plan is not None: + return (height, width), blend + return None + + return normalize + + +def plans_for_vae(vae, height, width, world_size): + """Select the default plans for a concrete VAE.""" + overlap = vae_api.tile_overlap(vae) + if overlap is None: + raise ValueError(f"{type(vae).__name__} has no tile overlap") + sample_shape = (height, width) + return select_plans( + sample_shape, + overlap, + world_size, + normalizer_for_vae(vae, sample_shape, world_size), + ) + + +def default_suite(plans, height, width, frames, diagnostics=False): + """Build the bounded suite from the selected tile plans. + + By default only the compositions an orchestrator can actually select: the two untiled + baselines and whole-tile distribution at each plan. `local` tiles without distributing and + `row-tiled` shards rows beneath the tiling, and callers reach neither - xFuser, for one, + branches straight between marking a VAE for tile parallelism and parallelizing its decoder, + with nothing in between. They are also the slow ones, together about 60% of the suite's + compute at 1024x1024 on four ranks, which is a poor trade for a number nobody can act on. + + `diagnostics` puts them back. They earn it when characterising a new geometry rather than + comparing plans: `local` is the only case with no collectives at all, so it separates what + tiling does to the decode from what the collectives cost, and its peak is the true floor for + a window - 651 MB against tile-runs' 806 MB on that sample, the difference being assembly + rather than tile. + """ + suite = baseline_suite(height, width, frames) + modes = ("local", "tile-runs") if diagnostics else ("tile-runs",) + for mode in modes: + for plan in plans: + window, overlap, profile = ( + plan["window"], + plan["overlap"], + plan["profile"], + ) + suite.append( + _cell( + f"{mode}-{profile}", + mode, + height, + width, + frames, + window, + overlap, + profile=profile, + plan_selection=plan, + ) + ) + if diagnostics: + # Row sharding beneath the tiling, at the finest plan the sample offers. Lightest by + # predicted window area, which is a model's opinion rather than a measurement, and one + # more reason this belongs with the diagnostics. + lightest = min(plans, key=lambda plan: plan["objectives"]["window_area"]) + suite.append( + _cell( + f"row-tiled-{lightest['profile']}", + "row-tiled", + height, + width, + frames, + lightest["window"], + lightest["overlap"], + profile=lightest["profile"], + plan_selection=lightest, + ) + ) + return suite + + +def baseline_suite(height, width, frames): + """Build the two untiled cases supported by encoders and decoders.""" + return [ + _cell("unsharded", "unsharded", height, width, frames), + _cell("row", "row", height, width, frames), + ] diff --git a/bench/harness/catalog.py b/bench/harness/catalog.py new file mode 100644 index 0000000..543461f --- /dev/null +++ b/bench/harness/catalog.py @@ -0,0 +1,282 @@ +"""VAE family specifications, construction, sampling, and adapter descriptions.""" + +from contextlib import nullcontext + +import torch + +# `shapes` is the family's canonical matrix, as (height, width, frames), and `--matrix` runs it. +# It lives here beside the architecture rather than in a caller's script so that pinning a commit +# pins the shapes too: two runs of the same SHA measured the same thing, on whatever machine, and +# a result that cannot say what it measured is one nobody can reproduce. +# +# Frames are carried even where they are ignored, so every entry reads the same. A family with no +# temporal axis discards them in `sample_for`; one with a temporal axis needs 1 plus a multiple of +# it. Qwen-Image is a 3D VAE that ships as a single-image model, which is why it asks for one +# frame rather than the video-shaped default. +FAMILIES = { + "flux2": { + "cls": "AutoencoderKLFlux2", + "config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 32, + "block_out_channels": [128, 256, 512, 512], + "layers_per_block": 2, + "norm_num_groups": 32, + "down_block_types": ["DownEncoderBlock2D"] * 4, + "up_block_types": ["UpDecoderBlock2D"] * 4, + "patch_size": [2, 2], + "mid_block_add_attention": True, + "use_quant_conv": True, + "use_post_quant_conv": True, + }, + "latent_channels": 32, + "spatial": 8, + "temporal": None, + "shapes": ((1024, 1024, 1), (2048, 2048, 1)), + "note": "FLUX.2 checkpoints", + }, + "kl": { + "cls": "AutoencoderKL", + "config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 16, + "block_out_channels": [128, 256, 512, 512], + "layers_per_block": 2, + "norm_num_groups": 32, + "down_block_types": ["DownEncoderBlock2D"] * 4, + "up_block_types": ["UpDecoderBlock2D"] * 4, + "sample_size": 1024, + }, + "latent_channels": 16, + "spatial": 8, + "temporal": None, + "shapes": ((1024, 1024, 1), (2048, 2048, 1)), + "note": "plain 2D KL autoencoders", + }, + "wan": { + "cls": "AutoencoderKLWan", + "config": { + "base_dim": 96, + "z_dim": 16, + "dim_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_scales": [], + "temperal_downsample": [False, True, True], + }, + "latent_channels": 16, + "spatial": 8, + "temporal": 4, + # Portrait 480p and 720p at the production length. 81 frames is 21 latent ones, which is + # enough that the unsharded case may not fit at 720p. That failure is recorded per cell, + # and the tiled configurations continue. + "shapes": ((832, 480, 81), (1280, 720, 81)), + "note": "Wan video autoencoders", + }, + "qwen_image": { + "cls": "AutoencoderKLQwenImage", + "config": { + "base_dim": 96, + "z_dim": 16, + "dim_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_scales": [], + "temperal_downsample": [False, True, True], + }, + "latent_channels": 16, + "spatial": 8, + "temporal": 4, + "shapes": ((1024, 1024, 1), (2048, 2048, 1)), + "note": "Qwen Image autoencoders", + }, + "hunyuan_video": { + "cls": "AutoencoderKLHunyuanVideo", + "config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 16, + "block_out_channels": [128, 256, 512, 512], + "layers_per_block": 2, + "norm_num_groups": 32, + "mid_block_add_attention": True, + "spatial_compression_ratio": 8, + "temporal_compression_ratio": 4, + }, + "latent_channels": 16, + "spatial": 8, + "temporal": 4, + "shapes": ((832, 480, 129), (1280, 720, 129)), + "note": "Hunyuan Video autoencoders", + }, + "hunyuan_video_15": { + "cls": "AutoencoderKLHunyuanVideo15", + "config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 32, + "block_out_channels": [128, 256, 512, 1024, 1024], + "layers_per_block": 2, + "downsample_match_channel": True, + "upsample_match_channel": True, + "spatial_compression_ratio": 16, + "temporal_compression_ratio": 4, + }, + "latent_channels": 32, + "spatial": 16, + "temporal": 4, + "shapes": ((832, 480, 129), (1280, 720, 129)), + "note": "Hunyuan Video 1.5 autoencoders", + }, + "ltx2": { + "cls": "AutoencoderKLLTX2Video", + "config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": [256, 512, 1024, 2048], + "decoder_block_out_channels": [256, 512, 1024], + "layers_per_block": [4, 6, 6, 2, 2], + "decoder_layers_per_block": [5, 5, 5, 5], + "spatio_temporal_scaling": [True, True, True, True], + "decoder_spatio_temporal_scaling": [True, True, True], + "decoder_inject_noise": [False, False, False, False], + "downsample_type": [ + "spatial", + "temporal", + "spatiotemporal", + "spatiotemporal", + ], + "upsample_factor": [2, 2, 2], + "upsample_residual": [True, True, True], + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + "latent_channels": 128, + "spatial": 32, + "temporal": 8, + "shapes": ((1536, 1024, 121), (1920, 1280, 121)), + "note": "LTX-2 autoencoders", + }, +} + + +def _dtype(value): + return getattr(torch, value) if isinstance(value, str) else value + + +def matrix_for(family): + """Return a family's canonical shapes, checked against what its VAE can accept. + + Checked here rather than left to `sample_for` because a matrix is meant to be run unattended + across machines: an axis that does not divide, or a frame count the temporal ratio rejects, + should fail while the pod is still starting rather than partway through the third shape. + """ + spec = FAMILIES[family] + shapes = spec.get("shapes") + if not shapes: + raise ValueError( + f"--family {family} has no canonical shapes; ask for --shape explicitly" + ) + ratio, temporal = spec["spatial"], spec["temporal"] + for height, width, frames in shapes: + if height % ratio or width % ratio: + raise ValueError( + f"{family} shape {height}x{width} is not divisible by " + f"compression ratio {ratio}" + ) + if temporal and (frames - 1) % temporal: + raise ValueError( + f"{family} shape {height}x{width}x{frames} needs 1 plus a multiple " + f"of {temporal} frames" + ) + return tuple(shapes) + + +def build_vae(family, dtype, device): + """Build a deterministic architecture with random weights.""" + import diffusers + + spec = FAMILIES[family] + cls = getattr(diffusers, spec["cls"], None) + if cls is None: + raise ValueError( + f"diffusers {diffusers.__version__} does not provide {spec['cls']} " + f"required by --family {family}" + ) + device = torch.device(device) + torch.manual_seed(0) + context = torch.device("meta") if device.type == "meta" else nullcontext() + with context: + vae = cls(**spec["config"]).eval() + if device.type != "meta": + vae = vae.to(device=device, dtype=_dtype(dtype)) + return vae + + +def sample_for(spec, half, height, width, dtype, device, batch=1, frames=1): + """Create the input tensor for one encoder or decoder call.""" + ratio = spec["spatial"] + if height % ratio or width % ratio: + raise ValueError( + f"{height}x{width} is not divisible by compression ratio {ratio}" + ) + temporal = spec["temporal"] + if temporal and (frames - 1) % temporal: + raise ValueError(f"--frames {frames} must be 1 plus a multiple of {temporal}") + if half == "decoder": + channels = spec["latent_channels"] + rows, columns = height // ratio, width // ratio + depth = 1 + (frames - 1) // temporal if temporal else None + else: + channels = spec["config"].get("in_channels", 3) + rows, columns = height, width + depth = frames if temporal else None + shape = (batch, channels, rows, columns) + if depth is not None: + shape = (batch, channels, depth, rows, columns) + torch.manual_seed(1) + return torch.randn(*shape, dtype=_dtype(dtype), device=device) + + +def run_half(vae, half, sample): + """Run one VAE half and return the comparable tensor.""" + if half == "decoder": + return vae.decode(sample).sample + encoded = vae.encode(sample) + distribution = getattr(encoded, "latent_dist", None) + return distribution.mean if distribution is not None else encoded.latent + + +def describe_vae(vae, half): + """Describe the intact half and DistVAE's selected public adapter.""" + from distvae import vae as vae_api + + part = getattr(vae, half) + blocks = tuple( + getattr(part, "up_blocks" if half == "decoder" else "down_blocks", None) or () + ) + + def named(obj): + cls = type(obj) + return f"{cls.__module__}.{cls.__name__}" + + choose = ( + vae_api.decoder_adapter_name + if half == "decoder" + else vae_api.encoder_adapter_name + ) + return { + "class": named(part), + "blocks": sorted({named(block) for block in blocks}), + "mid_block": named(getattr(part, "mid_block", None)), + "conv_norm_out": named(getattr(part, "conv_norm_out", None)), + "adapter": choose(vae), + } diff --git a/bench/harness/cli.py b/bench/harness/cli.py new file mode 100644 index 0000000..d146a4b --- /dev/null +++ b/bench/harness/cli.py @@ -0,0 +1,318 @@ +"""Command-line parsing and orchestration for the DistVAE benchmark.""" + +import argparse + +import torch.distributed as dist + +from . import cases, catalog, measure, report, shape_costs +from .distributed import ( + Runtime, + aggregate_rank_errors, + exception_record, + gather_rank_errors, + ranks_diverged, +) + + +def parser(): + """Build the compatibility command-line parser.""" + value = argparse.ArgumentParser( + description=( + "Measure native DistVAE VAE sharding, tiling, and tile distribution." + ) + ) + value.add_argument("--family", default="flux2", choices=sorted(catalog.FAMILIES)) + value.add_argument("--half", default="decoder", choices=["decoder", "encoder"]) + value.add_argument("--height", type=int, default=2048) + value.add_argument("--width", type=int, default=2048) + value.add_argument("--frames", type=int, default=17) + value.add_argument( + "--shape", + action="append", + help="explicit HxW or HxWxFRAMES input shape; repeat to request more", + ) + value.add_argument( + "--matrix", + action="store_true", + help=( + "run the family's canonical shapes, so a pinned commit fixes what was " + "measured; overridden by --shape" + ), + ) + value.add_argument("--dtype", default="bfloat16", choices=sorted(measure.MAX_REL)) + value.add_argument("--warmup", type=int, default=2) + value.add_argument("--iters", type=int, default=5) + value.add_argument("--batch", type=int, default=1) + value.add_argument( + "--case", + action="append", + help=( + "exact case; repeat unsharded, row, or " + "MODE:WINDOW_HxW@OVERLAP_HxW where MODE is local, tile-runs, " + "or row-tiled. Omit for the bounded default suite" + ), + ) + value.add_argument( + "--diagnostics", + action="store_true", + help=( + "add the local and row-tiled compositions, which no orchestrator selects " + "but which isolate tiling from its collectives" + ), + ) + value.add_argument( + "--phase-timing", + action="store_true", + help="measure decoder calls separately from tiled decode overhead", + ) + value.add_argument( + "--profile", + action="store_true", + help="profile one selected VAE-half call outside the timed iterations", + ) + value.add_argument( + "--profile-trace", + action="store_true", + help="export a harness-named Chrome trace; implies --profile", + ) + value.add_argument( + "--profile-memory", + action="store_true", + help="export a harness-named memory timeline; implies --profile", + ) + value.add_argument( + "--profile-dir", + default="bench-profile", + help="directory for requested profiler artifacts", + ) + value.add_argument( + "--tile-shape-costs", + action="store_true", + help="measure decoder latency and memory across tile shapes and batch sizes", + ) + value.add_argument( + "--tile-shape-batch", + type=int, + default=1, + help="largest power-of-two tile batch to measure", + ) + value.add_argument( + "--tile-shape-windows", + default="", + help="comma-separated latent HEIGHTxWIDTH windows to measure", + ) + value.add_argument("--max-rel", type=float) + value.add_argument("--skip-reference", action="store_true") + value.add_argument("--reference-max-latent-elems", type=int, default=16384) + value.add_argument( + "--describe-only", + action="store_true", + help="describe native adapter selection on the meta device and stop", + ) + value.add_argument("--timeout-min", type=int, default=30) + value.add_argument("--out", help="write versioned JSON here") + return value + + +def _shape(spec, cell): + return { + "height": cell["height"], + "width": cell["width"], + "frames": cell["frames"] if spec["temporal"] else None, + } + + +def _describe(args, cells, provenance_data=None): + spec = catalog.FAMILIES[args.family] + if not cells: + cells = [ + cases.parse_case("unsharded", height, width, frames) + for height, width, frames in cases.shapes_from_args(args) + ] + records = [] + for cell in cells: + vae = catalog.build_vae(args.family, args.dtype, "meta") + description = catalog.describe_vae(vae, args.half) + composition = { + **cell, + "execution": "describe-only", + "adapter": description["adapter"], + } + records.append( + report.make_record( + args.family, + args.half, + _shape(spec, cell), + composition, + {"description": description}, + dtype=args.dtype, + world_size=1, + provenance_data=provenance_data, + ) + ) + return records + + +def _measure(args, cells, runtime, provenance_data=None): + spec = catalog.FAMILIES[args.family] + if args.tile_shape_costs: + error = None + costs = {"frames": args.frames if spec["temporal"] else None} + try: + costs = shape_costs.tile_shape_costs( + args, + spec, + runtime, + lambda *parts: print(*parts, flush=True) if runtime.rank == 0 else None, + ) + measurement = {"tile_shape_costs": costs} + except (Exception, SystemExit) as caught: + error = exception_record(caught, runtime.rank) + measurement = {} + aggregate_error = aggregate_rank_errors(gather_rank_errors(error, runtime)) + composition = { + "name": "tile-shape-costs", + "execution": "tile-shape-costs", + "sharding": "unsharded", + "window": None, + "overlap": None, + "tile_distribution": None, + } + record = report.make_record( + args.family, + "decoder", + {"height": None, "width": None, "frames": costs.get("frames")}, + composition, + measurement, + aggregate_error, + dtype=args.dtype, + world_size=runtime.world_size, + provenance_data=provenance_data, + ) + if runtime.rank == 0: + report.render(record, "decoder") + return [record] + + if not cells: + cells = [] + selector = ( + catalog.build_vae(args.family, args.dtype, "meta") + if args.half == "decoder" + else None + ) + for height, width, frames in cases.shapes_from_args(args): + if args.half == "encoder": + cells.extend(cases.baseline_suite(height, width, frames)) + else: + plans = cases.plans_for_vae( + selector, height, width, runtime.world_size + ) + cells.extend( + cases.default_suite( + plans, height, width, frames, diagnostics=args.diagnostics + ) + ) + + references = {} + records = [] + + def say(*parts): + if runtime.rank == 0: + print(*parts, flush=True) + + for cell in cells: + error = None + try: + composition, measurement = measure.measure_cell( + args, spec, cell, runtime, references, say + ) + except (Exception, SystemExit) as caught: + error = exception_record(caught, runtime.rank) + print( + f"[rank {runtime.rank}] cell {cell['name']} failed: " + f"{error['type']}: {error['message']}", + flush=True, + ) + composition, measurement = dict(cell), {} + runtime.device_api.empty_cache() + + failures = gather_rank_errors(error, runtime) + aggregate_error = aggregate_rank_errors(failures) + record = report.make_record( + args.family, + args.half, + _shape(spec, cell), + composition, + measurement, + aggregate_error, + dtype=args.dtype, + world_size=runtime.world_size, + provenance_data=provenance_data, + ) + records.append(record) + if runtime.rank == 0: + report.render(record, args.half) + # Written after every cell rather than once at the end. A sweep spends hours + # reaching its later cells, and a failure there used to discard every cell before + # it along with itself - the measurements were already paid for, and losing them + # means running the whole matrix again to recover what was already known. + out = getattr(args, "out", None) + if out: + report.write_json(out, records) + if ranks_diverged(failures): + say( + f"stopping after {cell['name']}: the ranks have diverged and no further " + "measurement from this group would mean anything" + ) + break + return records + + +def main(argv=None): + """Run describe-only or accelerator measurement mode and return an exit status.""" + command = parser() + args = command.parse_args(argv) + if args.tile_shape_costs and args.half != "decoder": + command.error("--tile-shape-costs requires --half decoder") + if args.tile_shape_costs and args.describe_only: + command.error("--tile-shape-costs cannot be combined with --describe-only") + if args.tile_shape_costs: + cells = [] + else: + try: + cells = cases.cells_from_args(args) + except ValueError as error: + command.error(str(error)) + provenance_data = report.provenance() + + if args.describe_only: + records = _describe(args, cells, provenance_data) + for record in records: + report.render(record, args.half) + if args.out: + report.write_json(args.out, records) + return report.report_status(records) + + runtime = Runtime.start(args.timeout_min) + try: + records = _measure(args, cells, runtime, provenance_data) + if runtime.rank == 0 and args.out: + report.write_json(args.out, records) + status = report.report_status(records) + statuses = [None] * runtime.world_size + # Agreeing on an exit status is itself a collective, and a run that stopped because its + # ranks diverged is in no position to complete one. The records are already on disk by + # here, so fall back to this rank's own status rather than fail on the way out and lose + # the status of a run that otherwise finished. + try: + dist.all_gather_object(statuses, status, group=runtime.group) + except Exception: + return status + agreed = [value for value in statuses if isinstance(value, int)] + return max(agreed) if agreed else status + finally: + runtime.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/harness/distributed.py b/bench/harness/distributed.py new file mode 100644 index 0000000..d062a9a --- /dev/null +++ b/bench/harness/distributed.py @@ -0,0 +1,281 @@ +"""Distributed process lifecycle and exact collective accounting.""" + +import importlib +import os +import sys +from collections import defaultdict +from dataclasses import dataclass +from datetime import timedelta + +import torch +import torch.distributed as dist + + +def exception_record(error, rank): + """Represent a local exception without losing its originating rank or type.""" + preserved = getattr(error, "rank_error", None) + if preserved is not None: + return preserved + return {"type": type(error).__name__, "message": str(error), "rank": int(rank)} + + +def gather_rank_errors(local_error, runtime): + """Collect one optional error from every rank in collective order.""" + failures = [None] * runtime.world_size + dist.all_gather_object(failures, local_error, group=runtime.group) + return failures + + +DESYNCHRONIZED = "DesynchronizedRanks" + + +def _is_failure_record(failure): + return isinstance(failure, dict) and "rank" in failure + + +def ranks_diverged(failures): + """Return whether the group can still be trusted to run collectives together. + + A cell that fails on EVERY rank leaves the group in step - an out-of-memory on a decode + nobody can fit is the ordinary way a matrix run reports "this shape does not fit", and the + next cell measures normally afterwards. A cell that fails on SOME ranks does not: the ranks + that failed stopped issuing collectives while the others carried on, so from that point the + two are matching up different calls and nothing the group produces means anything. + + The distinction is the whole point of the check. Stopping on the first kind would end most + sweeps at their first unsharded cell; continuing through the second kind produces numbers + that look ordinary and are not. + """ + if any( + failure is not None and not _is_failure_record(failure) for failure in failures + ): + return True + reported = [failure is not None for failure in failures] + return any(reported) and not all(reported) + + +def aggregate_rank_errors(failures): + """Combine rank errors while preserving each original failure record. + + An entry that is not a failure record is reported as one rather than raising. Once the ranks + diverge, the gather that collects the errors pairs with whatever call the other ranks are + still inside, so what comes back can be another call site's payload - and reaching into it + for a failure record used to raise an AttributeError that both hid the failure underneath + and took the rest of the run down with it. + """ + details = [] + for rank, failure in enumerate(failures): + if failure is None: + continue + if not _is_failure_record(failure): + details.append( + { + "type": DESYNCHRONIZED, + "message": ( + "the ranks are no longer running the same sequence of " + f"collectives: the failure gathered for rank {rank} came back " + f"as {type(failure).__name__}, which is another call's payload " + "rather than a failure record" + ), + "rank": rank, + } + ) + continue + nested = failure.get("failures") + details.extend(nested if nested is not None else [failure]) + by_rank = {} + for failure in details: + by_rank.setdefault(failure["rank"], failure) + details = list(by_rank.values()) + if not details: + return None + return { + **details[0], + "failed_ranks": [failure["rank"] for failure in details], + "failures": details, + } + + +class RankError(RuntimeError): + """Propagate an aggregated rank failure without wrapping its identity.""" + + def __init__(self, error, context): + self.rank_error = error + super().__init__( + f"{context} failed on rank {error['rank']}: " + f"{error['type']}: {error['message']}" + ) + + +def accelerator_backend(): + """Return the available accelerator API and its distributed backend.""" + if torch.cuda.is_available(): + return "cuda", torch.cuda, "nccl" + try: + importlib.import_module("torch_musa") + except ModuleNotFoundError as error: + raise RuntimeError("measurement requires CUDA or MUSA") from error + musa = getattr(torch, "musa", None) + if musa is None or not musa.is_available(): + raise RuntimeError("measurement requires CUDA or MUSA") + return "musa", musa, "mccl" + + +class CollectiveLog: + """Count collective calls and tensor bytes by operation and call site.""" + + WRAPPED = ( + "all_reduce", + "all_gather", + "all_gather_into_tensor", + "broadcast", + "isend", + "irecv", + "recv", + "send", + "barrier", + "batch_isend_irecv", + ) + + def __init__(self): + self.enabled = False + self.by_call = defaultdict(lambda: {"calls": 0, "bytes": 0}) + self.by_site = defaultdict(lambda: {"calls": 0, "bytes": 0}) + self._originals = {} + + @staticmethod + def _nbytes(args): + total = 0 + for arg in args: + values = arg if isinstance(arg, (list, tuple)) else (arg,) + total += sum( + value.numel() * value.element_size() + for value in values + if isinstance(value, torch.Tensor) + ) + return total + + def _wrap(self, name, original): + def wrapper(*args, **kwargs): + if self.enabled: + frame = sys._getframe(1) + site = f"{os.path.basename(frame.f_code.co_filename)}:{frame.f_lineno}" + nested = ( + os.path.basename(frame.f_code.co_filename) == "distributed_c10d.py" + ) + label = f"{name} (batched)" if nested else name + size = self._nbytes(args) + for entry in (self.by_call[label], self.by_site[f"{name} @ {site}"]): + entry["calls"] += 1 + entry["bytes"] += size + return original(*args, **kwargs) + + return wrapper + + def install(self): + from torch.distributed import distributed_c10d + + for name in self.WRAPPED: + original = getattr(dist, name, None) + if original is None: + continue + self._originals[name] = original + wrapper = self._wrap(name, original) + setattr(dist, name, wrapper) + if getattr(distributed_c10d, name, None) is original: + setattr(distributed_c10d, name, wrapper) + + def uninstall(self): + from torch.distributed import distributed_c10d + + for name, original in self._originals.items(): + setattr(dist, name, original) + setattr(distributed_c10d, name, original) + self._originals.clear() + + def reset(self): + self.by_call.clear() + self.by_site.clear() + + def report(self): + return { + "by_call": { + name: dict(value) for name, value in sorted(self.by_call.items()) + }, + "by_site": { + name: dict(value) + for name, value in sorted( + self.by_site.items(), key=lambda item: -item[1]["calls"] + ) + }, + "total_calls": sum( + value["calls"] + for name, value in self.by_call.items() + if "(batched)" not in name + ), + "total_bytes": sum(value["bytes"] for value in self.by_call.values()), + } + + +def across_ranks(by_call, world_size, group): + """Report operation counts for every rank and the busiest rank.""" + gathered = [None] * world_size + counts = {name: entry["calls"] for name, entry in by_call.items()} + dist.all_gather_object(gathered, counts, group=group) + + def total(values): + return sum(calls for name, calls in values.items() if "(batched)" not in name) + + names = sorted({name for values in gathered for name in values}) + return { + "by_call_max": { + name: max(values.get(name, 0) for values in gathered) for name in names + }, + "total_calls_max": max(total(values) for values in gathered), + "total_calls_by_rank": [total(values) for values in gathered], + } + + +@dataclass +class Runtime: + rank: int + world_size: int + local_rank: int + device: torch.device + group: object + log: CollectiveLog + device_api: object + + @classmethod + def start(cls, timeout_min): + """Initialize the accelerator process group used by measurements.""" + device_type, device_api, backend = accelerator_backend() + missing = [name for name in ("RANK", "WORLD_SIZE") if name not in os.environ] + if missing: + raise RuntimeError( + "measurement requires a distributed launch environment; missing " + + ", ".join(missing) + ) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + device_api.set_device(local_rank) + device = torch.device(device_type, local_rank) + dist.init_process_group( + backend=backend, + init_method="env://", + timeout=timedelta(minutes=timeout_min), + ) + log = CollectiveLog() + log.install() + group = dist.group.WORLD + dist.all_reduce(torch.zeros(1, device=device), group=group) + return cls(rank, world_size, local_rank, device, group, log, device_api) + + def close(self): + """Synchronize, restore wrapped calls, and destroy the process group.""" + try: + dist.barrier(group=self.group) + finally: + self.log.uninstall() + dist.destroy_process_group() diff --git a/bench/harness/measure.py b/bench/harness/measure.py new file mode 100644 index 0000000..af7dfd0 --- /dev/null +++ b/bench/harness/measure.py @@ -0,0 +1,324 @@ +"""Benchmark execution, timing, memory, phase timing, and output agreement.""" + +import time +from collections import Counter + +import torch +import torch.distributed as dist +import torch.nn as nn + +from distvae import vae as vae_api +from distvae.vae.tiling import _latent_shape, latent_rows + +from . import catalog, profile +from .distributed import across_ranks +from .report import set_agreement_policy + +MAX_REL = {"float32": 1e-4, "float16": 2e-2, "bfloat16": 5e-2} + + +def _device_api(runtime): + return runtime.device_api + + +def _tile_latent_area(vae): + shape = _latent_shape(vae) + return shape[0] * shape[1] if shape is not None else None + + +def configure_tiling(vae, cell, runtime, half, say): + """Apply the requested tile window, overlap, and whole-tile distribution.""" + if cell["window"] is None: + return {"enabled": False} + if half != "decoder": + raise ValueError("tiling is a decode-side feature and requires --half decoder") + + vae_api.require_vae_support(vae, "tiling", "--case") + vae.enable_tiling() + native = vae_api.tile_shape(vae) + native_window = tuple(native) if native is not None else None + native_overlap = vae_api.tile_overlap(vae) + facts = { + "enabled": True, + "requested_window_px": tuple(cell["window"]), + "native_window_px": native_window, + "window_px": tuple(cell["window"]), + "native_overlap_px": native_overlap, + } + + requested = tuple(cell["window"]) + plan = vae_api.tile_shape_plan(vae, *requested) + if plan is None: + raise ValueError( + f"tile shape {requested} is invalid for {type(vae).__name__}" + ) + rows = latent_rows(vae, plan) + if cell["sharding"] == "row" and rows is not None and rows < runtime.world_size: + raise ValueError( + f"a {requested[0]}x{requested[1]}px tile has {rows} latent rows " + f"for {runtime.world_size} row shards" + ) + vae_api.apply_tile_plan(vae, plan) + facts["tile_latent_rows"] = rows + + overlap = cell.get("overlap") + if overlap is not None: + plan = vae_api.tile_overlap_plan( + vae, + *overlap, + sample_shape=(cell["height"], cell["width"]), + ) + if plan is None: + raise ValueError( + f"tile overlap {overlap} is unavailable for {type(vae).__name__}" + ) + vae_api.apply_tile_plan(vae, plan) + tiled_decode = vae_api.tiled_decode_for(vae) + if tiled_decode is not None: + vae.tiled_decode = tiled_decode + facts.update( + overlap=vae_api.tile_overlap(vae), + tile_latent_area=_tile_latent_area(vae), + ) + + if cell["tile_distribution"] is not None: + if not vae_api.supports_tile_parallel(vae): + raise ValueError( + f"{type(vae).__name__} does not support whole-tile distribution" + ) + dispatch, assemble = vae_api.sharing(runtime.group) + tiled_decode = vae_api.tiled_decode_for(vae, dispatch, assemble) + if tiled_decode is None: + raise ValueError(f"{type(vae).__name__} has no distributable tiled decode") + vae.tiled_decode = tiled_decode + facts["distribution"] = cell["tile_distribution"] + else: + facts["distribution"] = None + return facts + + +def configure_sharding(vae, cell, runtime, half): + """Install row sharding or leave the decoder whole.""" + if cell["sharding"] != "row": + return None + install = ( + vae_api.parallelize_decoder + if half == "decoder" + else vae_api.parallelize_encoder + ) + return install(vae, runtime.group) + + +class PhaseTimer(nn.Module): + """Measure decoder calls separately from the full tiled decode.""" + + def __init__(self, decoder, runtime, counters): + super().__init__() + self.decoder = decoder + self.runtime = runtime + self.counters = counters + + def forward(self, *args, **kwargs): + _device_api(self.runtime).synchronize(self.runtime.device) + start = time.perf_counter() + output = self.decoder(*args, **kwargs) + _device_api(self.runtime).synchronize(self.runtime.device) + self.counters["decoder_s"] += time.perf_counter() - start + self.counters["calls"] += 1 + return output + + +def install_phase_timing(vae, runtime): + """Wrap decoder and tiled decode calls for optional phase accounting.""" + counters = Counter() + vae.decoder = PhaseTimer(vae.decoder, runtime, counters) + tiled_decode = vae.tiled_decode + + def timed_decode(*args, **kwargs): + _device_api(runtime).synchronize(runtime.device) + start = time.perf_counter() + output = tiled_decode(*args, **kwargs) + _device_api(runtime).synchronize(runtime.device) + counters["total_s"] += time.perf_counter() - start + counters["decodes"] += 1 + return output + + vae.tiled_decode = timed_decode + return counters + + +def phase_report(counters, runtime): + """Summarize phase time per decode and load spread across ranks.""" + decodes = counters.get("decodes", 0) + if not decodes: + return None + total = counters["total_s"] / decodes + decoder = counters["decoder_s"] / decodes + result = { + "total_ms": total * 1e3, + "decoder_ms": decoder * 1e3, + "rest_ms": (total - decoder) * 1e3, + "calls_per_decode": counters["calls"] / decodes, + } + if runtime.world_size > 1: + gathered = [None] * runtime.world_size + dist.all_gather_object( + gathered, + (decoder, counters["calls"] / decodes), + group=runtime.group, + ) + result["decoder_ms_by_rank"] = [value * 1e3 for value, _ in gathered] + result["calls_by_rank"] = [calls for _, calls in gathered] + slowest = max(value for value, _ in gathered) + result["idle_share"] = sum(slowest - value for value, _ in gathered) / ( + runtime.world_size * slowest or 1 + ) + return result + + +def timed(run, iters, runtime): + """Measure synchronized latency samples.""" + samples = [] + for _ in range(iters): + dist.barrier(group=runtime.group) + _device_api(runtime).synchronize(runtime.device) + start = time.perf_counter() + run() + _device_api(runtime).synchronize(runtime.device) + samples.append(time.perf_counter() - start) + return _timing_report(samples) + + +def _timing_report(samples): + """Summarize a non-empty sequence of latency samples.""" + samples.sort() + return { + "median_s": samples[len(samples) // 2], + "mean_s": sum(samples) / len(samples), + "min_s": samples[0], + "max_s": samples[-1], + "samples_s": samples, + } + + +def agreement_with(actual, reference, dtype, max_rel): + """Measure raw error against an unsharded, untiled reference.""" + if tuple(actual.shape) != tuple(reference.shape): + agreement = { + "ok": False, + "disagreement_type": "shape", + "why": f"shape {tuple(actual.shape)} != {tuple(reference.shape)}", + } + else: + diff = (actual.float().cpu() - reference).abs() + scale = reference.abs().max().item() + tolerance = max_rel if max_rel is not None else MAX_REL[dtype] + relative = diff.max().item() / scale if scale else 0.0 + agreement = { + "ok": bool(relative <= tolerance), + "disagreement_type": "numerical", + "max_abs": diff.max().item(), + "mean_abs": diff.mean().item(), + "reference_max_abs": scale, + "max_rel_to_scale": relative, + "mean_rel_to_scale": diff.mean().item() / scale if scale else 0.0, + "share_off_by_1pc": ( + (diff > 0.01 * scale).float().mean().item() if scale else 0.0 + ), + "max_rel_allowed": tolerance, + } + return agreement + + +def measure_cell(args, spec, cell, runtime, references, say): + """Build and measure one normalized composition cell.""" + vae = catalog.build_vae(args.family, args.dtype, runtime.device) + sample = catalog.sample_for( + spec, + args.half, + cell["height"], + cell["width"], + args.dtype, + runtime.device, + args.batch, + cell["frames"], + ) + description = catalog.describe_vae(vae, args.half) + if cell["sharding"] == "row" and description["adapter"] is None: + raise ValueError(f"DistVAE has no adapter for {type(vae).__name__} {args.half}") + + latent_area = sample.shape[0] * sample.shape[-2] * sample.shape[-1] + if sample.ndim == 5: + latent_area *= sample.shape[2] + if args.half == "encoder": + latent_area //= spec["spatial"] ** 2 + key = (cell["height"], cell["width"], cell["frames"], args.half) + take_reference = ( + not args.skip_reference and latent_area <= args.reference_max_latent_elems + ) + if take_reference and key not in references: + with torch.no_grad(): + references[key] = catalog.run_half(vae, args.half, sample).float().cpu() + + adapter = configure_sharding(vae, cell, runtime, args.half) + tiling = configure_tiling(vae, cell, runtime, args.half, say) + counters = ( + install_phase_timing(vae, runtime) + if args.phase_timing and tiling["enabled"] + else Counter() + ) + + def once(): + with torch.no_grad(): + return catalog.run_half(vae, args.half, sample) + + for _ in range(args.warmup): + once() + _device_api(runtime).synchronize(runtime.device) + + runtime.log.reset() + runtime.log.enabled = True + try: + output = once() + finally: + runtime.log.enabled = False + collectives = runtime.log.report() + collectives.update( + across_ranks(runtime.log.by_call, runtime.world_size, runtime.group) + ) + counters.clear() + _device_api(runtime).reset_peak_memory_stats(runtime.device) + timing = timed(once, args.iters, runtime) + peak_mb = _device_api(runtime).max_memory_allocated(runtime.device) / (1024 * 1024) + phases = phase_report(counters, runtime) + profile_result = profile.profile_once(once, args, cell, runtime) + reference = references.get(key) + agreement = ( + agreement_with( + output, + reference, + args.dtype, + args.max_rel, + ) + if reference is not None + else None + ) + if agreement is not None: + set_agreement_policy(agreement, tiling["enabled"]) + composition = { + **cell, + "execution": "measurement", + "adapter": adapter, + "tiling_effective": tiling, + } + measurement = { + "description": description, + "latent_shape": list(sample.shape), + "collectives": collectives, + "timing": timing, + "phases": phases, + "profile": profile_result, + "peak_vram_mb": peak_mb, + "agreement": agreement, + } + return composition, measurement diff --git a/bench/harness/profile.py b/bench/harness/profile.py new file mode 100644 index 0000000..047522b --- /dev/null +++ b/bench/harness/profile.py @@ -0,0 +1,112 @@ +"""Optional profiler execution and artifact export.""" + +import importlib +from contextlib import ExitStack +from pathlib import Path + +import torch + +from .distributed import ( + RankError, + aggregate_rank_errors, + exception_record, + gather_rank_errors, +) + +PROFILE_SUMMARY_LIMIT = 16_000 + + +def _profiler_backend(device_type): + activity_name = device_type.upper() + if device_type == "musa": + try: + importlib.import_module("torch_musa") + except ModuleNotFoundError as error: + raise RuntimeError( + "MUSA profiling requires the optional torch_musa package" + ) from error + activity = getattr(torch.profiler.ProfilerActivity, activity_name, None) + if device_type != "cpu" and activity is None: + raise RuntimeError(f"torch.profiler has no {activity_name} activity") + device_api = getattr(torch, device_type, None) + memory = getattr(device_api, "memory", None) + recorder = getattr(memory, "_record_memory_history", None) + return activity, recorder, f"self_{device_type}_time_total" + + +def _artifact_stem(output_dir, stem, suffixes): + candidate = stem + occurrence = 1 + while any((output_dir / f"{candidate}.{suffix}").exists() for suffix in suffixes): + occurrence += 1 + candidate = f"{stem}-{occurrence}" + return candidate + + +def profile_once(run, args, cell=None, runtime=None): + """Profile one VAE-half call and export only explicitly requested artifacts.""" + if not (args.profile or args.profile_trace or args.profile_memory): + return None + + with ExitStack() as memory_setup: + with ExitStack() as profiler_setup: + local_error = None + try: + output_dir = Path(args.profile_dir) + shape = f"{cell['height']}x{cell['width']}x{cell['frames']}" + base = ( + f"{args.family}-{args.half}-{cell['name']}-{shape}" + f"-rank{runtime.rank}" + ) + suffixes = [] + if args.profile_trace: + suffixes.append("trace.json") + if args.profile_memory: + suffixes.append("memory.html") + stem = _artifact_stem(output_dir, base, suffixes) + artifacts = {} + if args.profile_trace: + artifacts["trace"] = str(output_dir / f"{stem}.trace.json") + if args.profile_memory: + artifacts["memory"] = str(output_dir / f"{stem}.memory.html") + if artifacts: + output_dir.mkdir(parents=True, exist_ok=True) + + accelerator, memory_recorder, sort_by = _profiler_backend( + runtime.device.type + ) + activities = [torch.profiler.ProfilerActivity.CPU] + if accelerator is not None: + activities.append(accelerator) + + if args.profile_memory and memory_recorder is not None: + memory_recorder(enabled="all") + memory_setup.callback(memory_recorder, enabled=None) + profiler = profiler_setup.enter_context( + torch.profiler.profile( + activities=activities, + profile_memory=args.profile_memory, + record_shapes=args.profile_memory, + with_stack=args.profile_memory, + ) + ) + except (Exception, SystemExit) as error: + local_error = exception_record(error, runtime.rank) + failure = ( + aggregate_rank_errors(gather_rank_errors(local_error, runtime)) + if getattr(runtime, "world_size", 1) > 1 + else aggregate_rank_errors([local_error]) + ) + if failure is not None: + raise RankError(failure, "profiler setup") + + run() + + summary = profiler.key_averages().table(sort_by=sort_by, row_limit=20)[ + :PROFILE_SUMMARY_LIMIT + ] + if args.profile_trace: + profiler.export_chrome_trace(artifacts["trace"]) + if args.profile_memory: + profiler.export_memory_timeline(artifacts["memory"]) + return {"summary": summary, "artifacts": artifacts} diff --git a/bench/harness/report.py b/bench/harness/report.py new file mode 100644 index 0000000..882dace --- /dev/null +++ b/bench/harness/report.py @@ -0,0 +1,241 @@ +"""Versioned benchmark records, provenance, rendering, and exit policy.""" + +import hashlib +import importlib.metadata +import json +import os +import platform +import socket +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +import torch + +SCHEMA_VERSION = 7 + + +def _version(distribution, module=None): + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return ( + getattr(module, "__version__", "unknown") + if module is not None + else "unknown" + ) + + +def _distvae_revision(): + try: + import distvae + + root = Path(distvae.__file__).resolve().parents[1] + result = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + capture_output=True, + check=True, + text=True, + timeout=2, + ) + return result.stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _git(root, *arguments): + try: + result = subprocess.run( + ["git", "-C", str(root), *arguments], + capture_output=True, + check=True, + text=True, + timeout=2, + ) + return result.stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _source_checkout(module): + location = getattr(module, "__file__", None) + if location is None: + return None + start = Path(location).resolve().parent + for root in (start, *start.parents): + if not (root / ".git").exists(): + continue + return { + "branch": _git(root, "rev-parse", "--abbrev-ref", "HEAD"), + "commit": _git(root, "rev-parse", "HEAD"), + "dirty": bool(_git(root, "status", "--porcelain")), + } + return None + + +def _benchmark_identity(): + try: + launcher = Path(sys.argv[0]).resolve() + harness = Path(__file__).resolve().parent + sources = sorted(harness.glob("*.py")) + digest = hashlib.sha256() + if launcher.is_file() and launcher not in sources: + sources.append(launcher) + for source in sources: + try: + label = source.relative_to(harness.parent) + except ValueError: + label = Path(source.name) + digest.update(str(label).encode()) + digest.update(b"\0") + digest.update(source.read_bytes()) + return { + "path": str(launcher), + "sha256": digest.hexdigest(), + "implementation": [str(source) for source in sources], + } + except OSError: + return None + + +def _device_identity(): + """Return the accelerator properties that make measurements comparable.""" + if not torch.cuda.is_available(): + return None + properties = torch.cuda.get_device_properties(torch.cuda.current_device()) + return { + "name": properties.name, + "arch": getattr(properties, "gcnArchName", None), + "total_memory": properties.total_memory, + "count": torch.cuda.device_count(), + } + + +def provenance(): + """Return library versions and the DistVAE source revision when available.""" + import diffusers + import distvae + + return { + "versions": { + "torch": torch.__version__, + "diffusers": _version("diffusers", diffusers), + "distvae": _version("distvae", distvae), + }, + "provenance": { + "recorded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "host": socket.gethostname(), + "hardware_family": os.environ.get("HW_FAMILY"), + "device": _device_identity(), + "python": platform.python_version(), + "argv": list(sys.argv), + "benchmark": _benchmark_identity(), + "distvae_git_revision": _distvae_revision(), + "distvae_checkout": _source_checkout(distvae), + }, + } + + +def make_record( + family, + half, + shape, + composition, + measurement=None, + error=None, + *, + dtype, + world_size, + provenance_data=None, +): + """Build one self-contained schema-versioned result.""" + dtype_name = str(dtype).removeprefix("torch.") + record = { + "schema_version": SCHEMA_VERSION, + **(provenance_data if provenance_data is not None else provenance()), + "family": family, + "half": half, + "shape": shape, + "composition": composition, + "runtime": {"dtype": dtype_name, "world_size": int(world_size)}, + "measurement": measurement or {}, + } + if error is not None: + record["error"] = error + return record + + +def set_agreement_policy(agreement, tiling_enabled): + """Record whether the raw agreement verdict controls process success.""" + numerical_difference = agreement.get("disagreement_type") == "numerical" + agreement["enforced"] = not (tiling_enabled and numerical_difference) + if tiling_enabled and numerical_difference: + agreement["measured_not_enforced"] = ( + "tiling changes arithmetic; the measured difference remains reported" + ) + + +def report_status(records): + """Return failure only after all records are ready to be written.""" + if any(record is None or "error" in record for record in records): + return 1 + for record in records: + agreement = record.get("measurement", {}).get("agreement") + if agreement and agreement.get("enforced", True) and not agreement["ok"]: + return 1 + return 0 + + +def write_json(path, records): + """Write one record as an object and a grid as an array.""" + payload = records[0] if len(records) == 1 else records + Path(path).write_text(json.dumps(payload, indent=2) + "\n") + + +def render(record, half): + """Render the compact human-readable view of one record.""" + measurement = record.get("measurement", {}) + mode = record.get("composition", {}).get("execution") + if "error" in record: + print( + f"{record['composition'].get('name', 'cell')} failed: " + f"{record['error']['type']}: {record['error']['message']}", + flush=True, + ) + return + if mode == "describe-only": + print(json.dumps(measurement["description"], sort_keys=True), flush=True) + return + if mode == "tile-shape-costs": + costs = measurement["tile_shape_costs"] + print( + json.dumps( + { + "latent_window": costs["latent_window"], + "analysis": costs["analysis"], + }, + sort_keys=True, + ), + flush=True, + ) + return + collectives = measurement["collectives"] + timing = measurement["timing"] + print(f"\n--- collectives per {half} call ---", flush=True) + for name, entry in collectives["by_call"].items(): + maximum = collectives["by_call_max"].get(name, entry["calls"]) + print( + f" {name:<24} {entry['calls']:>6} calls " + f"{maximum:>6} max {entry['bytes'] / 1e6:>10.2f} MB", + flush=True, + ) + print( + f"median {timing['median_s'] * 1000:.1f} ms " + f"peak {measurement['peak_vram_mb']:.0f} MB", + flush=True, + ) + agreement = measurement.get("agreement") + if agreement is not None: + verdict = "matches" if agreement["ok"] else "differs from" + print(f"output {verdict} the unsharded reference: {agreement}", flush=True) diff --git a/bench/harness/shape_costs.py b/bench/harness/shape_costs.py new file mode 100644 index 0000000..3cf94de --- /dev/null +++ b/bench/harness/shape_costs.py @@ -0,0 +1,232 @@ +"""Decoder cost measurements across tile shapes and batch sizes.""" + +import time + +import torch +import torch.distributed as dist + +from distvae import vae as vae_api + +from . import cases, catalog +from .distributed import ( + RankError, + aggregate_rank_errors, + exception_record, + gather_rank_errors, +) + + +def _device_api(runtime): + return runtime.device_api + + +def _timing_report(samples): + samples.sort() + return { + "median_s": samples[len(samples) // 2], + "mean_s": sum(samples) / len(samples), + "min_s": samples[0], + "max_s": samples[-1], + "samples_s": samples, + } + + +def _synchronize_failure(local_error, runtime): + """Share a rank-local case failure before any rank enters the next case.""" + failures = gather_rank_errors(local_error, runtime) + return aggregate_rank_errors(failures) + + +def _all_out_of_memory(failure): + """Return whether every underlying rank failure is an OOM.""" + failures = failure.get("failures", [failure]) + return all(item["type"] == "OutOfMemoryError" for item in failures) + + +def _shape_iterations(run, iterations, runtime): + """Run decoder iterations with a rank verdict exchange after each call.""" + samples = [] + for _ in range(iterations): + dist.barrier(group=runtime.group) + local_error = None + elapsed = None + try: + _device_api(runtime).synchronize(runtime.device) + start = time.perf_counter() + run() + _device_api(runtime).synchronize(runtime.device) + elapsed = time.perf_counter() - start + except Exception as error: + local_error = exception_record(error, runtime.rank) + failure = _synchronize_failure(local_error, runtime) + if failure is not None: + return None, failure + samples.append(elapsed) + return samples, None + + +def tile_shape_costs(args, spec, runtime, say): + """Measure decoder cost across representative tile shapes and batch sizes.""" + local_error = None + try: + vae = catalog.build_vae(args.family, args.dtype, runtime.device) + window = vae_api.tile_shape(vae) + if window is None: + raise ValueError( + f"{type(vae).__name__} has no native tile shape for shape analysis" + ) + latent_window = tuple(value // spec["spatial"] for value in window) + depth = 1 + (args.frames - 1) // spec["temporal"] if spec["temporal"] else None + + if args.tile_shape_windows: + shapes = [ + cases.parse_pair(value, "latent tile window") + for value in args.tile_shape_windows.split(",") + ] + if any(min(shape) <= 0 for shape in shapes): + raise ValueError("--tile-shape-windows axes must be positive") + else: + plans = cases.plans_for_vae( + vae, args.height, args.width, runtime.world_size + ) + shapes = [ + tuple(axis // spec["spatial"] for axis in plan["window"]) + for plan in plans + ] + if not shapes: + raise ValueError( + f"tile window produces no representative shapes at {latent_window}" + ) + if args.tile_shape_batch < 1: + raise ValueError("--tile-shape-batch must be positive") + + counts = [] + count = 1 + while count <= args.tile_shape_batch: + counts.append(count) + count *= 2 + dtype = getattr(torch, args.dtype) + effective_frames = args.frames if spec["temporal"] else None + except (Exception, SystemExit) as error: + local_error = exception_record(error, runtime.rank) + failure = _synchronize_failure(local_error, runtime) + if failure is not None: + raise RankError(failure, "tile shape setup") + + measured = [] + baseline = None + alone = {} + for rows, columns in shapes: + for count in counts: + shape = (count, spec["latent_channels"], rows, columns) + if depth is not None: + shape = (count, spec["latent_channels"], depth, rows, columns) + latent = None + local_error = None + try: + torch.manual_seed(1) + latent = torch.randn(*shape, dtype=dtype, device=runtime.device) + _device_api(runtime).reset_peak_memory_stats(runtime.device) + except Exception as error: + local_error = exception_record(error, runtime.rank) + failure = _synchronize_failure(local_error, runtime) + if failure is not None: + if not _all_out_of_memory(failure): + raise RankError(failure, "tile shape setup") + measured.append( + { + "rows": rows, + "columns": columns, + "tiles_in_the_call": count, + "latent_area": rows * columns, + "out_of_memory": True, + "failure_phase": "allocation", + "failed_ranks": failure["failed_ranks"], + } + ) + say(f"{rows}x{columns} x{count}: out of memory during allocation") + latent = None + _device_api(runtime).empty_cache() + break + + def once(): + with torch.no_grad(): + return catalog.run_half(vae, "decoder", latent) + + _, failure = _shape_iterations(once, args.warmup, runtime) + failure_phase = "warmup" + if failure is None: + samples, failure = _shape_iterations(once, args.iters, runtime) + failure_phase = "measurement" + if failure is not None: + if not _all_out_of_memory(failure): + raise RankError(failure, f"tile shape {failure_phase}") + measured.append( + { + "rows": rows, + "columns": columns, + "tiles_in_the_call": count, + "latent_area": rows * columns, + "out_of_memory": True, + "failure_phase": failure_phase, + "failed_ranks": failure["failed_ranks"], + } + ) + say(f"{rows}x{columns} x{count}: out of memory during {failure_phase}") + latent = None + _device_api(runtime).empty_cache() + break + timing = _timing_report(samples) + + area = rows * columns + median_ms = timing["median_s"] * 1000 + per_tile_ms = median_ms / count + if count == 1: + alone[(rows, columns)] = per_tile_ms + if baseline is None: + baseline = (per_tile_ms, area) + predicted_ms = baseline[0] * area / baseline[1] + entry = { + "rows": rows, + "columns": columns, + "tiles_in_the_call": count, + "latent_area": area, + "timing": timing, + "median_ms": median_ms, + "ms_per_tile": per_tile_ms, + "peak_vram_mb": _device_api(runtime).max_memory_allocated( + runtime.device + ) + / (1024 * 1024), + "ms_per_1k_latent_area": per_tile_ms / area * 1000, + "against_area_prediction": per_tile_ms / predicted_ms, + "against_single_tile": per_tile_ms / alone[(rows, columns)], + } + measured.append(entry) + say( + f"{rows}x{columns} x{count}: {median_ms:.1f} ms, " + f"{entry['peak_vram_mb']:.0f} MB" + ) + latent = None + _device_api(runtime).empty_cache() + + fitted = [entry for entry in measured if not entry.get("out_of_memory")] + batched = [entry for entry in fitted if entry["tiles_in_the_call"] > 1] + return { + "family": args.family, + "latent_window": latent_window, + "frames": effective_frames, + "shapes": measured, + "analysis": { + "highest_area_cost": ( + max(fitted, key=lambda entry: entry["against_area_prediction"]) + if fitted + else None + ), + "worst_batch_scaling": ( + max(batched, key=lambda entry: entry["against_single_tile"]) + if batched + else None + ), + }, + } diff --git a/bench/smoke_families.py b/bench/smoke_families.py new file mode 100644 index 0000000..9aeb447 --- /dev/null +++ b/bench/smoke_families.py @@ -0,0 +1,42 @@ +"""Build every catalog family on the meta device without weights or an accelerator.""" + +import torch + +if __package__: + from .harness.catalog import FAMILIES, sample_for +else: + from harness.catalog import FAMILIES, sample_for + + +def main(): + import diffusers + + print(f"diffusers {diffusers.__version__}") + failures = 0 + for name, spec in sorted(FAMILIES.items()): + cls = getattr(diffusers, spec["cls"], None) + if cls is None: + print(f" {name:<18} SKIP {spec['cls']} is not in this diffusers") + continue + try: + with torch.device("meta"): + vae = cls(**spec["config"]).eval() + latent = sample_for( + spec, "decoder", 512, 512, torch.bfloat16, "meta", frames=17 + ) + pixels = sample_for( + spec, "encoder", 512, 512, torch.bfloat16, "meta", frames=17 + ) + params = sum(p.numel() for p in vae.parameters()) + print( + f" {name:<18} OK {params / 1e6:>7.1f}M params " + f"latent {tuple(latent.shape)} input {tuple(pixels.shape)}" + ) + except Exception as error: + failures += 1 + print(f" {name:<18} FAIL {type(error).__name__}: {error}") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/distvae/__version__.py b/distvae/__version__.py index 02bbad0..d8a21cd 100644 --- a/distvae/__version__.py +++ b/distvae/__version__.py @@ -1 +1 @@ -__version__ = "0.0.0beta5" +__version__ = "0.0.0beta9" diff --git a/distvae/models/layers/asymmetric_zero_pad_conv2d.py b/distvae/models/layers/asymmetric_zero_pad_conv2d.py new file mode 100644 index 0000000..d886fde --- /dev/null +++ b/distvae/models/layers/asymmetric_zero_pad_conv2d.py @@ -0,0 +1,207 @@ +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn import functional as F + +from distvae.models.layers.conv_mixin import PatchConvMixin +from distvae.models.layers.conv_utils import ( + chunk_bounds, + get_world_size_and_rank, +) +from distvae.utils import ParallelContext, normalize_patch_dim + + +Size2 = Union[int, Tuple[int, int]] +Size4 = Union[int, Tuple[int, int, int, int]] + + +class AsymmetricZeroPadConv2d(nn.Conv2d, PatchConvMixin): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Size2 = 3, + stride: Size2 = 2, + dilation: Size2 = 1, + groups: int = 1, + bias: bool = True, + device=None, + dtype=None, + reversed_zero_padding: Size4 = 0, + block_size: Union[int, Tuple[int, int]] = 0, + parallel_context: ParallelContext = None, + ) -> None: + if isinstance(dilation, int): + assert dilation == 1, "dilation is not supported in AsymmetricZeroPadConv2d" + else: + for value in dilation: + assert value == 1, ( + "dilation is not supported in AsymmetricZeroPadConv2d" + ) + if not isinstance(parallel_context, ParallelContext): + raise TypeError("AsymmetricZeroPadConv2d requires a ParallelContext") + if isinstance(reversed_zero_padding, int): + reversed_zero_padding = ( + reversed_zero_padding, + reversed_zero_padding, + reversed_zero_padding, + reversed_zero_padding, + ) + elif isinstance(reversed_zero_padding, tuple): + assert len(reversed_zero_padding) == 4, ( + "reversed_zero_padding must be a tuple of 4 integers" + ) + else: + raise ValueError( + f"Unsupported reversed_zero_padding: {type(reversed_zero_padding)}" + ) + if ( + reversed_zero_padding[0] != 0 + or reversed_zero_padding[1] != 1 + or reversed_zero_padding[2] != 0 + or reversed_zero_padding[3] != 1 + ): + raise ValueError( + f"Unsupported reversed_zero_padding: {reversed_zero_padding}" + ) + if ( + isinstance(kernel_size, int) + and kernel_size != 3 + or isinstance(kernel_size, tuple) + and (kernel_size[0] != 3 or kernel_size[1] != 3) + ): + raise ValueError(f"Unsupported kernel_size: {kernel_size}") + if ( + isinstance(stride, int) + and stride != 2 + or isinstance(stride, tuple) + and (stride[0] != 2 or stride[1] != 2) + ): + raise ValueError(f"Unsupported stride: {stride}") + + self.reversed_zero_padding = reversed_zero_padding + self.block_size = block_size + self.parallel_context = parallel_context + self.patch_dim = normalize_patch_dim( + parallel_context.patch_dim, 4, spatial_only=True + ) + self.halo_buffer = {} + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + 0, + dilation, + groups, + bias, + "zeros", + device, + dtype, + ) + + def _patch_ndim(self) -> int: + """Return 4 for 2D (N, C, H, W).""" + return 4 + + def _conv_forward( + self, input: Tensor, weight: Tensor, bias: Optional[Tensor] + ) -> Tensor: + group_world_size, rank_in_group = get_world_size_and_rank( + self.parallel_context + ) + + reversed_zero_padding = tuple(self.reversed_zero_padding) + patch_dim = input.ndim + normalize_patch_dim(self.patch_dim, input.ndim) + # The pad-then-stride-2 arithmetic below assumes each band halves cleanly. Bands are cut + # in multiples of what the whole encoder narrows by, so they are still even here. + assert input.shape[patch_dim] % 2 == 0, ( + "input.shape[patch_dim] must be even" + ) + + if group_world_size == 1: + return F.conv2d( + F.pad( + input, + reversed_zero_padding, + mode="constant", + value=0, + ), + weight, + bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + ( + input, + patch_dim, + _patch_size, + _halo_width, + _kernel_size_patch_dim, + _padding_patch_dim, + _stride_patch_dim, + _global_start, + group_world_size, + rank_in_group, + ) = self._multi_rank_metadata_and_halo(input, self.halo_buffer) + + if rank_in_group == 0: + padding = list(reversed_zero_padding) + padding[2 * (2 - patch_dim + 1) + 1] = 0 + elif rank_in_group == group_world_size - 1: + padding = list(reversed_zero_padding) + padding[2 * (2 - patch_dim + 1)] = 0 + else: + padding = list(reversed_zero_padding) + padding[2 * (2 - patch_dim + 1)] = 0 + padding[2 * (2 - patch_dim + 1) + 1] = 0 + input = F.pad(input, tuple(padding), mode="constant", value=0) + + _, _, height, width = input.shape + if self._use_direct_path(input): + return F.conv2d( + input, + weight, + bias, + self.stride, + (0, 0), + self.dilation, + self.groups, + ) + + block_h, block_w = ( + (self.block_size, self.block_size) + if isinstance(self.block_size, int) + else self.block_size + ) + kernel_h, kernel_w = self.kernel_size + stride_h, stride_w = self.stride + rows = chunk_bounds(height, block_h, kernel_h, stride_h) + columns = chunk_bounds(width, block_w, kernel_w, stride_w) + + return torch.cat( + [ + torch.cat( + [ + F.conv2d( + input[:, :, top:bottom, left:right], + weight, + bias, + self.stride, + 0, + self.dilation, + self.groups, + ) + for left, right in columns + ], + dim=-1, + ) + for top, bottom in rows + ], + dim=-2, + ) diff --git a/distvae/models/layers/conv2d.py b/distvae/models/layers/conv2d.py index 8cadcfb..e16a8eb 100644 --- a/distvae/models/layers/conv2d.py +++ b/distvae/models/layers/conv2d.py @@ -4,16 +4,21 @@ import torch.nn as nn from torch import Tensor from torch.nn import functional as F -from torch.nn.modules.utils import _pair -from torch.nn.common_types import _size_2_t from distvae.models.layers.conv_utils import ( get_world_size_and_rank, - correct_end, - correct_start, + chunk_bounds, build_crop_slice, ) from distvae.models.layers.conv_mixin import PatchConvMixin +from distvae.utils import ParallelContext, normalize_patch_dim + + +Size2 = Union[int, Tuple[int, int]] + + +def _pair(value: Size2) -> Tuple[int, int]: + return value if isinstance(value, tuple) else (value, value) class PatchConv2d(nn.Conv2d, PatchConvMixin): @@ -21,18 +26,17 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: _size_2_t, - stride: _size_2_t = 1, - padding: Union[str, _size_2_t] = 0, - dilation: _size_2_t = 1, + kernel_size: Size2, + stride: Size2 = 1, + padding: Union[str, Size2] = 0, + dilation: Size2 = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', # TODO: refine this type device=None, dtype=None, block_size: Union[int, Tuple[int, int]] = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ) -> None: if isinstance(dilation, int): @@ -40,12 +44,13 @@ def __init__( else: for i in dilation: assert i == 1, "dilation is not supported in PatchConv2d" - assert patch_dim in (-2, -1, 2, 3), ( - "PatchConv2d patch_dim must be H (-2 or 2) or W (-1 or 3)" - ) + if not isinstance(parallel_context, ParallelContext): + raise TypeError("PatchConv2d requires a ParallelContext") self.block_size = block_size - self.patch_dim = patch_dim - self.use_uniform_patch = use_uniform_patch + self.parallel_context = parallel_context + self.patch_dim = normalize_patch_dim( + parallel_context.patch_dim, 4, spatial_only=True + ) self.halo_buffer = {} super().__init__( in_channels, out_channels, kernel_size, stride, padding, dilation, @@ -57,7 +62,7 @@ def _patch_ndim(self) -> int: def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): bs, channels, h, w = input.shape - group_world_size, global_rank, rank_in_group, local_rank = get_world_size_and_rank() + group_world_size, rank_in_group = get_world_size_and_rank(self.parallel_context) if (group_world_size == 1): if self.padding_mode != 'zeros': @@ -68,6 +73,7 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): self.padding, self.dilation, self.groups) else: + self._check_padding_mode(group_world_size) ( input, patch_dim, @@ -76,11 +82,10 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): kernel_size_patch_dim, padding_patch_dim, stride_patch_dim, - patch_index, + global_start, group_world_size, rank_in_group, - stride_shift, - ) = self._multi_rank_metadata_and_halo(input, self.use_uniform_patch, self.halo_buffer) + ) = self._multi_rank_metadata_and_halo(input, self.halo_buffer) conv_res: Tensor padding = self._adjust_padding_for_patch( self._reversed_padding_repeated_twice, @@ -108,28 +113,21 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): _pair(0), self.dilation, self.groups) # Always apply cropping when halos are present to remove halo regions from output - # This prevents rank boundary artifacts for all convolution configurations + # This prevents rank boundary artifacts for all convolution configurations. + # build_crop_slice also recognises the output that is already patch-sized, which + # is what the branches above that pad only the outer edges produce: there the + # halo stands in for the padding those branches dropped, so nothing is left over + # to crop and cropping anyway would eat into the patch itself. if halo_width[0] > 0 or halo_width[1] > 0: - if stride_patch_dim > 1: - # For stride > 1, use global position-based cropping - global_start = patch_index[rank_in_group] - crop_slice = build_crop_slice( - patch_dim, patch_size, halo_width, conv_res.shape[patch_dim], ndim=4, - global_start=global_start, - kernel_size=kernel_size_patch_dim, - padding=padding_patch_dim, - stride=stride_patch_dim, - input_halo_width=halo_width, - ) - conv_res = conv_res[tuple(crop_slice)].contiguous() - else: - # For stride=1, use simple halo-based cropping - crop_slice = 4 * [slice(None),] - if halo_width[1] == 0: - crop_slice[patch_dim] = slice(halo_width[0], None) - else: - crop_slice[patch_dim] = slice(halo_width[0], -halo_width[1]) - conv_res = conv_res[tuple(crop_slice)].contiguous() + crop_slice = build_crop_slice( + patch_dim, patch_size, halo_width, conv_res.shape[patch_dim], ndim=4, + global_start=global_start, + kernel_size=kernel_size_patch_dim, + padding=padding_patch_dim, + stride=stride_patch_dim, + input_halo_width=halo_width, + ) + conv_res = conv_res[tuple(crop_slice)].contiguous() return conv_res else: @@ -139,73 +137,30 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): input = F.pad(input, padding, mode="constant") _, _, h, w = input.shape - num_chunks_in_h = 0 - num_chunks_in_w = 0 - if isinstance(self.block_size, int): - num_chunks_in_h = (h + self.block_size - 1) // self.block_size - num_chunks_in_w = (w + self.block_size - 1) // self.block_size - elif isinstance(self.block_size, tuple): - num_chunks_in_h = (h + self.block_size[0] - 1) // self.block_size[0] - num_chunks_in_w = (w + self.block_size[1] - 1) // self.block_size[1] - unit_chunk_size_h = h // num_chunks_in_h - unit_chunk_size_w = w // num_chunks_in_w - if isinstance(self.kernel_size, int): - kernel_size_h, kernel_size_w = self.kernel_size, self.kernel_size - elif isinstance(self.kernel_size, tuple): - kernel_size_h, kernel_size_w = self.kernel_size - else: - raise ValueError( - f"kernel_size should be int or tuple, type:{type(self.kernel_size)}" - ) - - if isinstance(self.stride, int): - stride_h, stride_w = self.stride, self.stride - elif isinstance(self.stride, tuple): - stride_h, stride_w = self.stride - else: - raise ValueError( - f"stride should be int or tuple, type: {type(self.stride)}" - ) - - outputs = [] - for idx_h in range(num_chunks_in_h): - inner_output = [] - for idx_w in range(num_chunks_in_w): - start_w = idx_w * unit_chunk_size_w - start_h = idx_h * unit_chunk_size_h - end_w = (idx_w + 1) * unit_chunk_size_w - end_h = (idx_h + 1) * unit_chunk_size_h - if idx_w + 1 < num_chunks_in_w: - end_w = correct_end(end_w, kernel_size_w, stride_w) - else: - end_w = w - if idx_h + 1 < num_chunks_in_h: - end_h = correct_end(end_h, kernel_size_h, stride_h) - else: - end_h = h - - if idx_w > 0: - start_w = correct_start(start_w, stride_w) - if idx_h > 0: - start_h = correct_start(start_h, stride_h) - - inner_output.append( - F.conv2d( - input[:, :, start_h:end_h, start_w:end_w], - weight, - bias, - self.stride, - 0, - self.dilation, - self.groups, - ) + # nn.Conv2d normalises all three of these to pairs in its own __init__, so they + # are read as pairs rather than tested for which they are. + block_h, block_w = _pair(self.block_size) + kernel_h, kernel_w = _pair(self.kernel_size) + stride_h, stride_w = _pair(self.stride) + rows = chunk_bounds(h, block_h, kernel_h, stride_h) + columns = chunk_bounds(w, block_w, kernel_w, stride_w) + + outputs = torch.cat([ + torch.cat([ + F.conv2d( + input[:, :, top:bottom, left:right], + weight, + bias, + self.stride, + 0, + self.dilation, + self.groups, ) - outputs.append(torch.cat(inner_output, dim=-1)) - outputs = torch.cat(outputs, dim=-2) - # Get global position for precise output cropping when stride > 1 - global_start = patch_index[rank_in_group] - # Note: patch_size here is the LOCAL patch size (before halo exchange) - # but after stride_shift trimming + for left, right in columns + ], dim=-1) + for top, bottom in rows + ], dim=-2) + # patch_size here is this rank's own, read before the halo was exchanged. crop_slice = build_crop_slice( patch_dim, patch_size, halo_width, outputs.shape[patch_dim], ndim=4, global_start=global_start, diff --git a/distvae/models/layers/conv3d.py b/distvae/models/layers/conv3d.py index 4438cc8..f7456bf 100644 --- a/distvae/models/layers/conv3d.py +++ b/distvae/models/layers/conv3d.py @@ -1,10 +1,10 @@ -"""PatchConv3d: 5D convolution with patch-dim parallelism for distributed VAE. +"""PatchConv3d: 5D convolution with H/W patch parallelism for distributed VAE. When world size is 1, behaves as nn.Conv3d. When world size > 1, gathers patch -sizes, exchanges halos along the patch dimension (F, H, or W), then either runs a +sizes, exchanges halos along the patch dimension (H or W), then either runs a single conv and crops (direct path) or splits the padded input into overlapping chunks, convs each chunk, concatenates, and crops (chunked path). Supports -patch_dim in {-3, -2, -1, 2, 3, 4} for F, H, W. Dilation is not supported. +patch_dim in {-2, -1, 3, 4} for H and W. Dilation is not supported. """ from typing import Optional, Tuple, Union @@ -13,56 +13,66 @@ import torch.nn as nn from torch import Tensor from torch.nn import functional as F -from torch.nn.modules.utils import _triple -from torch.nn.common_types import _size_3_t from distvae.models.layers.conv_utils import ( get_world_size_and_rank, - correct_end, - correct_start, + chunk_bounds, build_crop_slice, ) from distvae.models.layers.conv_mixin import PatchConvMixin +from distvae.utils import ParallelContext, normalize_patch_dim + + +Size3 = Union[int, Tuple[int, int, int]] + + +def _triple(value: Size3) -> Tuple[int, int, int]: + return value if isinstance(value, tuple) else (value, value, value) class PatchConv3d(nn.Conv3d, PatchConvMixin): - """3D convolution with patch-dim parallelism; subclasses nn.Conv3d and PatchConvMixin. + """3D convolution with H/W patch parallelism. - patch_dim selects which spatial dimension is split across ranks (F=frame, H=height, - W=width). block_size controls when the chunked path is used: 0 or all spatial - sizes <= block_size => direct path (one conv + crop); otherwise chunked path. - Dilation must be 1. + ``patch_dim`` selects height or width for splitting across ranks. ``block_size`` + controls local convolution chunking across all three convolution dimensions: + an integer applies one limit to F, H, and W, while a tuple is ordered (F, H, W). + Zero, or all dimensions fitting their limits, selects the direct path. Dilation + must be 1. """ def __init__( self, in_channels: int, out_channels: int, - kernel_size: _size_3_t, - stride: _size_3_t = 1, - padding: Union[str, _size_3_t] = 0, - dilation: _size_3_t = 1, + kernel_size: Size3, + stride: Size3 = 1, + padding: Union[str, Size3] = 0, + dilation: Size3 = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', # TODO: refine this type device=None, dtype=None, block_size: Union[int, Tuple[int, int, int]] = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ) -> None: - """patch_dim: which spatial dim is split (F=-3/3, H=-2/2, W=-1/4). block_size: 0 => prefer direct path; int or (F,H,W) => chunked when any spatial > block_size.""" + """Initialize H/W sharding and optional local (F, H, W) chunk limits. + + ``patch_dim`` accepts H (-2 or 3) or W (-1 or 4). ``block_size`` is zero + for the direct path, an integer shared by F/H/W, or an (F, H, W) tuple. + """ if isinstance(dilation, int): assert dilation == 1, "dilation is not supported in PatchConv3d" else: for i in dilation: assert i == 1, "dilation is not supported in PatchConv3d" - assert patch_dim in (-3, -2, -1, 2, 3, 4), ( - "PatchConv3d patch_dim must be F (-3 or 3) or H (-2 or 2) or W (-1 or 4)" - ) + if not isinstance(parallel_context, ParallelContext): + raise TypeError("PatchConv3d requires a ParallelContext") self.block_size = block_size - self.patch_dim = patch_dim - self.use_uniform_patch = use_uniform_patch + self.parallel_context = parallel_context + self.patch_dim = normalize_patch_dim( + parallel_context.patch_dim, 5, spatial_only=True + ) self.halo_buffer = {} super().__init__( in_channels, out_channels, kernel_size, stride, padding, dilation, @@ -75,7 +85,7 @@ def _patch_ndim(self) -> int: def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): bs, channels, f, h, w = input.shape - group_world_size, global_rank, rank_in_group, local_rank = get_world_size_and_rank() + group_world_size, rank_in_group = get_world_size_and_rank(self.parallel_context) # Single rank: use standard F.conv3d (with optional padding_mode). if (group_world_size == 1): @@ -85,8 +95,9 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): _triple(0), self.dilation, self.groups) return F.conv3d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups) - # Multi-rank: get extended input and metadata from mixin (patch_index, halo_width, etc.), then choose direct or chunked path. + # Multi-rank: get extended input and metadata from mixin (halo_width, global_start, etc.), then choose direct or chunked path. else: + self._check_padding_mode(group_world_size) ( input, patch_dim, @@ -95,11 +106,10 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): kernel_size_patch_dim, padding_patch_dim, stride_patch_dim, - patch_index, + global_start, group_world_size, rank_in_group, - stride_shift, - ) = self._multi_rank_metadata_and_halo(input, self.use_uniform_patch, self.halo_buffer) + ) = self._multi_rank_metadata_and_halo(input, self.halo_buffer) conv_res: Tensor padding = self._adjust_padding_for_patch( self._reversed_padding_repeated_twice, @@ -129,28 +139,21 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): _triple(0), self.dilation, self.groups) # Always apply cropping when halos are present to remove halo regions from output - # This prevents rank boundary artifacts for all convolution configurations + # This prevents rank boundary artifacts for all convolution configurations. + # build_crop_slice also recognises the output that is already patch-sized, which + # is what the branches above that pad only the outer edges produce: there the + # halo stands in for the padding those branches dropped, so nothing is left over + # to crop and cropping anyway would eat into the patch itself. if halo_width[0] > 0 or halo_width[1] > 0: - if stride_patch_dim > 1: - # For stride > 1, use global position-based cropping - global_start = patch_index[rank_in_group] - crop_slice = build_crop_slice( - patch_dim, patch_size, halo_width, conv_res.shape[patch_dim], ndim=5, - global_start=global_start, - kernel_size=kernel_size_patch_dim, - padding=padding_patch_dim, - stride=stride_patch_dim, - input_halo_width=halo_width, - ) - conv_res = conv_res[tuple(crop_slice)].contiguous() - else: - # For stride=1, use simple halo-based cropping - crop_slice = [slice(None)] * 5 - if halo_width[1] == 0: - crop_slice[patch_dim] = slice(halo_width[0], None) - else: - crop_slice[patch_dim] = slice(halo_width[0], -halo_width[1]) - conv_res = conv_res[tuple(crop_slice)].contiguous() + crop_slice = build_crop_slice( + patch_dim, patch_size, halo_width, conv_res.shape[patch_dim], ndim=5, + global_start=global_start, + kernel_size=kernel_size_patch_dim, + padding=padding_patch_dim, + stride=stride_patch_dim, + input_halo_width=halo_width, + ) + conv_res = conv_res[tuple(crop_slice)].contiguous() return conv_res # Chunked path: pad input, split into overlapping chunks along F, H, W; conv each chunk with padding=0; concat outputs; crop to this rank's patch. @@ -161,74 +164,33 @@ def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): input = F.pad(input, padding, mode="constant") _, _, f, h, w = input.shape - if isinstance(self.block_size, int): - num_chunks_in_f = (f + self.block_size - 1) // self.block_size - num_chunks_in_h = (h + self.block_size - 1) // self.block_size - num_chunks_in_w = (w + self.block_size - 1) // self.block_size - else: - num_chunks_in_f = (f + self.block_size[0] - 1) // self.block_size[0] - num_chunks_in_h = (h + self.block_size[1] - 1) // self.block_size[1] - num_chunks_in_w = (w + self.block_size[2] - 1) // self.block_size[2] - unit_chunk_size_f = f // num_chunks_in_f - unit_chunk_size_h = h // num_chunks_in_h - unit_chunk_size_w = w // num_chunks_in_w - if isinstance(self.kernel_size, int): - kernel_size_f, kernel_size_h, kernel_size_w = self.kernel_size, self.kernel_size, self.kernel_size - else: - kernel_size_f, kernel_size_h, kernel_size_w = self.kernel_size - if isinstance(self.stride, int): - stride_f, stride_h, stride_w = self.stride, self.stride, self.stride - else: - stride_f, stride_h, stride_w = self.stride - - # Chunk boundaries aligned via correct_end/correct_start so conv outputs line up when concatenated. - outputs = [] - for idx_f in range(num_chunks_in_f): - outer_output = [] - for idx_h in range(num_chunks_in_h): - inner_output = [] - for idx_w in range(num_chunks_in_w): - start_f = idx_f * unit_chunk_size_f - start_w = idx_w * unit_chunk_size_w - start_h = idx_h * unit_chunk_size_h - end_f = (idx_f + 1) * unit_chunk_size_f - end_w = (idx_w + 1) * unit_chunk_size_w - end_h = (idx_h + 1) * unit_chunk_size_h - if idx_f + 1 < num_chunks_in_f: - end_f = correct_end(end_f, kernel_size_f, stride_f) - else: - end_f = f - if idx_w + 1 < num_chunks_in_w: - end_w = correct_end(end_w, kernel_size_w, stride_w) - else: - end_w = w - if idx_h + 1 < num_chunks_in_h: - end_h = correct_end(end_h, kernel_size_h, stride_h) - else: - end_h = h - if idx_f > 0: - start_f = correct_start(start_f, stride_f) - if idx_w > 0: - start_w = correct_start(start_w, stride_w) - if idx_h > 0: - start_h = correct_start(start_h, stride_h) - - inner_output.append( - F.conv3d( - input[:, :, start_f:end_f, start_h:end_h, start_w:end_w], - weight, - bias, - self.stride, - 0, - self.dilation, - self.groups, - ) + # nn.Conv3d normalises all three of these to triples in its own __init__, so they + # are read as triples rather than tested for which they are. + block_f, block_h, block_w = _triple(self.block_size) + kernel_f, kernel_h, kernel_w = _triple(self.kernel_size) + stride_f, stride_h, stride_w = _triple(self.stride) + frames = chunk_bounds(f, block_f, kernel_f, stride_f) + rows = chunk_bounds(h, block_h, kernel_h, stride_h) + columns = chunk_bounds(w, block_w, kernel_w, stride_w) + + outputs = torch.cat([ + torch.cat([ + torch.cat([ + F.conv3d( + input[:, :, first:last, top:bottom, left:right], + weight, + bias, + self.stride, + 0, + self.dilation, + self.groups, ) - outer_output.append(torch.cat(inner_output, dim=-1)) - outputs.append(torch.cat(outer_output, dim=-2)) - outputs = torch.cat(outputs, dim=-3) - # Get global position for precise output cropping when stride > 1 - global_start = patch_index[rank_in_group] + for left, right in columns + ], dim=-1) + for top, bottom in rows + ], dim=-2) + for first, last in frames + ], dim=-3) crop_slice = build_crop_slice( patch_dim, patch_size, halo_width, outputs.shape[patch_dim], ndim=5, global_start=global_start, diff --git a/distvae/models/layers/conv_mixin.py b/distvae/models/layers/conv_mixin.py index 1c32604..276b363 100644 --- a/distvae/models/layers/conv_mixin.py +++ b/distvae/models/layers/conv_mixin.py @@ -8,11 +8,12 @@ import torch.distributed as dist from torch import Tensor -from distvae.utils import DistributedEnv +from distvae.utils import normalize_patch_dim from distvae.models.layers.conv_utils import ( get_world_size_and_rank, calc_patch_index, calc_halo_width, + calc_halo_width_unit_stride, calc_bottom_halo_width, calc_top_halo_width, exchange_halo, @@ -29,7 +30,7 @@ class PatchConvMixin: Methods: _patch_ndim (return 4 or 5); _adjust_padding_for_patch (delegate to conv_utils); _use_direct_path (True if single rank or all spatial sizes <= block_size); - _multi_rank_metadata_and_halo (all_gather patch sizes, compute halo, exchange, return extended input + metadata). + _multi_rank_metadata_and_halo (compute halo, exchange it, return extended input + metadata). """ def _patch_ndim(self) -> int: @@ -42,6 +43,20 @@ def _adjust_padding_for_patch(self, padding, rank, world_size, patch_dim: int = padding, rank, world_size, patch_dim, ndim=self._patch_ndim() ) + def _check_padding_mode(self, group_world_size: int) -> None: + """Refuse a padding mode whose values a halo exchange cannot supply. + + Zero, replicate, and reflect padding need only local values or rows from neighboring + ranks. Circular padding reads the opposite image edge, which may belong to a + non-neighboring rank and cannot be supplied by the halo exchange. + """ + if group_world_size > 1 and self.padding_mode == "circular": + raise NotImplementedError( + f"{type(self).__name__} cannot shard a convolution padded circularly: its " + f"padding wraps to the opposite edge of the image, which is not on a " + f"neighbouring rank. Use a single rank for this VAE, or tile it instead." + ) + def _use_direct_path(self, input: Tensor) -> bool: """Return True if we can run a single conv and crop (no chunking). @@ -59,43 +74,26 @@ def _use_direct_path(self, input: Tensor) -> bool: spatial_sizes[i] <= block_size[i] for i in range(len(spatial_sizes)) ) - def _uniform_patch_index(self, t: torch.Tensor, group_world_size: int): - """Calculate the patch index for a uniform patch. - - Args: - patch_dim_size: The size of the patch dimension - group_world_size: The world size of the group - - Returns: - The patch index - """ - patch_dim = self.patch_dim if self.patch_dim >= 0 else t.ndim + self.patch_dim - patch_list = [ - torch.tensor( - [t.shape[patch_dim]], - dtype=torch.int64, - device=t.device - ) for _ in range(group_world_size) - ] - return calc_patch_index(patch_list) - def _multi_rank_metadata_and_halo( self, input: Tensor, - use_uniform_patch: bool = False, halo_buffer: dict = None ): - """All_gather patch sizes, compute patch_index and halo_width, exchange halos; return extended input and metadata. - - All-gathers each rank's patch size along the patch dimension, builds - patch_index (cumulative boundaries), computes halo_width for this rank and - prev_bottom_halo_width/next_top_halo_width for send sizes, exchanges halos - with neighbors via exchange_halo. Returns (input, patch_dim, patch_size, - halo_width, kernel_size_patch_dim, padding_patch_dim, stride_patch_dim, - patch_index, group_world_size, rank_in_group). + """Exchange this rank's halo and return the input extended with neighboring rows. + + A strided conv all-gathers each rank's patch size to build the cumulative boundaries + its halo widths and its output cropping both turn on. A unit-stride conv derives the + same widths from its kernel and skips the gather; it also has no use for the + boundaries, so it reports global_start as None. + + Returns (input, patch_dim, patch_size, halo_width, kernel_size_patch_dim, + padding_patch_dim, stride_patch_dim, global_start, group_world_size, rank_in_group). """ - group_world_size, global_rank, rank_in_group, local_rank = get_world_size_and_rank() - patch_dim = self.patch_dim if self.patch_dim >= 0 else input.ndim + self.patch_dim + context = getattr(self, "parallel_context", None) + group_world_size, rank_in_group = get_world_size_and_rank(context) + patch_dim = input.ndim + normalize_patch_dim( + self.patch_dim, input.ndim, spatial_only=True + ) patch_size = input.shape[patch_dim] spatial_idx = patch_dim - 2 kernel_size_patch_dim = ( @@ -113,17 +111,25 @@ def _multi_rank_metadata_and_halo( if isinstance(self.stride, tuple) else self.stride ) - if use_uniform_patch: - if halo_buffer is None: - patch_index = self._uniform_patch_index(input, group_world_size) - else: - key = ("patch_index", input.shape[patch_dim], torch.int64, input.device) - if key in halo_buffer: - patch_index = halo_buffer[key] - else: - patch_index = self._uniform_patch_index(input, group_world_size) - halo_buffer[key] = patch_index + prev_bottom_halo_width: int = 0 + next_top_halo_width: int = 0 + if stride_patch_dim == 1: + # At unit stride the halo depends on the kernel alone, so no rank has to be told + # where the others' patches begin and the gather below can be skipped. A rank one + # along is neither first nor last from this rank's point of view, which is why the + # widths it wants are the plain kernel halves. + patch_index = None + halo_width = calc_halo_width_unit_stride( + rank_in_group, group_world_size, kernel_size_patch_dim + ) + if rank_in_group != 0: + prev_bottom_halo_width = kernel_size_patch_dim // 2 + if rank_in_group != group_world_size - 1: + next_top_halo_width = (kernel_size_patch_dim - 1) // 2 else: + # Patchify cuts bands that differ in size wherever the row count does not divide by + # the rank count, and a strided conv's halo turns on where in the global stride grid + # a patch starts, so a rank cannot read this off its own patch and has to be told. patch_list = [ torch.zeros(1, dtype=torch.int64, device=input.device) for _ in range(group_world_size) @@ -135,36 +141,39 @@ def _multi_rank_metadata_and_halo( dtype=torch.int64, device=input.device, ), - group=DistributedEnv.get_vae_group(), + group=context.group, ) patch_index = calc_patch_index(patch_list) - halo_width = calc_halo_width( - rank_in_group, - patch_index, - kernel_size_patch_dim, - padding_patch_dim, - stride_patch_dim, - ) - prev_bottom_halo_width: int = 0 - next_top_halo_width: int = 0 - if rank_in_group != 0: - prev_bottom_halo_width = calc_bottom_halo_width( - rank_in_group - 1, + halo_width = calc_halo_width( + rank_in_group, patch_index, kernel_size_patch_dim, padding_patch_dim, stride_patch_dim, ) - if rank_in_group != group_world_size - 1: - next_top_halo_width = calc_top_halo_width( - rank_in_group + 1, - patch_index, - kernel_size_patch_dim, - padding_patch_dim, - stride_patch_dim, - ) - next_top_halo_width = max(0, next_top_halo_width) + if rank_in_group != 0: + prev_bottom_halo_width = calc_bottom_halo_width( + rank_in_group - 1, + patch_index, + kernel_size_patch_dim, + padding_patch_dim, + stride_patch_dim, + ) + if rank_in_group != group_world_size - 1: + next_top_halo_width = calc_top_halo_width( + rank_in_group + 1, + patch_index, + kernel_size_patch_dim, + padding_patch_dim, + stride_patch_dim, + ) + next_top_halo_width = max(0, next_top_halo_width) if self._patch_ndim() == 4: + # Backstop, not the guard. Bands differ by a unit, so this can be true on one rank and + # false on its neighbour, and a rank that stops here stops on its way into the + # exchange below - leaving the others waiting on rows that will not come. Patchify + # refuses the same case up front, where every rank works it out from the same numbers + # and they all refuse together. Anything reaching here came in already split. assert halo_width[0] <= patch_size and halo_width[1] <= patch_size, ( "halo width is larger than the patch dimension of input tensor" ) @@ -176,24 +185,19 @@ def _multi_rank_metadata_and_halo( halo_width, prev_bottom_halo_width, next_top_halo_width, - group_world_size, - rank_in_group, + context, halo_buffer, ) - # Stride alignment: when stride > 1, we need to align input to global stride grid - # to ensure output indices match across ranks (prevents border artifacts) - stride_shift = 0 - if halo_width[0] > 0 and stride_patch_dim > 1: - global_start = patch_index[rank_in_group] - shift = (global_start - halo_width[0] + padding_patch_dim) % stride_patch_dim - if shift != 0: - stride_shift = shift - # Trim `shift` pixels from the top to align to stride grid - trim_slice = [slice(None)] * input.ndim - trim_slice[patch_dim] = slice(shift, None) - input = input[tuple(trim_slice)] - halo_width = (max(0, halo_width[0] - shift), halo_width[1]) + # Where this rank's patch begins in the whole image. Only a strided conv needs it, and + # only a strided conv paid to find it out, so at unit stride there is nothing to report. + global_start = None if patch_index is None else patch_index[rank_in_group] + + # A block trimming the input back onto the global stride grid used to sit here. It never + # trimmed anything: the top halo is defined as the distance from the patch start back to + # the last output step before it, so start - halo + padding is that step's position, which + # is a whole number of strides by construction and leaves nothing to shift by. What it + # reported was therefore always zero, and both callers unpacked it and never read it. return ( input, patch_dim, @@ -202,8 +206,7 @@ def _multi_rank_metadata_and_halo( kernel_size_patch_dim, padding_patch_dim, stride_patch_dim, - patch_index, + global_start, group_world_size, rank_in_group, - stride_shift, ) diff --git a/distvae/models/layers/conv_utils.py b/distvae/models/layers/conv_utils.py index 368d931..ea2f2de 100644 --- a/distvae/models/layers/conv_utils.py +++ b/distvae/models/layers/conv_utils.py @@ -13,20 +13,18 @@ import torch.distributed as dist from torch import Tensor -from distvae.utils import DistributedEnv +from distvae.utils import ParallelContext -def get_world_size_and_rank(): - """Return distributed group and rank info from DistributedEnv. +def get_world_size_and_rank(parallel_context: ParallelContext): + """Return rank metadata captured by an immutable parallel context. Returns: - Tuple of (group_world_size, global_rank, rank_in_group, local_rank). + Tuple of (group_world_size, rank_in_group). """ - group_world_size = DistributedEnv.get_group_world_size() - global_rank = DistributedEnv.get_global_rank() - rank_in_group = DistributedEnv.get_rank_in_vae_group() - local_rank = DistributedEnv.get_local_rank() - return group_world_size, global_rank, rank_in_group, local_rank + if not isinstance(parallel_context, ParallelContext): + raise TypeError("patch convolution requires a ParallelContext") + return parallel_context.world_size, parallel_context.rank def calc_patch_index(patch_list: List[Tensor]): @@ -114,8 +112,7 @@ def calc_halo_width(rank, height_index, kernel_size, padding=0, stride=1): The halo is the region used for convolution but not included in this rank's output. The first rank forces top to 0; the last rank (world_size - 1, inferred - from len(height_index) - 1 or DistributedEnv.get_group_world_size()) forces - bottom to 0. + from len(height_index) - 1) forces bottom to 0. Returns: Tuple (top_halo_width, bottom_halo_width) in patch-dim elements. @@ -126,11 +123,36 @@ def calc_halo_width(rank, height_index, kernel_size, padding=0, stride=1): ] if rank == 0: halo_width[0] = 0 - elif rank == DistributedEnv.get_group_world_size() - 1: + elif rank == len(height_index) - 2: halo_width[1] = 0 return tuple(halo_width) +def calc_halo_width_unit_stride(rank, world_size, kernel_size): + """Compute (top, bottom) halo widths for a stride-1 conv, asking no other rank anything. + + Under unit stride every term that mentions where a patch sits cancels out of + calc_top_halo_width and calc_bottom_halo_width, and the halo comes down to the kernel: + a rank needs the (kernel_size - 1) // 2 rows above it that its first output row reads, + and kernel_size // 2 rows below it for its last. Padding cancels too, because it shifts + the output grid and the patch start by the same amount. + + That matters because the alternative is an all_gather of one integer per convolution, + and on a Wan decode those gathers are half of every collective the model makes. + + Args: + rank: This rank's index within the VAE group. + world_size: Size of the VAE group. + kernel_size: Kernel size along the patch dimension. + + Returns: + Tuple (top_halo_width, bottom_halo_width), matching calc_halo_width at stride 1. + """ + top = 0 if rank == 0 else (kernel_size - 1) // 2 + bottom = 0 if rank == world_size - 1 else kernel_size // 2 + return top, bottom + + def correct_end(end, kernel_size, stride): """Adjust chunk end so conv output at that boundary aligns with stride. @@ -157,6 +179,46 @@ def correct_start(start, stride): return ((start + stride - 1) // stride) * stride +def chunk_bounds(extent, block, kernel_size, stride) -> List[Tuple[int, int]]: + """Where each chunk of one axis begins and ends, so that convolving them separately and + concatenating gives what convolving the whole axis would + + The chunks divide the axis evenly and are then grown at each cut: every chunk but the last + runs on to the last input its final output position reads, and every chunk but the first + begins at the first input the next output step needs. Both are the same two corrections the + 2D and the 3D path each used to write out per axis, which is two of them and five copies. + + Never more chunks than leave every one of them at least a kernel long. Asking for more cuts + an axis into pieces no convolution can be run on at all, which raises out of torch rather + than costing accuracy - and the count that does it is not obvious from the block size: a + frame axis of 4 padded to 6, chunked by 4 at stride 2, ends on a chunk of 2. Fewer chunks + only means a larger intermediate, which is the knob's own currency, and the output is the + same however the axis is divided. + + Args: + extent: Length of the axis, after any padding. + block: Requested chunk length; the count is the ceiling of extent over it. + kernel_size, stride: Conv parameters along this axis. + + Returns: + List of (start, end) input-space bounds, one per chunk. + """ + chunks = (extent + block - 1) // block + # A chunk spans its share of the axis less what the corrections at either end move it, which + # is a stride at most, so a share of kernel + stride - 1 is what keeps the shortest of them + # at a kernel. At stride 1 that is the kernel itself. + chunks = max(1, min(chunks, extent // (kernel_size + stride - 1))) + unit = extent // chunks + bounds = [] + for idx in range(chunks): + start = idx * unit + bounds.append(( + correct_start(start, stride) if idx else start, + extent if idx + 1 == chunks else correct_end(start + unit, kernel_size, stride), + )) + return bounds + + def build_crop_slice( patch_dim: int, patch_size: int, @@ -283,8 +345,7 @@ def exchange_halo( halo_width: tuple, prev_bottom_halo_width: int, next_top_halo_width: int, - group_world_size: int, - rank_in_group: int, + parallel_context: ParallelContext, halo_buffer: dict = None, ) -> Tensor: """Exchange halo regions with previous and next ranks; return extended local tensor. @@ -292,10 +353,13 @@ def exchange_halo( Send: bottom halo to next rank (size next_top_halo_width), top halo to prev (size prev_bottom_halo_width). Receive: top halo from prev (halo_width[0]), bottom halo from next (halo_width[1]). Concatenate [top_halo_recv, input, - bottom_halo_recv] along patch_dim and return. Uses non-blocking isend and - blocking recv, then wait on sends. + bottom_halo_recv] along patch_dim and return. All four are issued as one + batch and waited on together. Args: + patch_index: Cumulative patch boundaries, or None when the caller never gathered them. + They only serve the bounds checks here, which are skipped in that case rather than + paid for with a collective. halo_buffer: Optional dict to cache/reuse comms buffers for better performance """ ndim = input.ndim @@ -304,77 +368,59 @@ def exchange_halo( indices_start = [slice(None)] * ndim indices_start[patch_dim] = slice(0, prev_bottom_halo_width) - to_next = None - to_prev = None + if not isinstance(parallel_context, ParallelContext): + raise TypeError("exchange_halo requires a ParallelContext") + vae_group = parallel_context.group + rank_in_group = parallel_context.rank + ops = [] top_halo_recv = None bottom_halo_recv = None global_rank_of_next = None global_rank_of_prev = None - if next_top_halo_width > 0: - global_rank_of_next = DistributedEnv.get_global_rank_from_group_rank(rank_in_group + 1) - bottom_halo_send = input[tuple(indices_end)].contiguous() - to_next = dist.isend( - bottom_halo_send, - global_rank_of_next, - group=DistributedEnv.get_vae_group(), - ) - if halo_width[0] > 0: - assert patch_index[rank_in_group] - halo_width[0] >= patch_index[rank_in_group - 1], ( - "width of top halo region is larger than the input tensor of prev rank" - ) + def recv_buffer(name: str, width: int) -> Tensor: recv_shape = list(input.shape) - recv_shape[patch_dim] = halo_width[0] + recv_shape[patch_dim] = width if halo_buffer is None: - top_halo_recv = torch.empty( + return torch.empty(recv_shape, dtype=input.dtype, device=input.device) + key = (name, tuple(recv_shape), input.dtype, input.device) + if key not in halo_buffer: + halo_buffer[key] = torch.empty( recv_shape, dtype=input.dtype, device=input.device ) - else: - key = ("top_recv", tuple(recv_shape), input.dtype, input.device) - if key in halo_buffer: - top_halo_recv = halo_buffer[key] - else: - top_halo_recv = torch.empty( - recv_shape, dtype=input.dtype, device=input.device - ) - halo_buffer[key] = top_halo_recv - global_rank_of_prev = DistributedEnv.get_global_rank_from_group_rank(rank_in_group - 1) - dist.recv(top_halo_recv, global_rank_of_prev, group=DistributedEnv.get_vae_group()) + return halo_buffer[key] + + if next_top_halo_width > 0: + global_rank_of_next = parallel_context.global_rank(rank_in_group + 1) + bottom_halo_send = input[tuple(indices_end)].contiguous() + ops.append(dist.P2POp(dist.isend, bottom_halo_send, global_rank_of_next, group=vae_group)) + if halo_width[0] > 0: + assert patch_index is None or ( + patch_index[rank_in_group] - halo_width[0] >= patch_index[rank_in_group - 1] + ), "width of top halo region is larger than the input tensor of prev rank" + top_halo_recv = recv_buffer("top_recv", halo_width[0]) + global_rank_of_prev = parallel_context.global_rank(rank_in_group - 1) + ops.append(dist.P2POp(dist.irecv, top_halo_recv, global_rank_of_prev, group=vae_group)) if prev_bottom_halo_width > 0: top_halo_send = input[tuple(indices_start)].contiguous() if global_rank_of_prev is None: - global_rank_of_prev = DistributedEnv.get_global_rank_from_group_rank(rank_in_group - 1) - to_prev = dist.isend( - top_halo_send, - global_rank_of_prev, - group=DistributedEnv.get_vae_group(), - ) + global_rank_of_prev = parallel_context.global_rank(rank_in_group - 1) + ops.append(dist.P2POp(dist.isend, top_halo_send, global_rank_of_prev, group=vae_group)) if halo_width[1] > 0: - assert patch_index[rank_in_group + 1] + halo_width[1] <= patch_index[rank_in_group + 2], ( - "width of bottom halo region is larger than the input tensor of next rank" - ) - recv_shape = list(input.shape) - recv_shape[patch_dim] = halo_width[1] - if halo_buffer is None: - bottom_halo_recv = torch.empty( - recv_shape, dtype=input.dtype, device=input.device - ) - else: - key = ("bottom_recv", tuple(recv_shape), input.dtype, input.device) - if key in halo_buffer: - bottom_halo_recv = halo_buffer[key] - else: - bottom_halo_recv = torch.empty( - recv_shape, dtype=input.dtype, device=input.device - ) - halo_buffer[key] = bottom_halo_recv + assert patch_index is None or ( + patch_index[rank_in_group + 1] + halo_width[1] <= patch_index[rank_in_group + 2] + ), "width of bottom halo region is larger than the input tensor of next rank" + bottom_halo_recv = recv_buffer("bottom_recv", halo_width[1]) if global_rank_of_next is None: - global_rank_of_next = DistributedEnv.get_global_rank_from_group_rank(rank_in_group + 1) - dist.recv( - bottom_halo_recv, - global_rank_of_next, - group=DistributedEnv.get_vae_group(), - ) + global_rank_of_next = parallel_context.global_rank(rank_in_group + 1) + ops.append(dist.P2POp(dist.irecv, bottom_halo_recv, global_rank_of_next, group=vae_group)) + + # Batching exposes both independent directions at once and lets NCCL reuse the wider group's + # communicator instead of constructing one for each point-to-point operation. + if ops: + for work in dist.batch_isend_irecv(ops): + work.wait() + if halo_width[0] < 0: trim_slice = [slice(None)] * ndim trim_slice[patch_dim] = slice(-halo_width[0], None) @@ -383,9 +429,5 @@ def exchange_halo( input = torch.cat([top_halo_recv, input], dim=patch_dim) if bottom_halo_recv is not None: input = torch.cat([input, bottom_halo_recv], dim=patch_dim) - if to_next is not None: - to_next.wait() - if to_prev is not None: - to_prev.wait() return input diff --git a/distvae/models/layers/normalization.py b/distvae/models/layers/normalization.py index cd2a2dc..2579217 100644 --- a/distvae/models/layers/normalization.py +++ b/distvae/models/layers/normalization.py @@ -1,129 +1,23 @@ import math -import numbers -from typing import Optional - import torch import torch.nn as nn import torch.distributed as dist from torch import Tensor -from diffusers.models.activations import get_activation -from distvae.utils import DistributedEnv - - -class PatchAdaGroupNorm(nn.Module): - def __init__( - self, - embedding_dim: int, - out_dim: int, - num_groups: int, - act_fn: Optional[str] = None, - eps: float = 1e-5, - patch_dim: int = -2, - ): - super().__init__() - self.patch_dim = patch_dim - self.num_groups = num_groups - self.eps = eps - - if act_fn is None: - self.act = None - else: - self.act = get_activation(act_fn) - - self.linear = nn.Linear(embedding_dim, out_dim * 2) - - def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: - patch_dim = self.patch_dim if self.patch_dim >= 0 else x.ndim + self.patch_dim - - if self.act: - emb = self.act(emb) - emb = self.linear(emb) - # Support 4D (N,C,H,W) and 5D (N,C,F,H,W); patch dim is always first spatial (index 2). - emb = emb[:, :, None, None, None] if x.ndim == 5 else emb[:, :, None, None] - scale, shift = emb.chunk(2, dim=1) - - world_size = DistributedEnv.get_world_size() - patch_size_list = [torch.empty([1], dtype=torch.int64) for _ in range(world_size)] - dist.all_gather( - patch_size_list, - torch.tensor([x.shape[patch_dim]], dtype=torch.int64), - group=DistributedEnv.get_vae_group() - ) - patch_size = torch.tensor(patch_size_list).sum().item() - - channels_per_group = x.shape[1] // self.num_groups - nelements = ( - channels_per_group * - math.prod(x.shape[2: patch_dim]) * patch_size * math.prod(x.shape[patch_dim + 1:]) - ) - partial_sum = x.sum_to_size(x.shape[0], self.num_groups) - partial_sum_list = [ - torch.empty([x.shape[0], self.num_groups], dtype=x.dtype, device=x.device) - for _ in range(world_size) - ] - dist.all_gather(partial_sum_list, partial_sum, group=DistributedEnv.get_vae_group()) - group_sum = torch.tensor(partial_sum_list, device=x.device).sum(dim=0) - E = group_sum / nelements - partial_var = ((x - E) ** 2).sum_to_size(x.shape[0], self.num_groups) - partial_var_list = [ - torch.empty([x.shape[0], self.num_groups], dtype=x.dtype, device=x.device) - for _ in range(world_size) - ] - dist.all_gather(partial_var_list, partial_var, group=DistributedEnv.get_vae_group()) - group_var = torch.tensor(partial_var_list, device=x.device).sum(dim=0) - var = group_var / nelements - - x = (x - E) / torch.sqrt(var + self.eps) - x = x * (1 + scale) + shift - - return x +from distvae.utils import ParallelContext, normalize_patch_dim class PatchGroupNorm(nn.GroupNorm): - r"""Applies Group Normalization over a mini-batch of inputs. - - This layer implements the operation as described in - the paper `Group Normalization `__ - - .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta - - The input channels are separated into :attr:`num_groups` groups, each containing - ``num_channels / num_groups`` channels. :attr:`num_channels` must be divisible by - :attr:`num_groups`. The mean and standard-deviation are calculated - separately over the each group. :math:`\gamma` and :math:`\beta` are learnable - per-channel affine transform parameter vectors of size :attr:`num_channels` if - :attr:`affine` is ``True``. - The standard-deviation is calculated via the biased estimator, equivalent to - `torch.var(input, unbiased=False)`. + """Inference-only GroupNorm over spatial shards held by a VAE process group. - This layer uses statistics computed from input data in both training and - evaluation modes. + Each rank supplies its local, potentially uneven shard. Group sums and squared + deviations are reduced across ``parallel_context.group``, so every rank normalizes + its shard with the statistics of the complete unsharded tensor. The biased variance + estimator and affine transform match :class:`torch.nn.GroupNorm`. - Args: - num_groups (int): number of groups to separate the channels into - num_channels (int): number of channels expected in input - eps: a value added to the denominator for numerical stability. Default: 1e-5 - affine: a boolean value that when set to ``True``, this module - has learnable per-channel affine parameters initialized to ones (for weights) - and zeros (for biases). Default: ``True``. - - Shape: - - Input: :math:`(N, C, *)` where :math:`C=\text{num\_channels}` - - Output: :math:`(N, C, *)` (same shape as input) - - Examples:: - - >>> input = torch.randn(20, 6, 10, 10) - >>> # Separate 6 channels into 3 groups - >>> m = nn.GroupNorm(3, 6) - >>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm) - >>> m = nn.GroupNorm(6, 6) - >>> # Put all 6 channels into a single group (equivalent with LayerNorm) - >>> m = nn.GroupNorm(1, 6) - >>> # Activating the module - >>> output = m(input) + ``parallel_context`` identifies the process group and spatial patch dimension. + ``forward`` runs without gradient tracking and returns a tensor with the same local + shape as its input. """ def __init__( @@ -134,9 +28,12 @@ def __init__( affine: bool = True, device=None, dtype=None, - patch_dim: int = -2, + parallel_context: ParallelContext = None, ) -> None: - self.patch_dim = patch_dim + if not isinstance(parallel_context, ParallelContext): + raise TypeError("PatchGroupNorm requires a ParallelContext") + self.parallel_context = parallel_context + self.patch_dim = parallel_context.patch_dim super().__init__( num_groups=num_groups, num_channels=num_channels, @@ -149,77 +46,57 @@ def __init__( def forward(self, x: Tensor) -> Tensor: ndim = x.ndim shape = x.shape - patch_dim = self.patch_dim if self.patch_dim >= 0 else ndim + self.patch_dim - + patch_dim = ndim + normalize_patch_dim(self.patch_dim, ndim, spatial_only=True) + vae_group = self.parallel_context.group + group_world_size = self.parallel_context.world_size x = x.detach() - # Support 4D (N,C,H,W) and 5D (N,C,F,H,W); patch dim is first spatial (index 2). - patch_size = torch.tensor(shape[patch_dim], dtype=torch.int64, device=x.device) - dist.all_reduce(patch_size, group=DistributedEnv.get_vae_group()) channels_per_group = shape[1] // self.num_groups + + x = x.view(shape[0], self.num_groups, -1, *shape[2: ]) + reduced = tuple(range(2, x.ndim)) + # [bs, num_groups, 1, 1, 1] for 4D input, one more 1 for 5D. + per_group = (shape[0], self.num_groups, *([1] * (x.ndim - 2))) + + # This rank's row count travels with its group sums. Both are sums over the same group of + # ranks, so combining them changes no arithmetic, and sent alone the row count costs a + # whole round trip to move one number. Float32 holds a row count exactly either way. + # Support 4D (N,C,H,W) and 5D (N,C,F,H,W); patch dim is first spatial (index 2). + totals = torch.empty( + 1 + shape[0] * self.num_groups, dtype=torch.float32, device=x.device + ) + totals[0] = shape[patch_dim] + totals[1:] = x.sum(dim=reduced, dtype=torch.float32).flatten() + # Summing one rank's numbers across one rank returns them unchanged, so on a single-rank + # group both reductions here are the identity. They are still real collectives, though: + # a decode of a VAE with twenty-five group norms issued fifty of them to talk to nobody. + if group_world_size > 1: + dist.all_reduce(totals, group=vae_group) + + patch_size = totals[0] nelements = ( channels_per_group * math.prod(shape[2: patch_dim]) * patch_size * math.prod(shape[patch_dim + 1: ]) ) - nelements_rank = (nelements // patch_size) * shape[patch_dim] - - x = x.view(shape[0], self.num_groups, -1, *shape[2: ]) - group_sum = x.mean(dim=tuple(range(2, x.ndim)), dtype=torch.float32) - group_sum = group_sum * nelements_rank - dist.all_reduce(group_sum, group=DistributedEnv.get_vae_group()) - # shape: [bs, num_groups, 1, 1, 1] or [bs, num_groups, 1, 1, 1, 1] - E = (group_sum / nelements)[:, :, None, None, None].to(x.dtype) - group_var_sum = torch.empty( - (x.shape[0], self.num_groups), - dtype=torch.float32, - device=x.device - ) - torch.var(x, dim=tuple(range(2, x.ndim)), out=group_var_sum) - group_var_sum = group_var_sum * nelements_rank - dist.all_reduce(group_var_sum, group=DistributedEnv.get_vae_group()) - var = (group_var_sum / nelements)[:, :, None, None, None].to(x.dtype) - if ndim == 5: - E = E.unsqueeze(-1) - var = var.unsqueeze(-1) + group_sum = totals[1:].view(shape[0], self.num_groups) + E = (group_sum / nelements).view(per_group).to(x.dtype) + + # Squared about the mean of the whole group rather than this rank's share of it. A rank + # holding a brighter patch has a mean of its own, and deviations measured from that one + # leave out how far the patch itself sits from the middle, so the summed variance comes + # out short of the variance the unsharded norm computes. + group_square_sum = ((x - E) ** 2).sum(dim=reduced, dtype=torch.float32) + if group_world_size > 1: + dist.all_reduce(group_square_sum, group=vae_group) + # Divided by the count, not one less than it, which is the estimator nn.GroupNorm uses. + var = (group_square_sum / nelements).view(per_group).to(x.dtype) x = (x - E) / torch.sqrt(var + self.eps) - x = x.view(x.shape[0], -1, *shape[2: ]) + x = x.view(shape[0], -1, *shape[2: ]) if self.weight is not None and self.bias is not None: weight = self.weight.view(1, -1, *([1] * (ndim - 2))) bias = self.bias.view(1, -1, *([1] * (ndim - 2))) x = x * weight + bias - return x - - -class RMSNorm(nn.Module): - def __init__(self, dim, eps: float, elementwise_affine: bool = True): - super().__init__() - - self.eps = eps - - if isinstance(dim, numbers.Integral): - dim = (dim,) - - self.dim = torch.Size(dim) - - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim)) - else: - self.weight = None - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - - if self.weight is not None: - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - hidden_states = hidden_states * self.weight - else: - hidden_states = hidden_states.to(input_dtype) - - return hidden_states \ No newline at end of file + return x \ No newline at end of file diff --git a/distvae/models/layers/wan/__init__.py b/distvae/models/layers/wan/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/distvae/models/layers/wan/zeropadconv2d.py b/distvae/models/layers/wan/zeropadconv2d.py deleted file mode 100644 index 4a5a405..0000000 --- a/distvae/models/layers/wan/zeropadconv2d.py +++ /dev/null @@ -1,224 +0,0 @@ -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn -from torch import Tensor -from torch.nn import functional as F -from torch.nn.modules.utils import _pair -from torch.nn.common_types import _size_2_t,_size_4_t - -from distvae.models.layers.conv_utils import ( - get_world_size_and_rank, - correct_end, - correct_start, -) -from distvae.models.layers.conv_mixin import PatchConvMixin - - -class WanZeroPadConv2d(nn.Conv2d, PatchConvMixin): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: _size_2_t = 3, - stride: _size_2_t = 2, - dilation: _size_2_t = 1, - groups: int = 1, - bias: bool = True, - device=None, - dtype=None, - reversed_zero_padding: Union[int, _size_4_t] = 0, - block_size: Union[int, Tuple[int, int, int]] = 0, - patch_dim: int = -2, - use_uniform_patch: bool = True, - ) -> None: - if not use_uniform_patch: - raise NotImplementedError("WanZeroPadConv2d not implemented for use_uniform_patch=False") - if isinstance(dilation, int): - assert dilation == 1, "dilation is not supported in WanZeroPadConv2d" - else: - for i in dilation: - assert i == 1, "dilation is not supported in WanZeroPadConv2d" - assert patch_dim in (-2, -1), ( - "WanZeroPadConv2d patch_dim must be H (-2) or W (-1)" - ) - if isinstance(reversed_zero_padding, int): - reversed_zero_padding = ( - reversed_zero_padding, reversed_zero_padding, reversed_zero_padding, reversed_zero_padding - ) - elif isinstance(reversed_zero_padding, tuple): - assert len(reversed_zero_padding) == 4, "reversed_zero_padding must be a tuple of 4 integers" - else: - raise ValueError(f"Unsupported reversed_zero_padding: {type(reversed_zero_padding)}") - if ( - reversed_zero_padding[0] != 0 or - reversed_zero_padding[1] != 1 or - reversed_zero_padding[2] != 0 or - reversed_zero_padding[3] != 1 - ): - raise ValueError(f"Unsupported reversed_zero_padding: {reversed_zero_padding}") - # Validate kernel_size and stride - if ( - isinstance(kernel_size, int) and kernel_size != 3 or - isinstance(kernel_size, tuple) and (kernel_size[0] != 3 or kernel_size[1] != 3) - ): - raise ValueError(f"Unsupported kernel_size: {kernel_size}") - if ( - isinstance(stride, int) and stride != 2 or - isinstance(stride, tuple) and (stride[0] != 2 or stride[1] != 2) - ): - raise ValueError(f"Unsupported stride: {stride}") - - self.reversed_zero_padding = reversed_zero_padding - self.block_size = block_size - self.patch_dim = patch_dim - self.use_uniform_patch = use_uniform_patch - self.halo_buffer = {} - super().__init__( - in_channels, - out_channels, - kernel_size, - stride, - 0, - dilation, - groups, - bias, - "zeros", - device, - dtype - ) - - def _patch_ndim(self) -> int: - """Return 4 for 2D (N, C, H, W).""" - return 4 - - def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): - group_world_size, global_rank, rank_in_group, local_rank = get_world_size_and_rank() - - bs, channels, h, w = input.shape - reversed_zero_padding = tuple(self.reversed_zero_padding) - - patch_dim = self.patch_dim if self.patch_dim >= 0 else input.ndim + self.patch_dim - assert input.shape[patch_dim] % 2 == 0, "input.shape[patch_dim] must be even" - - # Single rank: use standard F.conv2d - if group_world_size == 1: - output = F.conv2d( - F.pad( - input, - reversed_zero_padding, - mode="constant", - value=0 - ), - weight, - bias, - self.stride, - self.padding, - self.dilation, - self.groups - ) - - return output - # Multi-rank: get extended input and metadata from mixin (patch_index, halo_width, etc.), then choose direct or chunked path. - else: - # Metadata and halo exchange - ( - input, - patch_dim, - patch_size, - halo_width, - kernel_size_patch_dim, - padding_patch_dim, - stride_patch_dim, - patch_index, - group_world_size, - rank_in_group, - _, - ) = self._multi_rank_metadata_and_halo(input, self.use_uniform_patch, self.halo_buffer) - - # ZeroPad2d - if rank_in_group == 0: - padding = list(reversed_zero_padding) - padding[2 * (2 - patch_dim + 1) + 1] = 0 - elif rank_in_group == group_world_size - 1: - padding = list(reversed_zero_padding) - padding[2 * (2 - patch_dim + 1)] = 0 - else: - padding = list(reversed_zero_padding) - padding[2 * (2 - patch_dim + 1)] = 0 - padding[2 * (2 - patch_dim + 1) + 1] = 0 - input = F.pad(input, tuple(padding), mode="constant", value=0) - - # Conv2d - output: Tensor - _, channels, h, w = input.shape - # Direct path: one conv over the extended (halo-padded) input - if self._use_direct_path(input): - output = F.conv2d( - input, - weight, - bias, - self.stride, - _pair(0), - self.dilation, - self.groups - ) - - return output - # Chunked path: pad input, split into overlapping chunks along F, H, W; conv each chunk with padding=0; concat outputs - else: - _, channels, h, w = input.shape - if isinstance(self.block_size, int): - num_chunks_in_h = (h + self.block_size - 1) // self.block_size - num_chunks_in_w = (w + self.block_size - 1) // self.block_size - else: - num_chunks_in_h = (h + self.block_size[0] - 1) // self.block_size[0] - num_chunks_in_w = (w + self.block_size[1] - 1) // self.block_size[1] - unit_chunk_size_h = h // num_chunks_in_h - unit_chunk_size_w = w // num_chunks_in_w - if isinstance(self.kernel_size, int): - kernel_size_h, kernel_size_w = self.kernel_size, self.kernel_size - else: - kernel_size_h, kernel_size_w = self.kernel_size - if isinstance(self.stride, int): - stride_h, stride_w = self.stride, self.stride - else: - stride_h, stride_w = self.stride - - # Chunk boundaries aligned via correct_end/correct_start so conv outputs line up when concatenated. - output = [] - for idx_h in range(num_chunks_in_h): - inner_output = [] - for idx_w in range(num_chunks_in_w): - start_w = idx_w * unit_chunk_size_w - start_h = idx_h * unit_chunk_size_h - end_w = (idx_w + 1) * unit_chunk_size_w - end_h = (idx_h + 1) * unit_chunk_size_h - if idx_w + 1 < num_chunks_in_w: - end_w = correct_end(end_w, kernel_size_w, stride_w) - else: - end_w = w - if idx_h + 1 < num_chunks_in_h: - end_h = correct_end(end_h, kernel_size_h, stride_h) - else: - end_h = h - if idx_w > 0: - start_w = correct_start(start_w, stride_w) - if idx_h > 0: - start_h = correct_start(start_h, stride_h) - - inner_output.append( - F.conv2d( - input[:, :, start_h:end_h, start_w:end_w], - weight, - bias, - self.stride, - 0, - self.dilation, - self.groups, - ) - ) - output.append(torch.cat(inner_output, dim=-1)) - output = torch.cat(output, dim=2) - - return output diff --git a/distvae/models/resnet.py b/distvae/models/resnet.py deleted file mode 100644 index 7298fd4..0000000 --- a/distvae/models/resnet.py +++ /dev/null @@ -1,370 +0,0 @@ -from typing import Optional - -import torch -import torch.distributed -import torch.nn as nn - -from diffusers.utils import deprecate -from diffusers.models.resnet import ResnetBlock2D -from diffusers.models.activations import get_activation -from diffusers.models.normalization import AdaGroupNorm -from diffusers.models.attention_processor import SpatialNorm -from diffusers.models.downsampling import Downsample2D -from diffusers.models.upsampling import Upsample2D - -from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter - -# class ResnetBlockCondNorm2D(nn.Module): -# r""" -# A Resnet block that use normalization layer that incorporate conditioning information. - -# Parameters: -# in_channels (`int`): The number of channels in the input. -# out_channels (`int`, *optional*, default to be `None`): -# The number of output channels for the first conv2d layer. If None, same as `in_channels`. -# dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. -# temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. -# groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. -# groups_out (`int`, *optional*, default to None): -# The number of groups to use for the second normalization layer. if set to None, same as `groups`. -# eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. -# non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. -# time_embedding_norm (`str`, *optional*, default to `"ada_group"` ): -# The normalization layer for time embedding `temb`. Currently only support "ada_group" or "spatial". -# kernel (`torch.FloatTensor`, optional, default to None): FIR filter, see -# [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. -# output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. -# use_in_shortcut (`bool`, *optional*, default to `True`): -# If `True`, add a 1x1 nn.conv2d layer for skip-connection. -# up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. -# down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. -# conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the -# `conv_shortcut` output. -# conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. -# If None, same as `out_channels`. -# """ - -# def __init__( -# self, -# *, -# in_channels: int, -# out_channels: Optional[int] = None, -# conv_shortcut: bool = False, -# dropout: float = 0.0, -# temb_channels: int = 512, -# groups: int = 32, -# groups_out: Optional[int] = None, -# eps: float = 1e-6, -# non_linearity: str = "swish", -# time_embedding_norm: str = "ada_group", # ada_group, spatial -# output_scale_factor: float = 1.0, -# use_in_shortcut: Optional[bool] = None, -# up: bool = False, -# down: bool = False, -# conv_shortcut_bias: bool = True, -# conv_2d_out_channels: Optional[int] = None, -# ): -# super().__init__() -# self.in_channels = in_channels -# out_channels = in_channels if out_channels is None else out_channels -# self.out_channels = out_channels -# self.use_conv_shortcut = conv_shortcut -# self.up = up -# self.down = down -# self.output_scale_factor = output_scale_factor -# self.time_embedding_norm = time_embedding_norm - -# conv_cls = nn.Conv2d - -# if groups_out is None: -# groups_out = groups - -# if self.time_embedding_norm == "ada_group": # ada_group -# self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps) -# elif self.time_embedding_norm == "spatial": -# self.norm1 = SpatialNorm(in_channels, temb_channels) -# else: -# raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}") - -# self.conv1 = conv_cls(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - -# if self.time_embedding_norm == "ada_group": # ada_group -# self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps) -# elif self.time_embedding_norm == "spatial": # spatial -# self.norm2 = SpatialNorm(out_channels, temb_channels) -# else: -# raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}") - -# self.dropout = torch.nn.Dropout(dropout) - -# conv_2d_out_channels = conv_2d_out_channels or out_channels -# self.conv2 = conv_cls(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1) - -# self.nonlinearity = get_activation(non_linearity) - -# self.upsample = self.downsample = None -# if self.up: -# self.upsample = Upsample2D(in_channels, use_conv=False) -# elif self.down: -# self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op") - -# self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut - -# self.conv_shortcut = None -# if self.use_in_shortcut: -# self.conv_shortcut = conv_cls( -# in_channels, -# conv_2d_out_channels, -# kernel_size=1, -# stride=1, -# padding=0, -# bias=conv_shortcut_bias, -# ) - -# def forward(self, input_tensor: torch.FloatTensor, temb: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: -# if len(args) > 0 or kwargs.get("scale", None) is not None: -# deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." -# deprecate("scale", "1.0.0", deprecation_message) - -# hidden_states = input_tensor - -# hidden_states = self.norm1(hidden_states, temb) - -# hidden_states = self.nonlinearity(hidden_states) - -# if self.upsample is not None: -# # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 -# if hidden_states.shape[0] >= 64: -# input_tensor = input_tensor.contiguous() -# hidden_states = hidden_states.contiguous() -# input_tensor = self.upsample(input_tensor) -# hidden_states = self.upsample(hidden_states) - -# elif self.downsample is not None: -# input_tensor = self.downsample(input_tensor) -# hidden_states = self.downsample(hidden_states) - -# hidden_states = self.conv1(hidden_states) - -# hidden_states = self.norm2(hidden_states, temb) - -# hidden_states = self.nonlinearity(hidden_states) - -# hidden_states = self.dropout(hidden_states) -# hidden_states = self.conv2(hidden_states) - -# if self.conv_shortcut is not None: -# input_tensor = self.conv_shortcut(input_tensor) - -# output_tensor = (input_tensor + hidden_states) / self.output_scale_factor - -# return output_tensor - - -class PatchResnetBlock2D(nn.Module): - r""" - A Resnet block. - - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv2d layer. If None, same as `in_channels`. - dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. - temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. - groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. - groups_out (`int`, *optional*, default to None): - The number of groups to use for the second normalization layer. if set to None, same as `groups`. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. - non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. - time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config. - By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" - for a stronger conditioning with scale and shift. - kernel (`torch.FloatTensor`, optional, default to None): FIR filter, see - [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. - output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. - use_in_shortcut (`bool`, *optional*, default to `True`): - If `True`, add a 1x1 nn.conv2d layer for skip-connection. - up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. - down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. - conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the - `conv_shortcut` output. - conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. - If None, same as `out_channels`. - """ - - def __init__( - self, - *, - in_channels: int, - out_channels: Optional[int] = None, - conv_shortcut: bool = False, - dropout: float = 0.0, - temb_channels: int = 512, - groups: int = 32, - groups_out: Optional[int] = None, - pre_norm: bool = True, - eps: float = 1e-6, - non_linearity: str = "swish", - skip_time_act: bool = False, - time_embedding_norm: str = "default", # default, scale_shift, - kernel: Optional[torch.FloatTensor] = None, - output_scale_factor: float = 1.0, - use_in_shortcut: Optional[bool] = None, - up: bool = False, - down: bool = False, - conv_shortcut_bias: bool = True, - conv_2d_out_channels: Optional[int] = None, - conv_block_size = 0, - ): - assert temb_channels is None, "temb_channels is not supported currently." - assert up is False, "Upsampling is not supported currently." - assert down is False, "Downsampling is not supported currently." - - super().__init__() - if time_embedding_norm == "ada_group": - raise ValueError( - "This class cannot be used with `time_embedding_norm==ada_group`, please use `ResnetBlockCondNorm2D` instead", - ) - if time_embedding_norm == "spatial": - raise ValueError( - "This class cannot be used with `time_embedding_norm==spatial`, please use `ResnetBlockCondNorm2D` instead", - ) - - self.pre_norm = True - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - self.up = up - self.down = down - self.output_scale_factor = output_scale_factor - self.time_embedding_norm = time_embedding_norm - self.skip_time_act = skip_time_act - - linear_cls = nn.Linear - conv_cls = nn.Conv2d - - if groups_out is None: - groups_out = groups - - self.norm1 = GroupNormAdapter(torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)) - - self.conv1 = Conv2dAdapter(conv_cls(in_channels, out_channels, kernel_size=3, stride=1, padding=1), block_size=conv_block_size) - - #TODO: Add support for temb_channels - assert temb_channels is None, "temb_channels is not supported currently." - self.time_emb_proj = None - # if temb_channels is not None: - # if self.time_embedding_norm == "default": - # self.time_emb_proj = linear_cls(temb_channels, out_channels) - # elif self.time_embedding_norm == "scale_shift": - # self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels) - # else: - # raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ") - # else: - # self.time_emb_proj = None - - self.norm2 = GroupNormAdapter(torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)) - - self.dropout = torch.nn.Dropout(dropout) - conv_2d_out_channels = conv_2d_out_channels or out_channels - self.conv2 = Conv2dAdapter(conv_cls(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1), block_size=conv_block_size) - - self.nonlinearity = get_activation(non_linearity) - - self.upsample = self.downsample = None - - #TODO: Add support for upsample and downsample - assert self.up is False, "Upsampling is not supported currently." - assert self.down is False, "Downsampling is not supported currently." - # if self.up: - # if kernel == "fir": - # fir_kernel = (1, 3, 3, 1) - # self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel) - # elif kernel == "sde_vp": - # self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest") - # else: - # self.upsample = Upsample2D(in_channels, use_conv=False) - # elif self.down: - # if kernel == "fir": - # fir_kernel = (1, 3, 3, 1) - # self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel) - # elif kernel == "sde_vp": - # self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2) - # else: - # self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op") - - self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut - - self.conv_shortcut = None - if self.use_in_shortcut: - self.conv_shortcut = Conv2dAdapter( - conv_cls( - in_channels, - conv_2d_out_channels, - kernel_size=1, - stride=1, - padding=0, - bias=conv_shortcut_bias, - ), - block_size=conv_block_size - ) - - def forward(self, input_tensor: torch.FloatTensor, temb: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: - #TODO: Add support for temb - assert temb is None, "temb is not supported currently." - - if len(args) > 0 or kwargs.get("scale", None) is not None: - deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." - deprecate("scale", "1.0.0", deprecation_message) - - hidden_states = input_tensor - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 - if hidden_states.shape[0] >= 64: - input_tensor = input_tensor.contiguous() - hidden_states = hidden_states.contiguous() - input_tensor = self.upsample(input_tensor) - hidden_states = self.upsample(hidden_states) - elif self.downsample is not None: - input_tensor = self.downsample(input_tensor) - hidden_states = self.downsample(hidden_states) - - hidden_states = self.conv1(hidden_states) - - if self.time_emb_proj is not None: - if not self.skip_time_act: - temb = self.nonlinearity(temb) - temb = self.time_emb_proj(temb)[:, :, None, None] - - if self.time_embedding_norm == "default": - if temb is not None: - hidden_states = hidden_states + temb - hidden_states = self.norm2(hidden_states) - elif self.time_embedding_norm == "scale_shift": - if temb is None: - raise ValueError( - f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}" - ) - time_scale, time_shift = torch.chunk(temb, 2, dim=1) - hidden_states = self.norm2(hidden_states) - hidden_states = hidden_states * (1 + time_scale) + time_shift - else: - hidden_states = self.norm2(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = (input_tensor + hidden_states) / self.output_scale_factor - - return output_tensor \ No newline at end of file diff --git a/distvae/models/unets/unet_2d_blocks.py b/distvae/models/unets/unet_2d_blocks.py deleted file mode 100644 index b3026d8..0000000 --- a/distvae/models/unets/unet_2d_blocks.py +++ /dev/null @@ -1,320 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from diffusers.utils import deprecate, is_torch_version, logging -from diffusers.utils.torch_utils import apply_freeu -from diffusers.models.activations import get_activation -from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0 -from diffusers.models.normalization import AdaGroupNorm -from diffusers.models.resnet import ( - Downsample2D, - FirDownsample2D, - FirUpsample2D, - KDownsample2D, - KUpsample2D, - ResnetBlock2D, - ResnetBlockCondNorm2D, - Upsample2D, -) -from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel -from diffusers.models.transformers.transformer_2d import Transformer2DModel -from diffusers.models.unets.unet_2d_blocks import ( - AttnSkipUpBlock2D, - AttnUpBlock2D, - AttnUpDecoderBlock2D, - CrossAttnUpBlock2D, - KCrossAttnUpBlock2D, - KUpBlock2D, - ResnetUpsampleBlock2D, - SimpleCrossAttnUpBlock2D, - SkipUpBlock2D, - UpBlock2D, - UpDecoderBlock2D -) - -from distvae.modules.adapters.resnet_adapters import ResnetBlock2DAdapter -from distvae.modules.adapters.upsampling_adapters import Upsample2DAdapter - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def get_up_block( - up_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - add_upsample: bool, - resnet_eps: float, - resnet_act_fn: str, - resolution_idx: Optional[int] = None, - transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = None, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - resnet_time_scale_shift: str = "default", - attention_type: str = "default", - resnet_skip_time_act: bool = False, - resnet_out_scale_factor: float = 1.0, - cross_attention_norm: Optional[str] = None, - attention_head_dim: Optional[int] = None, - upsample_type: Optional[str] = None, - dropout: float = 0.0, - conv_block_size = 0, -) -> nn.Module: - # If attn head dim is not defined, we default it to the number of heads - if attention_head_dim is None: - logger.warning( - f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." - ) - attention_head_dim = num_attention_heads - - up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type - if up_block_type == "UpBlock2D": - return UpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "ResnetUpsampleBlock2D": - return ResnetUpsampleBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - ) - elif up_block_type == "CrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnUpBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - ) - elif up_block_type == "SimpleCrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D") - return SimpleCrossAttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - ) - elif up_block_type == "AttnUpBlock2D": - if add_upsample is False: - upsample_type = None - else: - upsample_type = upsample_type or "conv" # default to 'conv' - - return AttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - upsample_type=upsample_type, - ) - elif up_block_type == "SkipUpBlock2D": - return SkipUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "AttnSkipUpBlock2D": - return AttnSkipUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "UpDecoderBlock2D": - return PatchUpDecoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - temb_channels=temb_channels, - conv_block_size=conv_block_size, - ) - elif up_block_type == "AttnUpDecoderBlock2D": - return AttnUpDecoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - temb_channels=temb_channels, - ) - elif up_block_type == "KUpBlock2D": - return KUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - ) - elif up_block_type == "KCrossAttnUpBlock2D": - return KCrossAttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - ) - - raise ValueError(f"{up_block_type} does not exist.") - - - -class PatchUpDecoderBlock2D(UpDecoderBlock2D): - def __init__( - self, - in_channels: int, - out_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - temb_channels: Optional[int] = None, - conv_block_size = 0, - ): - #TODO: Add support for spatial time embedding - assert resnet_time_scale_shift != "spatial", "'spatial' has not been supported for UpDecoderBlock2D yet." - super().__init__(in_channels, out_channels, resolution_idx, - dropout, num_layers, resnet_eps, resnet_time_scale_shift, - resnet_act_fn, resnet_groups, resnet_pre_norm, output_scale_factor, - add_upsample, temb_channels) - patched_resnet = [] - for resnet in self.resnets: - patched_resnet.append(ResnetBlock2DAdapter(resnet, conv_block_size=conv_block_size)) - self.resnets = nn.ModuleList(patched_resnet) - - if add_upsample: - patched_upsamplers = [] - for upsampler in self.upsamplers: - patched_upsamplers.append(Upsample2DAdapter(upsampler, conv_block_size=conv_block_size)) - self.upsamplers = nn.ModuleList(patched_upsamplers) diff --git a/distvae/models/upsampling.py b/distvae/models/upsampling.py deleted file mode 100644 index d832d62..0000000 --- a/distvae/models/upsampling.py +++ /dev/null @@ -1,53 +0,0 @@ -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusers.utils import deprecate -from diffusers.models.upsampling import Upsample2D - -from distvae.models.layers.normalization import RMSNorm -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter - - -class PatchUpsample2D(Upsample2D): - """A 2D upsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - use_conv_transpose (`bool`, default `False`): - option to use a convolution transpose. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - name (`str`, default `conv`): - name of the upsampling 2D layer. - """ - - def __init__( - self, - channels: int, - use_conv: bool = False, - use_conv_transpose: bool = False, - out_channels: Optional[int] = None, - name: str = "conv", - kernel_size: Optional[int] = None, - padding=1, - norm_type=None, - eps=None, - elementwise_affine=None, - bias=True, - interpolate=True, - conv_block_size = 0, - ): - assert norm_type is None, "norm_type has not been supported for PatchUpsample2D yat." - assert use_conv_transpose is False, "use_conv_transpose has not been supported for PatchUpsample2D yet." - super().__init__(channels, use_conv, use_conv_transpose, out_channels, name, - kernel_size, padding, norm_type, eps, elementwise_affine, - bias, interpolate) - if name == "conv": - self.conv = Conv2dAdapter(self.conv, block_size=conv_block_size) - else: - self.Conv2d_0 = Conv2dAdapter(self.Conv2d_0, block_size=conv_block_size) \ No newline at end of file diff --git a/distvae/models/vae.py b/distvae/models/vae.py deleted file mode 100644 index 99ae494..0000000 --- a/distvae/models/vae.py +++ /dev/null @@ -1,1001 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from dataclasses import dataclass -from typing import Optional, Tuple - -import numpy as np -import torch -import torch.distributed -import torch.nn as nn - -from diffusers.utils import BaseOutput, is_torch_version -from diffusers.utils.torch_utils import randn_tensor -from diffusers.models.activations import get_activation -from diffusers.models.attention_processor import SpatialNorm -from diffusers.models.unets.unet_2d_blocks import ( - AutoencoderTinyBlock, - UNetMidBlock2D, - get_down_block -) -from distvae.models.unets.unet_2d_blocks import ( - get_up_block, -) -from distvae.models.layers.conv2d import PatchConv2d -from distvae.models.layers.normalization import PatchGroupNorm -from distvae.modules.patch_utils import Patchify, DePatchify - - -@dataclass -class DecoderOutput(BaseOutput): - r""" - Output of decoding method. - - Args: - sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): - The decoded output sample from the last layer of the model. - """ - - sample: torch.FloatTensor - - -class Encoder(nn.Module): - r""" - The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. - - Args: - in_channels (`int`, *optional*, defaults to 3): - The number of input channels. - out_channels (`int`, *optional*, defaults to 3): - The number of output channels. - down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available - options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): - The number of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - norm_num_groups (`int`, *optional*, defaults to 32): - The number of groups for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - double_z (`bool`, *optional*, defaults to `True`): - Whether to double the number of output channels for the last block. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - double_z: bool = True, - mid_block_add_attention=True, - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[0], - kernel_size=3, - stride=1, - padding=1, - ) - - self.mid_block = None - self.down_blocks = nn.ModuleList([]) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=self.layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - add_downsample=not is_final_block, - resnet_eps=1e-6, - downsample_padding=0, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=None, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default", - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=None, - add_attention=mid_block_add_attention, - ) - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - - conv_out_channels = 2 * out_channels if double_z else out_channels - self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: - r"""The forward method of the `Encoder` class.""" - - sample = self.conv_in(sample) - - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - # down - if is_torch_version(">=", "1.11.0"): - for down_block in self.down_blocks: - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(down_block), sample, use_reentrant=False - ) - # middle - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.mid_block), sample, use_reentrant=False - ) - else: - for down_block in self.down_blocks: - sample = torch.utils.checkpoint.checkpoint(create_custom_forward(down_block), sample) - # middle - sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample) - - else: - # down - for down_block in self.down_blocks: - sample = down_block(sample) - - # middle - sample = self.mid_block(sample) - - # post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample - - -class PatchDecoder(nn.Module): - r""" - The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. - - Args: - in_channels (`int`, *optional*, defaults to 3): - The number of input channels. - out_channels (`int`, *optional*, defaults to 3): - The number of output channels. - up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): - The number of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - norm_num_groups (`int`, *optional*, defaults to 32): - The number of groups for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - norm_type (`str`, *optional*, defaults to `"group"`): - The normalization type to use. Can be either `"group"` or `"spatial"`. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - norm_type: str = "group", # group, spatial - mid_block_add_attention=True, - conv_block_size = 0, - ): - assert norm_type == "group", "Only group normalization is supported in PatchDecoder. Please use Decoder instead." - super().__init__() - for up_block in up_block_types: - assert up_block in ["UpDecoderBlock2D"], "Only UpDecoderBlock2D is supported in PatchDecoder. Please use Decoder instead." - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1, - padding=1, - ) - - self.mid_block = None - self.up_blocks = nn.ModuleList([]) - - temb_channels = in_channels if norm_type == "spatial" else None - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default" if norm_type == "group" else norm_type, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - add_attention=mid_block_add_attention, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - prev_output_channel=None, - add_upsample=not is_final_block, - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=temb_channels, - resnet_time_scale_shift=norm_type, - conv_block_size=conv_block_size - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # patchify - self.patch = Patchify() - # unpatchify - self.depatch = DePatchify() - - - # out - if norm_type == "spatial": - self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) - else: - self.conv_norm_out = PatchGroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = PatchConv2d(block_out_channels[0], out_channels, 3, padding=1, block_size=conv_block_size) - - self.gradient_checkpointing = False - - def forward( - self, - sample: torch.FloatTensor, - latent_embeds: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - r"""The forward method of the `Decoder` class.""" - - sample = self.conv_in(sample) - - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - if is_torch_version(">=", "1.11.0"): - # middle - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.mid_block), - sample, - latent_embeds, - use_reentrant=False, - ) - sample = sample.to(upscale_dtype) - sample = self.patch(sample) - # up - for up_block in self.up_blocks: - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(up_block), - sample, - latent_embeds, - use_reentrant=False, - ) - else: - # middle - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.mid_block), sample, latent_embeds - ) - sample = sample.to(upscale_dtype) - sample = self.patch(sample) - # up - for up_block in self.up_blocks: - sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds) - else: - # middle - sample = self.mid_block(sample, latent_embeds) - sample = sample.to(upscale_dtype) - # up - sample = self.patch(sample) - for up_block in self.up_blocks: - sample = up_block(sample, latent_embeds) - - # post-process - if latent_embeds is None: - sample = self.conv_norm_out(sample) - else: - sample = self.conv_norm_out(sample, latent_embeds) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - sample = self.depatch(sample) - - return sample - - -class UpSample(nn.Module): - r""" - The `UpSample` layer of a variational autoencoder that upsamples its input. - - Args: - in_channels (`int`, *optional*, defaults to 3): - The number of input channels. - out_channels (`int`, *optional*, defaults to 3): - The number of output channels. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - ) -> None: - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1) - - def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: - r"""The forward method of the `UpSample` class.""" - x = torch.relu(x) - x = self.deconv(x) - return x - - -class MaskConditionEncoder(nn.Module): - """ - used in AsymmetricAutoencoderKL - """ - - def __init__( - self, - in_ch: int, - out_ch: int = 192, - res_ch: int = 768, - stride: int = 16, - ) -> None: - super().__init__() - - channels = [] - while stride > 1: - stride = stride // 2 - in_ch_ = out_ch * 2 - if out_ch > res_ch: - out_ch = res_ch - if stride == 1: - in_ch_ = res_ch - channels.append((in_ch_, out_ch)) - out_ch *= 2 - - out_channels = [] - for _in_ch, _out_ch in channels: - out_channels.append(_out_ch) - out_channels.append(channels[-1][0]) - - layers = [] - in_ch_ = in_ch - for l in range(len(out_channels)): - out_ch_ = out_channels[l] - if l == 0 or l == 1: - layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=3, stride=1, padding=1)) - else: - layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=4, stride=2, padding=1)) - in_ch_ = out_ch_ - - self.layers = nn.Sequential(*layers) - - def forward(self, x: torch.FloatTensor, mask=None) -> torch.FloatTensor: - r"""The forward method of the `MaskConditionEncoder` class.""" - out = {} - for l in range(len(self.layers)): - layer = self.layers[l] - x = layer(x) - out[str(tuple(x.shape))] = x - x = torch.relu(x) - return out - - -class MaskConditionDecoder(nn.Module): - r"""The `MaskConditionDecoder` should be used in combination with [`AsymmetricAutoencoderKL`] to enhance the model's - decoder with a conditioner on the mask and masked image. - - Args: - in_channels (`int`, *optional*, defaults to 3): - The number of input channels. - out_channels (`int`, *optional*, defaults to 3): - The number of output channels. - up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): - The number of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - norm_num_groups (`int`, *optional*, defaults to 32): - The number of groups for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - norm_type (`str`, *optional*, defaults to `"group"`): - The normalization type to use. Can be either `"group"` or `"spatial"`. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - norm_type: str = "group", # group, spatial - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1, - padding=1, - ) - - self.mid_block = None - self.up_blocks = nn.ModuleList([]) - - temb_channels = in_channels if norm_type == "spatial" else None - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default" if norm_type == "group" else norm_type, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - prev_output_channel=None, - add_upsample=not is_final_block, - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=temb_channels, - resnet_time_scale_shift=norm_type, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # condition encoder - self.condition_encoder = MaskConditionEncoder( - in_ch=out_channels, - out_ch=block_out_channels[0], - res_ch=block_out_channels[-1], - ) - - # out - if norm_type == "spatial": - self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) - else: - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward( - self, - z: torch.FloatTensor, - image: Optional[torch.FloatTensor] = None, - mask: Optional[torch.FloatTensor] = None, - latent_embeds: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - r"""The forward method of the `MaskConditionDecoder` class.""" - sample = z - sample = self.conv_in(sample) - - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - if is_torch_version(">=", "1.11.0"): - # middle - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.mid_block), - sample, - latent_embeds, - use_reentrant=False, - ) - sample = sample.to(upscale_dtype) - - # condition encoder - if image is not None and mask is not None: - masked_image = (1 - mask) * image - im_x = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.condition_encoder), - masked_image, - mask, - use_reentrant=False, - ) - - # up - for up_block in self.up_blocks: - if image is not None and mask is not None: - sample_ = im_x[str(tuple(sample.shape))] - mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest") - sample = sample * mask_ + sample_ * (1 - mask_) - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(up_block), - sample, - latent_embeds, - use_reentrant=False, - ) - if image is not None and mask is not None: - sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask) - else: - # middle - sample = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.mid_block), sample, latent_embeds - ) - sample = sample.to(upscale_dtype) - - # condition encoder - if image is not None and mask is not None: - masked_image = (1 - mask) * image - im_x = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.condition_encoder), - masked_image, - mask, - ) - - # up - for up_block in self.up_blocks: - if image is not None and mask is not None: - sample_ = im_x[str(tuple(sample.shape))] - mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest") - sample = sample * mask_ + sample_ * (1 - mask_) - sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds) - if image is not None and mask is not None: - sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask) - else: - # middle - sample = self.mid_block(sample, latent_embeds) - sample = sample.to(upscale_dtype) - - # condition encoder - if image is not None and mask is not None: - masked_image = (1 - mask) * image - im_x = self.condition_encoder(masked_image, mask) - - # up - for up_block in self.up_blocks: - if image is not None and mask is not None: - sample_ = im_x[str(tuple(sample.shape))] - mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest") - sample = sample * mask_ + sample_ * (1 - mask_) - sample = up_block(sample, latent_embeds) - if image is not None and mask is not None: - sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask) - - # post-process - if latent_embeds is None: - sample = self.conv_norm_out(sample) - else: - sample = self.conv_norm_out(sample, latent_embeds) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample - - -class VectorQuantizer(nn.Module): - """ - Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly avoids costly matrix - multiplications and allows for post-hoc remapping of indices. - """ - - # NOTE: due to a bug the beta term was applied to the wrong term. for - # backwards compatibility we use the buggy version by default, but you can - # specify legacy=False to fix it. - def __init__( - self, - n_e: int, - vq_embed_dim: int, - beta: float, - remap=None, - unknown_index: str = "random", - sane_index_shape: bool = False, - legacy: bool = True, - ): - super().__init__() - self.n_e = n_e - self.vq_embed_dim = vq_embed_dim - self.beta = beta - self.legacy = legacy - - self.embedding = nn.Embedding(self.n_e, self.vq_embed_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.used: torch.Tensor - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed + 1 - print( - f"Remapping {self.n_e} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices." - ) - else: - self.re_embed = n_e - - self.sane_index_shape = sane_index_shape - - def remap_to_used(self, inds: torch.LongTensor) -> torch.LongTensor: - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - match = (inds[:, :, None] == used[None, None, ...]).long() - new = match.argmax(-1) - unknown = match.sum(2) < 1 - if self.unknown_index == "random": - new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds: torch.LongTensor) -> torch.LongTensor: - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds >= self.used.shape[0]] = 0 # simply set to zero - back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) - return back.reshape(ishape) - - def forward(self, z: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor, Tuple]: - # reshape z -> (batch, height, width, channel) and flatten - z = z.permute(0, 2, 3, 1).contiguous() - z_flattened = z.view(-1, self.vq_embed_dim) - - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - min_encoding_indices = torch.argmin(torch.cdist(z_flattened, self.embedding.weight), dim=1) - - z_q = self.embedding(min_encoding_indices).view(z.shape) - perplexity = None - min_encodings = None - - # compute loss for embedding - if not self.legacy: - loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2) - else: - loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2) - - # preserve gradients - z_q: torch.FloatTensor = z + (z_q - z).detach() - - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - if self.remap is not None: - min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis - min_encoding_indices = self.remap_to_used(min_encoding_indices) - min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten - - if self.sane_index_shape: - min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3]) - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - - def get_codebook_entry(self, indices: torch.LongTensor, shape: Tuple[int, ...]) -> torch.FloatTensor: - # shape specifying (batch, height, width, channel) - if self.remap is not None: - indices = indices.reshape(shape[0], -1) # add batch axis - indices = self.unmap_to_all(indices) - indices = indices.reshape(-1) # flatten again - - # get quantized latent vectors - z_q: torch.FloatTensor = self.embedding(indices) - - if shape is not None: - z_q = z_q.view(shape) - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q - - -class DiagonalGaussianDistribution(object): - def __init__(self, parameters: torch.Tensor, deterministic: bool = False): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, device=self.parameters.device, dtype=self.parameters.dtype - ) - - def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor: - # make sure sample is on the same device as the parameters and has same dtype - sample = randn_tensor( - self.mean.shape, - generator=generator, - device=self.parameters.device, - dtype=self.parameters.dtype, - ) - x = self.mean + self.std * sample - return x - - def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3], - ) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=[1, 2, 3], - ) - - def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: - if self.deterministic: - return torch.Tensor([0.0]) - logtwopi = np.log(2.0 * np.pi) - return 0.5 * torch.sum( - logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, - dim=dims, - ) - - def mode(self) -> torch.Tensor: - return self.mean - - -class EncoderTiny(nn.Module): - r""" - The `EncoderTiny` layer is a simpler version of the `Encoder` layer. - - Args: - in_channels (`int`): - The number of input channels. - out_channels (`int`): - The number of output channels. - num_blocks (`Tuple[int, ...]`): - Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to - use. - block_out_channels (`Tuple[int, ...]`): - The number of output channels for each block. - act_fn (`str`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], - act_fn: str, - ): - super().__init__() - - layers = [] - for i, num_block in enumerate(num_blocks): - num_channels = block_out_channels[i] - - if i == 0: - layers.append(nn.Conv2d(in_channels, num_channels, kernel_size=3, padding=1)) - else: - layers.append( - nn.Conv2d( - num_channels, - num_channels, - kernel_size=3, - padding=1, - stride=2, - bias=False, - ) - ) - - for _ in range(num_block): - layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) - - layers.append(nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, padding=1)) - - self.layers = nn.Sequential(*layers) - self.gradient_checkpointing = False - - def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: - r"""The forward method of the `EncoderTiny` class.""" - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - if is_torch_version(">=", "1.11.0"): - x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x, use_reentrant=False) - else: - x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x) - - else: - # scale image from [-1, 1] to [0, 1] to match TAESD convention - x = self.layers(x.add(1).div(2)) - - return x - - -class DecoderTiny(nn.Module): - r""" - The `DecoderTiny` layer is a simpler version of the `Decoder` layer. - - Args: - in_channels (`int`): - The number of input channels. - out_channels (`int`): - The number of output channels. - num_blocks (`Tuple[int, ...]`): - Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to - use. - block_out_channels (`Tuple[int, ...]`): - The number of output channels for each block. - upsampling_scaling_factor (`int`): - The scaling factor to use for upsampling. - act_fn (`str`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], - upsampling_scaling_factor: int, - act_fn: str, - ): - super().__init__() - - layers = [ - nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=1), - get_activation(act_fn), - ] - - for i, num_block in enumerate(num_blocks): - is_final_block = i == (len(num_blocks) - 1) - num_channels = block_out_channels[i] - - for _ in range(num_block): - layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) - - if not is_final_block: - layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor)) - - conv_out_channel = num_channels if not is_final_block else out_channels - layers.append( - nn.Conv2d( - num_channels, - conv_out_channel, - kernel_size=3, - padding=1, - bias=is_final_block, - ) - ) - - self.layers = nn.Sequential(*layers) - self.gradient_checkpointing = False - - def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: - r"""The forward method of the `DecoderTiny` class.""" - # Clamp. - x = torch.tanh(x / 3) * 3 - - if self.training and self.gradient_checkpointing: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - - return custom_forward - - if is_torch_version(">=", "1.11.0"): - x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x, use_reentrant=False) - else: - x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x) - - else: - x = self.layers(x) - - # scale image from [0, 1] to [-1, 1] to match diffusers convention - return x.mul(2).sub(1) diff --git a/distvae/modules/adapters/__init__.py b/distvae/modules/adapters/__init__.py index 00ed0f3..08defe4 100644 --- a/distvae/modules/adapters/__init__.py +++ b/distvae/modules/adapters/__init__.py @@ -1,31 +1,71 @@ -# Export downsampling adapters -from .downsampling_adapters import ( - WanResampleDownAdapter, - WanResidualDownBlockAdapter, -) +"""Public adapter exports, loaded only when requested.""" -# Export upsampling adapters -from .upsampling_adapters import ( - Upsample2DAdapter, - WanResampleAdapter, - WanResidualUpBlockAdapter, - WanUpBlockAdapter, -) +from importlib import import_module -# Export other adapters -from .midblock_adapters import WanMidBlockAdapter -from .resnet_adapters import WanResidualBlockAdapter -__all__ = [ - # Downsampling +_DOWNSAMPLING = ( + "Downsample2DAdapter", + "HunyuanVideo15DownBlockAdapter", + "HunyuanVideo15DownsampleAdapter", + "HunyuanVideoDownBlockAdapter", + "HunyuanVideoDownsampleAdapter", + "LTX2VideoDownBlockAdapter", + "LTX2VideoDownsamplerAdapter", + "QwenImageResampleDownAdapter", "WanResampleDownAdapter", "WanResidualDownBlockAdapter", - # Upsampling +) +_UPSAMPLING = ( + "HunyuanVideo15UpBlockAdapter", + "HunyuanVideo15UpsampleAdapter", + "HunyuanVideoUpBlockAdapter", + "HunyuanVideoUpsampleAdapter", + "LTX2VideoUpBlockAdapter", + "LTX2VideoUpsamplerAdapter", + "QwenImageResampleAdapter", + "QwenImageUpBlockAdapter", "Upsample2DAdapter", "WanResampleAdapter", "WanResidualUpBlockAdapter", "WanUpBlockAdapter", - # Other +) +_MIDBLOCK = ( + "HunyuanVideo15MidBlockAdapter", + "HunyuanVideoMidBlockAdapter", + "LTX2VideoMidBlockAdapter", + "QwenImageMidBlockAdapter", "WanMidBlockAdapter", +) +_RESNET = ( + "HunyuanVideo15ResnetBlockAdapter", + "HunyuanVideoResnetBlockAdapter", + "LTX2VideoResnetBlockAdapter", + "QwenImageResidualBlockAdapter", "WanResidualBlockAdapter", +) +_EXPORTS = { + **{name: "downsampling_adapters" for name in _DOWNSAMPLING}, + **{name: "upsampling_adapters" for name in _UPSAMPLING}, + **{name: "midblock_adapters" for name in _MIDBLOCK}, + **{name: "resnet_adapters" for name in _RESNET}, +} + +__all__ = [ + *_DOWNSAMPLING, + *_UPSAMPLING, + *_MIDBLOCK, + *_RESNET, ] + + +def __getattr__(name): + module_name = _EXPORTS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f"{__name__}.{module_name}"), name) + globals()[name] = value + return value + + +def __dir__(): + return sorted((*globals(), *__all__)) diff --git a/distvae/modules/adapters/adapter_utils.py b/distvae/modules/adapters/adapter_utils.py new file mode 100644 index 0000000..c1cd10b --- /dev/null +++ b/distvae/modules/adapters/adapter_utils.py @@ -0,0 +1,24 @@ +def adopt_convolution_parameters(target, original): + """Make a replacement convolution reuse the original Parameters.""" + target.weight = original.weight + target.bias = original.bias + return target + + +def replace_child_convolution( + module, + adapter, + *, + child="conv", + conv_block_size=0, + parallel_context=None, +): + """Replace a child convolution while giving its weights to the adapter.""" + convolution = getattr(module, child) + adapted = adapter( + convolution, + block_size=conv_block_size, + parallel_context=parallel_context, + ) + setattr(module, child, adapted) + return adapted diff --git a/distvae/modules/adapters/diffusers_blocks.py b/distvae/modules/adapters/diffusers_blocks.py new file mode 100644 index 0000000..6022749 --- /dev/null +++ b/distvae/modules/adapters/diffusers_blocks.py @@ -0,0 +1,46 @@ +"""The diffusers VAE block classes the adapters are written against, resolved optionally. + +DistVAE adapts several VAE families whose blocks arrived across a range of diffusers releases, +and an install new enough for one need not carry another. Importing them all eagerly would let +one missing family break every other family's adapter at import time, so each is resolved to +None instead and the adapter that wanted it says which class the installed diffusers is short +of, at the point someone tries to use it. +""" + +import importlib +from typing import Optional, Tuple + +WAN = "diffusers.models.autoencoders.autoencoder_kl_wan" +QWEN_IMAGE = "diffusers.models.autoencoders.autoencoder_kl_qwenimage" +HUNYUAN_VIDEO = "diffusers.models.autoencoders.autoencoder_kl_hunyuan_video" +HUNYUAN_VIDEO_15 = "diffusers.models.autoencoders.autoencoder_kl_hunyuanvideo15" +LTX2_VIDEO = "diffusers.models.autoencoders.autoencoder_kl_ltx2" + + +def block(module: str, name: str) -> Optional[type]: + """The named class from a diffusers module, or None where this release has neither""" + try: + found = getattr(importlib.import_module(module), name, None) + except ImportError: + return None + return found if isinstance(found, type) else None + + +def resolved(*blocks: Optional[type]) -> Tuple[type, ...]: + """The blocks that were found, for an isinstance check the rest simply cannot pass""" + return tuple(found for found in blocks if found is not None) + + +def require(supported: Tuple[type, ...], adapter: str, requires: str) -> None: + """Refuse an adapter whose diffusers classes are not in this release + + Without this the isinstance check against an empty tuple would report that the block passed + in was the wrong type, when the truth is that the right type does not exist here. + """ + if not supported: + import diffusers + + raise ImportError( + f"{adapter} needs {requires}, which diffusers {diffusers.__version__} does not " + f"provide. A newer diffusers is required to shard this VAE." + ) diff --git a/distvae/modules/adapters/downsampling_adapters.py b/distvae/modules/adapters/downsampling_adapters.py index ebfb0c6..cefb561 100644 --- a/distvae/modules/adapters/downsampling_adapters.py +++ b/distvae/modules/adapters/downsampling_adapters.py @@ -1,98 +1,398 @@ +from typing import Tuple + import torch.nn as nn -from distvae.models.layers.wan.zeropadconv2d import WanZeroPadConv2d -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter, WanCausalConv3dAdapter -from distvae.modules.adapters.resnet_adapters import WanResidualBlockAdapter -from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample, WanResidualDownBlock +from distvae.models.layers.asymmetric_zero_pad_conv2d import ( + AsymmetricZeroPadConv2d, +) +from distvae.modules.adapters.adapter_utils import ( + adopt_convolution_parameters, + replace_child_convolution, +) +from distvae.utils import ParallelContext, cache_cursor +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, + require, + resolved, +) +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + HunyuanVideo15CausalConv3dAdapter, + HunyuanVideoCausalConv3dAdapter, + LTX2VideoCausalConv3dAdapter, + QwenImageCausalConv3dAdapter, + WanCausalConv3dAdapter, +) +from distvae.modules.adapters.resnet_adapters import ( + HunyuanVideo15ResnetBlockAdapter, + HunyuanVideoResnetBlockAdapter, + LTX2VideoResnetBlockAdapter, + WanResidualBlockAdapter, +) +from diffusers.models.downsampling import Downsample2D + +WanResample = block(WAN, "WanResample") +WanResidualDownBlock = block(WAN, "WanResidualDownBlock") +QwenImageResample = block(QWEN_IMAGE, "QwenImageResample") +HunyuanVideoDownsampleCausal3D = block(HUNYUAN_VIDEO, "HunyuanVideoDownsampleCausal3D") +HunyuanVideoDownBlock3D = block(HUNYUAN_VIDEO, "HunyuanVideoDownBlock3D") +HunyuanVideo15Downsample = block(HUNYUAN_VIDEO_15, "HunyuanVideo15Downsample") +HunyuanVideo15DownBlock3D = block(HUNYUAN_VIDEO_15, "HunyuanVideo15DownBlock3D") +LTX2VideoCausalConv3d = block(LTX2_VIDEO, "LTX2VideoCausalConv3d") +LTX2VideoDownsampler3d = block(LTX2_VIDEO, "LTX2VideoDownsampler3d") +LTX2VideoDownBlock3D = block(LTX2_VIDEO, "LTX2VideoDownBlock3D") + + +def _zero_pad_strided_conv(conv, conv_block_size, parallel_context): + """A sharded stand-in for a (0, 1, 0, 1) zero pad followed by a stride-2 convolution + + The pair cannot be split as written, because a rank's bottom row is padding only if it is the + bottom row of the whole image. One module pads outside edges and exchanges halos across rank + boundaries. Wan resampling and diffusers Downsample2D both use this operation. + """ + padding = conv.padding + if (isinstance(padding, int) and padding != 0) or ( + isinstance(padding, tuple) and sum(padding) != 0 + ): + raise ValueError(f"Unsupported padding: {padding}") + sharded = AsymmetricZeroPadConv2d( + in_channels=conv.in_channels, + out_channels=conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + dilation=conv.dilation, + groups=conv.groups, + bias=conv.bias is not None, + device=conv.weight.device, + dtype=conv.weight.dtype, + reversed_zero_padding=(0, 1, 0, 1), + block_size=conv_block_size, + parallel_context=parallel_context, + ) + return adopt_convolution_parameters(sharded, conv) + + +class Downsample2DAdapter(nn.Module): + """Shard the convolution in the 2D downsampler used by AutoencoderKL and Flux.2. -class WanResampleDownAdapter(nn.Module): + When ``padding == 0``, preserve the module's explicit ``(0, 1, 0, 1)`` padding and replace + only the convolution. Otherwise, adapt the existing strided convolution directly. + Average-pooling configurations need no cross-rank data because Patchify assigns complete row + pairs. Any normalization reduces over channels and therefore remains unsharded. """ - Adapter for WanResample used in downsampling operations. - Handles temporal convolution and spatial downsampling with distributed patching. + + def __init__( + self, + downsampler: Downsample2D, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + assert isinstance(downsampler, Downsample2D), ( + "Downsample2DAdapter does not support downsampler except Downsample2D" + ) + self.downsampler = downsampler + self.pads_by_hand = downsampler.use_conv and downsampler.padding == 0 + if not downsampler.use_conv: + return + conv = downsampler.conv + if self.pads_by_hand: + sharded = _zero_pad_strided_conv( + conv, conv_block_size, parallel_context + ) + else: + sharded = Conv2dAdapter( + conv, + block_size=conv_block_size, + parallel_context=parallel_context, + ) + downsampler.conv = sharded + # Some configurations name the same convolution twice. Both have to move, or the original + # stays alive holding a second copy of the weights. + if getattr(downsampler, "Conv2d_0", None) is conv: + downsampler.Conv2d_0 = sharded + + def forward(self, hidden_states, *args, **kwargs): + if not self.pads_by_hand: + return self.downsampler(hidden_states, *args, **kwargs) + if self.downsampler.norm is not None: + hidden_states = self.downsampler.norm( + hidden_states.permute(0, 2, 3, 1) + ).permute(0, 3, 1, 2) + return self.downsampler.conv(hidden_states) + + +class _CausalResampleDownAdapter(nn.Module): + """Shards a downsampling resample containing temporal and strided spatial convolutions. + + The spatial half is a zero pad of (0, 1, 0, 1) followed by a stride-2 convolution with no + padding of its own. Splitting that needs the pad and the convolution taken together, since a + rank's bottom row is padding only if it is the bottom row of the whole image, so the pair is + replaced by one module that pads the outside edges and exchanges halos on the inside ones. """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + def __init__( self, - wan_resample: WanResample, + resample: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = True, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_resample, WanResample), ( - "WanResampleDownAdapter does not support resample except WanResample" + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(resample, self._supported), ( + f"{adapter} does not support resample except {self._requires}" ) - self.resample = wan_resample - if patch_dim == -3: - raise ValueError("WanResampleDownAdapter does not support patch_dim F (-3); use H (-2) or W (-1).") + self.resample = resample - # Adapt time_conv if present - if hasattr(wan_resample, "time_conv") and wan_resample.time_conv is not None: - wan_resample.time_conv = WanCausalConv3dAdapter( - wan_resample.time_conv, + if getattr(resample, "time_conv", None) is not None: + resample.time_conv = self._conv_adapter( + resample.time_conv, block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=parallel_context, ) - # Adapt the resample layers - if isinstance(wan_resample.resample, nn.Sequential): - count = 0 - for layer in wan_resample.resample: - count += 1 - if isinstance(layer, nn.ZeroPad2d): - continue - elif isinstance(layer, nn.Conv2d): - in_channels = layer.in_channels - out_channels = layer.out_channels - kernel_size = layer.kernel_size - stride = layer.stride - if ( - isinstance(layer.padding, int) and layer.padding != 0 or - isinstance(layer.padding, tuple) and (sum(layer.padding) != 0) - ): - raise ValueError(f"Unsupported padding: {layer.padding}") - dilation = layer.dilation - groups = layer.groups - bias = layer.bias is not None - device = layer.weight.device - dtype = layer.weight.dtype - _weight = layer.weight - _bias = layer.bias - else: - raise ValueError(f"Unsupported layer type: {type(layer)}") - if count != 2: - raise ValueError(f"WanResampleDownAdapter expects 2 layers, got {count}") - - self.resample.resample = WanZeroPadConv2d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - groups=groups, - bias=bias, - device=device, - dtype=dtype, - reversed_zero_padding=(0, 1, 0, 1), + if isinstance(resample.resample, nn.Sequential): + layers = list(resample.resample) + convs = [layer for layer in layers if isinstance(layer, nn.Conv2d)] + pads = [layer for layer in layers if isinstance(layer, nn.ZeroPad2d)] + if len(layers) != 2 or len(convs) != 1 or len(pads) != 1: + raise ValueError( + f"{adapter} expects a zero pad and one convolution, got " + f"{[type(layer).__name__ for layer in layers]}" + ) + resample.resample = _zero_pad_strided_conv( + convs[0], conv_block_size, parallel_context + ) + elif isinstance(resample.resample, nn.Conv2d): + resample.resample = Conv2dAdapter( + resample.resample, block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=parallel_context, ) - self.resample.resample.weight.data = _weight.data - if _bias is not None: - self.resample.resample.bias.data = _bias.data - else: - # Single conv layer - if isinstance(wan_resample.resample, nn.Conv2d): - self.resample.resample = Conv2dAdapter( - wan_resample.resample, - block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + + def forward(self, x, feat_cache=None, feat_idx=None): + return self.resample(x, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx)) + + +class WanResampleDownAdapter(_CausalResampleDownAdapter): + _supported = resolved(WanResample) + _requires = "WanResample" + _conv_adapter = WanCausalConv3dAdapter + + +class QwenImageResampleDownAdapter(_CausalResampleDownAdapter): + _supported = resolved(QwenImageResample) + _requires = "QwenImageResample" + _conv_adapter = QwenImageCausalConv3dAdapter + + +class _PaddedCausalDownsampleAdapter(nn.Module): + """Shards a HunyuanVideo downsampler, of which the convolution is the only sharded part + + Whatever the downsampler does after that convolution reads one input position per output + one: HunyuanVideo strides, and 1.5 folds each pair of rows and columns into channels. Either + way a rank can do it to its own rows, provided it holds whole pairs of them, which the bands + Patchify cuts guarantee by being whole multiples of what the encoder narrows by. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + + def __init__( + self, + downsampler: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(downsampler, self._supported), ( + f"{adapter} does not support downsampler except {self._requires}" + ) + self.downsampler = downsampler + replace_child_convolution( + downsampler, + self._conv_adapter, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, hidden_states): + return self.downsampler(hidden_states) + + +class HunyuanVideoDownsampleAdapter(_PaddedCausalDownsampleAdapter): + _supported = resolved(HunyuanVideoDownsampleCausal3D) + _requires = "HunyuanVideoDownsampleCausal3D" + _conv_adapter = HunyuanVideoCausalConv3dAdapter + + +class HunyuanVideo15DownsampleAdapter(_PaddedCausalDownsampleAdapter): + _supported = resolved(HunyuanVideo15Downsample) + _requires = "HunyuanVideo15Downsample" + _conv_adapter = HunyuanVideo15CausalConv3dAdapter + + +class _PaddedCausalDownBlockAdapter(nn.Module): + """Shards a HunyuanVideo down block: its residual blocks and its downsampler""" + + _supported: Tuple[type, ...] = () + _requires: str = "" + _resnet_adapter = None + _downsample_adapter = None + + def __init__( + self, + down_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(down_block, self._supported), ( + f"{adapter} does not support down block except {self._requires}" + ) + options = dict( + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + self.down_block = down_block + down_block.resnets = nn.ModuleList( + [self._resnet_adapter(resnet, **options) for resnet in down_block.resnets] + ) + if down_block.downsamplers is not None: + down_block.downsamplers = nn.ModuleList( + [self._downsample_adapter(down, **options) for down in down_block.downsamplers] + ) + + def forward(self, hidden_states): + return self.down_block(hidden_states) + + +class HunyuanVideoDownBlockAdapter(_PaddedCausalDownBlockAdapter): + _supported = resolved(HunyuanVideoDownBlock3D) + _requires = "HunyuanVideoDownBlock3D" + _resnet_adapter = HunyuanVideoResnetBlockAdapter + _downsample_adapter = HunyuanVideoDownsampleAdapter + + +class HunyuanVideo15DownBlockAdapter(_PaddedCausalDownBlockAdapter): + _supported = resolved(HunyuanVideo15DownBlock3D) + _requires = "HunyuanVideo15DownBlock3D" + _resnet_adapter = HunyuanVideo15ResnetBlockAdapter + _downsample_adapter = HunyuanVideo15DownsampleAdapter + + +class LTX2VideoDownsamplerAdapter(nn.Module): + """Shards an LTX-2 downsampler, which is its convolution + + What follows the convolution moves space into channels and averages the input the same way + for the residual, reading one input position per output one, so a rank can do it to its own + rows alone. + """ + + _supported = resolved(LTX2VideoDownsampler3d) + _requires = "LTX2VideoDownsampler3d" + + def __init__( + self, + downsampler: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(downsampler, self._supported), ( + f"{adapter} does not support downsampler except {self._requires}" + ) + self.downsampler = downsampler + replace_child_convolution( + downsampler, + LTX2VideoCausalConv3dAdapter, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, hidden_states, causal: bool = True): + return self.downsampler(hidden_states, causal=causal) + + +class LTX2VideoDownBlockAdapter(nn.Module): + """Shards an LTX-2 down block: its residual blocks and its downsampler + + Which downsampler that is depends on how the stage was configured: a strided causal + convolution where it downsamples by striding, or the space-to-channel downsampler where it + does so by folding. Both are handled because a checkpoint may hold either. + """ + + _supported = resolved(LTX2VideoDownBlock3D) + _requires = "LTX2VideoDownBlock3D" + + def __init__( + self, + down_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(down_block, self._supported), ( + f"{adapter} does not support down block except {self._requires}" + ) + options = dict( + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + self.down_block = down_block + down_block.resnets = nn.ModuleList( + [LTX2VideoResnetBlockAdapter(resnet, **options) for resnet in down_block.resnets] + ) + if down_block.downsamplers is not None: + down_block.downsamplers = nn.ModuleList( + [self._adapt_downsampler( + down, adapter, conv_block_size, parallel_context ) + for down in down_block.downsamplers] + ) - def forward(self, x, feat_cache=None, feat_idx=[0]): - return self.resample(x, feat_cache=feat_cache, feat_idx=feat_idx) + @staticmethod + def _adapt_downsampler( + downsampler, adapter, conv_block_size, parallel_context + ): + if LTX2VideoDownsampler3d is not None and isinstance(downsampler, LTX2VideoDownsampler3d): + return LTX2VideoDownsamplerAdapter( + downsampler, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + if LTX2VideoCausalConv3d is not None and isinstance(downsampler, LTX2VideoCausalConv3d): + return LTX2VideoCausalConv3dAdapter( + downsampler, + block_size=conv_block_size, + parallel_context=parallel_context, + ) + raise TypeError( + f"{adapter} cannot shard a downsampler of type {type(downsampler).__name__}. It " + f"handles LTX2VideoDownsampler3d and LTX2VideoCausalConv3d." + ) + + def forward(self, hidden_states, temb=None, generator=None, causal: bool = True): + return self.down_block(hidden_states, temb, generator, causal=causal) class WanResidualDownBlockAdapter(nn.Module): @@ -100,20 +400,23 @@ class WanResidualDownBlockAdapter(nn.Module): Adapter for WanResidualDownBlock used in the encoder (Wan2.2). Patches residual blocks and downsampler with distributed processing support. """ + _supported = resolved(WanResidualDownBlock) + def __init__( self, - wan_residual_down_block: WanResidualDownBlock, + wan_residual_down_block: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = True, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_residual_down_block, WanResidualDownBlock), ( + require( + self._supported, + type(self).__name__, + "WanResidualDownBlock", + ) + assert isinstance(wan_residual_down_block, self._supported), ( "WanResidualDownBlockAdapter only supports WanResidualDownBlock" ) - if patch_dim == -3: - raise ValueError("WanResidualDownBlockAdapter does not support patch_dim F (-3); use H (-2) or W (-1).") - self.down_block = wan_residual_down_block if hasattr(wan_residual_down_block, "resnets"): adapted_resnets = [] @@ -122,19 +425,21 @@ def __init__( WanResidualBlockAdapter( resnet, conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + parallel_context=parallel_context, ) ) self.down_block.resnets = nn.ModuleList(adapted_resnets) if hasattr(wan_residual_down_block, "downsampler") and wan_residual_down_block.downsampler is not None: - if isinstance(wan_residual_down_block.downsampler, WanResample): + if WanResample is not None and isinstance( + wan_residual_down_block.downsampler, WanResample + ): self.down_block.downsampler = WanResampleDownAdapter( wan_residual_down_block.downsampler, conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + parallel_context=parallel_context, ) - - def forward(self, hidden_states, feat_cache=None, feat_idx=[0]): - return self.down_block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx) + + def forward(self, hidden_states, feat_cache=None, feat_idx=None): + return self.down_block( + hidden_states, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx) + ) diff --git a/distvae/modules/adapters/layers/attn_adapters.py b/distvae/modules/adapters/layers/attn_adapters.py index 6e1850c..c5a6717 100644 --- a/distvae/modules/adapters/layers/attn_adapters.py +++ b/distvae/modules/adapters/layers/attn_adapters.py @@ -1,54 +1,44 @@ from typing import Any import torch -import torch.distributed as dist import torch.nn as nn -from distvae.utils import DistributedEnv +from distvae.modules.patch_utils import gather_patches +from distvae.utils import ParallelContext, normalize_patch_dim -class WanAttentionBlockAdapter(torch.nn.Module): +class GatheredAttentionAdapter(torch.nn.Module): """Runs attention on the full sequence by gathering along the patch dim, then narrows back to the local patch. - Supports unequal patch sizes across ranks (e.g. after Patchify without padding). + Attention is the one layer in a VAE that relates every position to every other, so unlike a + convolution it cannot be satisfied with a halo. Nothing here reads the wrapped module, only + calls it, so this covers whichever attention block a family happens to use. + + Patches need not be the same size across ranks: the gather below pads for transport only. """ def __init__( self, module: nn.Module, - patch_dim: int = -2, + parallel_context: ParallelContext = None, ) -> None: super().__init__() + if not isinstance(parallel_context, ParallelContext): + raise TypeError("GatheredAttentionAdapter requires a ParallelContext") self.module = module - self.patch_dim = patch_dim + self.parallel_context = parallel_context + self.patch_dim = parallel_context.patch_dim def forward(self, hidden_states: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: - patch_dim = self.patch_dim if self.patch_dim >= 0 else hidden_states.ndim + self.patch_dim - rank = DistributedEnv.get_rank_in_vae_group() - world_size = DistributedEnv.get_group_world_size() - device = hidden_states.device - - # Gather chunk sizes from all ranks - size_list = [torch.empty(1, dtype=torch.int64, device=device) for _ in range(world_size)] - dist.all_gather( - size_list, - torch.tensor([hidden_states.shape[patch_dim]], dtype=torch.int64, device=device), - group=DistributedEnv.get_vae_group(), - ) - chunk_sizes = [size_list[i].item() for i in range(world_size)] - - base_shape = list(hidden_states.shape) - gathered_tensors = [] - for i in range(world_size): - shape = base_shape.copy() - shape[patch_dim] = chunk_sizes[i] - gathered_tensors.append(torch.empty(shape, dtype=hidden_states.dtype, device=device)) - dist.all_gather(gathered_tensors, hidden_states.contiguous(), group=DistributedEnv.get_vae_group()) - - combined_tensor = torch.cat(gathered_tensors, dim=patch_dim) - forward_output = self.module(combined_tensor, *args, **kwargs) - start_idx = sum(chunk_sizes[:rank]) - local_output = torch.narrow( - forward_output, patch_dim, start_idx, chunk_sizes[rank] + patch_dim = hidden_states.ndim + normalize_patch_dim( + self.patch_dim, hidden_states.ndim, spatial_only=True ) - return local_output \ No newline at end of file + rank = self.parallel_context.rank + + patches, sizes = gather_patches(hidden_states, self.parallel_context) + whole = self.module(torch.cat(patches, dim=patch_dim), *args, **kwargs) + return torch.narrow(whole, patch_dim, sum(sizes[:rank]), sizes[rank]) + + +# The name this was introduced under, before other families turned out to need the same thing. +WanAttentionBlockAdapter = GatheredAttentionAdapter \ No newline at end of file diff --git a/distvae/modules/adapters/layers/conv_adapters.py b/distvae/modules/adapters/layers/conv_adapters.py index d44b601..6234f39 100644 --- a/distvae/modules/adapters/layers/conv_adapters.py +++ b/distvae/modules/adapters/layers/conv_adapters.py @@ -1,10 +1,29 @@ +from typing import Tuple + import torch import torch.nn as nn import torch.nn.functional as F -from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d from distvae.models.layers.conv2d import PatchConv2d from distvae.models.layers.conv3d import PatchConv3d +from distvae.modules.adapters.adapter_utils import adopt_convolution_parameters +from distvae.utils import ParallelContext +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, + require, + resolved, +) + +WanCausalConv3d = block(WAN, "WanCausalConv3d") +QwenImageCausalConv3d = block(QWEN_IMAGE, "QwenImageCausalConv3d") +HunyuanVideoCausalConv3d = block(HUNYUAN_VIDEO, "HunyuanVideoCausalConv3d") +HunyuanVideo15CausalConv3d = block(HUNYUAN_VIDEO_15, "HunyuanVideo15CausalConv3d") +LTX2VideoCausalConv3d = block(LTX2_VIDEO, "LTX2VideoCausalConv3d") class Conv2dAdapter(nn.Module): @@ -13,8 +32,7 @@ def __init__( conv2d: nn.Conv2d, *, block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() for i in conv2d.dilation: @@ -32,12 +50,9 @@ def __init__( device=conv2d.weight.device, dtype=conv2d.weight.dtype, block_size=block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=parallel_context, ) - self.conv2d.weight.data = conv2d.weight.data - if conv2d.bias is not None: - self.conv2d.bias.data = conv2d.bias.data + adopt_convolution_parameters(self.conv2d, conv2d) def forward(self, x): return self.conv2d(x) @@ -49,8 +64,7 @@ def __init__( conv3d: nn.Conv3d, *, block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() for i in conv3d.dilation: @@ -68,31 +82,40 @@ def __init__( device=conv3d.weight.device, dtype=conv3d.weight.dtype, block_size=block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=parallel_context, ) - self.conv3d.weight.data = conv3d.weight.data - if conv3d.bias is not None: - self.conv3d.bias.data = conv3d.bias.data + adopt_convolution_parameters(self.conv3d, conv3d) def forward(self, x): return self.conv3d(x) -class WanCausalConv3dAdapter(nn.Module): +class _CausalConv3dAdapter(nn.Module): + """Shards a causal 3D convolution that subclasses nn.Conv3d and holds its padding in _padding. + + Only the spatial half of that padding reaches PatchConv3d, which exchanges halos so a rank + pads where the image ends rather than where its own patch happens to. The temporal half is + applied here instead, before the convolution, because the frame axis is not the one split + across ranks and its causal padding has to stay one-sided. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + def __init__( - self, - causal_conv3d: WanCausalConv3d, + self, + causal_conv3d: nn.Conv3d, *, block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) for i in causal_conv3d.dilation: - assert i == 1, "dilation is not supported in WanCausalConv3dAdapter" - assert isinstance(causal_conv3d, WanCausalConv3d), ( - "WanCausalConv3dAdapter does not support causal_conv3d except WanCausalConv3d" + assert i == 1, f"dilation is not supported in {adapter}" + assert isinstance(causal_conv3d, self._supported), ( + f"{adapter} does not support causal_conv3d except {self._requires}" ) self.conv3d = PatchConv3d( in_channels=causal_conv3d.in_channels, @@ -107,12 +130,9 @@ def __init__( device=causal_conv3d.weight.device, dtype=causal_conv3d.weight.dtype, block_size=block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=parallel_context, ) - self.conv3d.weight.data = causal_conv3d.weight.data - if causal_conv3d.bias is not None: - self.conv3d.bias.data = causal_conv3d.bias.data + adopt_convolution_parameters(self.conv3d, causal_conv3d) self._padding = (0, 0, 0, 0, causal_conv3d._padding[4], causal_conv3d._padding[5]) def forward(self, x, cache_x=None): @@ -122,4 +142,138 @@ def forward(self, x, cache_x=None): x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) - return self.conv3d(x) \ No newline at end of file + return self.conv3d(x) + + +class WanCausalConv3dAdapter(_CausalConv3dAdapter): + _supported = resolved(WanCausalConv3d) + _requires = "WanCausalConv3d" + + +class QwenImageCausalConv3dAdapter(_CausalConv3dAdapter): + """Qwen-Image's causal convolution, which is WanCausalConv3d under a different name""" + + _supported = resolved(QwenImageCausalConv3d) + _requires = "QwenImageCausalConv3d" + + +class _PaddedCausalConv3dAdapter(nn.Module): + """Shards a causal 3D convolution that holds a plain nn.Conv3d and pads in its own forward. + + Unlike Wan's, these pad by replication rather than with zeros, which left alone would have + each rank repeat its own top and bottom rows where it ought to be reading its neighbour's. + Moving the spatial half of the padding into PatchConv3d settles that: it exchanges halos and + replicates only at the edges of the real image. The temporal half is applied here, because + the frame axis is not the one being split and its padding has to stay one-sided. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + + def __init__( + self, + causal_conv3d: nn.Module, + *, + block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(causal_conv3d, self._supported), ( + f"{adapter} does not support causal_conv3d except {self._requires}" + ) + conv = causal_conv3d.conv + for i in conv.dilation: + assert i == 1, f"dilation is not supported in {adapter}" + assert tuple(conv.padding) == (0, 0, 0), ( + f"{adapter} expects all padding to live in time_causal_padding, but the " + f"convolution also pads by {tuple(conv.padding)}" + ) + # F.pad orders its argument (W, W, H, H, F, F). + pad_w, _, pad_h, _, pad_front, pad_back = causal_conv3d.time_causal_padding + self.conv3d = PatchConv3d( + in_channels=conv.in_channels, + out_channels=conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + padding=(0, pad_h, pad_w), + dilation=conv.dilation, + groups=conv.groups, + bias=conv.bias is not None, + padding_mode=causal_conv3d.pad_mode, + device=conv.weight.device, + dtype=conv.weight.dtype, + block_size=block_size, + parallel_context=parallel_context, + ) + adopt_convolution_parameters(self.conv3d, conv) + self.pad_mode = causal_conv3d.pad_mode + self._padding = (0, 0, 0, 0, pad_front, pad_back) + + def forward(self, hidden_states): + # Padding one axis and then the other reaches the same place as padding both at once: + # replication reads a clamped index per axis, and clamping them in turn is the same. + hidden_states = F.pad(hidden_states, self._padding, mode=self.pad_mode) + return self.conv3d(hidden_states) + + +class HunyuanVideoCausalConv3dAdapter(_PaddedCausalConv3dAdapter): + _supported = resolved(HunyuanVideoCausalConv3d) + _requires = "HunyuanVideoCausalConv3d" + + +class HunyuanVideo15CausalConv3dAdapter(_PaddedCausalConv3dAdapter): + _supported = resolved(HunyuanVideo15CausalConv3d) + _requires = "HunyuanVideo15CausalConv3d" + + +class LTX2VideoCausalConv3dAdapter(nn.Module): + """Shards LTX-2's causal convolution, which needs less rearranging than the others. + + Its spatial padding already sits inside the nn.Conv3d rather than being applied around it, + so swapping that convolution for a PatchConv3d built from the same arguments is the whole of + it. The temporal padding repeats the first and last frames along an axis nobody splits, and + happens in the wrapped module's own forward, which is left to run as it is. + """ + + _supported = resolved(LTX2VideoCausalConv3d) + _requires = "LTX2VideoCausalConv3d" + + def __init__( + self, + causal_conv3d: nn.Module, + *, + block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(causal_conv3d, self._supported), ( + f"{adapter} does not support causal_conv3d except {self._requires}" + ) + conv = causal_conv3d.conv + for i in conv.dilation: + assert i == 1, f"dilation is not supported in {adapter}" + self.causal_conv3d = causal_conv3d + sharded = PatchConv3d( + in_channels=conv.in_channels, + out_channels=conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + padding=conv.padding, + dilation=conv.dilation, + groups=conv.groups, + bias=conv.bias is not None, + padding_mode=conv.padding_mode, + device=conv.weight.device, + dtype=conv.weight.dtype, + block_size=block_size, + parallel_context=parallel_context, + ) + adopt_convolution_parameters(sharded, conv) + causal_conv3d.conv = sharded + + def forward(self, hidden_states, causal: bool = True): + return self.causal_conv3d(hidden_states, causal=causal) diff --git a/distvae/modules/adapters/layers/norm_adapters.py b/distvae/modules/adapters/layers/norm_adapters.py index 6ea645e..9e3f228 100644 --- a/distvae/modules/adapters/layers/norm_adapters.py +++ b/distvae/modules/adapters/layers/norm_adapters.py @@ -1,16 +1,24 @@ -import torch import torch.nn as nn + from distvae.models.layers.normalization import PatchGroupNorm +from distvae.utils import ParallelContext class GroupNormAdapter(nn.Module): - def __init__(self, group_norm: nn.GroupNorm): + """A GroupNorm whose statistics use this adapter's immutable group and split axis.""" + + def __init__( + self, + group_norm: nn.GroupNorm, + parallel_context: ParallelContext = None, + ): super().__init__() self.group_norm = PatchGroupNorm( num_groups=group_norm.num_groups, num_channels=group_norm.num_channels, eps=group_norm.eps, - affine=group_norm.affine + affine=group_norm.affine, + parallel_context=parallel_context, ) if group_norm.affine: self.group_norm.weight = group_norm.weight diff --git a/distvae/modules/adapters/midblock_adapters.py b/distvae/modules/adapters/midblock_adapters.py index cb2836a..36e47ec 100644 --- a/distvae/modules/adapters/midblock_adapters.py +++ b/distvae/modules/adapters/midblock_adapters.py @@ -1,33 +1,166 @@ +from typing import Tuple + import torch.nn as nn -from diffusers.models.autoencoders.autoencoder_kl_wan import WanMidBlock -from distvae.modules.adapters.layers.attn_adapters import WanAttentionBlockAdapter -from distvae.modules.adapters.resnet_adapters import WanResidualBlockAdapter +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, + require, + resolved, +) +from distvae.modules.adapters.layers.attn_adapters import GatheredAttentionAdapter +from distvae.utils import ParallelContext, cache_cursor +from distvae.modules.adapters.resnet_adapters import ( + HunyuanVideo15ResnetBlockAdapter, + HunyuanVideoResnetBlockAdapter, + LTX2VideoResnetBlockAdapter, + QwenImageResidualBlockAdapter, + WanResidualBlockAdapter, +) + +WanMidBlock = block(WAN, "WanMidBlock") +QwenImageMidBlock = block(QWEN_IMAGE, "QwenImageMidBlock") +HunyuanVideoMidBlock3D = block(HUNYUAN_VIDEO, "HunyuanVideoMidBlock3D") +HunyuanVideo15MidBlock = block(HUNYUAN_VIDEO_15, "HunyuanVideo15MidBlock") +LTX2VideoMidBlock3d = block(LTX2_VIDEO, "LTX2VideoMidBlock3d") + + +class _CausalMidBlockAdapter(nn.Module): + """Shards a mid block: its residual blocks stay local, its attentions have to gather""" + + _supported: Tuple[type, ...] = () + _requires: str = "" + _resnet_adapter = None -class WanMidBlockAdapter(nn.Module): def __init__( self, - wan_mid_block: WanMidBlock, + mid_block: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_mid_block, WanMidBlock), "WanMidBlockAdapter does not support mid block except WanMidBlock" - self.mid_block = wan_mid_block + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(mid_block, self._supported), ( + f"{adapter} does not support mid block except {self._requires}" + ) + self.mid_block = mid_block self.mid_block.resnets = nn.ModuleList([ - WanResidualBlockAdapter( + self._resnet_adapter( resnet, conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) for resnet in wan_mid_block.resnets + parallel_context=parallel_context, + ) for resnet in mid_block.resnets ]) self.mid_block.attentions = nn.ModuleList([ - WanAttentionBlockAdapter(attn, patch_dim=patch_dim) - for attn in wan_mid_block.attentions + GatheredAttentionAdapter(attn, parallel_context=parallel_context) + if attn is not None else attn + for attn in mid_block.attentions + ]) + + def forward(self, x, feat_cache=None, feat_idx=None): + return self.mid_block(x, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx)) + + +class WanMidBlockAdapter(_CausalMidBlockAdapter): + _supported = resolved(WanMidBlock) + _requires = "WanMidBlock" + _resnet_adapter = WanResidualBlockAdapter + + +class QwenImageMidBlockAdapter(_CausalMidBlockAdapter): + _supported = resolved(QwenImageMidBlock) + _requires = "QwenImageMidBlock" + _resnet_adapter = QwenImageResidualBlockAdapter + + +class HunyuanVideo15MidBlockAdapter(_CausalMidBlockAdapter): + """Shards HunyuanVideo 1.5's mid block: residual blocks stay local, attentions gather""" + + _supported = resolved(HunyuanVideo15MidBlock) + _requires = "HunyuanVideo15MidBlock" + _resnet_adapter = HunyuanVideo15ResnetBlockAdapter + + def forward(self, hidden_states): + return self.mid_block(hidden_states) + + +class HunyuanVideoMidBlockAdapter(nn.Module): + """Shards HunyuanVideo's mid block, or gathers around the whole of it when it has attention. + + Its attention cannot be wrapped on its own the way every other family's can. The mid block + flattens (F, H, W) into a sequence and builds the causal mask itself, both from the height + it can see, so an attention handed a patch would also be handed a mask cut for a patch and + would quietly attend over the wrong span. Gathering around the entire block avoids + reimplementing that forward, and costs little: the mid block runs at the latent resolution, + which is the cheapest point in the decoder, and it is the up blocks after it that hold the + activations worth splitting. + """ + + def __init__( + self, + mid_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + supported = resolved(HunyuanVideoMidBlock3D) + require(supported, adapter, "HunyuanVideoMidBlock3D") + assert isinstance(mid_block, supported), ( + f"{adapter} does not support mid block except HunyuanVideoMidBlock3D" + ) + if any(attn is not None for attn in mid_block.attentions): + self.mid_block = GatheredAttentionAdapter( + mid_block, parallel_context=parallel_context + ) + else: + mid_block.resnets = nn.ModuleList([ + HunyuanVideoResnetBlockAdapter( + resnet, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) for resnet in mid_block.resnets + ]) + self.mid_block = mid_block + + def forward(self, hidden_states): + return self.mid_block(hidden_states) + + +class LTX2VideoMidBlockAdapter(nn.Module): + """Shards an LTX-2 mid block, which is only residual blocks + + Alone among these families LTX-2 puts no attention in its mid block, so nothing here needs + to see the whole image and every rank can stay on its own patch throughout. + """ + + def __init__( + self, + mid_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + supported = resolved(LTX2VideoMidBlock3d) + require(supported, adapter, "LTX2VideoMidBlock3d") + assert isinstance(mid_block, supported), ( + f"{adapter} does not support mid block except LTX2VideoMidBlock3d" + ) + self.mid_block = mid_block + mid_block.resnets = nn.ModuleList([ + LTX2VideoResnetBlockAdapter( + resnet, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) for resnet in mid_block.resnets ]) - def forward(self, x, feat_cache=None, feat_idx=[0]): - return self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + def forward(self, hidden_states, temb=None, generator=None, causal: bool = True): + return self.mid_block(hidden_states, temb, generator, causal=causal) diff --git a/distvae/modules/adapters/resnet_adapters.py b/distvae/modules/adapters/resnet_adapters.py index 645a0ac..e5e3095 100644 --- a/distvae/modules/adapters/resnet_adapters.py +++ b/distvae/modules/adapters/resnet_adapters.py @@ -1,84 +1,247 @@ +from typing import Tuple + import torch import torch.nn as nn -from distvae.models.resnet import PatchResnetBlock2D -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter, WanCausalConv3dAdapter +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, + require, + resolved, +) +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + HunyuanVideo15CausalConv3dAdapter, + HunyuanVideoCausalConv3dAdapter, + LTX2VideoCausalConv3dAdapter, + QwenImageCausalConv3dAdapter, + WanCausalConv3dAdapter, +) from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter +from distvae.utils import ParallelContext, cache_cursor from diffusers.models.resnet import ResnetBlock2D -from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanResidualBlock + +WanResidualBlock = block(WAN, "WanResidualBlock") +QwenImageResidualBlock = block(QWEN_IMAGE, "QwenImageResidualBlock") +HunyuanVideoResnetBlockCausal3D = block(HUNYUAN_VIDEO, "HunyuanVideoResnetBlockCausal3D") +HunyuanVideo15ResnetBlock = block(HUNYUAN_VIDEO_15, "HunyuanVideo15ResnetBlock") +LTX2VideoResnetBlock3d = block(LTX2_VIDEO, "LTX2VideoResnetBlock3d") class ResnetBlock2DAdapter(nn.Module): + """Shards a 2D residual block: its two convolutions, its two group norms, and any shortcut + + Wrap the existing residual block in place. Reconstructing it as ``PatchResnetBlock2D`` would + allocate unused parameters and replace configuration not represented in that constructor. + """ + def __init__( - self, - resnet: ResnetBlock2D, - *, + self, + resnet: ResnetBlock2D, + *, conv_block_size = 0, + parallel_context: ParallelContext = None, ): super().__init__() assert resnet.time_emb_proj is None, "temb_channels is not supported in ResnetBlock2DAdapter currently" assert resnet.up is False, "up sample is not supported in ResnetBlock2DAdapter currently" assert resnet.down is False, "ResnetBlock2DAdapter does not support down sample currently" - self.resnet = PatchResnetBlock2D( - in_channels=resnet.in_channels, - out_channels=resnet.out_channels, - conv_shortcut=resnet.use_conv_shortcut, - dropout=0, - temb_channels=None, - groups=1, - groups_out=None, - pre_norm=resnet.pre_norm, - skip_time_act=resnet.skip_time_act, - time_embedding_norm=resnet.time_embedding_norm, - output_scale_factor=resnet.output_scale_factor, - use_in_shortcut=resnet.use_in_shortcut, - up=resnet.up, - down=resnet.down, + self.resnet = resnet + options = dict(parallel_context=parallel_context) + resnet.conv1 = Conv2dAdapter( + resnet.conv1, block_size=conv_block_size, **options + ) + resnet.norm1 = GroupNormAdapter(resnet.norm1, **options) + resnet.conv2 = Conv2dAdapter( + resnet.conv2, block_size=conv_block_size, **options ) - self.resnet.use_in_shortcut = resnet.use_in_shortcut - self.resnet.conv1 = Conv2dAdapter(resnet.conv1, block_size=conv_block_size) - self.resnet.norm1 = GroupNormAdapter(resnet.norm1) - self.resnet.conv2 = Conv2dAdapter(resnet.conv2, block_size=conv_block_size) - self.resnet.norm2 = GroupNormAdapter(resnet.norm2) - self.resnet.dropout = resnet.dropout - self.resnet.nonlinearity = resnet.nonlinearity - self.resnet.conv_shortcut = Conv2dAdapter(resnet.conv_shortcut, block_size=conv_block_size) if resnet.conv_shortcut is not None else None - + resnet.norm2 = GroupNormAdapter(resnet.norm2, **options) + if resnet.conv_shortcut is not None: + resnet.conv_shortcut = Conv2dAdapter( + resnet.conv_shortcut, block_size=conv_block_size, **options + ) def forward(self, x, temb: torch.FloatTensor = None, *args, **kwargs): return self.resnet(x, temb, *args, **kwargs) -class WanResidualBlockAdapter(nn.Module): +class _CausalResidualBlockAdapter(nn.Module): + """Shards a residual block built from two causal 3D convolutions and an optional shortcut. + + The norms either side of them are RMS, which reduces over channels and so needs nothing from + the other ranks; only the convolutions reach across the split. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + def __init__( self, - wan_residual_block: WanResidualBlock, + residual_block: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_residual_block, WanResidualBlock), ( - "WanResidualBlockAdapter does not support resnet except WanResidualBlock" + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(residual_block, self._supported), ( + f"{adapter} does not support resnet except {self._requires}" + ) + self.residual_block = residual_block + for name in ("conv1", "conv2"): + setattr( + self.residual_block, + name, + self._conv_adapter( + getattr(residual_block, name), + block_size=conv_block_size, + parallel_context=parallel_context, + ), + ) + # Adapt conv_shortcut if it's not nn.Identity + if not isinstance(residual_block.conv_shortcut, nn.Identity): + self.residual_block.conv_shortcut = self._conv_adapter( + residual_block.conv_shortcut, + block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + return self.residual_block( + x, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx) ) - self.residual_block = wan_residual_block - self.residual_block.conv1 = WanCausalConv3dAdapter( - wan_residual_block.conv1, - block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + + +class WanResidualBlockAdapter(_CausalResidualBlockAdapter): + _supported = resolved(WanResidualBlock) + _requires = "WanResidualBlock" + _conv_adapter = WanCausalConv3dAdapter + + +class QwenImageResidualBlockAdapter(_CausalResidualBlockAdapter): + _supported = resolved(QwenImageResidualBlock) + _requires = "QwenImageResidualBlock" + _conv_adapter = QwenImageCausalConv3dAdapter + + +class _PaddedCausalResnetBlockAdapter(nn.Module): + """Shards a HunyuanVideo residual block: its two causal convolutions, and any GroupNorms + + HunyuanVideo normalises with GroupNorm, whose statistics span the axis being split and so + have to be summed across ranks. HunyuanVideo 1.5 replaced those with RMS, which reduces over + channels and needs nothing from anyone; the isinstance check below is what tells them apart. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + + def __init__( + self, + resnet: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(resnet, self._supported), ( + f"{adapter} does not support resnet except {self._requires}" ) - self.residual_block.conv2 = WanCausalConv3dAdapter( - wan_residual_block.conv2, - block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + self.resnet = resnet + for name in ("conv1", "conv2"): + setattr( + resnet, + name, + self._conv_adapter( + getattr(resnet, name), + block_size=conv_block_size, + parallel_context=parallel_context, + ), + ) + for name in ("norm1", "norm2"): + norm = getattr(resnet, name) + if isinstance(norm, nn.GroupNorm): + setattr( + resnet, + name, + GroupNormAdapter( + norm, + parallel_context=parallel_context, + ), + ) + # Where the shortcut is a causal convolution it needs the same treatment; where it is a + # bare 1x1x1 it reads one position per output and is already right on a patch. + if isinstance(resnet.conv_shortcut, self._conv_adapter._supported): + resnet.conv_shortcut = self._conv_adapter( + resnet.conv_shortcut, + block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, hidden_states): + return self.resnet(hidden_states) + + +class HunyuanVideoResnetBlockAdapter(_PaddedCausalResnetBlockAdapter): + _supported = resolved(HunyuanVideoResnetBlockCausal3D) + _requires = "HunyuanVideoResnetBlockCausal3D" + _conv_adapter = HunyuanVideoCausalConv3dAdapter + + +class HunyuanVideo15ResnetBlockAdapter(_PaddedCausalResnetBlockAdapter): + _supported = resolved(HunyuanVideo15ResnetBlock) + _requires = "HunyuanVideo15ResnetBlock" + _conv_adapter = HunyuanVideo15CausalConv3dAdapter + + +class LTX2VideoResnetBlockAdapter(nn.Module): + """Shards an LTX-2 residual block, which is its two convolutions and nothing else. + + Both its norms reduce over channels, its shortcut is a 1x1x1 convolution reading one + position per output, and its timestep conditioning arrives shaped to broadcast over space. + """ + + _supported = resolved(LTX2VideoResnetBlock3d) + _requires = "LTX2VideoResnetBlock3d" + + def __init__( + self, + resnet: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(resnet, self._supported), ( + f"{adapter} does not support resnet except {self._requires}" ) - # Adapt conv_shortcut if it's not nn.Identity - if not isinstance(wan_residual_block.conv_shortcut, nn.Identity): - self.residual_block.conv_shortcut = WanCausalConv3dAdapter( - wan_residual_block.conv_shortcut, block_size=conv_block_size, patch_dim=patch_dim + if resnet.per_channel_scale1 is not None or resnet.per_channel_scale2 is not None: + # Each rank would draw its own noise for its own rows, and the ranks together would + # not reconstruct the field a single one draws, so a sharded decode could not match + # an unsharded one at all. No shipped LTX-2 or LTX-2.3 config turns this on. + raise NotImplementedError( + f"{adapter} cannot shard a residual block with inject_noise enabled: the noise " + f"is drawn per rank and would not add up to the noise one rank draws. Decode " + f"this VAE on a single rank, or tile it instead." + ) + self.resnet = resnet + for name in ("conv1", "conv2"): + setattr( + resnet, + name, + LTX2VideoCausalConv3dAdapter( + getattr(resnet, name), + block_size=conv_block_size, + parallel_context=parallel_context, + ), ) - def forward(self, x, feat_cache=None, feat_idx=[0]): - return self.residual_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + def forward(self, inputs, temb=None, generator=None, causal: bool = True): + return self.resnet(inputs, temb, generator, causal=causal) diff --git a/distvae/modules/adapters/unets/unet_2d_blocks_adapters.py b/distvae/modules/adapters/unets/unet_2d_blocks_adapters.py index b234e46..ecc0276 100644 --- a/distvae/modules/adapters/unets/unet_2d_blocks_adapters.py +++ b/distvae/modules/adapters/unets/unet_2d_blocks_adapters.py @@ -2,13 +2,12 @@ import torch import torch.nn as nn +from distvae.modules.adapters.downsampling_adapters import Downsample2DAdapter from distvae.modules.adapters.resnet_adapters import ResnetBlock2DAdapter from distvae.modules.adapters.upsampling_adapters import Upsample2DAdapter -from distvae.models.unets.unet_2d_blocks import PatchUpDecoderBlock2D -from diffusers.models.unets.unet_2d_blocks import UpDecoderBlock2D -from diffusers.models.resnet import ResnetBlock2D -from diffusers.models.upsampling import Upsample2D +from diffusers.models.unets.unet_2d_blocks import DownEncoderBlock2D, UpDecoderBlock2D +from distvae.utils import ParallelContext class UpDecoderBlock2DAdapter(nn.Module): @@ -17,26 +16,67 @@ def __init__( up_block: UpDecoderBlock2D, *, conv_block_size = 0, + parallel_context: ParallelContext = None, ): super().__init__() assert up_block is not None and isinstance(up_block, UpDecoderBlock2D), "up_block must be a UpDecoderBlock2D instance" - self.up_block = PatchUpDecoderBlock2D( - in_channels=32, - out_channels=32, - add_upsample=False, - conv_block_size=conv_block_size - ) - self.up_block.resolution_idx = up_block.resolution_idx - self.up_block.resnets = nn.ModuleList([ - ResnetBlock2DAdapter(resnet, conv_block_size=conv_block_size) for resnet in up_block.resnets if isinstance(resnet, ResnetBlock2D) + self.up_block = up_block + up_block.resnets = nn.ModuleList([ + ResnetBlock2DAdapter( + resnet, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) for resnet in up_block.resnets ]) if up_block.upsamplers is not None: - self.up_block.upsamplers = nn.ModuleList([ - Upsample2DAdapter(upsampler, conv_block_size=conv_block_size) for upsampler in up_block.upsamplers if isinstance(upsampler, Upsample2D) + up_block.upsamplers = nn.ModuleList([ + Upsample2DAdapter( + upsampler, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) for upsampler in up_block.upsamplers ]) - assert len(self.up_block.upsamplers) == len(up_block.upsamplers), "Number of upsamplers in the adapter must match the number of upsamplers in the original block" - - assert len(self.up_block.resnets) == len(up_block.resnets), "Number of resnets in the adapter must match the number of resnets in the original block" def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None): - return self.up_block(hidden_states, temb) \ No newline at end of file + return self.up_block(hidden_states, temb) + + +class DownEncoderBlock2DAdapter(nn.Module): + """Shards the 2D down block AutoencoderKL and Flux.2 encode with: its resnets and downsampler + + Unlike the up block, this block is wrapped in place rather than rebuilt. Its forward method + runs both components in order without additional patch metadata. + """ + + def __init__( + self, + down_block: DownEncoderBlock2D, + *, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + assert isinstance(down_block, DownEncoderBlock2D), ( + "down_block must be a DownEncoderBlock2D instance" + ) + self.down_block = down_block + down_block.resnets = nn.ModuleList([ + ResnetBlock2DAdapter( + resnet, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + for resnet in down_block.resnets + ]) + if down_block.downsamplers is not None: + down_block.downsamplers = nn.ModuleList([ + Downsample2DAdapter( + downsampler, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + for downsampler in down_block.downsamplers + ]) + + def forward(self, hidden_states: torch.FloatTensor, *args, **kwargs): + return self.down_block(hidden_states, *args, **kwargs) \ No newline at end of file diff --git a/distvae/modules/adapters/upsampling_adapters.py b/distvae/modules/adapters/upsampling_adapters.py index 411a137..7b28f70 100644 --- a/distvae/modules/adapters/upsampling_adapters.py +++ b/distvae/modules/adapters/upsampling_adapters.py @@ -1,14 +1,48 @@ -from typing import Optional +from typing import Optional, Tuple import torch import torch.nn as nn -from distvae.utils import DistributedEnv -from distvae.models.upsampling import PatchUpsample2D -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter, WanCausalConv3dAdapter -from distvae.modules.adapters.resnet_adapters import WanResidualBlockAdapter +from distvae.modules.adapters.adapter_utils import replace_child_convolution +from distvae.utils import ParallelContext, cache_cursor +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, + require, + resolved, +) +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + HunyuanVideo15CausalConv3dAdapter, + HunyuanVideoCausalConv3dAdapter, + LTX2VideoCausalConv3dAdapter, + QwenImageCausalConv3dAdapter, + WanCausalConv3dAdapter, +) +from distvae.modules.adapters.resnet_adapters import ( + HunyuanVideo15ResnetBlockAdapter, + HunyuanVideoResnetBlockAdapter, + LTX2VideoResnetBlockAdapter, + QwenImageResidualBlockAdapter, + WanResidualBlockAdapter, +) from diffusers.models.upsampling import Upsample2D -from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample, WanResidualUpBlock, WanUpBlock + +WanResample = block(WAN, "WanResample") +WanResidualUpBlock = block(WAN, "WanResidualUpBlock") +WanUpBlock = block(WAN, "WanUpBlock") +QwenImageResample = block(QWEN_IMAGE, "QwenImageResample") +QwenImageUpBlock = block(QWEN_IMAGE, "QwenImageUpBlock") +HunyuanVideoUpsampleCausal3D = block(HUNYUAN_VIDEO, "HunyuanVideoUpsampleCausal3D") +HunyuanVideoUpBlock3D = block(HUNYUAN_VIDEO, "HunyuanVideoUpBlock3D") +HunyuanVideo15Upsample = block(HUNYUAN_VIDEO_15, "HunyuanVideo15Upsample") +HunyuanVideo15UpBlock3D = block(HUNYUAN_VIDEO_15, "HunyuanVideo15UpBlock3D") +LTX2VideoUpsampler3d = block(LTX2_VIDEO, "LTX2VideoUpsampler3d") +LTX2VideoUpBlock3d = block(LTX2_VIDEO, "LTX2VideoUpBlock3d") class Upsample2DAdapter(nn.Module): @@ -17,6 +51,7 @@ def __init__( upsample2d: Upsample2D, *, conv_block_size = 0, + parallel_context: ParallelContext = None, ): super().__init__() assert upsample2d.norm is None, "upsample2dBlock2DAdapter does not support normalization" @@ -24,20 +59,19 @@ def __init__( assert not isinstance(upsample2d.conv, nn.ConvTranspose2d), "upsample2dBlock2DAdapter does not support transpose conv" else: assert not isinstance(upsample2d.Conv2d_0, nn.ConvTranspose2d), "upsample2dBlock2DAdapter does not support transpose conv" - self.upsample2d = PatchUpsample2D( - channels=upsample2d.channels, - use_conv=upsample2d.use_conv, - use_conv_transpose=upsample2d.use_conv_transpose, - out_channels=upsample2d.out_channels, - name=upsample2d.name, - kernel_size=None, - padding=1, - interpolate=upsample2d.interpolate - ) + self.upsample2d = upsample2d if upsample2d.name == "conv": - self.upsample2d.conv = Conv2dAdapter(upsample2d.conv, block_size=conv_block_size) + upsample2d.conv = Conv2dAdapter( + upsample2d.conv, + block_size=conv_block_size, + parallel_context=parallel_context, + ) else: - self.upsample2d.Conv2d_0 = Conv2dAdapter(upsample2d.Conv2d_0, block_size=conv_block_size) + upsample2d.Conv2d_0 = Conv2dAdapter( + upsample2d.Conv2d_0, + block_size=conv_block_size, + parallel_context=parallel_context, + ) def forward( @@ -46,139 +80,312 @@ def forward( return self.upsample2d(hidden_states, output_size, *args, **kwargs) -class WanResampleAdapter(nn.Module): +class _CausalResampleAdapter(nn.Module): + """Shards a resample block: the 2D convolution it upsamples with, and its temporal one. + + The interpolation between them is nearest-neighbour, which reads a single input pixel per + output pixel, so a rank can upsample its own rows without hearing from its neighbours. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + def __init__( self, - wan_resample: WanResample, + resample: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_resample, WanResample), ( - "WanResampleAdapter does not support resample except WanResample" + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(resample, self._supported), ( + f"{adapter} does not support resample except {self._requires}" ) - self.resample = wan_resample - if patch_dim == -3: - raise ValueError("WanResampleAdapter does not support patch_dim F (-3); use H (-2) or W (-1).") - if hasattr(wan_resample, "time_conv"): - wan_resample.time_conv = WanCausalConv3dAdapter( - wan_resample.time_conv, + self.resample = resample + if hasattr(resample, "time_conv"): + resample.time_conv = self._conv_adapter( + resample.time_conv, block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - if isinstance(wan_resample.resample, nn.Sequential): - resample = [] - for layer in wan_resample.resample: - if isinstance(layer, nn.Conv2d): - resample.append( - Conv2dAdapter( - layer, - block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - ) - else: - resample.append(layer) - self.resample.resample = nn.Sequential(*resample) - else: - self.resample.resample = wan_resample.resample + parallel_context=parallel_context, + ) + if isinstance(resample.resample, nn.Sequential): + self.resample.resample = nn.Sequential(*[ + Conv2dAdapter( + layer, + block_size=conv_block_size, + parallel_context=parallel_context, + ) if isinstance(layer, nn.Conv2d) else layer + for layer in resample.resample + ]) + + def forward(self, x, feat_cache=None, feat_idx=None): + return self.resample(x, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx)) - def forward(self, x, feat_cache=None, feat_idx=[0]): - return self.resample(x, feat_cache=feat_cache, feat_idx=feat_idx) +class WanResampleAdapter(_CausalResampleAdapter): + _supported = resolved(WanResample) + _requires = "WanResample" + _conv_adapter = WanCausalConv3dAdapter + + +class QwenImageResampleAdapter(_CausalResampleAdapter): + _supported = resolved(QwenImageResample) + _requires = "QwenImageResample" + _conv_adapter = QwenImageCausalConv3dAdapter + + +class _CausalUpBlockAdapter(nn.Module): + """Shards an up block: its residual blocks, and whichever resample it upsamples with""" + + _supported: Tuple[type, ...] = () + _requires: str = "" + _resnet_adapter = None + _resample_adapter = None + _resample_types: Tuple[type, ...] = () + # Which attribute the wrapped block is kept under, since a decoder reaches back through it. + _attr = "up_block" + # Wan threads first_chunk through its up blocks to tell the temporal cache it is starting + # over. The families forked from it dropped that argument. + _takes_first_chunk = True -class WanResidualUpBlockAdapter(nn.Module): def __init__( self, - wan_residual_up_block: WanResidualUpBlock, + up_block: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_residual_up_block, WanResidualUpBlock), ( - "WanResidualUpBlockAdapter does not support up block except WanResidualUpBlock" + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(up_block, self._supported), ( + f"{adapter} does not support up block except {self._requires}" ) - self.residual_up_block = wan_residual_up_block - self.residual_up_block.resnets = nn.ModuleList([ - WanResidualBlockAdapter( - resnet, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) for resnet in wan_residual_up_block.resnets - ]) - if hasattr(wan_residual_up_block, "upsamplers"): - if wan_residual_up_block.upsamplers is not None: - self.residual_up_block.upsamplers = nn.ModuleList([ - WanResampleAdapter( - upsampler, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) if isinstance(upsampler, WanResample) else upsampler - for upsampler in wan_residual_up_block.upsamplers + options = dict( + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + up_block.resnets = nn.ModuleList( + [self._resnet_adapter(resnet, **options) for resnet in up_block.resnets] + ) + if hasattr(up_block, "upsamplers"): + if up_block.upsamplers is not None: + up_block.upsamplers = nn.ModuleList([ + self._resample_adapter(upsampler, **options) + if isinstance(upsampler, self._resample_types) else upsampler + for upsampler in up_block.upsamplers ]) - elif hasattr(wan_residual_up_block, "upsampler"): - if wan_residual_up_block.upsampler is not None: - upsampler = wan_residual_up_block.upsampler - if isinstance(upsampler, WanResample): - self.residual_up_block.upsampler = WanResampleAdapter( - upsampler, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) - - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): - return self.residual_up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) - - -class WanUpBlockAdapter(nn.Module): + elif hasattr(up_block, "upsampler"): + if isinstance(up_block.upsampler, self._resample_types): + up_block.upsampler = self._resample_adapter(up_block.upsampler, **options) + setattr(self, self._attr, up_block) + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + up_block = getattr(self, self._attr) + feat_idx = cache_cursor(feat_idx) + if self._takes_first_chunk: + return up_block( + x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk + ) + return up_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + +class WanResidualUpBlockAdapter(_CausalUpBlockAdapter): + _supported = resolved(WanResidualUpBlock) + _requires = "WanResidualUpBlock" + _resnet_adapter = WanResidualBlockAdapter + _resample_adapter = WanResampleAdapter + _resample_types = resolved(WanResample) + _attr = "residual_up_block" + + +class WanUpBlockAdapter(_CausalUpBlockAdapter): + _supported = resolved(WanUpBlock) + _requires = "WanUpBlock" + _resnet_adapter = WanResidualBlockAdapter + _resample_adapter = WanResampleAdapter + _resample_types = resolved(WanResample) + + +class QwenImageUpBlockAdapter(_CausalUpBlockAdapter): + _supported = resolved(QwenImageUpBlock) + _requires = "QwenImageUpBlock" + _resnet_adapter = QwenImageResidualBlockAdapter + _resample_adapter = QwenImageResampleAdapter + _resample_types = resolved(QwenImageResample) + _takes_first_chunk = False + + +class _PaddedCausalUpsampleAdapter(nn.Module): + """Shards a HunyuanVideo upsampler, which is only its convolution + + What surrounds that convolution is nearest-neighbour interpolation in one family and a + channel-to-space shuffle in the other. Both read a single input position per output one, so + a rank can upsample its own rows knowing nothing about anyone else's. + """ + + _supported: Tuple[type, ...] = () + _requires: str = "" + _conv_adapter = None + def __init__( self, - wan_up_block: WanUpBlock, + upsampler: nn.Module, conv_block_size = 0, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext = None, ): super().__init__() - assert isinstance(wan_up_block, WanUpBlock), ( - "WanUpBlockAdapter does not support up block except WanUpBlock" + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(upsampler, self._supported), ( + f"{adapter} does not support upsampler except {self._requires}" ) - self.up_block = wan_up_block - self.up_block.resnets = nn.ModuleList([ - WanResidualBlockAdapter( - resnet, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) for resnet in wan_up_block.resnets - ]) - if hasattr(wan_up_block, "upsamplers"): - if wan_up_block.upsamplers is not None: - self.up_block.upsamplers = nn.ModuleList([ - WanResampleAdapter( - upsampler, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) if isinstance(upsampler, WanResample) else upsampler - for upsampler in wan_up_block.upsamplers - ]) - elif hasattr(wan_up_block, "upsampler"): - if wan_up_block.upsampler is not None: - upsampler = wan_up_block.upsampler - if isinstance(upsampler, WanResample): - self.up_block.upsampler = WanResampleAdapter( - upsampler, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) - - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): - return self.up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + self.upsampler = upsampler + replace_child_convolution( + upsampler, + self._conv_adapter, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, hidden_states): + return self.upsampler(hidden_states) + + +class HunyuanVideoUpsampleAdapter(_PaddedCausalUpsampleAdapter): + _supported = resolved(HunyuanVideoUpsampleCausal3D) + _requires = "HunyuanVideoUpsampleCausal3D" + _conv_adapter = HunyuanVideoCausalConv3dAdapter + + +class HunyuanVideo15UpsampleAdapter(_PaddedCausalUpsampleAdapter): + _supported = resolved(HunyuanVideo15Upsample) + _requires = "HunyuanVideo15Upsample" + _conv_adapter = HunyuanVideo15CausalConv3dAdapter + + +class _PaddedCausalUpBlockAdapter(nn.Module): + """Shards a HunyuanVideo up block: its residual blocks and its upsampler""" + + _supported: Tuple[type, ...] = () + _requires: str = "" + _resnet_adapter = None + _upsample_adapter = None + + def __init__( + self, + up_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(up_block, self._supported), ( + f"{adapter} does not support up block except {self._requires}" + ) + options = dict( + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + self.up_block = up_block + up_block.resnets = nn.ModuleList( + [self._resnet_adapter(resnet, **options) for resnet in up_block.resnets] + ) + if up_block.upsamplers is not None: + up_block.upsamplers = nn.ModuleList( + [self._upsample_adapter(up, **options) for up in up_block.upsamplers] + ) + + def forward(self, hidden_states): + return self.up_block(hidden_states) + + +class HunyuanVideoUpBlockAdapter(_PaddedCausalUpBlockAdapter): + _supported = resolved(HunyuanVideoUpBlock3D) + _requires = "HunyuanVideoUpBlock3D" + _resnet_adapter = HunyuanVideoResnetBlockAdapter + _upsample_adapter = HunyuanVideoUpsampleAdapter + + +class HunyuanVideo15UpBlockAdapter(_PaddedCausalUpBlockAdapter): + _supported = resolved(HunyuanVideo15UpBlock3D) + _requires = "HunyuanVideo15UpBlock3D" + _resnet_adapter = HunyuanVideo15ResnetBlockAdapter + _upsample_adapter = HunyuanVideo15UpsampleAdapter + + +class LTX2VideoUpsamplerAdapter(nn.Module): + """Shards an LTX-2 upsampler, which is its convolution + + What follows the convolution moves channels into space, reading one input position per + output one, so a rank can do it to its own rows alone. + """ + + _supported = resolved(LTX2VideoUpsampler3d) + _requires = "LTX2VideoUpsampler3d" + + def __init__( + self, + upsampler: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(upsampler, self._supported), ( + f"{adapter} does not support upsampler except {self._requires}" + ) + self.upsampler = upsampler + replace_child_convolution( + upsampler, + LTX2VideoCausalConv3dAdapter, + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + + def forward(self, hidden_states, causal: bool = True): + return self.upsampler(hidden_states, causal=causal) + + +class LTX2VideoUpBlockAdapter(nn.Module): + """Shards an LTX-2 up block: an optional leading residual block, an upsampler, and resnets + + Unlike the other families the upsampler comes before the residual blocks rather than after, + which changes nothing about what has to be sharded, only the order it runs in. + """ + + _supported = resolved(LTX2VideoUpBlock3d) + _requires = "LTX2VideoUpBlock3d" + + def __init__( + self, + up_block: nn.Module, + conv_block_size = 0, + parallel_context: ParallelContext = None, + ): + super().__init__() + adapter = type(self).__name__ + require(self._supported, adapter, self._requires) + assert isinstance(up_block, self._supported), ( + f"{adapter} does not support up block except {self._requires}" + ) + options = dict( + conv_block_size=conv_block_size, + parallel_context=parallel_context, + ) + self.up_block = up_block + if up_block.conv_in is not None: + up_block.conv_in = LTX2VideoResnetBlockAdapter(up_block.conv_in, **options) + if up_block.upsamplers is not None: + up_block.upsamplers = nn.ModuleList( + [LTX2VideoUpsamplerAdapter(up, **options) for up in up_block.upsamplers] + ) + up_block.resnets = nn.ModuleList( + [LTX2VideoResnetBlockAdapter(resnet, **options) for resnet in up_block.resnets] + ) + + def forward(self, hidden_states, temb=None, generator=None, causal: bool = True): + return self.up_block(hidden_states, temb, generator, causal=causal) diff --git a/distvae/modules/adapters/vae/__init__.py b/distvae/modules/adapters/vae/__init__.py index d211178..2fb339f 100644 --- a/distvae/modules/adapters/vae/__init__.py +++ b/distvae/modules/adapters/vae/__init__.py @@ -1,11 +1,40 @@ -# Export decoder adapters -from .decoder_adapters import DecoderAdapter, WanDecoderAdapter +"""Public VAE adapter exports, loaded only when requested.""" -# Export encoder adapters -from .encoder_adapters import WanEncoderAdapter +from importlib import import_module -__all__ = [ + +_DECODERS = ( "DecoderAdapter", + "HunyuanVideo15DecoderAdapter", + "HunyuanVideoDecoderAdapter", + "LTX2VideoDecoderAdapter", + "QwenImageDecoderAdapter", "WanDecoderAdapter", +) +_ENCODERS = ( + "EncoderAdapter", + "HunyuanVideo15EncoderAdapter", + "HunyuanVideoEncoderAdapter", + "LTX2VideoEncoderAdapter", + "QwenImageEncoderAdapter", "WanEncoderAdapter", -] +) +_EXPORTS = { + **{name: "decoder_adapters" for name in _DECODERS}, + **{name: "encoder_adapters" for name in _ENCODERS}, +} + +__all__ = [*_DECODERS, *_ENCODERS] + + +def __getattr__(name): + module_name = _EXPORTS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f"{__name__}.{module_name}"), name) + globals()[name] = value + return value + + +def __dir__(): + return sorted((*globals(), *__all__)) diff --git a/distvae/modules/adapters/vae/causal_setup.py b/distvae/modules/adapters/vae/causal_setup.py new file mode 100644 index 0000000..6851527 --- /dev/null +++ b/distvae/modules/adapters/vae/causal_setup.py @@ -0,0 +1,90 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch.nn as nn +from torch.distributed import ProcessGroup + +from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter +from distvae.modules.patch_utils import DePatchify, Patchify, widest_halo +from distvae.utils import ( + ParallelContext, + normalize_patch_dim, + parallel_context, +) + + +@dataclass(frozen=True) +class CausalVAEAdapterSetup: + """Immutable setup shared by causal encoder and decoder halves.""" + + adapter: str + conv_adapter: object + block_adapters: Tuple[Tuple[Optional[type], object], ...] + conv_block_size: object + patch_dim: int + parallel_context: ParallelContext + + @classmethod + def create( + cls, + *, + adapter, + conv_adapter, + block_adapters, + conv_block_size, + patch_dim, + vae_group: ProcessGroup, + ): + patch_dim = normalize_patch_dim(patch_dim, 5, spatial_only=True) + return cls( + adapter=adapter, + conv_adapter=conv_adapter, + block_adapters=block_adapters, + conv_block_size=conv_block_size, + patch_dim=patch_dim, + parallel_context=parallel_context(vae_group, patch_dim, ndim=5), + ) + + @property + def options(self): + return {"parallel_context": self.parallel_context} + + def adapt_convolution(self, convolution): + return self.conv_adapter( + convolution, block_size=self.conv_block_size, **self.options + ) + + def adapt_blocks(self, blocks, kind): + return nn.ModuleList([self.adapt_block(one, kind) for one in blocks]) + + def adapt_block(self, block, kind): + for block_type, block_adapter in self.block_adapters: + if block_type is not None and isinstance(block, block_type): + return block_adapter( + block, conv_block_size=self.conv_block_size, **self.options + ) + handled = ", ".join( + block_type.__name__ + for block_type, _ in self.block_adapters + if block_type is not None + ) + raise TypeError( + f"{self.adapter} cannot shard a {kind} block of type " + f"{type(block).__name__}. It handles " + f"{handled or 'no block type the installed diffusers provides'}." + ) + + def adapt_group_norm(self, norm): + if isinstance(norm, nn.GroupNorm): + return GroupNormAdapter(norm, **self.options) + return norm + + def patchers(self, module, scale_factor=1): + return ( + Patchify( + scale_factor=scale_factor, + halo=widest_halo(module), + **self.options, + ), + DePatchify(**self.options), + ) diff --git a/distvae/modules/adapters/vae/decoder_adapters.py b/distvae/modules/adapters/vae/decoder_adapters.py index 61f9f90..783eb8c 100644 --- a/distvae/modules/adapters/vae/decoder_adapters.py +++ b/distvae/modules/adapters/vae/decoder_adapters.py @@ -1,30 +1,67 @@ -import time -from typing import Optional +from typing import List, Optional, Tuple import torch import torch.nn as nn from torch.distributed import ProcessGroup -from torch.profiler import profile, ProfilerActivity from diffusers.models.autoencoders.vae import Decoder from diffusers.models.unets.unet_2d_blocks import UpDecoderBlock2D -from diffusers.models.autoencoders.autoencoder_kl_wan import ( - WanUpBlock, - WanResidualUpBlock, -) -from distvae.models.vae import PatchDecoder -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter, WanCausalConv3dAdapter +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, +) +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + HunyuanVideo15CausalConv3dAdapter, + HunyuanVideoCausalConv3dAdapter, + LTX2VideoCausalConv3dAdapter, + QwenImageCausalConv3dAdapter, + WanCausalConv3dAdapter, +) from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter from distvae.modules.adapters.unets.unet_2d_blocks_adapters import UpDecoderBlock2DAdapter -from distvae.modules.adapters.upsampling_adapters import WanResidualUpBlockAdapter, WanUpBlockAdapter -from distvae.modules.adapters.midblock_adapters import WanMidBlockAdapter +from distvae.modules.adapters.vae.causal_setup import CausalVAEAdapterSetup +from distvae.modules.adapters.upsampling_adapters import ( + HunyuanVideo15UpBlockAdapter, + HunyuanVideoUpBlockAdapter, + LTX2VideoUpBlockAdapter, + QwenImageUpBlockAdapter, + WanResidualUpBlockAdapter, + WanUpBlockAdapter, +) +from distvae.modules.adapters.midblock_adapters import ( + HunyuanVideo15MidBlockAdapter, + HunyuanVideoMidBlockAdapter, + LTX2VideoMidBlockAdapter, + QwenImageMidBlockAdapter, + WanMidBlockAdapter, +) from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.utils import DistributedEnv +from distvae.utils import ( + cache_cursor, + normalize_patch_dim, + parallel_context, +) + +WanUpBlock = block(WAN, "WanUpBlock") +WanResidualUpBlock = block(WAN, "WanResidualUpBlock") +QwenImageUpBlock = block(QWEN_IMAGE, "QwenImageUpBlock") +HunyuanVideoUpBlock3D = block(HUNYUAN_VIDEO, "HunyuanVideoUpBlock3D") +HunyuanVideo15UpBlock3D = block(HUNYUAN_VIDEO_15, "HunyuanVideo15UpBlock3D") +LTX2VideoUpBlock3d = block(LTX2_VIDEO, "LTX2VideoUpBlock3d") + + +def _reject_benchmark_options(use_profiler: bool, verbose: bool): + if use_profiler or verbose: + raise ValueError( + "Decoder adapter profiling and verbose timing moved to the bench harness; " + "run bench/distvae_bench.py for benchmark instrumentation." + ) -try: - import torch_musa -except ModuleNotFoundError: - pass class DecoderAdapter(nn.Module): def __init__( @@ -35,193 +72,218 @@ def __init__( use_profiler: bool = False, verbose: bool = False, conv_block_size = 0, + patch_dim: int = -2, ): super().__init__() - assert isinstance(decoder.conv_norm_out, nn.GroupNorm), "DecoderAdapter does not support normalization method except GroupNorm" + _reject_benchmark_options(use_profiler, verbose) + assert isinstance(decoder.conv_norm_out, nn.GroupNorm), ( + "DecoderAdapter requires conv_norm_out to be nn.GroupNorm" + ) for up_block in decoder.up_blocks: - assert isinstance(up_block, UpDecoderBlock2D), "DecoderAdapter does not support up block except UpDecoderBlock2D" - DistributedEnv.initialize(vae_group) - self.decoder = PatchDecoder() - self.decoder.layers_per_block = decoder.layers_per_block - self.decoder.conv_in = decoder.conv_in - self.decoder.mid_block = decoder.mid_block + assert isinstance(up_block, UpDecoderBlock2D), ( + "DecoderAdapter requires every up block to be UpDecoderBlock2D" + ) + patch_dim = normalize_patch_dim(patch_dim, 4, spatial_only=True) + self.patch_dim = patch_dim + self.parallel_context = parallel_context(vae_group, patch_dim, ndim=4) + options = dict(parallel_context=self.parallel_context) + self.decoder = decoder self.decoder.up_blocks = nn.ModuleList([ - UpDecoderBlock2DAdapter(up_block, conv_block_size=conv_block_size) for up_block in decoder.up_blocks + UpDecoderBlock2DAdapter( + up_block, conv_block_size=conv_block_size, **options + ) for up_block in decoder.up_blocks ]) - self.decoder.conv_norm_out = GroupNormAdapter(decoder.conv_norm_out) + self.decoder.conv_norm_out = GroupNormAdapter(decoder.conv_norm_out, **options) self.decoder.conv_act = decoder.conv_act - self.decoder.conv_out = Conv2dAdapter(decoder.conv_out, block_size=conv_block_size) - self.use_profiler = use_profiler - self.verbose = verbose + self.decoder.conv_out = Conv2dAdapter( + decoder.conv_out, block_size=conv_block_size, **options + ) + self.patch = Patchify(**options) + self.depatch = DePatchify(**options) self.vae_group = vae_group + self.train(decoder.training) def forward( self, sample: torch.FloatTensor, latent_embeds: Optional[torch.FloatTensor] = None, ): - rank = DistributedEnv.get_global_rank() - device_type = DistributedEnv.get_device_type() - start_time = time.time() - elapsed_time = 0 - if self.use_profiler: - if device_type == "musa": - torch.musa.memory._record_memory_history(enabled=None) - activities=[ProfilerActivity.CPU,ProfilerActivity.MUSA] - else: - torch.cuda.memory._record_memory_history(enabled=None) - activities=[ProfilerActivity.CPU,ProfilerActivity.CUDA] - - with profile( - activities=activities, - on_trace_ready=torch.profiler.tensorboard_trace_handler( - f"./profile/patch_vae_{rank}" - ), - profile_memory=True, - with_stack=True, - record_shapes=True, - ) as prof: - output = self.decoder(sample, latent_embeds) - prof.export_memory_timeline(f"patch_vae_profiler_mem_{rank}.html") + if torch.is_grad_enabled(): + raise RuntimeError( + "DecoderAdapter is inference-only; use torch.no_grad() or inference mode " + "(torch.inference_mode())." + ) + + decoder = self.decoder + sample = decoder.conv_in(sample) + upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype + + sample = decoder.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + sample = self.patch(sample) + for up_block in decoder.up_blocks: + sample = up_block(sample, latent_embeds) + + if latent_embeds is None: + sample = decoder.conv_norm_out(sample) else: - output = self.decoder(sample, latent_embeds) + sample = decoder.conv_norm_out(sample, latent_embeds) + sample = decoder.conv_act(sample) + sample = decoder.conv_out(sample) + return self.depatch(sample) - end_time = time.time() - elapsed_time = end_time - start_time - peak_memory = DistributedEnv.get_peak_memory(device_type) - if self.verbose and rank == 0: - print( - f"Decoder: [elapsed_time: {elapsed_time:.2f} sec," - f"peak_memory: {peak_memory/1e9} GB]" - ) - return output +class _CausalDecoderAdapter(nn.Module): + """Shards a causal 3D video decoder across ranks along one spatial axis. + + These decoders share a skeleton: a causal convolution in, a mid block, a run of up blocks, a + normalisation, and a causal convolution out. Where that norm is RMS it needs no sharding, + reducing over channels rather than over the axis being split; GroupNorm requires the + distributed wrapper. The families also use different classes for the remaining layers and + pass different temporal-cache state through their forward methods. + """ + _label = "Decoder" + _conv_adapter = None + _mid_adapter = None + _up_block_adapters: Tuple[Tuple[Optional[type], type], ...] = () + # Wan and the families forked from it thread a temporal cache through every forward so a + # decode can be split into chunks of frames. The HunyuanVideo decoders take a tensor and + # nothing else. + _takes_feature_cache = True + # Of those that do, Wan alone also passes first_chunk, to tell the cache it is starting over. + _takes_first_chunk = True + _setup_type = CausalVAEAdapterSetup -class WanDecoderAdapter(nn.Module): def __init__( - self, - decoder: Decoder, + self, + decoder: nn.Module, vae_group: ProcessGroup = None, *, - use_uniform_patch: bool = True, use_profiler: bool = False, verbose: bool = False, conv_block_size = 0, patch_dim: int = -2, ): super().__init__() - if patch_dim == -3: - raise ValueError("WanDecoderAdapter does not support patch_dim F (-3); use H (-2) or W (-1).") - DistributedEnv.initialize(vae_group) - self.patch_dim = patch_dim - DistributedEnv.set_patch_dim(patch_dim) - self.decoder = decoder - self.decoder.conv_in = WanCausalConv3dAdapter( - decoder.conv_in, block_size=conv_block_size, patch_dim=patch_dim, use_uniform_patch=use_uniform_patch - ) - self.decoder.mid_block = WanMidBlockAdapter( - decoder.mid_block, conv_block_size=conv_block_size, patch_dim=patch_dim, use_uniform_patch=use_uniform_patch + _reject_benchmark_options(use_profiler, verbose) + setup = self._setup_type.create( + adapter=type(self).__name__, + conv_adapter=self._conv_adapter, + block_adapters=self._up_block_adapters, + conv_block_size=conv_block_size, + patch_dim=patch_dim, + vae_group=vae_group, ) - up_blocks = [] - for up_block in decoder.up_blocks: - if isinstance(up_block, WanUpBlock): - up_blocks.append( - WanUpBlockAdapter( - up_block, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - ) - elif isinstance(up_block, WanResidualUpBlock): - up_blocks.append( - WanResidualUpBlockAdapter( - up_block, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - ) - self.decoder.up_blocks = nn.ModuleList(up_blocks) - self.decoder.conv_out = WanCausalConv3dAdapter( - decoder.conv_out, block_size=conv_block_size, patch_dim=patch_dim, use_uniform_patch=use_uniform_patch + self._setup = setup + self.patch_dim = setup.patch_dim + self.parallel_context = setup.parallel_context + self.decoder = decoder + self.decoder.conv_in = setup.adapt_convolution(decoder.conv_in) + self.decoder.mid_block = self._mid_adapter( + decoder.mid_block, conv_block_size=conv_block_size, **setup.options ) - self.patchify = Patchify(patch_dim=patch_dim, use_uniform_patch=use_uniform_patch) - self.depatchify = DePatchify(patch_dim=patch_dim, use_uniform_patch=use_uniform_patch) - self.use_uniform_patch = use_uniform_patch - self.use_profiler = use_profiler - self.verbose = verbose + self.decoder.up_blocks = setup.adapt_blocks(decoder.up_blocks, "up") + self.decoder.conv_out = setup.adapt_convolution(decoder.conv_out) + # HunyuanVideo ends on a GroupNorm, whose statistics span the axis being split. The RMS + # norms the other families end on do not, and are left as they are. + if hasattr(decoder, "conv_norm_out"): + self.decoder.conv_norm_out = setup.adapt_group_norm(decoder.conv_norm_out) + # Read after the whole stack is adapted, so it sees every convolution that will exchange. + self.patchify, self.depatchify = setup.patchers(self.decoder) self.vae_group = vae_group - def _forward( - self, - sample: torch.FloatTensor, - feat_cache: Optional[torch.FloatTensor] = None, - feat_idx: Optional[int] = 0, - first_chunk: bool = False, - patchify: bool = True - ): - if self.use_uniform_patch and not patchify: - raise ValueError("WanDecoderAdapter does not support use_uniform_patch for already patchified inputs.") + def _run_decoder(self, sample, feat_cache, feat_idx, first_chunk): + if not self._takes_feature_cache: + return self.decoder(sample) + feat_idx = cache_cursor(feat_idx) + if self._takes_first_chunk: + return self.decoder( + sample, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk + ) + return self.decoder(sample, feat_cache=feat_cache, feat_idx=feat_idx) - if self.use_uniform_patch: - patch_dim = self.patch_dim if self.patch_dim >= 0 else sample.ndim + self.patch_dim - patch_dim_size = sample.shape[patch_dim] + def _sharded_decode(self, sample: torch.FloatTensor, patchify: bool, run): + """Split the sample across ranks, run the decoder on this rank's share, and reassemble + Kept apart from forward because the families do not agree on what a decoder call looks + like: some thread a temporal cache through it, LTX-2 a timestep embedding. Splitting and + reassembling is the same either way. + """ if patchify: sample = self.patchify(sample) - output = self.decoder(sample, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) - output = self.depatchify(output) - - if self.use_uniform_patch: - group_world_size = DistributedEnv.get_group_world_size() - upsampling_factor = output.shape[patch_dim] // (sample.shape[patch_dim] * group_world_size) - output = output.narrow(patch_dim, 0, patch_dim_size * upsampling_factor) - - return output + return self.depatchify(run(sample)) def forward( self, sample: torch.FloatTensor, feat_cache: Optional[torch.FloatTensor] = None, - feat_idx: Optional[int] = 0, + feat_idx: Optional[List[int]] = None, first_chunk: bool = False, patchify: bool = True, ): - rank = DistributedEnv.get_global_rank() - device_type = DistributedEnv.get_device_type() - start_time = time.time() - elapsed_time = 0 - if self.use_profiler: - if device_type == "musa": - torch.musa.memory._record_memory_history(enabled=None) - activities=[ProfilerActivity.CPU,ProfilerActivity.MUSA] - else: - torch.cuda.memory._record_memory_history(enabled=None) - activities=[ProfilerActivity.CPU,ProfilerActivity.CUDA] - - with profile( - activities=activities, - on_trace_ready=torch.profiler.tensorboard_trace_handler( - f"./profile/patch_vae_{rank}" - ), - profile_memory=True, - with_stack=True, - record_shapes=True, - ) as prof: - output = self._forward(sample, feat_cache, feat_idx, first_chunk, patchify) - prof.export_memory_timeline(f"patch_vae_profiler_mem_{rank}.html") - else: - output = self._forward(sample, feat_cache, feat_idx, first_chunk, patchify) + return self._sharded_decode( + sample, + patchify, + lambda x: self._run_decoder(x, feat_cache, feat_idx, first_chunk), + ) - end_time = time.time() - elapsed_time = end_time - start_time - peak_memory = DistributedEnv.get_peak_memory(device_type) - if self.verbose and rank == 0: - print( - f"WanDecoder: [elapsed_time: {elapsed_time:.2f} sec, " - f"peak_memory: {peak_memory/1e9} GB]" - ) - return output \ No newline at end of file +class WanDecoderAdapter(_CausalDecoderAdapter): + _label = "WanDecoder" + _conv_adapter = WanCausalConv3dAdapter + _mid_adapter = WanMidBlockAdapter + _up_block_adapters = ( + (WanUpBlock, WanUpBlockAdapter), + (WanResidualUpBlock, WanResidualUpBlockAdapter), + ) + + +class QwenImageDecoderAdapter(_CausalDecoderAdapter): + """Qwen-Image's decoder, which is Wan's without the first_chunk argument""" + + _label = "QwenImageDecoder" + _conv_adapter = QwenImageCausalConv3dAdapter + _mid_adapter = QwenImageMidBlockAdapter + _up_block_adapters = ((QwenImageUpBlock, QwenImageUpBlockAdapter),) + _takes_first_chunk = False + + +class HunyuanVideoDecoderAdapter(_CausalDecoderAdapter): + _label = "HunyuanVideoDecoder" + _conv_adapter = HunyuanVideoCausalConv3dAdapter + _mid_adapter = HunyuanVideoMidBlockAdapter + _up_block_adapters = ((HunyuanVideoUpBlock3D, HunyuanVideoUpBlockAdapter),) + _takes_feature_cache = False + + +class HunyuanVideo15DecoderAdapter(_CausalDecoderAdapter): + _label = "HunyuanVideo15Decoder" + _conv_adapter = HunyuanVideo15CausalConv3dAdapter + _mid_adapter = HunyuanVideo15MidBlockAdapter + _up_block_adapters = ((HunyuanVideo15UpBlock3D, HunyuanVideo15UpBlockAdapter),) + _takes_feature_cache = False + + +class LTX2VideoDecoderAdapter(_CausalDecoderAdapter): + """LTX-2's decoder, which takes a timestep embedding where the others take a temporal cache + + The embedding arrives shaped to broadcast over space, so it needs no sharding of its own, + and neither does the channel-to-space shuffle this decoder ends on. + """ + + _label = "LTX2VideoDecoder" + _conv_adapter = LTX2VideoCausalConv3dAdapter + _mid_adapter = LTX2VideoMidBlockAdapter + _up_block_adapters = ((LTX2VideoUpBlock3d, LTX2VideoUpBlockAdapter),) + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + causal: Optional[bool] = None, + patchify: bool = True, + ): + return self._sharded_decode( + hidden_states, patchify, lambda x: self.decoder(x, temb, causal) + ) diff --git a/distvae/modules/adapters/vae/encoder_adapters.py b/distvae/modules/adapters/vae/encoder_adapters.py index 8fcb186..05ef870 100644 --- a/distvae/modules/adapters/vae/encoder_adapters.py +++ b/distvae/modules/adapters/vae/encoder_adapters.py @@ -1,167 +1,336 @@ -from typing import Optional +from typing import List, Optional, Tuple import torch import torch.nn as nn from torch.distributed import ProcessGroup -from distvae.modules.adapters.layers.conv_adapters import WanCausalConv3dAdapter -from distvae.modules.adapters.midblock_adapters import WanMidBlockAdapter +from distvae.modules.adapters.diffusers_blocks import ( + HUNYUAN_VIDEO, + HUNYUAN_VIDEO_15, + LTX2_VIDEO, + QWEN_IMAGE, + WAN, + block, +) from distvae.modules.adapters.downsampling_adapters import ( + HunyuanVideo15DownBlockAdapter, + HunyuanVideoDownBlockAdapter, + LTX2VideoDownBlockAdapter, + QwenImageResampleDownAdapter, + WanResampleDownAdapter, WanResidualDownBlockAdapter, - WanResampleDownAdapter ) -from distvae.modules.adapters.resnet_adapters import WanResidualBlockAdapter -from distvae.modules.adapters.layers.attn_adapters import WanAttentionBlockAdapter -from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.utils import DistributedEnv - -from diffusers.models.autoencoders.autoencoder_kl_wan import ( - WanResidualDownBlock, - WanResidualBlock, - WanResample, - WanAttentionBlock, +from distvae.modules.adapters.layers.attn_adapters import GatheredAttentionAdapter +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + HunyuanVideo15CausalConv3dAdapter, + HunyuanVideoCausalConv3dAdapter, + LTX2VideoCausalConv3dAdapter, + QwenImageCausalConv3dAdapter, + WanCausalConv3dAdapter, +) +from distvae.modules.adapters.midblock_adapters import ( + HunyuanVideo15MidBlockAdapter, + HunyuanVideoMidBlockAdapter, + LTX2VideoMidBlockAdapter, + QwenImageMidBlockAdapter, + WanMidBlockAdapter, +) +from distvae.modules.adapters.resnet_adapters import ( + QwenImageResidualBlockAdapter, + WanResidualBlockAdapter, ) +from distvae.modules.adapters.unets.unet_2d_blocks_adapters import DownEncoderBlock2DAdapter +from distvae.modules.adapters.vae.causal_setup import CausalVAEAdapterSetup +from distvae.modules.patch_utils import Patchify, DePatchify, widest_halo +from distvae.utils import ( + cache_cursor, + normalize_patch_dim, + parallel_context, +) + +from diffusers.models.autoencoders.vae import Encoder +from diffusers.models.unets.unet_2d_blocks import DownEncoderBlock2D + +WanAttentionBlock = block(WAN, "WanAttentionBlock") +WanResample = block(WAN, "WanResample") +WanResidualBlock = block(WAN, "WanResidualBlock") +WanResidualDownBlock = block(WAN, "WanResidualDownBlock") +QwenImageAttentionBlock = block(QWEN_IMAGE, "QwenImageAttentionBlock") +QwenImageResample = block(QWEN_IMAGE, "QwenImageResample") +QwenImageResidualBlock = block(QWEN_IMAGE, "QwenImageResidualBlock") +HunyuanVideoDownBlock3D = block(HUNYUAN_VIDEO, "HunyuanVideoDownBlock3D") +HunyuanVideo15DownBlock3D = block(HUNYUAN_VIDEO_15, "HunyuanVideo15DownBlock3D") +LTX2VideoDownBlock3D = block(LTX2_VIDEO, "LTX2VideoDownBlock3D") + + +class EncoderAdapter(nn.Module): + """Shards the 2D encoder AutoencoderKL and Flux.2 use, over its down blocks alone. + + The mirror of the 2D decoder adapter, which splits after its mid block rather than before. + Here the split is undone before the mid block, so the attention in it and the GroupNorm after + it see the whole feature map and need no sharding of their own. What that costs is running the + narrowest part of the encoder on every rank, and what it buys is that the down blocks, which + carry the image at full size and are the reason to encode in parallel at all, are the part + that gets split. + """ -class WanEncoderAdapter(nn.Module): def __init__( self, - encoder, + encoder: Encoder, vae_group: ProcessGroup = None, *, - use_uniform_patch: bool = True, vae_scale_factor: int = 8, conv_block_size = 0, patch_dim: int = -2, ): super().__init__() - if patch_dim == -3: - raise ValueError("WanEncoderAdapter does not support patch_dim F (-3); use H (-2) or W (-1).") - - DistributedEnv.initialize(vae_group) + adapter = type(self).__name__ + patch_dim = normalize_patch_dim(patch_dim, 4, spatial_only=True) + if patch_dim != -2: + # The resnet adapter this reaches through splits H and says nothing about which axis. + raise ValueError(f"{adapter} only supports patch_dim H (-2).") + for down_block in encoder.down_blocks: + assert isinstance(down_block, DownEncoderBlock2D), ( + f"{adapter} does not support down block except DownEncoderBlock2D" + ) + # A band has to be a whole multiple of what the encoder narrows by, and here that can be + # counted rather than taken on trust: one halving per stage that carries a downsampler. + # A caller working from a config default rather than from the blocks would otherwise cut + # bands that a later stage halves into a row it does not own. + counted = 2 ** sum( + 1 for down_block in encoder.down_blocks if down_block.downsamplers + ) + if vae_scale_factor != counted: + raise ValueError( + f"{adapter} was told this encoder narrows by {vae_scale_factor}, but its " + f"down blocks narrow by {counted}." + ) self.patch_dim = patch_dim - DistributedEnv.set_patch_dim(patch_dim) - self.vae_scale_factor = vae_scale_factor + self.parallel_context = parallel_context(vae_group, patch_dim, ndim=4) self.encoder = encoder - - # Patch the conv_in layer - self.encoder.conv_in = WanCausalConv3dAdapter( + encoder.conv_in = Conv2dAdapter( encoder.conv_in, block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + parallel_context=self.parallel_context, ) - # Patch the down_blocks - down_blocks = [] - for i, down_block in enumerate(encoder.down_blocks): - if isinstance(down_block, WanResidualDownBlock): - # Wan2.2 style: wrapped in WanResidualDownBlock - down_blocks.append( - WanResidualDownBlockAdapter( - down_block, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - ) - elif isinstance(down_block, WanResidualBlock): - # Wan2.1 style: individual residual block - down_blocks.append( - WanResidualBlockAdapter( - down_block, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - ) - ) - elif isinstance(down_block, WanResample): - # Wan2.1 style: individual downsample block - down_blocks.append( - WanResampleDownAdapter( - down_block, - conv_block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch - ) - ) - elif isinstance(down_block, WanAttentionBlock): - # Attention blocks need to see full spatial context, so wrap with adapter - down_blocks.append( - WanAttentionBlockAdapter(down_block, patch_dim=patch_dim) - ) - else: - # Unknown block type - keep as-is and log warning - import warnings - warnings.warn( - f"Unsupported down_block type {type(down_block).__name__} at index {i} in encoder, " - f"keeping original. This may cause issues with parallel VAE." - ) - down_blocks.append(down_block) - self.encoder.down_blocks = nn.ModuleList(down_blocks) - # Patch the mid_block - self.encoder.mid_block = WanMidBlockAdapter( - encoder.mid_block, + encoder.down_blocks = nn.ModuleList([ + DownEncoderBlock2DAdapter( + down_block, + conv_block_size=conv_block_size, + parallel_context=self.parallel_context, + ) + for down_block in encoder.down_blocks + ]) + self.patchify = Patchify( + scale_factor=vae_scale_factor, + parallel_context=self.parallel_context, + halo=widest_halo(self.encoder), + ) + self.depatchify = DePatchify(parallel_context=self.parallel_context) + self.vae_group = vae_group + + def forward(self, sample: torch.FloatTensor): + sample = self.encoder.conv_in(self.patchify(sample)) + for down_block in self.encoder.down_blocks: + sample = down_block(sample) + sample = self.encoder.mid_block(self.depatchify(sample)) + sample = self.encoder.conv_act(self.encoder.conv_norm_out(sample)) + return self.encoder.conv_out(sample) + + +def _gathered(attention: nn.Module, **options) -> nn.Module: + """Adapt an attention block, which needs the whole image rather than a patch of it + + Written as a function so it can sit in a down block table beside the adapters that shard a + convolution, none of whose sizing options a gather has any use for. + """ + return GatheredAttentionAdapter( + attention, + parallel_context=options["parallel_context"], + ) + + +class _CausalEncoderAdapter(nn.Module): + """Shards a causal 3D video encoder across ranks along one spatial axis. + + The mirror of _CausalDecoderAdapter, over the same skeleton read the other way: a causal + convolution in, a run of down blocks, a mid block, a normalisation, a causal convolution + out. What differs is what a band has to be a multiple of. An encoder narrows what it is + handed, so a band is cut in whole multiples of the VAE's spatial ratio and the latent rows it + produces are its own, where a decoder cuts latent rows and multiplies. + """ + + _label = "Encoder" + _conv_adapter = None + _mid_adapter = None + # Which adapter fits which down block class. A family whose blocks are not in the installed + # diffusers leaves None in the type slot, which no block can match. + _down_block_adapters: Tuple[Tuple[Optional[type], object], ...] = () + # Wan and the family forked from it thread a temporal cache through every forward. The + # HunyuanVideo and LTX-2 encoders take a tensor and nothing else. + _takes_feature_cache = True + _setup_type = CausalVAEAdapterSetup + + def __init__( + self, + encoder: nn.Module, + vae_group: ProcessGroup = None, + *, + vae_scale_factor: int = 8, + conv_block_size = 0, + patch_dim: int = -2, + ): + super().__init__() + setup = self._setup_type.create( + adapter=type(self).__name__, + conv_adapter=self._conv_adapter, + block_adapters=self._down_block_adapters, conv_block_size=conv_block_size, patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, + vae_group=vae_group, ) - # Patch the conv_out layer - self.encoder.conv_out = WanCausalConv3dAdapter( - encoder.conv_out, - block_size=conv_block_size, - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch + self._setup = setup + self.patch_dim = setup.patch_dim + self.parallel_context = setup.parallel_context + self.vae_scale_factor = vae_scale_factor + self.encoder = encoder + self.encoder.conv_in = setup.adapt_convolution(encoder.conv_in) + self.encoder.down_blocks = setup.adapt_blocks(encoder.down_blocks, "down") + self.encoder.mid_block = self._mid_adapter( + encoder.mid_block, conv_block_size=conv_block_size, **setup.options ) - self.use_uniform_patch = use_uniform_patch - self.patchify = Patchify( - patch_dim=patch_dim, - use_uniform_patch=use_uniform_patch, - scale_factor=vae_scale_factor, + self.encoder.conv_out = setup.adapt_convolution(encoder.conv_out) + # HunyuanVideo ends on a GroupNorm, whose statistics span the axis being split. The RMS + # norms the other families end on do not, and are left as they are. + if hasattr(encoder, "conv_norm_out"): + self.encoder.conv_norm_out = setup.adapt_group_norm(encoder.conv_norm_out) + # Each band is a whole multiple of what the encoder narrows by, so it starts on the grid + # the strided convolutions step along and the latent rows it produces are its own. + # The scale factor comes from the public VAE orchestration, because some families narrow + # by folding space into channels rather than by convolution stride. Read the halo only + # after the complete stack has been adapted. + self.patchify, self.depatchify = setup.patchers( + self.encoder, vae_scale_factor + ) + self.vae_group = vae_group + + def _run_encoder(self, sample, feat_cache, feat_idx): + if not self._takes_feature_cache: + return self.encoder(sample) + return self.encoder( + sample, feat_cache=feat_cache, feat_idx=cache_cursor(feat_idx) ) - self.depatchify = DePatchify(patch_dim=patch_dim, use_uniform_patch=use_uniform_patch) - def _forward( + def _sharded_encode(self, sample: torch.FloatTensor, patchify: bool, run): + """Split the sample across ranks, encode this rank's share, and reassemble + + Kept apart from forward because the families do not agree on what an encoder call looks + like: some thread a temporal cache through it, LTX-2 takes a causal flag. Splitting and + reassembling is the same either way. + """ + if patchify: + sample = self.patchify(sample) + return self.depatchify(run(sample)) + + def forward( self, sample: torch.FloatTensor, feat_cache: Optional[torch.FloatTensor] = None, - feat_idx: Optional[int] = 0, + feat_idx: Optional[List[int]] = None, patchify: bool = True, ): - """Internal forward with optional patchify.""" - if self.use_uniform_patch and not patchify: - raise ValueError("WanEncoderAdapter does not support use_uniform_patch for already patchified inputs.") + # A one-element list the causal blocks advance in place; see the decoder's forward for + # why it is neither a mutable default nor the bare 0 this used to take. + feat_idx = cache_cursor(feat_idx) + return self._sharded_encode( + sample, patchify, lambda x: self._run_encoder(x, feat_cache, feat_idx) + ) - if self.use_uniform_patch: - patch_dim = self.patch_dim if self.patch_dim >= 0 else sample.ndim + self.patch_dim - patch_dim_size = sample.shape[patch_dim] - if patchify: - sample = self.patchify(sample) - output = self.encoder(sample, feat_cache=feat_cache, feat_idx=feat_idx) - output = self.depatchify(output) +class WanEncoderAdapter(_CausalEncoderAdapter): + """Wan's encoder, whose down blocks come either grouped or one layer at a time + + Wan 2.2 wraps each stage in a WanResidualDownBlock; Wan 2.1 lays the same residual blocks, + attentions and resamples out flat in one list. Both ship, and the encoder class alone does + not say which, so both shapes are handled. + """ + + _label = "WanEncoder" + _conv_adapter = WanCausalConv3dAdapter + _mid_adapter = WanMidBlockAdapter + _down_block_adapters = ( + (WanResidualDownBlock, WanResidualDownBlockAdapter), + (WanResidualBlock, WanResidualBlockAdapter), + (WanResample, WanResampleDownAdapter), + (WanAttentionBlock, _gathered), + ) + + +class QwenImageEncoderAdapter(_CausalEncoderAdapter): + """Qwen-Image's encoder, which is Wan 2.1's laid out flat and renamed - if self.use_uniform_patch: - downsampling_factor = self.vae_scale_factor - output = output.narrow(patch_dim, 0, patch_dim_size // downsampling_factor) + Its resample carries the same zero-pad-then-strided-convolution downsample as Wan's, so the + one thing it does not inherit outright is the residual down block Wan 2.2 groups its stages + into, which Qwen-Image has no equivalent of. + """ - return output + _label = "QwenImageEncoder" + _conv_adapter = QwenImageCausalConv3dAdapter + _mid_adapter = QwenImageMidBlockAdapter + _down_block_adapters = ( + (QwenImageResidualBlock, QwenImageResidualBlockAdapter), + (QwenImageResample, QwenImageResampleDownAdapter), + (QwenImageAttentionBlock, _gathered), + ) + + +class HunyuanVideoEncoderAdapter(_CausalEncoderAdapter): + """HunyuanVideo's encoder, which groups its stages and ends on a GroupNorm + + That norm reduces over the axis being split, so the base wraps it. Its mid block holds + diffusers' own attention, flattened over frames and rows and columns together, which the mid + block adapter gathers around rather than trying to shard. + """ + + _label = "HunyuanVideoEncoder" + _conv_adapter = HunyuanVideoCausalConv3dAdapter + _mid_adapter = HunyuanVideoMidBlockAdapter + _down_block_adapters = ((HunyuanVideoDownBlock3D, HunyuanVideoDownBlockAdapter),) + _takes_feature_cache = False + + +class HunyuanVideo15EncoderAdapter(_CausalEncoderAdapter): + """HunyuanVideo 1.5's encoder, which downsamples by folding space into channels + + It ends on an RMS norm, which reduces over channels and so needs no sharding. Its downsampler + packs each pair of rows and columns into channels, which reads one input position per output + one so long as a rank holds whole pairs of rows, and the bands Patchify cuts do. + """ + + _label = "HunyuanVideo15Encoder" + _conv_adapter = HunyuanVideo15CausalConv3dAdapter + _mid_adapter = HunyuanVideo15MidBlockAdapter + _down_block_adapters = ((HunyuanVideo15DownBlock3D, HunyuanVideo15DownBlockAdapter),) + _takes_feature_cache = False + + +class LTX2VideoEncoderAdapter(_CausalEncoderAdapter): + """LTX-2's encoder, which takes a causal flag where the others take a temporal cache + + Its mid block holds no attention, so nothing here has to be gathered: every layer is a + convolution or a norm that reduces over channels. + """ + + _label = "LTX2VideoEncoder" + _conv_adapter = LTX2VideoCausalConv3dAdapter + _mid_adapter = LTX2VideoMidBlockAdapter + _down_block_adapters = ((LTX2VideoDownBlock3D, LTX2VideoDownBlockAdapter),) def forward( self, - sample: torch.FloatTensor, - feat_cache: Optional[torch.FloatTensor] = None, - feat_idx: Optional[int] = 0, + hidden_states: torch.FloatTensor, + causal: Optional[bool] = None, patchify: bool = True, ): - """ - Forward pass through the encoder. - - Args: - sample: Input tensor to encode - feat_cache: Optional feature cache for temporal consistency - feat_idx: Feature index for caching - patchify: Whether to apply patchify/depatchify (default: True) - - Returns: - Encoded latent tensor - """ - return self._forward(sample, feat_cache, feat_idx, patchify) + return self._sharded_encode(hidden_states, patchify, lambda x: self.encoder(x, causal)) diff --git a/distvae/modules/patch_utils.py b/distvae/modules/patch_utils.py index fdf9050..7155d1d 100644 --- a/distvae/modules/patch_utils.py +++ b/distvae/modules/patch_utils.py @@ -1,89 +1,188 @@ +from typing import List, Tuple + import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist -from distvae.utils import DistributedEnv +from distvae.models.layers.conv_mixin import PatchConvMixin +from distvae.utils import ParallelContext, normalize_patch_dim + + +class VAERowSplitError(ValueError): + """A row-sharded VAE cannot divide this axis into complete processing units.""" + + def __init__(self, rows: int, factor: int): + self.rows = rows + self.factor = factor + super().__init__( + f"Cannot split {rows} rows into multiples of {factor}: the VAE narrows this " + f"axis by {factor}, so every band must contain a whole multiple of {factor} rows." + ) + + +def _patch_axis(conv) -> int: + """Which entry of a convolution's per-axis tuples describes the axis being split""" + patch_dim = conv.patch_dim + if patch_dim < 0: + patch_dim += conv._patch_ndim() + return patch_dim - 2 + + +def widest_halo(module: nn.Module) -> int: + """Return the maximum halo width required by any convolution in ``module``. + + For the split axis, neither halo exceeds ``kernel_size // 2``; stride and padding cancel from + the bound. Therefore the largest kernel in the module determines the maximum required halo. + """ + widest = 0 + # Every patched convolution, by the mixin that gives them their halo rather than by the two + # plain subclasses: AsymmetricZeroPadConv2d exchanges a halo like the others and is neither + # of them, so naming the subclasses left its kernel out of the bound this guard is built from. + for conv in module.modules(): + if not isinstance(conv, PatchConvMixin): + continue + kernel = conv.kernel_size + if isinstance(kernel, tuple): + kernel = kernel[_patch_axis(conv)] + widest = max(widest, kernel // 2) + return widest + + +def gather_patches( + patch: torch.Tensor, + parallel_context: ParallelContext, +) -> Tuple[List[torch.Tensor], List[int]]: + """All-gather patches that need not be the same size along patch_dim + + dist.all_gather insists every rank contributes the same shape, so a rank holding fewer rows + than its neighbours cannot take part directly. Each rank pads its patch out to the widest + before the collective and the padding is sliced off on the far side, so it exists only for + the length of the transfer and never reaches a convolution. + + Returns each rank's patch in rank order, and the sizes, which callers need to locate their + own rows within the whole. + """ + if not isinstance(parallel_context, ParallelContext): + raise TypeError("gather_patches requires a ParallelContext") + patch_dim = patch.ndim + normalize_patch_dim( + parallel_context.patch_dim, patch.ndim, spatial_only=True + ) + group = parallel_context.group + world_size = parallel_context.world_size + + # With one rank, no collection or size discovery is required. The gather operations would + # return the input unchanged. Callers concatenate the returned list, and cat copies the input, + # so returning the original tensor does not introduce aliasing. + if world_size == 1: + return [patch], [patch.shape[patch_dim]] + + gathered_sizes = [ + torch.empty(1, dtype=torch.int64, device=patch.device) for _ in range(world_size) + ] + dist.all_gather( + gathered_sizes, + torch.tensor([patch.shape[patch_dim]], dtype=torch.int64, device=patch.device), + group=group, + ) + sizes = [int(size.item()) for size in gathered_sizes] + widest = max(sizes) + + padded = patch + if patch.shape[patch_dim] < widest: + # torch.nn.functional.pad counts its pairs from the last dimension backwards. + pad = [0] * (2 * patch.ndim) + pad[2 * (patch.ndim - patch_dim - 1) + 1] = widest - patch.shape[patch_dim] + padded = F.pad(patch, tuple(pad)) + + buffers = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(buffers, padded.contiguous(), group=group) + + return [ + buffer.narrow(patch_dim, 0, size) for buffer, size in zip(buffers, sizes) + ], sizes class Patchify(nn.Module): + """Assign each rank one contiguous band along the patch dimension. + + Bands are cut in whole multiples of scale_factor, the amount the VAE narrows or widens this + axis by, so every band begins on the downstream strided-convolution grid. Bands may differ + in size when the axis does not divide evenly, so gathers pad them during transport. + + Padding to an even split changes the computation: convolution and attention propagate the + network's response to padded values into retained rows before any final crop. + + Validate that every band can supply the required halo before any rank enters convolution + communication. This prevents one rank from exiting while other ranks wait for its rows. + """ + def __init__( self, - patch_dim: int = -2, - use_uniform_patch: bool = False, + parallel_context: ParallelContext, scale_factor: int = 1, + halo: int = 0, ): super().__init__() - self.group_world_size = DistributedEnv.get_group_world_size() - self.rank_in_vae_group = DistributedEnv.get_rank_in_vae_group() - self.patch_dim = patch_dim - self.use_uniform_patch = use_uniform_patch + if not isinstance(parallel_context, ParallelContext): + raise TypeError("Patchify requires a ParallelContext") + self.parallel_context = parallel_context + self.group_world_size = parallel_context.world_size + self.rank_in_vae_group = parallel_context.rank + self.patch_dim = parallel_context.patch_dim self.scale_factor = scale_factor + self.halo = halo def forward(self, hidden_state): - patch_dim = self.patch_dim if self.patch_dim >= 0 else hidden_state.ndim + self.patch_dim - if self.use_uniform_patch: - factor = self.scale_factor * self.group_world_size - patch_dim_size = hidden_state.shape[patch_dim] - pad_size = 2 * [0] * hidden_state.ndim - # remember that torch.pad operates on the last dimension first - # pad_size for patch_dim is the number of elements to pad to the next multiple of factor - pad_size[2 * (hidden_state.ndim - patch_dim - 1) + 1] = ( - factor - patch_dim_size % factor - ) % factor - hidden_state = F.pad(hidden_state, tuple(pad_size), mode='constant', value=0) - chunks = torch.chunk(hidden_state, self.group_world_size, dim=patch_dim) - return chunks[self.rank_in_vae_group].clone() + patch_dim = hidden_state.ndim + normalize_patch_dim( + self.patch_dim, hidden_state.ndim, spatial_only=True + ) + size = hidden_state.shape[patch_dim] + factor = max(1, self.scale_factor) + if size % factor: + raise VAERowSplitError(size, factor) + units = size // factor + if units < self.group_world_size: + raise ValueError( + f"Cannot split {size} rows across {self.group_world_size} ranks: that leaves " + f"{units} band{'' if units == 1 else 's'} of {factor} rows to go round. Use at " + f"most {units} rank{'' if units == 1 else 's'} for this VAE." + ) + # The ranks that come first each take one extra band where the count does not divide. + band, remainder = divmod(units, self.group_world_size) + # A unit is the narrowest a band gets: an encoder is on its way down to one row per unit + # and a decoder is on its way up from one. So the thinnest band anyone will hold at any + # point in the run is `band` rows, and a halo wider than that is a rank reaching past its + # neighbour into a rank it does not border. Erring towards refusal for an encoder whose + # widest kernel sits early, where the rows have not been spent yet. + if self.halo > band: + fits = units // self.halo + raise ValueError( + f"Cannot split {size} rows across {self.group_world_size} ranks: that leaves " + f"{band} row{'' if band == 1 else 's'} per rank at the narrowest, and this VAE " + f"has a convolution reaching {self.halo} rows past a band into its neighbour's. " + f"Use at most {fits} rank{'' if fits == 1 else 's'} for this VAE, or tile it " + f"instead." + ) + rank = self.rank_in_vae_group + start = (rank * band + min(rank, remainder)) * factor + length = (band + (1 if rank < remainder else 0)) * factor + # `narrow` alone would retain the complete input storage and a width slice would be + # non-contiguous. Materialize an independent contiguous rank-local shard. + return hidden_state.narrow(patch_dim, start, length).clone() class DePatchify(nn.Module): - def __init__(self, patch_dim: int = -2, use_uniform_patch: bool = False): + def __init__(self, parallel_context: ParallelContext): super().__init__() - self.group_world_size = DistributedEnv.get_group_world_size() - self.rank_in_vae_group = DistributedEnv.get_rank_in_vae_group() - self.local_rank = DistributedEnv.get_local_rank() - self.patch_dim = patch_dim - self.use_uniform_patch = use_uniform_patch + if not isinstance(parallel_context, ParallelContext): + raise TypeError("DePatchify requires a ParallelContext") + self.parallel_context = parallel_context + self.patch_dim = parallel_context.patch_dim def forward(self, patch_hidden_state): - patch_dim = self.patch_dim if self.patch_dim >= 0 else patch_hidden_state.ndim + self.patch_dim - if self.use_uniform_patch: - patch_size_list = [ - torch.tensor( - [patch_hidden_state.shape[patch_dim]], - dtype=torch.int64, - device=patch_hidden_state.device - ) - for _ in range(self.group_world_size) - ] - else: - patch_size_list = [ - torch.empty([1], dtype=torch.int64, device=patch_hidden_state.device) - for _ in range(self.group_world_size) - ] - dist.all_gather( - patch_size_list, - torch.tensor( - [patch_hidden_state.shape[patch_dim]], - dtype=torch.int64, - device=patch_hidden_state.device - ), - group=DistributedEnv.get_vae_group() - ) - hidden_state_shape = list(patch_hidden_state.shape) - patch_hidden_state_list = [] - for i in range(self.group_world_size): - hidden_state_shape[patch_dim] = patch_size_list[i].item() - patch_hidden_state_list.append( - torch.empty( - hidden_state_shape, - dtype=patch_hidden_state.dtype, - device=patch_hidden_state.device - ) - ) - dist.all_gather( - patch_hidden_state_list, - patch_hidden_state.contiguous(), - group=DistributedEnv.get_vae_group() + patch_dim = patch_hidden_state.ndim + normalize_patch_dim( + self.patch_dim, patch_hidden_state.ndim, spatial_only=True ) - return torch.cat(patch_hidden_state_list, dim=patch_dim) + patches, _ = gather_patches(patch_hidden_state, self.parallel_context) + return torch.cat(patches, dim=patch_dim) diff --git a/distvae/utils.py b/distvae/utils.py index 023f7d9..0f9e7a2 100644 --- a/distvae/utils.py +++ b/distvae/utils.py @@ -2,87 +2,78 @@ import torch.distributed as dist from torch.distributed import ProcessGroup import os +from dataclasses import dataclass +from typing import List, Optional, Tuple -try: - import torch_musa -except ModuleNotFoundError: - pass -class DistributedEnv: - _vae_group = None - _local_rank = None - _world_size = None # 添加新的类变量 - _patch_dim = -2 # -3=F, -2=H, -1=W; same for 2D/3D +def cache_cursor(feat_idx: Optional[List[int]]) -> List[int]: + """The caller's position in the feature cache, or a fresh one at the start of it - @classmethod - def initialize(cls, vae_group: ProcessGroup): - if vae_group is None: - cls._vae_group = dist.group.WORLD - else: - cls._vae_group = vae_group - cls._local_rank = int(os.environ.get('LOCAL_RANK', 0)) # FIXME: in ray all local_rank is 0 - cls._rank_mapping = None - cls._init_rank_mapping() - - @classmethod - def get_vae_group(cls) -> ProcessGroup: - if cls._vae_group is None: - raise RuntimeError("DistributedEnv not initialized. Call initialize() first.") - return cls._vae_group + The causal video decoders walk their feature cache with a one-element list, advancing it as + each layer takes its slot. That cursor cannot be a default argument: Python binds one list + per function at definition, so every call omitting it would share the same one, and a second + decode would carry on reading from wherever the first one stopped. What that gives is not an + error but a video conditioned on the tail of the previous decode. + """ + return [0] if feat_idx is None else feat_idx - @classmethod - def get_global_rank(cls) -> int: - return dist.get_rank() - - @classmethod - def _init_rank_mapping(cls): - """Initialize the mapping between group ranks and global ranks""" - if cls._rank_mapping is None: - # Get all ranks in the group - ranks = [None] * cls.get_group_world_size() - dist.all_gather_object(ranks, cls.get_global_rank(), group=cls.get_vae_group()) - cls._rank_mapping = ranks - @classmethod - def get_global_rank_from_group_rank(cls, group_rank: int) -> int: - """Convert a rank in VAE group to global rank using cached mapping. - - Args: - group_rank: The rank in VAE group - - Returns: - The corresponding global rank - - Raises: - RuntimeError: If the group_rank is invalid - """ - if cls._rank_mapping is None: - cls._init_rank_mapping() - - if group_rank < 0 or group_rank >= cls.get_group_world_size(): - raise RuntimeError(f"Invalid group rank: {group_rank}. Must be in range [0, {cls.get_group_world_size()-1}]") - - return cls._rank_mapping[group_rank] - - @classmethod - def get_rank_in_vae_group(cls) -> int: - return dist.get_rank(cls.get_vae_group()) +def normalize_patch_dim(patch_dim: int, ndim: int, *, spatial_only: bool = False) -> int: + """Return a canonical negative patch axis after validating it for the tensor rank.""" + if not isinstance(patch_dim, int) or isinstance(patch_dim, bool): + raise ValueError(f"patch_dim must be an integer, got {patch_dim!r}") + if ndim not in (4, 5): + raise ValueError(f"patch_dim validation supports 4D or 5D tensors, got {ndim}D") + positive = patch_dim if patch_dim >= 0 else ndim + patch_dim + if positive < 2 or positive >= ndim: + raise ValueError(f"patch_dim {patch_dim} is not a data axis of a {ndim}D tensor") + if spatial_only and ndim == 5 and positive == 2: + raise ValueError( + f"patch_dim {patch_dim} selects the frame axis; only H (-2 or 3) and " + "W (-1 or 4) are supported" + ) + return positive - ndim - @classmethod - def get_group_world_size(cls) -> int: - return dist.get_world_size(cls.get_vae_group()) - @classmethod - def set_patch_dim(cls, dim: int): - cls._patch_dim = dim +@dataclass(frozen=True) +class ParallelContext: + """Immutable distributed settings owned by one adapted VAE.""" - @classmethod - def get_patch_dim(cls) -> int: - return cls._patch_dim + group: Optional[ProcessGroup] + rank: int + world_size: int + patch_dim: int + global_ranks: Tuple[int, ...] = () + + def global_rank(self, group_rank: int) -> int: + if self.global_ranks: + return self.global_ranks[group_rank] + if self.world_size == 1: + return dist.get_rank() if dist.is_initialized() else 0 + return dist.get_global_rank(self.group, group_rank) + +def parallel_context( + vae_group: Optional[ProcessGroup], patch_dim: int, *, ndim: int +) -> ParallelContext: + """Capture one adapter's group and axis without changing process-global state.""" + group = dist.group.WORLD if vae_group is None else vae_group + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + global_ranks = tuple(dist.get_global_rank(group, one) for one in range(world_size)) + return ParallelContext( + group=group, + rank=rank, + world_size=world_size, + patch_dim=normalize_patch_dim(patch_dim, ndim, spatial_only=True), + global_ranks=global_ranks, + ) + + +class DistributedEnv: @classmethod def get_local_rank(cls) -> int: - return cls._local_rank + return int(os.environ.get("LOCAL_RANK", 0)) @classmethod def get_device(cls) -> torch.device: @@ -109,7 +100,10 @@ def get_torch_distributed_backend(cls) -> str: elif hasattr(torch, "musa") and torch.musa.is_available(): return "mccl" else: - raise NotImplementedError("No Accelerators(NV/MTT GPU accelerators) available") + # Sharding is correctness-testable without an accelerator, and gloo is the only + # backend that gets there. Raising instead would make every distributed entry point + # unreachable on a CPU-only machine, tests included. + return "gloo" @classmethod def record_memory_history(cls): diff --git a/distvae/vae/__init__.py b/distvae/vae/__init__.py new file mode 100644 index 0000000..5e6b52e --- /dev/null +++ b/distvae/vae/__init__.py @@ -0,0 +1,52 @@ +"""Public VAE orchestration APIs for DistVAE.""" + +from distvae.utils import ParallelContext +from distvae.modules.patch_utils import VAERowSplitError + +from .parallel import ( + decoder_adapter_name, + encoder_adapter_name, + encoder_scale_factor, + parallelize_decoder, + parallelize_encoder, +) +from .tile_parallel import ( + context_of, + mark, + sharing, +) +from .tiling import ( + apply_tile_plan, + is_tile_padding_error, + latent_rows, + require_vae_support, + supports_tile_parallel, + tile_overlap, + tile_overlap_plan, + tile_shape, + tile_shape_plan, + tiled_decode_for, +) + +__all__ = [ + "ParallelContext", + "VAERowSplitError", + "apply_tile_plan", + "context_of", + "decoder_adapter_name", + "encoder_adapter_name", + "encoder_scale_factor", + "is_tile_padding_error", + "latent_rows", + "mark", + "parallelize_decoder", + "parallelize_encoder", + "require_vae_support", + "sharing", + "supports_tile_parallel", + "tile_overlap", + "tile_overlap_plan", + "tile_shape", + "tile_shape_plan", + "tiled_decode_for", +] diff --git a/distvae/vae/parallel.py b/distvae/vae/parallel.py new file mode 100644 index 0000000..b4ff924 --- /dev/null +++ b/distvae/vae/parallel.py @@ -0,0 +1,298 @@ +"""Which DistVAE adapter, if any, can shard a given diffusers VAE, and how to size it. + +DistVAE shards a VAE by rebuilding it out of sharded convolutions, norms and upsampling, so an +adapter only fits a decoder or encoder assembled from the blocks it was written against. Which +adapter fits which VAE class, and the numbers an adapter has to be told about the VAE, are knowledge +about those two libraries and nothing else. Keeping it here means an integration adds support by +declaring it rather than by carrying its own copy of the wiring. +""" + +import importlib +from typing import NamedTuple, Optional, Tuple + +import torch.nn as nn + +DECODER_MODULE = "distvae.modules.adapters.vae.decoder_adapters" +ENCODER_MODULE = "distvae.modules.adapters.vae.encoder_adapters" +# Adapters by name rather than import, so an installed DistVAE predating one of them fails naming +# the adapter it lacks instead of on importing this module. +TWO_D = "DecoderAdapter" +WAN = "WanDecoderAdapter" +QWEN_IMAGE = "QwenImageDecoderAdapter" +HUNYUAN_VIDEO = "HunyuanVideoDecoderAdapter" +HUNYUAN_VIDEO_15 = "HunyuanVideo15DecoderAdapter" +LTX2_VIDEO = "LTX2VideoDecoderAdapter" + +TWO_D_ENCODER = "EncoderAdapter" + + +class _Family(NamedTuple): + """One VAE family: the adapters that fit its two halves, and the blocks that identify them""" + + decoder: str + encoder: str + module: str + up_blocks: Tuple[str, ...] + down_blocks: Tuple[str, ...] + mid_block: Tuple[str, ...] + + +# Each half of a family is recognised by the classes its blocks are built from, named here rather +# than imported: a VAE that arrived after the installed diffusers leaves its entry empty and +# matches nothing, instead of failing this module's import for every other VAE. The two halves are +# recognised separately rather than one from the other, because sharding either one replaces its +# blocks with adapters, and the half already sharded would no longer answer to anything. +_FAMILIES = ( + _Family( + WAN, + "WanEncoderAdapter", + "autoencoder_kl_wan", + ("WanUpBlock", "WanResidualUpBlock"), + # Wan 2.2 groups each encoder stage into a WanResidualDownBlock; 2.1 lays the same + # residual blocks, attentions and resamples out flat in one list. + ( + "WanResidualDownBlock", + "WanResidualBlock", + "WanAttentionBlock", + "WanResample", + ), + ("WanMidBlock",), + ), + _Family( + QWEN_IMAGE, + "QwenImageEncoderAdapter", + "autoencoder_kl_qwenimage", + ("QwenImageUpBlock",), + ("QwenImageResidualBlock", "QwenImageAttentionBlock", "QwenImageResample"), + ("QwenImageMidBlock",), + ), + _Family( + HUNYUAN_VIDEO, + "HunyuanVideoEncoderAdapter", + "autoencoder_kl_hunyuan_video", + ("HunyuanVideoUpBlock3D",), + ("HunyuanVideoDownBlock3D",), + ("HunyuanVideoMidBlock3D",), + ), + _Family( + HUNYUAN_VIDEO_15, + "HunyuanVideo15EncoderAdapter", + "autoencoder_kl_hunyuanvideo15", + ("HunyuanVideo15UpBlock3D",), + ("HunyuanVideo15DownBlock3D",), + ("HunyuanVideo15MidBlock",), + ), + _Family( + LTX2_VIDEO, + "LTX2VideoEncoderAdapter", + "autoencoder_kl_ltx2", + ("LTX2VideoUpBlock3d",), + ("LTX2VideoDownBlock3D",), + ("LTX2VideoMidBlock3d",), + ), +) + + +def _blocks(module: str, names: Tuple[str, ...]) -> Tuple[type, ...]: + """Those of these block classes the installed diffusers has""" + try: + found = importlib.import_module(f"diffusers.models.autoencoders.{module}") + except ImportError: + return () + return tuple( + block + for block in (getattr(found, name, None) for name in names) + if isinstance(block, type) + ) + + +def _family_of(half, attr: str) -> Optional[_Family]: + """The family this half of a VAE belongs to, by the blocks it is assembled from + + Named for the attribute holding them, up_blocks or down_blocks, which is what _Family calls + the classes it expects to find there too. + """ + blocks = tuple(getattr(half, attr, None) or ()) + mid_block = getattr(half, "mid_block", None) + for family in _FAMILIES: + types = _blocks(family.module, getattr(family, attr)) + mid_types = _blocks(family.module, family.mid_block) + # The mid block is checked too because these families fork one another closely enough + # that the blocks either side of it would not tell two of them apart. + if types and mid_types and all(isinstance(block, types) for block in blocks): + if isinstance(mid_block, mid_types): + return family + return None + + +def decoder_adapter_name(vae) -> Optional[str]: + """The DistVAE adapter that fits this VAE's decoder, None when none does""" + # The adapters assert this themselves, from inside a half-built replacement decoder. Asking + # first keeps an unsupported VAE from reaching that point, and lets a model be told it is + # unsupported rather than shown an assertion from a library it did not name. + decoder = getattr(vae, "decoder", None) + up_blocks = tuple(getattr(decoder, "up_blocks", None) or ()) + if not up_blocks: + return None + + from diffusers.models.unets.unet_2d_blocks import UpDecoderBlock2D + + if all(isinstance(block, UpDecoderBlock2D) for block in up_blocks) and isinstance( + getattr(decoder, "conv_norm_out", None), nn.GroupNorm + ): + return TWO_D + + family = _family_of(decoder, "up_blocks") + if family is None: + return None + return None if _injects_noise(decoder) else family.decoder + + +def encoder_adapter_name(vae) -> Optional[str]: + """The DistVAE adapter that fits this VAE's encoder, None when none does""" + encoder = getattr(vae, "encoder", None) + down_blocks = tuple(getattr(encoder, "down_blocks", None) or ()) + if not down_blocks: + return None + + from diffusers.models.unets.unet_2d_blocks import DownEncoderBlock2D + + # The 2D encoder adapter asks only this of it: everything after the down blocks runs whole on + # every rank, so the norm it ends on is its own business in a way the decoder's is not. + if all(isinstance(block, DownEncoderBlock2D) for block in down_blocks): + return TWO_D_ENCODER + + family = _family_of(encoder, "down_blocks") + if family is None: + return None + # No LTX-2 encoder injects noise, since only its decoder is offered the option, but the + # residual block adapter refuses either half that does and this is what asks first. + return None if _injects_noise(encoder) else family.encoder + + +def _injects_noise(half) -> bool: + """Whether this half of an LTX-2 VAE adds noise inside its residual blocks""" + # DistVAE cannot shard one that does: each rank would draw noise for its own rows, and the + # ranks together would not reconstruct what one rank draws. It refuses from inside a half it + # has already half-replaced, so this asks first. No released LTX-2 checkpoint turns it on. + return any( + getattr(block, "per_channel_scale1", None) is not None + or getattr(block, "per_channel_scale2", None) is not None + for block in half.modules() + ) + + +def _patch_size(vae) -> Optional[int]: + """Return the VAE's spatial patching factor applied outside its convolution stack.""" + # Wan stores one scalar factor, which is the only layout either adapter can use. Flux 2 stores + # its boundary pixel-unshuffle factor as `(2, 2)`; that tuple does not describe adapter-level + # patching, so anything other than one number is treated as no patching. + patch_size = getattr(vae.config, "patch_size", None) + return patch_size if isinstance(patch_size, int) and patch_size > 1 else None + + +def _two_d_scale_factor(vae) -> Optional[int]: + """A 2D encoder's ratio, counted off its stages rather than read from its config""" + # These VAEs record no spatial ratio, and the 8 every shipped one comes to is a consequence of + # having four stages rather than a number stated anywhere. Counting the stages that downsample + # gets it right for a checkpoint with some other number of them. + from diffusers.models.unets.unet_2d_blocks import DownEncoderBlock2D + + blocks = tuple(getattr(getattr(vae, "encoder", None), "down_blocks", None) or ()) + if not blocks or not all(isinstance(block, DownEncoderBlock2D) for block in blocks): + return None + return 2 ** sum(1 for block in blocks if block.downsamplers) + + +def encoder_scale_factor(vae) -> int: + """The encoder's own spatial downsampling, which is what the encoder adapter shards by""" + counted = _two_d_scale_factor(vae) + if counted is not None: + return counted + # A VAE that patches folds that factor into its spatial ratio, and the adapter needs the conv + # stack's share of it alone: Cosmos 3's 16 is 8 from the encoder and 2 from patching. + factor = getattr(vae.config, "scale_factor_spatial", None) or 8 + patch_size = _patch_size(vae) + return factor // patch_size if patch_size else factor + + +def _adapter(module: str, name: str, vae) -> type: + try: + return getattr(importlib.import_module(module), name) + except (ImportError, AttributeError) as e: + raise ValueError( + f"The installed DistVAE does not provide {name}, which this VAE " + f"({type(vae).__name__}) needs. Try installing the latest DistVAE from " + f"https://github.com/xdit-project/DistVAE." + ) from e + + +def _keep_causal_cache_length(vae) -> None: + """Hold a recounting feature cache at the length the unsharded VAE would have given it + + A causal VAE threads a list of cached frames through a call, one entry per causal + convolution, and Qwen-Image sizes that list by counting those convolutions afresh every + time, asking isinstance against its own class. Sharding replaces every one of them with an + adapter, so the count comes to zero, the list is empty, and the first convolution to reach + for its entry indexes off the end of it. Wan counts once when it is built and never notices. + + Counting before the replacement and holding the answer is what the VAE would have done for + itself had it cached the count, and has to happen before either half is replaced, which is + why both entry points call it first and only the first call does anything. + """ + if getattr(vae, "_distvae_cache_length_kept", False) or not hasattr( + vae, "clear_cache" + ): + return + recount = vae.clear_cache + recount() + counts = { + name: value for name, value in vars(vae).items() if name.endswith("conv_num") + } + if not counts: + return + + def clear_cache(): + recount() + for name, count in counts.items(): + setattr(vae, name, count) + # _conv_num goes with _feat_map, _enc_conv_num with _enc_feat_map. + setattr(vae, f"{name[: -len('conv_num')]}feat_map", [None] * count) + + vae.clear_cache = clear_cache + vae._distvae_cache_length_kept = True + + +def parallelize_decoder(vae, vae_group) -> str: + """Replace this VAE's decoder with a sharded one, returning the adapter that did it""" + name = decoder_adapter_name(vae) + if name is None: + raise ValueError( + f"DistVAE cannot shard this VAE decoder ({type(vae).__name__}): no adapter " + f"matches its decoder blocks. Use Diffusers VAE tiling to lower decode memory instead." + ) + _keep_causal_cache_length(vae) + decoder = _adapter(DECODER_MODULE, name, vae)(vae.decoder, vae_group=vae_group) + patch_size = _patch_size(vae) + # The adapter crops its output by the ratio it upsamples, and its patchify assumes no patching + # because Wan does none. A VAE that patches upsamples by that much again. + if patch_size and hasattr(decoder, "patchify"): + decoder.patchify.scale_factor = patch_size + vae.decoder = decoder.to(vae.device) + return name + + +def parallelize_encoder(vae, vae_group) -> str: + """Replace this VAE's encoder with a sharded one, returning the adapter that did it""" + name = encoder_adapter_name(vae) + if name is None: + raise ValueError( + f"Parallel VAE encoding is not available for this VAE ({type(vae).__name__}): " + f"DistVAE has no adapter for its encoder blocks." + ) + _keep_causal_cache_length(vae) + adapter = _adapter(ENCODER_MODULE, name, vae) + vae.encoder = adapter( + vae.encoder, vae_group=vae_group, vae_scale_factor=encoder_scale_factor(vae) + ).to(vae.device) + return name diff --git a/distvae/vae/tile_parallel.py b/distvae/vae/tile_parallel.py new file mode 100644 index 0000000..4e15d99 --- /dev/null +++ b/distvae/vae/tile_parallel.py @@ -0,0 +1,590 @@ +"""Distribute complete VAE tiles among the ranks of a process group. + +Tiling and sharding both split a VAE decode, and composing them splits it twice. DistVAE shards +the rows of each tile independently, so every tile requires Patchify, a halo exchange per +convolution, a reduction per norm, and a gather. This communication cost is per tile rather than +per pixel, so it increases as the window narrows and the tile count grows. + +Tiles are independent, although rows within a tile are not. Distributing complete tiles requires +two exchanges for the full decode regardless of tile count, and each rank decodes its assigned +tiles without row sharding. + +The caller supplies one callable per decoder invocation and receives every result on every rank +in call order. +""" + +import functools +import math +import warnings +from typing import Callable, Dict, List, NamedTuple, Optional, Sequence, Set, Tuple + +import torch +import torch.distributed as dist + +from distvae.utils import ParallelContext + +# Recorded on the VAE itself, because the decision is made when the decoder would otherwise be +# sharded and acted on later, when the tile window is settled and the decode is installed. +GROUP_ATTR = "_distvae_tile_parallel_context" + +Call = Callable[[], torch.Tensor] +Dispatch = Callable[[Sequence[Call]], List[torch.Tensor]] + +Where = Tuple[int, int] +Decode = Callable[[Sequence[Where]], Dict[Where, torch.Tensor]] + +# Both diffusers tiling loops blend down the second-from-last axis and across the last one, on a +# 4D sample and a 5D one alike, so the assembly below needs no axis of its own to be told. +DOWN, ACROSS = -2, -1 + +# Every rank independently searches the same tile assignment, so the bound must be deterministic: +# a wall-clock deadline could leave ranks with different owners. This budget permits exhaustive +# levelling for small grids while large grids keep the already weight-balanced contiguous runs. +MAX_LEVEL_CANDIDATES = 100_000 + + +class Blend(NamedTuple): + """Functions and dimensions used by Diffusers tiling loops to combine adjacent tiles.""" + + down: Callable # blend_v: mixes a tile's first `deep_down` rows with the tile above's last + across: Callable # blend_h: mixes its first `deep_across` columns with the left tile's last + deep_down: int + deep_across: int + crop: Callable[ + [torch.Tensor], torch.Tensor + ] # the corner of a blended tile that is kept + # The configured window gives every rank the same whole-tile dimensions. A rank-local decoded + # tile may be clipped; using it could make one rank fall back while the others enter a gather. + tile_down: int + tile_across: int + + +def mark(vae, context: ParallelContext) -> None: + """Record this VAE's immutable tile-distribution context.""" + if not isinstance(context, ParallelContext): + raise TypeError("tile-parallel metadata requires a ParallelContext") + setattr(vae, GROUP_ATTR, context) + + +def context_of(vae) -> Optional[ParallelContext]: + """Return this VAE's tile-parallel context, if one was recorded.""" + return getattr(vae, GROUP_ATTR, None) + + +def group_of(vae): + """Return the process group recorded for this VAE's tile distribution.""" + context = context_of(vae) + return context.group if context is not None else None + + +def _distributed(context_or_group): + """Return group, rank, and size from a context or a direct group argument.""" + if isinstance(context_or_group, ParallelContext): + return ( + context_or_group.group, + context_or_group.rank, + context_or_group.world_size, + ) + group = context_or_group + return group, dist.get_rank(group), dist.get_world_size(group) + + +def in_order(calls: Sequence[Call]) -> List[torch.Tensor]: + """Execute all calls sequentially in input order.""" + return [call() for call in calls] + + +def dispatch_over(group) -> Dispatch: + """Distribute calls across `group` and return all results on every rank.""" + group, rank, world_size = _distributed(group) + if world_size < 2: + return in_order + + def dispatch(calls: Sequence[Call]) -> List[torch.Tensor]: + # Fewer calls than ranks and some rank contributes nothing to the exchange, with no + # tensor of its own to take a dtype and a device from. A decode that small is a tile or + # two, so every rank simply making every call costs less than arranging not to. + if len(calls) < world_size: + return in_order(calls) + made = [ + call() if n % world_size == rank else None for n, call in enumerate(calls) + ] + return _share(made, group, world_size) + + return dispatch + + +def sharing(group) -> Tuple[Dispatch, Callable]: + """Return contiguous-run assembly and per-call dispatch for a process group. + + Use contiguous-run assembly when every rank can receive a tile and each tile is large enough + to blend locally. Otherwise, distribute individual decoder calls and assemble the results on + every rank. + """ + return dispatch_over(group), functools.partial(assemble_in_runs, group) + + +def runs(weights: Sequence[int], world_size: int) -> List[Tuple[int, int]]: + """Split tiles into one contiguous, weight-balanced run per rank. + + Contiguous runs preserve tiling-loop order and keep most adjacent tiles on the same rank. + Tiles provide finer load balancing than complete grid rows. + + Balance by tile area rather than tile count. Boundary clipping makes tiles in the last grid + row and column smaller, so equal tile counts can assign less work to the last rank. + + Minimize the maximum run weight using binary search over feasible weight limits and a greedy + feasibility check. + """ + if world_size < 2: + return [(0, len(weights))] + low, high = max(weights), sum(weights) + while low < high: + middle = (low + high) // 2 + if len(_greedy(weights, middle)) <= world_size: + high = middle + else: + low = middle + 1 + return _widen(_greedy(weights, low), world_size) + + +def shares(weights: Sequence[int], world_size: int) -> List[int]: + """Return the cached tile assignment as a caller-owned list.""" + return list(_shares(tuple(weights), world_size)) + + +@functools.lru_cache(maxsize=128) +def _shares(weights: Tuple[int, ...], world_size: int) -> Tuple[int, ...]: + """Which rank decodes each tile: contiguous runs, levelled by moving or swapping a few tiles + + A run is the cheap shape to blend, since its tiles' neighbours are mostly its own, but it is + a coarse shape to balance. Nine tiles over four ranks split by weight as evenly as contiguity + allows still leaves the heaviest rank a quarter above the average, because the tiles are large + against the share and a run cannot skip one. No weighing fixes that; only a finer assignment. + + So the runs are a starting point rather than the answer. Within a fixed candidate budget, moves + and pairwise swaps are searched together across every rank pair. Each accepted change strictly + lowers the descending load vector, or keeps that vector while restoring a tile to its original + run. Among equally balanced choices, fewer tiles displaced from those runs win, followed by + tiles already beside their new owner. The budget and total tie-break are deterministic because + every rank computes this independently. + + A move never takes a rank's last tile. Swaps preserve every rank's tile count. + """ + owner: List[int] = [] + for rank, (start, stop) in enumerate(runs(weights, world_size)): + owner.extend([rank] * (stop - start)) + if world_size < 2: + return tuple(owner) + + load = [0] * world_size + for n, weight in enumerate(weights): + load[owner[n]] += weight + + original = owner.copy() + count = [owner.count(rank) for rank in range(world_size)] + displaced = 0 + + def objective(loads, moved): + return tuple(sorted(loads, reverse=True)), moved + + # Conservatively count every move to another rank and every tile pair. Some are skipped below, + # but charging for them makes this a simple hard ceiling independent of the current ownership. + tiles = len(weights) + candidates_per_round = ( + tiles * (world_size - 1) + tiles * (tiles - 1) // 2 + ) + rounds = MAX_LEVEL_CANDIDATES // candidates_per_round + + # Every accepted operation strictly lowers `objective`, so stopping at the budget can only + # leave the assignment no worse than the weighted runs it started from. + for _ in range(rounds): + current = objective(load, displaced) + best = None + + for moved, weight in enumerate(weights): + donor = owner[moved] + if count[donor] == 1: + continue + for receiver in range(world_size): + if receiver == donor: + continue + loads = load.copy() + loads[donor] -= weight + loads[receiver] += weight + next_displaced = displaced + next_displaced -= int(owner[moved] != original[moved]) + next_displaced += int(receiver != original[moved]) + candidate = objective(loads, next_displaced) + if candidate >= current: + continue + beside = any( + 0 <= neighbour < len(weights) and owner[neighbour] == receiver + for neighbour in (moved - 1, moved + 1) + ) + key = (candidate, 0 if beside else 1, 0, donor, receiver, moved) + if best is None or key < best[0]: + best = (key, "move", moved, receiver, loads, next_displaced) + + for first in range(len(weights)): + first_rank = owner[first] + for second in range(first + 1, len(weights)): + second_rank = owner[second] + if first_rank == second_rank: + continue + loads = load.copy() + loads[first_rank] += weights[second] - weights[first] + loads[second_rank] += weights[first] - weights[second] + next_displaced = displaced + next_displaced -= int(first_rank != original[first]) + next_displaced -= int(second_rank != original[second]) + next_displaced += int(second_rank != original[first]) + next_displaced += int(first_rank != original[second]) + candidate = objective(loads, next_displaced) + if candidate >= current: + continue + + def rank_after(tile): + if tile == first: + return second_rank + if tile == second: + return first_rank + return owner[tile] + + beside = 0 + for tile, receiver in ( + (first, second_rank), + (second, first_rank), + ): + beside += not any( + 0 <= neighbour < len(weights) + and rank_after(neighbour) == receiver + for neighbour in (tile - 1, tile + 1) + ) + key = ( + candidate, + beside, + 1, + first_rank, + second_rank, + first, + second, + ) + if best is None or key < best[0]: + best = ( + key, + "swap", + first, + second, + loads, + next_displaced, + ) + + if best is None: + break + _, operation, first, second, load, displaced = best + if operation == "move": + donor = owner[first] + owner[first] = second + count[donor] -= 1 + count[second] += 1 + else: + owner[first], owner[second] = owner[second], owner[first] + return tuple(owner) + + +def _greedy(weights: Sequence[int], ceiling: int) -> List[Tuple[int, int]]: + """The fewest contiguous runs none of which weighs more than `ceiling`""" + out, start, carried = [], 0, 0 + for at, weight in enumerate(weights): + if carried and carried + weight > ceiling: + out.append((start, at)) + start, carried = at, 0 + carried += weight + out.append((start, len(weights))) + return out + + +def _widen(split: List[Tuple[int, int]], world_size: int) -> List[Tuple[int, int]]: + """Enough runs for every rank, by halving the ones holding most tiles + + A ceiling that a few ranks can meet leaves the rest with nothing, and a rank holding no tile + has no tensor of its own to take a dtype and a device from. Halving cannot raise the heaviest + run, so nothing found above is given up here. + """ + while len(split) < world_size: + widest = max(range(len(split)), key=lambda n: split[n][1] - split[n][0]) + start, stop = split[widest] + if stop - start < 2: + break # fewer tiles than ranks, which the caller declines before asking + middle = (start + stop) // 2 + split[widest : widest + 1] = [(start, middle), (middle, stop)] + return split + + +def assemble_in_runs( + group, + rows: int, + columns: int, + decode: Decode, + blend: Blend, + weights: Sequence[int], +) -> Optional[torch.Tensor]: + """Assemble a tile grid with each rank decoding and blending its own run, None if it can't + + Dealing tiles out divides the decoding and leaves the blending on every rank, a cost that + does not shrink however many ranks join the group. Giving a rank a share of neighbouring + tiles lets it blend its own and send only the finished pieces, so the blending divides too. + + The reason a share can be blended alone is a property of the two blends. `blend_v` writes a + tile's *first* rows and `blend_h` its *first* columns, so neither ever writes the last rows or + the last columns - and those are the only parts of a tile that the tiles after it read. A + tile's edges are therefore final while it is still raw, and one exchange of raw edges lets + every rank blend its run exactly as a single rank walking the whole grid would, waiting on + nobody else's blending. + + What comes back is each rank's cropped tiles, which are disjoint and tile the image exactly, + so the gather carries the image once rather than every overlapping tile. + + Where a tile is smaller than twice the blend the argument fails, because the rows and columns + the blends write would reach into the ones their neighbours read. Every reason to decline is + one every rank reaches the same way, from the grid and the window rather than from the tiles + a rank happens to hold: a rank that fell back alone would leave the others waiting in a + gather it never joins, which hangs a decode rather than failing it. + """ + group, rank, world_size = _distributed(group) + if world_size < 2: + return None + order = [(i, j) for i in range(rows) for j in range(columns)] + # Fewer tiles than ranks and some rank would hold nothing, with no tensor of its own to take a + # dtype and a device from. A decode that small has nothing worth dividing anyway. + if len(order) < world_size: + if rank == 0: + warnings.warn( + f"VAE tile grid has {len(order)} tiles for {world_size} ranks; " + f"whole-tile distribution is disabled and every rank will decode all " + f"{len(order)} tiles locally. Use fewer VAE ranks or a smaller tile window.", + RuntimeWarning, + stacklevel=2, + ) + return None + if ( + blend.tile_down < 2 * blend.deep_down + or blend.tile_across < 2 * blend.deep_across + ): + return None + + owner = shares(weights, world_size) + mine = decode([at for n, at in enumerate(order) if owner[n] == rank]) + + # Exchanged before anything is blended, both because the edges are raw at that point and + # because a rank waiting on a neighbour's blending would serialise what this is dividing. + edges = _share_edges( + order, + mine, + owner, + rank, + _wanted(owner, columns, blend), + group, + world_size, + blend, + ) + + blended: Dict[Where, torch.Tensor] = {} + kept: List[Optional[torch.Tensor]] = [None] * len(order) + for n, (i, j) in enumerate(order): + if owner[n] != rank: + continue + tile = mine[(i, j)] + # A blend no rows deep is one the tiles do not overlap enough to need, which a wide + # enough stride leaves. Skipped rather than called with a zero depth, because a depth of + # zero reads as "the whole tile" everywhere an edge is sliced off the end of one. + if i > 0 and blend.deep_down: + # The neighbour itself where this rank blended it, and its edge rebuilt from the raw + # ones otherwise. Both carry the same values; only the cost differs. + above = blended.get((i - 1, j)) + tile = blend.down( + above if above is not None else _edge_above(edges, i, j, blend), + tile, + blend.deep_down, + ) + if j > 0 and blend.deep_across: + left = blended.get((i, j - 1)) + tile = blend.across( + left if left is not None else _edge_left(edges, i, j, blend), + tile, + blend.deep_across, + ) + blended[(i, j)] = tile + kept[n] = blend.crop(tile) + + shared = _share(kept, group, world_size) + return torch.cat( + [ + torch.cat(shared[i * columns : (i + 1) * columns], dim=ACROSS) + for i in range(rows) + ], + dim=DOWN, + ) + + +def assemble_here( + rows: int, columns: int, decode: Decode, blend: Blend +) -> torch.Tensor: + """Assemble the whole grid on this rank, which is what diffusers' own loop does""" + mine = decode([(i, j) for i in range(rows) for j in range(columns)]) + # Both blends write into the tile they are handed, so each tile is blended against neighbours + # that were themselves already blended, and the scan order that makes is part of the result. + made = [] + above: Optional[List[torch.Tensor]] = None + for i in range(rows): + row = [mine[(i, j)] for j in range(columns)] + kept = [] + for j, tile in enumerate(row): + if above is not None: + tile = blend.down(above[j], tile, blend.deep_down) + if j > 0: + tile = blend.across(row[j - 1], tile, blend.deep_across) + row[j] = tile + kept.append(blend.crop(tile)) + made.append(torch.cat(kept, dim=ACROSS)) + above = row + return torch.cat(made, dim=DOWN) + + +def _wanted(owner: Sequence[int], columns: int, blend: Blend) -> Set[int]: + """Return tiles whose unblended edges are needed by another rank. + + A tile blends with its upper and left neighbors. An edge must be transferred only when that + neighbor belongs to another rank. Contiguous assignments usually require one boundary row per + rank; load-balancing moves may add boundaries around the moved tile. + + An axis with zero overlap requires no edge transfer. + """ + wanted: Set[int] = set() + for n, rank in enumerate(owner): + row, column = divmod(n, columns) + if blend.deep_down and row > 0 and owner[n - columns] != rank: + wanted.add(n - columns) + if blend.deep_across and column > 0: + wanted.add(n - columns - 1) + if blend.deep_across and column > 0 and owner[n - 1] != rank: + wanted.add(n - 1) + if blend.deep_down and row > 0: + wanted.add(n - columns - 1) + return wanted + + +def _share_edges( + order: Sequence[Where], + mine: Dict[Where, torch.Tensor], + owner: Sequence[int], + rank: int, + wanted: Set[int], + group, + world_size: int, + blend: Blend, +) -> Dict[Where, Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]]: + """The last rows and last columns of the tiles another rank will read, raw, on every rank + + The edges rather than the tiles: what a neighbour reads is one blend deep, so this carries a + fraction of what dealing the tiles themselves round would have to. + """ + sending: List[Optional[torch.Tensor]] = [None] * (2 * len(order)) + for n in wanted: + if owner[n] != rank: + continue + tile = mine[order[n]] + # Cloned because the blending below writes into the tiles these came from, and an edge is + # only the edge a neighbour needs while it is still raw. Guarded on the depth because a + # slice from -0 is the whole tile rather than none of it, which would send the grid + # itself round in place of its seams. + if blend.deep_down: + sending[2 * n] = tile[..., -blend.deep_down :, :].clone() + if blend.deep_across: + sending[2 * n + 1] = tile[..., -blend.deep_across :].clone() + shared = _share(sending, group, world_size, next(iter(mine.values()))) + return {at: (shared[2 * n], shared[2 * n + 1]) for n, at in enumerate(order)} + + +def _edge_above(edges, i: int, j: int, blend: Blend) -> torch.Tensor: + """The last rows of the tile above, as its own rank would have blended them + + Only its blending across the columns reaches its last rows, and that reads its left + neighbour's last columns, which nothing writes. One blend of two raw edges rebuilds it. + """ + below, _ = edges[(i - 1, j)] + if j == 0 or not blend.deep_across: + return below + left, _ = edges[(i - 1, j - 1)] + return blend.across(left, below.clone(), blend.deep_across) + + +def _edge_left(edges, i: int, j: int, blend: Blend) -> torch.Tensor: + """The last columns of the tile to the left, as its own rank would have blended them + + Only its blending down the rows reaches its last columns, and that reads the corner where the + tile above it meets the tile above and to its left - raw on both counts. + """ + _, beside = edges[(i, j - 1)] + if i == 0 or not blend.deep_down: + return beside + above, _ = edges[(i - 1, j - 1)] + return blend.down(above[..., -blend.deep_across :], beside.clone(), blend.deep_down) + + +def _share( + made: List[Optional[torch.Tensor]], + group, + world_size: int, + like: Optional[torch.Tensor] = None, +) -> List[torch.Tensor]: + """Fill in the calls this rank did not make from the ranks that did. + + `like` supplies edge metadata when this rank has no local tensor of the required shape. The + last run needs it because no following run provides an edge shape. + """ + mine = [(n, tensor) for n, tensor in enumerate(made) if tensor is not None] + + # A rank cannot work out the shape of a call it did not make: tiles at the right and bottom + # edges are clipped by the latent bounds, and a rank can hold none of them. One object + # exchange settles that for the whole decode. + manifest: List = [None] * world_size + dist.all_gather_object( + manifest, [(n, tuple(t.shape)) for n, t in mine], group=group + ) + + # Nothing to send is a real answer here, not an empty group: the last run has no run after it + # to read its edges. Its rank still joins the exchange above, so nobody is left waiting. + width = max(sum(math.prod(shape) for _, shape in entries) for entries in manifest) + if width == 0: + return list(made) + + # Then one tensor exchange for the results themselves, flattened together and padded to the + # largest share, since all_gather wants every rank sending the same count. Ranks differ by at + # most one call, so the padding is at most one call's worth of the traffic. + # + # Filled a call at a time rather than concatenated into the buffer, which would hold a second + # copy of everything this rank decoded while the first was still alive. Left uninitialised + # past what this rank sends: the manifest bounds what each share is read back out of, so the + # padding is never looked at. + sample = mine[0][1] if mine else like + sending = torch.empty(width, dtype=sample.dtype, device=sample.device) + at = 0 + for _, tensor in mine: + sending[at : at + tensor.numel()] = tensor.reshape(-1) + at += tensor.numel() + received = [torch.empty_like(sending) for _ in range(world_size)] + dist.all_gather(received, sending, group=group) + + shared: List[torch.Tensor] = list(made) + for entries, buffer in zip(manifest, received): + at = 0 + for n, shape in entries: + size = math.prod(shape) + # What this rank decoded itself is kept as it decoded it, rather than read back out + # of its own copy in the buffer. + if shared[n] is None: + shared[n] = buffer[at : at + size].view(shape) + at += size + return shared diff --git a/distvae/vae/tiling.py b/distvae/vae/tiling.py new file mode 100644 index 0000000..da5936e --- /dev/null +++ b/distvae/vae/tiling.py @@ -0,0 +1,846 @@ +"""What the installed diffusers can tile or slice, and how wide a window it can decode. + +Everything here is knowledge about diffusers VAEs: which tiling attributes a class carries, how +they relate, and which releases have them. The calling integration decides whether to use them. +""" + +import functools +import inspect +import math +from typing import Callable, List, NamedTuple, Optional, Tuple + +import diffusers +import torch + +# Diffusers represents tiling windows with several attribute layouts: a latent/pixel pair +# (AutoencoderKL and friends), a pixel window plus a stride (Wan, Qwen-Image, the video VAEs), and +# either of those keyed by height and width. Frame tiling is left out on purpose, being unrelated +# to a spatial tile edge. +PIXEL_ATTRS = ( + "tile_sample_min_size", + "tile_sample_min_height", + "tile_sample_min_width", +) +LATENT_ATTRS = ( + "tile_latent_min_size", + "tile_latent_min_height", + "tile_latent_min_width", +) +STRIDE_ATTRS = ("tile_sample_stride_height", "tile_sample_stride_width") +SCALED_ATTRS = LATENT_ATTRS + STRIDE_ATTRS +OVERLAP_ATTRS = ( + "tile_overlap_factor", + "tile_overlap_factor_height", + "tile_overlap_factor_width", +) +def require_vae_support(vae, feature: str, flag: str) -> None: + """Raise unless the installed diffusers really implements `feature` for this VAE""" + # Diffusers hands every autoencoder the enable_tiling and enable_slicing methods through a + # shared mixin, implemented or not, so their presence proves nothing. The state flag the mixin + # itself checks does. Wan added support in Diffusers 0.34, later than the minimum supported + # Diffusers version. + if not hasattr(vae, f"use_{feature}"): + raise ValueError( + f"{flag} is not supported by this VAE ({type(vae).__name__}) in the installed " + f"diffusers {diffusers.__version__}." + ) + + +def is_tile_padding_error(error: BaseException) -> bool: + """Whether a decode failure is the padding error a too-narrow tile window causes""" + # Torch raises this from a pad deep inside the decoder, where a tile arrives thinner than the + # convolution's own padding: "Padding size should be less than the corresponding input + # dimension, but got: padding (1, 1) at dimension 4 of input [1, 8, 3, 4, 1]". Text is all + # there is to key on, and a rewording upstream only costs the hint, since anything unmatched + # reaches the caller as the decoder wrote it. + return "padding size should be less than" in str(error).lower() + + +def tile_shape(vae) -> Optional[Tuple[int, int]]: + """The VAE's pixel-space tile window as (height, width), if it carries one.""" + height = getattr(vae, "tile_sample_min_height", None) + width = getattr(vae, "tile_sample_min_width", None) + if all(isinstance(value, int) and value > 0 for value in (height, width)): + return height, width + square = getattr(vae, "tile_sample_min_size", None) + if isinstance(square, int) and square > 0: + return square, square + return None + + +def _tile_defaults(vae) -> dict: + """Every tiling attribute the VAE carries, as the reference to rescale from""" + defaults = {} + for attr in PIXEL_ATTRS + SCALED_ATTRS + OVERLAP_ATTRS: + value = getattr(vae, attr, None) + if ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and value > 0 + ): + defaults[attr] = value + return defaults + + +def spatial_ratio(vae) -> Optional[int]: + """Pixels per latent pixel, where the VAE says so; config first, since reading a config key + off the module is deprecated""" + for source in (getattr(vae, "config", None), vae): + ratio = ( + getattr(source, "spatial_compression_ratio", None) + if source is not None + else None + ) + if isinstance(ratio, int) and ratio > 0: + return ratio + return None + + +def _is_whole(value: float) -> bool: + """Whole within float error, so 30 x (1 - 1/3) counts as 20 and not 20.000000000000004""" + return abs(value - round(value)) < 1e-9 + + +def tile_shape_plan(vae, height: int, width: int) -> Optional[dict]: + """Tiling attributes rescaled independently to an exact (height, width) window. + + Scalar-window VAEs receive complete per-axis attributes for DistVAE's replacement overlap + loop. VAEs that already define per-axis windows retain their native attribute layout. + """ + if not all( + isinstance(value, int) and not isinstance(value, bool) and value > 0 + for value in (height, width) + ): + return None + + defaults = _tile_defaults(vae) + legacy_scalar = all( + attr in defaults for attr in ("tile_sample_min_size", "tile_latent_min_size") + ) + keyed = all( + attr in defaults + for attr in ("tile_sample_min_height", "tile_sample_min_width") + ) + if keyed: + source_pixels = ( + defaults["tile_sample_min_height"], + defaults["tile_sample_min_width"], + ) + pixel_attrs = ("tile_sample_min_height", "tile_sample_min_width") + latent_attrs = ("tile_latent_min_height", "tile_latent_min_width") + elif "tile_sample_min_size" in defaults: + source_pixels = (defaults["tile_sample_min_size"],) * 2 + pixel_attrs = ("tile_sample_min_height", "tile_sample_min_width") + latent_attrs = ("tile_latent_min_height", "tile_latent_min_width") + else: + return None + + targets = (height, width) + plan = dict(zip(pixel_attrs, targets)) + scalar_latent = defaults.get("tile_latent_min_size") + factors = ( + defaults.get("tile_overlap_factor_height", defaults.get("tile_overlap_factor")), + defaults.get("tile_overlap_factor_width", defaults.get("tile_overlap_factor")), + ) + + for axis, (target, source) in enumerate(zip(targets, source_pixels)): + latent_source = defaults.get(latent_attrs[axis], scalar_latent) + if latent_source is not None: + latent = target * latent_source / source + if latent < 1 or not _is_whole(latent): + return None + latent = round(latent) + factor = factors[axis] + if ( + isinstance(factor, float) + and factor < 1.0 + and not _overlap_lands(latent, target, factor) + ): + return None + plan[latent_attrs[axis]] = latent + + stride_attr = STRIDE_ATTRS[axis] + stride_source = defaults.get(stride_attr) + if stride_source is not None: + stride = target * stride_source / source + if stride < 1 or not _is_whole(stride): + return None + plan[stride_attr] = round(stride) + + if legacy_scalar: + # AutoencoderKL and Flux decide whether to enter tiled_decode with one scalar threshold. + # The smaller axis is conservative: crossing either requested window must cross it, while + # overlap_windows reads the exact keyed rectangle above once the loop is entered. + plan["tile_sample_min_size"] = min(targets) + plan["tile_latent_min_size"] = min( + plan["tile_latent_min_height"], plan["tile_latent_min_width"] + ) + + granularity = _stride_granularity(vae) if any( + attr in plan for attr in STRIDE_ATTRS + ) else None + if granularity is not None and any( + value % granularity + for value in targets + tuple(plan[attr] for attr in STRIDE_ATTRS if attr in plan) + ): + return None + + ratio = spatial_ratio(vae) + if ratio is not None and not any(attr in plan for attr in LATENT_ATTRS): + if any(target < ratio or target % ratio for target in targets): + return None + return plan + + +def apply_tile_plan(vae, plan: dict) -> None: + """Set a planned window on the VAE""" + # Newer VAE classes also take these through enable_tiling(), but only some of them, with a + # different signature each, and the body is a plain assignment either way. + for attr, value in plan.items(): + setattr(vae, attr, value) + + +def _latent_shape(vae, plan: Optional[dict] = None) -> Optional[Tuple[int, int]]: + """Latent tile height and width under `plan`, or None where the VAE does not say.""" + # Without a plan the VAE's own attributes are the plan, which is how a caller asks about a + # window that no flag set - a VAE tiling at its own default, or one a model turned on at + # load. + if plan is None: + plan = _tile_defaults(vae) + keyed = tuple( + plan.get(attr) + for attr in ("tile_latent_min_height", "tile_latent_min_width") + ) + if all(value is not None for value in keyed): + return keyed + scalar = plan.get("tile_latent_min_size") + if scalar is not None: + return scalar, scalar + ratio = spatial_ratio(vae) + if ratio is not None: + keyed = tuple( + plan.get(attr) + for attr in ("tile_sample_min_height", "tile_sample_min_width") + ) + if all(value is not None for value in keyed): + return tuple(value // ratio for value in keyed) + scalar = plan.get("tile_sample_min_size") + if scalar is not None: + edge = scalar // ratio + return edge, edge + return None + + +def latent_rows(vae, plan: Optional[dict] = None) -> Optional[int]: + """How many latent rows a tile holds, under `plan` or as the VAE stands.""" + shape = _latent_shape(vae, plan) + return shape[0] if shape is not None else None + + +def overlap_windows(vae) -> Optional[Tuple[Tuple[int, int], Tuple[int, int]]]: + """The latent and pixel tile windows as (down, across) pairs, None where the VAE has neither + + AutoencoderKL and FLUX.2 store one square edge; HunyuanVideo 1.5 stores one edge per axis. + Normalize both attribute layouts to a pair so one loop can support all three classes. + """ + keyed = [ + getattr(vae, attr, None) + for attr in ( + "tile_latent_min_height", + "tile_latent_min_width", + "tile_sample_min_height", + "tile_sample_min_width", + ) + ] + if all(isinstance(value, int) and value > 0 for value in keyed): + return (keyed[0], keyed[1]), (keyed[2], keyed[3]) + square = getattr(vae, "tile_latent_min_size", None) + if isinstance(square, int) and square > 0: + pixels = getattr(vae, "tile_sample_min_size", None) + return ( + ((square, square), (pixels, pixels)) + if isinstance(pixels, int) and pixels > 0 + else None + ) + return None + + +def _overlap_factors(vae) -> Optional[Tuple[float, float]]: + """Configured overlap factors by axis, preferring keyed values.""" + keyed = ( + getattr(vae, "tile_overlap_factor_height", None), + getattr(vae, "tile_overlap_factor_width", None), + ) + if all( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and 0.0 <= value < 1.0 + for value in keyed + ): + return keyed + scalar = getattr(vae, "tile_overlap_factor", None) + if ( + isinstance(scalar, (int, float)) + and not isinstance(scalar, bool) + and 0.0 <= scalar < 1.0 + ): + return scalar, scalar + return None + + +def tiles_by_overlap_factor(vae) -> bool: + """Whether this VAE tiles with the loop `overlap_tiled_decode` reimplements""" + # AutoencoderKL, AutoencoderKLFlux2 and HunyuanVideo 1.5 walk a latent window at a stride + # derived from an overlap fraction. Wan, Qwen-Image and the other video VAEs walk a stride + # they store outright, over a loop with different blending, and keep their own tiled_decode. + if any(getattr(vae, attr, None) for attr in STRIDE_ATTRS): + return False + if overlap_windows(vae) is None: + return False + # The scalar attribute identifies this family rather than CogVideoX, whose keyed-factor loop + # also tiles frames inside the spatial loop. DistVAE adds keyed values to this family when a + # caller requests rectangular overlap, but leaves the scalar marker in place. + return ( + isinstance(getattr(vae, "tile_overlap_factor", None), (int, float)) + and not isinstance(getattr(vae, "tile_overlap_factor", None), bool) + and _overlap_factors(vae) is not None + and callable(getattr(vae, "blend_v", None)) + and callable(getattr(vae, "blend_h", None)) + ) + + +WINDOW_ATTRS_FOR_STRIDE = ("tile_sample_min_height", "tile_sample_min_width") +"""The window each stride in STRIDE_ATTRS steps across, in the same order""" + + +def tile_overlap(vae) -> Optional[Tuple[int, int]]: + """Return absolute output-pixel overlap as (height, width) for any supported attribute layout.""" + strides = [getattr(vae, attr, None) for attr in STRIDE_ATTRS] + windows = [getattr(vae, attr, None) for attr in WINDOW_ATTRS_FOR_STRIDE] + if all(isinstance(value, int) and value > 0 for value in strides + windows): + overlap = tuple(window - stride for stride, window in zip(strides, windows)) + return ( + overlap + if all( + 0 <= value < window for value, window in zip(overlap, windows) + ) + else None + ) + factors = _overlap_factors(vae) + shape = tile_shape(vae) + if factors is not None and shape is not None: + return tuple(int(window * factor) for window, factor in zip(shape, factors)) + return None + + +def _stride_granularity(vae) -> Optional[int]: + """Return the required pixel-stride multiple for a consistent stored-stride tiling loop. + + That loop divides the stride it stores twice: by the compression ratio, to step the latent + grid, and - where the family decodes into a pixel unshuffle - by the patch size, to place the + crop. Both are integer divisions, so a stride that is not a multiple of each truncates in one + of them and the grid and the crop stop describing the same region. + """ + ratio = spatial_ratio(vae) + if ratio is None: + return None + loop = _STRIDE_LOOPS.get(type(vae).__name__) + patch = getattr(vae.config, "patch_size", None) if loop and loop.patches else None + if isinstance(patch, int) and patch > 1: + return math.lcm(ratio, patch) + return ratio + + +def _overlap_lands(latent: int, pixel: int, factor: float) -> bool: + """Whether the overlap-fraction loop's own arithmetic agrees with itself on this axis + + The loop derives the latent step by truncating `latent x (1 - factor)`, and crops each + decoded tile to `pixel - int(pixel x factor)`. Unless the second is the first in pixels, the + tiles step by one amount and are cropped by another, and the assembled image comes out a + different size than the decode was asked for - with nothing downstream to catch it. + + Checked by recomputing what the loop will compute, rather than by reasoning about the + algebra, because the factor is a float and the two truncations do not have to fall the same + way on both sides of it. + """ + stride = int(latent * (1.0 - factor)) + if stride < 1: + return False + ratio, remainder = divmod(pixel, latent) + return remainder == 0 and pixel - int(pixel * factor) == stride * ratio + + +def tile_overlap_plan( + vae, + overlap_height: int, + overlap_width: int, + sample_shape: Optional[Tuple[int, int]] = None, +) -> Optional[dict]: + """Plan an exact absolute output-pixel overlap, or None when it is not representable.""" + requested = (overlap_height, overlap_width) + if not all( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 + for value in requested + ): + return None + shape = tile_shape(vae) + if shape is None: + return None + active_axes = (True, True) + if sample_shape is not None: + if ( + not isinstance(sample_shape, tuple) + or len(sample_shape) != 2 + or not all( + isinstance(value, int) and not isinstance(value, bool) and value > 0 + for value in sample_shape + ) + ): + return None + active_axes = tuple(sample > window for sample, window in zip(sample_shape, shape)) + if any( + (active and overlap >= window) or (not active and overlap != 0) + for active, overlap, window in zip(active_axes, requested, shape) + ): + return None + + if tiles_by_stored_stride(vae): + step = _stride_granularity(vae) + if step is None: + return None + strides = ( + window - overlap for window, overlap in zip(shape, requested) + ) + plan = dict(zip(STRIDE_ATTRS, strides)) + for stride in plan.values(): + if stride <= 0 or stride % step: + return None + return plan + + if not tiles_by_overlap_factor(vae): + return None + windows = overlap_windows(vae) + if windows is None: + return None + (latent_down, latent_across), (pixel_down, pixel_across) = windows + axes = ((latent_down, pixel_down), (latent_across, pixel_across)) + factors = tuple(overlap / pixel for overlap, (_, pixel) in zip(requested, axes)) + if not all( + _overlap_lands(latent, pixel, factor) + for (latent, pixel), factor in zip(axes, factors) + ): + return None + plan = { + "tile_overlap_factor_height": factors[0], + "tile_overlap_factor_width": factors[1], + } + if hasattr(vae, "tile_overlap_factor") and factors[0] == factors[1]: + plan["tile_overlap_factor"] = factors[0] + return plan + + +def _returns_decoder_output(vae) -> bool: + """Return whether this class's tiled_decode returns DecoderOutput rather than a tensor. + + The replacement must preserve the return type expected by `_decode`. Most classes accept + `return_dict` and return DecoderOutput; HunyuanVideo 1.5 accepts no such argument and returns + a tensor directly. + + Read off the class rather than the instance, so that installing twice cannot end up reading + the first install's signature instead of the original. + """ + own = getattr(type(vae), "tiled_decode", None) + if own is None: + return True + try: + return "return_dict" in inspect.signature(own).parameters + except (TypeError, ValueError): + return True + + +class _StrideLoop(NamedTuple): + """Where the stride-walked tiling loops differ from one another""" + + patches: ( + bool # decodes into a pixel unshuffle, and unpatchifies the assembled sample + ) + clamps: bool # holds the assembled sample in [-1, 1] + first_chunk: bool # tells the decoder which frame starts the tile + frame_cache: ( + bool # decodes a tile frame by frame, threading the VAE's own feature cache + ) + post_quant: bool # puts a tile through post_quant_conv before decoding it + conditioned: ( + bool # carries a timestep embedding and a causality flag into the decoder + ) + + +# The VAEs whose stride-walked loop is reimplemented below, by name, because no attribute says +# which loop body a class has. All four walk the same grid and blend it the same way, and differ +# only in what a tile costs to turn into a decoder call. +# +# HunyuanVideo and LTX-2 keep no feature cache, so one decoder call handles all frames in a +# spatial tile. Their temporal loops call this spatial loop once per frame chunk. LTX-2 disables +# temporal tiling by default; HunyuanVideo enables it. +# +# CogVideoX is excluded because its spatial loop also tiles frames, so its tiles are not +# independent. HunyuanVideo 1.5 uses overlap-fraction tiling and is handled by +# overlap_tiled_decode. +_STRIDE_LOOPS = { + "AutoencoderKLWan": _StrideLoop( + patches=True, + clamps=True, + first_chunk=True, + frame_cache=True, + post_quant=True, + conditioned=False, + ), + "AutoencoderKLQwenImage": _StrideLoop( + patches=False, + clamps=False, + first_chunk=False, + frame_cache=True, + post_quant=True, + conditioned=False, + ), + "AutoencoderKLHunyuanVideo": _StrideLoop( + patches=False, + clamps=False, + first_chunk=False, + frame_cache=False, + post_quant=True, + conditioned=False, + ), + "AutoencoderKLLTX2Video": _StrideLoop( + patches=False, + clamps=False, + first_chunk=False, + frame_cache=False, + post_quant=False, + conditioned=True, + ), +} + + +def tiles_by_stored_stride(vae) -> bool: + """Whether this VAE tiles with the stride-walked loop `strided_tiled_decode` reimplements""" + loop = _STRIDE_LOOPS.get(type(vae).__name__) + if loop is None: + return False + # The class is the loop, but the pieces it walks are still checked, so that a VAE refactored + # out from under this fails the question rather than the decode. + parts = ["blend_v", "blend_h", "decoder"] + if loop.post_quant: + parts.append("post_quant_conv") + if loop.frame_cache: + parts.append("clear_cache") + return all( + isinstance(getattr(vae, attr, None), int) + for attr in ( + "tile_sample_min_height", + "tile_sample_min_width", + "tile_sample_stride_height", + "tile_sample_stride_width", + "spatial_compression_ratio", + ) + ) and all(callable(getattr(vae, attr, None)) for attr in parts) + + +def supports_tile_parallel(vae) -> bool: + """Whether this VAE's tiling loop is one of the ones reimplemented here + + Deciding which rank makes which decoder call means owning the loop that makes them, so this + is what a caller asks before planning to decode a VAE's tiles apart from one another. + """ + return tiles_by_overlap_factor(vae) or tiles_by_stored_stride(vae) + + +def tiled_decode_for( + vae, + dispatch: Optional[Callable] = None, + assemble: Optional[Callable] = None, +) -> Optional[Callable]: + """The tiled_decode to install on this VAE, None where its loop is not one reimplemented here""" + overlapping = overlap_tiled_decode(vae, dispatch, assemble) + if overlapping is not None: + return overlapping + # The stride-walked loop is reimplemented for one reason, which is to hand its tiles round; + # left to decode them all here it would only be diffusers' own loop with a second author. + if dispatch is None and assemble is None: + return None + return strided_tiled_decode(vae, dispatch, assemble) + + +def _latent_areas(down, across, window, bounds) -> List[int]: + """The latent area each tile of the grid covers, in the order the loop walks + + What a tile costs to decode follows the latent it is cut from, and the tiles on the last row + and the last column are cut short by the bounds. Which of them are short is the same on every + rank, being read off the grid rather than off a decoded tile. + """ + deep, wide = window if isinstance(window, tuple) else (window, window) + return [ + (min(top + deep, bounds[0]) - top) * (min(left + wide, bounds[1]) - left) + for top in down + for left in across + ] + + +def overlap_tiled_decode( + vae, + dispatch: Optional[Callable] = None, + assemble: Optional[Callable] = None, +) -> Optional[Callable]: + """A tiled_decode for the overlap-fraction family, None where the VAE is not one of them + + One tile per decoder call, as upstream does. This preserves exact decoder-call semantics while + allowing independent tiles to be dispatched in any order. + + `dispatch` decides which rank makes each call and defaults to this rank making all calls in + order. `distvae.vae.tile_parallel` supplies a dispatcher for distributing calls to a group. + + `assemble` goes further and divides the blending too, by giving each rank a run of + neighbouring tiles to decode and stitch by itself. Where it declines - too few tiles to give + every rank one, or tiles too small to blend against a neighbour's edge alone - the decode + falls back to `dispatch`, which divides the decoder calls and leaves the blending everywhere. + + AutoencoderKL, Flux.2, and HunyuanVideo 1.5 share this loop but use different window + attributes and return types. Normalize every window to ``(height, width)`` and preserve the + original method's return type. Height and width are always the final two dimensions. + """ + if not tiles_by_overlap_factor(vae): + return None + + from diffusers.models.autoencoders.vae import DecoderOutput + + from distvae.vae import tile_parallel as vae_tile_parallel + + # Treat either config.use_post_quant_conv or a non-null post_quant_conv as enabling + # post-quantization convolution. + use_post_quant_conv = getattr( + getattr(vae, "config", None), "use_post_quant_conv", None + ) + if use_post_quant_conv is None: + use_post_quant_conv = getattr(vae, "post_quant_conv", None) is not None + + def decode_tiles(z): + (latent_down, latent_across), (pixel_down, pixel_across) = overlap_windows(vae) + factor_down, factor_across = _overlap_factors(vae) + stride_down = int(latent_down * (1 - factor_down)) + stride_across = int(latent_across * (1 - factor_across)) + blend_down = int(pixel_down * factor_down) + blend_across = int(pixel_across * factor_across) + limit_down = pixel_down - blend_down + limit_across = pixel_across - blend_across + + down = range(0, z.shape[-2], stride_down) + across = range(0, z.shape[-1], stride_across) + + def latent_at(i, j): + tile = z[ + ..., + down[i] : down[i] + latent_down, + across[j] : across[j] + latent_across, + ] + return vae.post_quant_conv(tile) if use_post_quant_conv else tile + + def decode_with(share): + def decode(where): + # Every call is built before any is made, so that a dispatcher can see them all + # and hand them round. Each holds a latent tile, which is the small side of the + # decode; what they return is not held any longer than it was before. + at_order = list(where) + calls = [ + functools.partial(vae.decoder, latent_at(*at)) for at in at_order + ] + return dict(zip(at_order, share(calls))) + + return decode + + blend = vae_tile_parallel.Blend( + down=vae.blend_v, + across=vae.blend_h, + deep_down=blend_down, + deep_across=blend_across, + crop=lambda tile: tile[..., :limit_down, :limit_across], + tile_down=pixel_down, + tile_across=pixel_across, + ) + if assemble is not None: + # A run decodes its own tiles, so the calls stay here rather than going round again. + dec = assemble( + len(down), + len(across), + decode_with(vae_tile_parallel.in_order), + blend, + _latent_areas(down, across, (latent_down, latent_across), z.shape[-2:]), + ) + if dec is not None: + return dec + share = dispatch if dispatch is not None else vae_tile_parallel.in_order + return vae_tile_parallel.assemble_here( + len(down), len(across), decode_with(share), blend + ) + + def tiled_decode(z, return_dict: bool = True): + dec = decode_tiles(z) + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def bare_tiled_decode(z): + return decode_tiles(z) + + return tiled_decode if _returns_decoder_output(vae) else bare_tiled_decode + + +def strided_tiled_decode( + vae, dispatch: Optional[Callable] = None, assemble: Optional[Callable] = None +) -> Optional[Callable]: + """A tiled_decode for the video VAEs that walk a stride they store, None where it can't + + Upstream's loop, with the tiles built as calls rather than made where they are built, so that + `dispatch` can hand them round a group. Where the family keeps a feature cache a tile is a + frame loop threading it, cleared at the start of each tile, so a tile is independent of every + other tile in the way the frames inside it are not; where it keeps none, a tile is one call. + + A tile remains one call, preserving the VAE's cache and conditioning boundaries. + """ + if not tiles_by_stored_stride(vae): + return None + + from diffusers.models.autoencoders.vae import DecoderOutput + + from distvae.vae import tile_parallel as vae_tile_parallel + + loop = _STRIDE_LOOPS[type(vae).__name__] + patch_size = getattr(vae.config, "patch_size", None) if loop.patches else None + + # `temb` and `causal` are LTX-2's, which conditions its decoder on them and passes them + # through tiled_decode. Other families omit those arguments, so their defaults allow one + # replacement signature to support all four loops. + def tiled_decode(z, temb=None, causal=None, return_dict: bool = True): + _, _, num_frames, height, width = z.shape + ratio = vae.spatial_compression_ratio + sample_height = height * ratio + sample_width = width * ratio + latent_min_height = vae.tile_sample_min_height // ratio + latent_min_width = vae.tile_sample_min_width // ratio + latent_stride_height = vae.tile_sample_stride_height // ratio + latent_stride_width = vae.tile_sample_stride_width // ratio + sample_stride_height = vae.tile_sample_stride_height + sample_stride_width = vae.tile_sample_stride_width + if patch_size is not None: + sample_height //= patch_size + sample_width //= patch_size + sample_stride_height //= patch_size + sample_stride_width //= patch_size + blend_height = ( + vae.tile_sample_min_height // patch_size - sample_stride_height + ) + blend_width = vae.tile_sample_min_width // patch_size - sample_stride_width + else: + blend_height = vae.tile_sample_min_height - sample_stride_height + blend_width = vae.tile_sample_min_width - sample_stride_width + + down = range(0, height, latent_stride_height) + across = range(0, width, latent_stride_width) + + def tile_at(i, j): + def cut(frames=slice(None)): + return z[ + :, + :, + frames, + down[i] : down[i] + latent_min_height, + across[j] : across[j] + latent_min_width, + ] + + def decode_frame_by_frame(): + # The cache is per tile and threaded through the frames of one, which is why the + # frames cannot be handed round but the tiles can. + vae.clear_cache() + frames = [] + for k in range(num_frames): + vae._conv_idx = [0] + tile = vae.post_quant_conv(cut(slice(k, k + 1))) + extra = {"first_chunk": k == 0} if loop.first_chunk else {} + frames.append( + vae.decoder( + tile, + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + **extra, + ) + ) + return torch.cat(frames, dim=2) + + def decode_at_once(): + tile = vae.post_quant_conv(cut()) if loop.post_quant else cut() + if loop.conditioned: + return vae.decoder(tile, temb, causal=causal) + return vae.decoder(tile) + + return decode_frame_by_frame if loop.frame_cache else decode_at_once + + def decode_with(share): + def decode(where): + made = share([tile_at(*at) for at in where]) + if loop.frame_cache: + vae.clear_cache() + return dict(zip(where, made)) + + return decode + + blend = vae_tile_parallel.Blend( + down=vae.blend_v, + across=vae.blend_h, + deep_down=blend_height, + deep_across=blend_width, + crop=lambda tile: tile[ + :, :, :, :sample_stride_height, :sample_stride_width + ], + tile_down=( + vae.tile_sample_min_height // patch_size + if patch_size is not None + else vae.tile_sample_min_height + ), + tile_across=( + vae.tile_sample_min_width // patch_size + if patch_size is not None + else vae.tile_sample_min_width + ), + ) + dec = None + if assemble is not None: + dec = assemble( + len(down), + len(across), + decode_with(vae_tile_parallel.in_order), + blend, + _latent_areas( + down, + across, + (latent_min_height, latent_min_width), + (height, width), + ), + ) + if dec is None: + share = dispatch if dispatch is not None else vae_tile_parallel.in_order + dec = vae_tile_parallel.assemble_here( + len(down), len(across), decode_with(share), blend + ) + dec = dec[:, :, :, :sample_height, :sample_width] + + if patch_size is not None: + from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify + + dec = unpatchify(dec, patch_size=patch_size) + if loop.clamps: + dec = torch.clamp(dec, min=-1.0, max=1.0) + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + return tiled_decode diff --git a/docs/figure.png b/docs/figure.png new file mode 100644 index 0000000..9a82b07 Binary files /dev/null and b/docs/figure.png differ diff --git a/docs/figure.svg b/docs/figure.svg new file mode 100644 index 0000000..9a673dc --- /dev/null +++ b/docs/figure.svg @@ -0,0 +1,274 @@ + + + +DistVAE parallelism +The two modes below each decode a 1024 × 1024 image from a 128 × 128 latent on four GPUs. +The same five metrics compare both alternatives. + +collective: every rank waits + +halo swap: neighbours only + +held at once + +tile overlap + +1. Row sharding +one decoder call, split into four bands of 32 latent rows +Peak memory is rank-bound and every layer syncs, so the interconnect can be the bottleneck. +There is nothing to choose here, since the band is the latent divided by the GPU count. +The image is what a single GPU would produce, apart from floating-point rounding. + +rank 0 + +rank 1 + +rank 2 + +rank 3 + + +rank 0 +rank 1 +rank 2 +rank 3 + + + + +conv + + + + + + + +norm + + + + + +conv + + + + + + + +norm + + + + + +conv + + + + + + + + + + + + + + + + + + +image + +peak activations +25% +work +1.00× +seams +none +imbalance +0.0% +syncs +every layer +Split the rows +The bands never overlap. +Decode in lockstep +Convolutions swap edge rows, and norms reduce across all four ranks. + +2. Tile distribution +the same four GPUs, at two windows +Peak memory is tile-bound and the decode needs only two collectives, but it repeats more work. +Window and overlap are yours to set in output pixels, so tune them to your VAE and your GPUs. +The two windows below are worked examples, chosen to show the trade rather than to be copied. +Full-width strips overlap less and stay contiguous in row-major memory, while a grid holds less at once. +The image is close but not exact: a blend hides seams, but norms over too small a tile can leave the colour blocky. +Cut the rows only +344 px tall overlapping 88 px, full width +Four strips, one per rank, have the same shape as the row-sharded bands, but they overlap and nothing syncs until the end. + + + + + + + + +0 +1 +2 +3 + + + + + +rank 0 + +0 +rank 1 + +1 +rank 2 + +2 +rank 3 + +3 + +idle + + + +edges, image + +peak activations +34% +work +1.26× +seams +3 +imbalance +6.8% +syncs +twice +Cut and overlap the rows +One call each, and one rank waits +With one strip per rank, there is nothing for the scheduler to decide. +The last strip is 32 latent rows where the others are 43. +That shortfall is exactly the overlap, so closing it would thin the blend. +Cut both axes +432 × 296 px overlapping 72 px on both axes +Fifteen tiles across four ranks let the load be levelled, and a rank now holds a window rather than a strip. + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0 +1 +2 +3 +4 +5 +6 +7 +8 +9 +10 +11 +12 +13 +14 + + + + + + + + +rank 0 + +0 + +1 + +2 +rank 1 + +3 + +5 + +6 +rank 2 + +7 + +8 + +10 + +14 +rank 3 + +4 + +9 + +11 + +12 + +13 + + + + +edges, image + +peak activations +12% +work +1.46× +seams +22 +imbalance +0.4% +syncs +twice +Distribute the tiles +Each rank decodes its assigned tiles sequentially +Each rank starts with a contiguous run, then single tiles move to level it. +Rank 3 decodes five tiles to rank 0's three, and they still finish together. + \ No newline at end of file diff --git a/docs/make_figure.py b/docs/make_figure.py new file mode 100644 index 0000000..2073762 --- /dev/null +++ b/docs/make_figure.py @@ -0,0 +1,749 @@ +"""Draw the README's figure: row sharding, then tile distribution at two windows. + +The first row shows row sharding. The second uses four full-width tiles, matching the +row-sharding geometry while changing only the execution method. The third splits both spatial +axes. + +The five columns show per-rank data, redundant work, blend boundaries, maximum load imbalance, +and synchronization frequency. Row sharding uses the same formulas as both tiling configurations, +so the values are directly comparable. + +The right-hand panels are all the same axis, time, with one lane per rank. Each is scaled to +its own heaviest rank, so all three rows end at the same x and the lengths mean nothing +across rows; the work column is what to read for that. Within a row the blocks stay +proportional to the work in their tile, which is what makes the clipped tiles at a grid's +edges visibly cheap, and cheap is why dealing tiles out by area beats dealing them by count. + +Both windows are written as the pair of planner calls that would set them, a window and an +overlap in output pixels, so the figure cannot show a configuration the API could not be +asked for. The grids are drawn at the extents that pair leaves, so tiles overlap on the page +as they do in the loop. Neither window is a recommendation, and the heading says so: they +are the two ends of the trade, and the strips end has advantages no column here can show, +being one contiguous span of a row-major tensor where a grid's tile is a stride through +every row it touches. + +No row draws its output, because all three produce the same picture. What tiling changes is +a seam a good window renders invisible, so a panel of the result would be either blank or an +exaggeration of what a blend leaves behind. Each heading says what its mode costs the image +in words instead, beside what it costs memory and the interconnect. + +Written as plain SVG, so the figure rebuilds with no toolchain and stays legible in a diff. +The tile-to-rank assignment is asked of the scheduler rather than drawn by hand, so the +picture cannot drift from what a decode actually does. Nothing is positioned at an absolute +y: every block reports where it ended and the next starts from there, so a caption can be +added without re-tuning the page. + + python docs/make_figure.py +""" + +import importlib.util +import os +import sys +import types + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(HERE) +OUT = os.path.join(HERE, "figure.svg") +# The README points at the PNG, because a local SVG does not preview in every editor a README +# is read in. It is written only where cairosvg is installed, so rebuilding the SVG itself +# still needs nothing but a Python interpreter. +PNG = os.path.join(HERE, "figure.png") +RASTER = 2 + +RANKS = 4 +# A 1024 by 1024 image, and so 128 latent rows on an eight-fold VAE. +BOUND = 128 +SCALE_VAE = 8 + +NUMBERS = ("no", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", + "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen") + + +def word(n): + """A small count spelled out, since the captions are prose and the columns are not""" + return NUMBERS[n] if n < len(NUMBERS) else str(n) + + +def scheduler(): + """The tile-to-rank assignment the library ships, so the figure cannot invent one""" + try: + from distvae.vae.tile_parallel import shares + return shares + except ImportError: + pass + # shares() is pure integer arithmetic, so rebuilding a docs figure should not need a + # torch install. Load that one module, with the imports it never reaches stubbed out. + for name, attrs in (("torch", {"Tensor": object}), ("torch.distributed", {}), + ("distvae", {}), + ("distvae.utils", {"ParallelContext": type("Ctx", (), {})})): + sys.modules[name] = types.ModuleType(name) + sys.modules[name].__dict__.update(attrs) + spec = importlib.util.spec_from_file_location( + "_tile_parallel", os.path.join(ROOT, "distvae", "vae", "tile_parallel.py")) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.shares + + +SHARES = scheduler() + + +class Axis: + """One axis of the grid, in the two numbers the planners are actually given + + Both are absolute output pixels, because that is the interface: `tile_shape_plan` takes + a window and `tile_overlap_plan` takes an overlap, never a fraction of one. An axis the + sample already fits is inactive and must use zero overlap, which represents a full-width + strip. The stride is derived from the window and overlap rather than configured directly. + + Everything below is in latent units, since that is what the grid is drawn in. `at` is + where each tile starts and `extent` how far it reaches once the bound has clipped it. + `deep` is the overlap, and so both the band a second tile also covers and how far a + blend reaches into the tile on the far side of a join. + """ + + def __init__(self, window_px, overlap_px): + self.window_px, self.overlap_px = window_px, overlap_px + self.window, self.deep = window_px // SCALE_VAE, overlap_px // SCALE_VAE + stride = self.window - self.deep + self.at = list(range(0, BOUND, stride)) + self.extent = [min(o + self.window, BOUND) - o for o in self.at] + self.count = len(self.at) + + +class Split: + """Metrics for one way of dividing the latent. + + `load` is the latent area decoded by each rank. Row sharding and tiling use the same + calculations for peak activation area, total decoded work, seams, and load imbalance. + """ + + def __init__(self, weight, owner, seams): + self.weight, self.owner, self.seams = weight, owner, seams + self.run = [[n for n, who in enumerate(owner) if who == r] for r in range(RANKS)] + self.load = [sum(weight[n] for n in run) for run in self.run] + # Activations follow area, so the largest single call is what sets the memory a + # rank needs, whatever else it goes on to decode afterwards. + self.held = max(weight) / BOUND ** 2 + # The tiles together cover more latent than there is, and every unit over is a + # patch of image decoded twice. + self.work = sum(weight) / BOUND ** 2 + # The heaviest rank's excess over an even split determines how long other ranks wait. + self.imbalance = max(self.load) / (sum(self.load) / RANKS) - 1 + + +class Grid(Split): + """Metrics for the tile grid produced by a window and overlap.""" + + # Tile distribution synchronizes once during dispatch and once during assembly. + syncs = "twice" + + def __init__(self, down, across): + self.down, self.across = down, across + self.tiles = down.count * across.count + weight = [d * a for d in down.extent for a in across.extent] + # A seam is a join between two tiles, which is a pair of neighbours rather than a + # band of overlap: the grid has one per adjacency on each axis. Corners, where four + # tiles meet, are left out of the count for the same reason the legend leaves them + # out, so this is the number of places a blend has to work rather than of blends. + seams = down.count * (across.count - 1) + across.count * (down.count - 1) + super().__init__(weight, SHARES(weight, RANKS), seams) + self.biggest = max(range(self.tiles), key=lambda n: weight[n]) + + +class Bands(Split): + """Metrics for row sharding. + + Each rank receives one band. Total decoded work equals the latent area and there are no tile + seams. + """ + + syncs = "every layer" + + def __init__(self): + rows = [BOUND // RANKS + (r < BOUND % RANKS) for r in range(RANKS)] + super().__init__([r * BOUND for r in rows], list(range(RANKS)), 0) + + +# Both windows are written as the pair of planner calls that would set them, so the figure +# cannot describe a configuration the API could not be asked for: +# +# tile_shape_plan(vae, 352, 1408) +# tile_overlap_plan(vae, 88, 0, sample_shape=(1024, 1024)) +# +# A 344-pixel height is 43 latent rows. With an 88-pixel overlap, the 32-row stride produces +# four strips whose latent heights are 43, 43, 43, and 32. A 1408-pixel width exceeds the image, +# so width is untiled and uses zero overlap. A 352-pixel height would use a 33-row stride and +# shorten the final strip to 29 rows without increasing its blend. +STRIPS = Grid(Axis(344, 88), Axis(1408, 0)) + +# A 432 × 296-pixel window with 72-pixel overlap produces the rectangular grid. Compared with +# the squarest planner-selected grid, it holds 12.2% instead of 14.1% of the activations, decodes +# 1.46× instead of 1.47× the latent area, has 0.4% instead of 15.1% load imbalance, and creates +# 22 instead of 24 seams. Both axes use the same overlap so every seam has the same blend width. +TILED = Grid(Axis(432, 72), Axis(296, 72)) + +SHARDED = Bands() + +FONT = "system-ui, -apple-system, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif" +MONO = "'SF Mono', Menlo, Consolas, monospace" + +INK = "#1b1f24" +MUTED = "#6a737d" +# The rule that opens a mode. Darker and longer than the hairline over a costs row, because +# it has to read as the top of a section rather than as one more underline inside one. +EDGE = "#c2c8ce" +# Black is spoken for: it marks what a rank holds. Everything a rank says is this instead, +# light enough to sit under the blocks it crosses rather than on top of them. +SYNC = "#a7b0b8" +RULE = "#e3e6e8" + +RANK = [ + ("#dce9f7", "#3d6d99"), + ("#dff0da", "#4f8b3f"), + ("#fce8d0", "#bf7a2e"), + ("#f9d9dc", "#b04a52"), +] +# How far a rank's fill is carried towards white, so that the tiles panel can lay them on +# top of one another and have two deep still read as a fill rather than as ink. +TINT = 0.55 + +# Two columns: the latent, and the timeline it is decoded on. There is no third for the +# result, because the result is the same picture every way and a panel of it would either +# be blank or, drawn with anything visible on it, overstate what a blend leaves behind. +# What each mode does to the image is said in its heading instead, beside what it does to +# memory and to the interconnect, so all three read side by side. +COL1, PANEL = 40, 120 +TRACK = 216 +LABEL = 42 +# Where a heading's muted tail starts, the same for every row and every window under them, +# so one column of bold runs down the page and one column of grey runs beside it. +TAIL = COL1 + 150 +# Set to the longest line of writing, which is what bounds the page now that no panel +# reaches further than the headings do. Measured against a wide fallback rather than the +# font the README will pick, so a substitution loses margin instead of clipping a word. +# The height is not a constant: draw() adds it up from what the rows come to. +W = 680 + +LANE = 30 +# One layer of the sharded decode, and the room after it for whatever that layer syncs. The +# two are sized together so the timelines reach the width the headings above them set, +# rather than stopping short and leaving the right of the page empty. +LAYER = 48 +GAP = 18 +# The lanes stack to the same height the latent panels are drawn at, which is what lets a +# row read straight across. +TALL = RANKS * LANE +UNIT = PANEL / BOUND +# How far the busiest lane in a row runs, in pixels. Each row is scaled to its own heaviest +# rank rather than to the heaviest in the figure, so all three end their collectives at the +# same x. One clock across the rows would be the more informative drawing, but the rows are +# far apart and the difference between them is under a tenth: at that size a short row reads +# as a rendering fault rather than as a shorter decode, and the work column says it better. +# Within a row the blocks stay proportional, which is what makes the clipped tiles cheap. +SPAN = 5 * (LAYER + GAP) +# The gap the elided layers leave in the row-sharding timeline. Wide enough for a run of +# dots in every lane with the memory bracket's open edge clear of them, since that edge +# ends three units before the collective that closes the row. +ELIDED = 24 +# One monospace digit at size 9, the size the blocks are numbered, which is what decides +# whether a block is wide enough to hold its own number. +DIGIT = 5.4 + +# The columns every row is closed with. Four of them are what a way of splitting the latent +# costs; the fifth is what it buys, and it is here because without it the readout says only +# that tiling is worse, which is true of every column and beside the point. +COSTS = ( + ("peak activations", lambda s: f"{s.held:.0%}"), + ("work", lambda s: f"{s.work:.2f}×"), + # A count, except at zero, where the difference is not a small number of seams but a + # mode that never blends anything and so has none to hide. + ("seams", lambda s: str(s.seams) if s.seams else "none"), + ("imbalance", lambda s: f"{s.imbalance:.1%}"), + ("syncs", lambda s: s.syncs), +) +# Label over value rather than beside it, so a column is as wide as its widest single word +# and five of them fit where four sat before. The stack also stops the readout reading as +# another line of caption, which is what it looked like set on one line. +PITCH = 112 +STACK = 15 + +HATCH = ( + '' + f'' +) + + +out = [] + + +def add(s): + out.append(s) + + +# --------------------------------------------------------------------------- primitives + + +def rect(x, y, w, h, fill, stroke, rx=3, sw=1.2, opacity=None, fill_opacity=None): + o = f' opacity="{opacity}"' if opacity is not None else "" + f = f' fill-opacity="{fill_opacity}"' if fill_opacity is not None else "" + add( + f'' + ) + + +def text(x, y, s, size=11, fill=INK, anchor="start", weight="400", font=FONT): + add( + f'{s}' + ) + + +def note(x, y, s): + """A muted caption line, which is most of the writing on the page""" + text(x, y, s, size=10, fill=MUTED) + + +def tag(x, y, s): + """The small monospace label that names a mark rather than describes it""" + text(x, y, s, size=8.5, fill=MUTED, anchor="middle", font=MONO) + + +def line(x1, y1, x2, y2, colour, width=1.2, extra=""): + add( + f'' + ) + + +def flow(x1, y, x2): + line(x1, y, x2, y, MUTED, 1.4, ' marker-end="url(#fwd)"') + + +def seam(x, y, w, h): + """Where tiles overlap, and so where the blend writes: the same band, twice over""" + rect(x, y, w, h, "url(#seam)", "none", rx=0, sw=0) + + +def collective(x, top, bottom): + """A bar across every lane: a call no rank leaves before the others arrive""" + rect(x, top, 5, bottom - top, SYNC, "none", rx=2, sw=0) + + +def halo(x, boundaries): + """Short arrows across each internal lane boundary: neighbours only, not the group""" + for yb in boundaries: + line(x, yb - 9, x, yb + 9, SYNC, 1.3, + ' marker-end="url(#down)" marker-start="url(#up)"') + + +def peak(x, y, w, h, open_right=False): + """What one rank holds at once, which is what sets the memory it needs + + Left open where the layers it spans are themselves elided, so the edge is dotted for + the same reason the dots beside it are: the drawing stops there, the decode does not. + """ + if not open_right: + rect(x, y, w, h, "none", INK, rx=1, sw=2.6) + return + add( + f'' + ) + line(x + w, y, x + w, y + h, INK, 2.6, + ' stroke-linecap="round" stroke-dasharray="0.1 5"') + + +def head(mid, colour, width=5, back=False): + d = "M10,0 L0,5 L10,10 z" if back else "M0,0 L10,5 L0,10 z" + ref = 1 if back else 9 + return ( + f'' + f'' + ) + + +def lanes(y): + return [y + r * LANE for r in range(RANKS)] + + +def carries_on(x, y): + """The layers the drawing stops short of, marked in every lane rather than between them + + What row sharding costs is that the block-sync-block pattern to the left repeats for the + whole depth of the decoder, so the elision has to read as every rank going on doing that; + a single glyph between the lanes reads as one gap in the middle instead. + """ + for ly, (_, stroke) in zip(lanes(y), RANK): + for step in range(3): + add( + f'' + ) + + +# ------------------------------------------------------------------------------- blocks + + +def heading(y, n, title, tail, headline, *notes): + """A mode's header: what it is, what it costs, and the detail under that + + Numbered because the two modes are alternatives and a reader arriving at the top of a + long page can otherwise take them for the two halves of one pipeline. The number counts + the modes and nothing else, which is why the stages below are no longer numbered too. + + Returns the y the panels below it start at, so a row that grows a line pushes the page + down instead of needing every coordinate under it re-tuned. + """ + text(COL1, y, f"{n}. {title}", size=13, weight="700") + text(TAIL, y, tail, size=11, fill=MUTED) + text(COL1, y + 18, headline, size=11, weight="600") + for k, line_ in enumerate(notes): + note(COL1, y + 34 + k * 16, line_) + return y + 40 + len(notes) * 16 + + +def divider(y): + """The rule that opens a mode, and the only thing on the page drawn edge to edge + + The two modes get one and the windows inside tile distribution do not, which is what + keeps a section from reading as a row: "Row sharding" and "Cut the rows only" are set a + point and a half apart, and on their own that is not enough to rank them. + """ + line(COL1, y, W - COL1, y, EDGE, 1.4) + return y + + +def caption(x, y, title, *lines): + """A stage and what it says, under whichever of the two panels it belongs to + + Set at the panel's own left edge, and read in the order the arrow between the panels + already points, so neither stage needs a number to say where it comes. + """ + text(x, y, title, size=10.5, weight="600") + for k, line_ in enumerate(lines): + note(x, y + 16 + k * 14, line_) + return y + 16 + len(lines) * 14 + + +def costs(y, split): + """What a row costs and what it buys, in the columns every other row uses + + Drawn from the same attributes whichever way the latent was divided, so a reader + comparing the three rows is comparing arithmetic rather than prose. Returns the baseline + of the values, since that is what the captions below have to clear. + """ + line(COL1, y - 12, COL1 + (len(COSTS) - 1) * PITCH + 66, y - 12, RULE) + for k, (name, show) in enumerate(COSTS): + text(COL1 + k * PITCH, y, name, size=9.5, fill=MUTED) + text(COL1 + k * PITCH, y + STACK, show(split), size=11.5, weight="600") + return y + STACK + + +# --------------------------------------------------------------------------- the latent + + +def box(x0, y0, grid, n): + """Tile n's origin on the page and the window the bounds leave it""" + i, j = divmod(n, grid.across.count) + return (x0 + grid.across.at[j] * UNIT, y0 + grid.down.at[i] * UNIT, + grid.across.extent[j] * UNIT, grid.down.extent[i] * UNIT) + + +def tiles(x0, y0, grid): + """Every tile at its true extent, so the overlaps are the drawing's own + + Fills are transparent and lie on top of one another, so a band two tiles cover comes out + twice as deep, and the four-way corners deeper still. That is the redundant work. + """ + for n, who in enumerate(grid.owner): + rect(*box(x0, y0, grid, n), RANK[who][0], "none", rx=0, sw=0, fill_opacity=TINT) + for n, who in enumerate(grid.owner): + rect(*box(x0, y0, grid, n), "none", RANK[who][1], rx=1, sw=1) + for n in range(grid.tiles): + bx, by, bw, bh = box(x0, y0, grid, n) + text(bx + bw / 2, by + bh / 2 + 3, str(n), anchor="middle", size=8.5, font=MONO) + + +def overlaps(x0, y0, grid): + """The bands a second tile also covers, and so where the blend writes + + Every band is the overlap deep, because only the last tile on an axis is ever clipped + and nothing starts after it. An axis the window already spans is not cut at all, and so + contributes nothing here, which is what the strip row's three clean joins come from. + """ + for k in range(1, grid.down.count): + seam(x0, y0 + grid.down.at[k] * UNIT, PANEL, grid.down.deep * UNIT) + for k in range(1, grid.across.count): + seam(x0 + grid.across.at[k] * UNIT, y0, grid.across.deep * UNIT, PANEL) + + +# --------------------------------------------------------------------------------- rows + + +def sharding(y): + """Row sharding: one decoder call, split into bands, syncing all the way down""" + y = heading( + y, 1, "Row sharding", + f"one decoder call, split into {word(RANKS)} bands of {BOUND // RANKS} latent rows", + # A sync rather than a collective, because the convolutions swap halos with their + # neighbours and only the norms reduce across the group, and the legend draws those + # as two different things. + "Peak memory is rank-bound and every layer syncs, so the interconnect can be the " + "bottleneck.", + # Nothing to pick here, which is the contrast the tiling header is written against. + "There is nothing to choose here, since the band is the latent divided by the GPU " + "count.", + # The third thing a reader is choosing between, said where the other two are said. + # The caveat is only that a reduction adds its terms in a different order, so the + # last bits move. Naming that mechanism costs a clause and buys nothing: what the + # reader is weighing is this line against the tiling row's, and the contrast is + # between rounding they will never see and blocky colour they might. + "The image is what a single GPU would produce, apart from floating-point rounding.", + ) + bottom = y + TALL + + for r, ly in enumerate(lanes(y)): + rect(COL1, ly, PANEL, LANE, *RANK[r], rx=0) + text(COL1 + PANEL / 2, ly + LANE / 2 + 4, f"rank {r}", anchor="middle", size=10.5) + peak(COL1, y, PANEL, LANE) + + flow(COL1 + PANEL + 12, y + TALL / 2, TRACK - 10) + + for r, ly in enumerate(lanes(y)): + note(TRACK, ly + LANE / 2 + 4, f"rank {r}") + first = x = TRACK + LABEL + # Every rank runs the same layer at the same moment, so one header names them all. What + # follows a layer is the layer's own business: a convolution wants rows from its + # neighbours, a norm wants a statistic from everybody. + for layer in ("conv", "norm", "conv", "norm", "conv"): + for r, ly in enumerate(lanes(y)): + rect(x, ly + 4, LAYER, LANE - 8, *RANK[r], rx=2) + tag(x + LAYER / 2, y - 6, layer) + x += LAYER + if layer == "conv": + halo(x + GAP / 2, lanes(y)[1:]) + else: + collective(x + GAP / 2 - 2.5, y, bottom) + x += GAP + carries_on(x + 3, y) + x += ELIDED + collective(x, y, bottom) + # The band is one allocation and it is live the whole way down: the axis here is time, + # so what the mark spans is how long a rank holds it, not how much it is holding. + peak(first - 3, y + 1.5, x - first, LANE - 3, open_right=True) + # Named for what it carries, as the tile rows' are: all the bars are the same gather, + # and the count is the difference worth reading. + tag(x + 2.5, bottom + 12, "image") + + cap = costs(bottom + 32, SHARDED) + 26 + return max( + # The one fact the strip row below is written against: bands meet, tiles overlap. + # Short because the left caption column ends where the right one starts, at TRACK. + caption(COL1, cap, "Split the rows", "The bands never overlap."), + # One line, not two: how often it syncs is a column now, so this is left to say + # only what a sync is, which the legend then splits into its two marks. + caption(TRACK, cap, "Decode in lockstep", + "Convolutions swap edge rows, and norms reduce across all four ranks."), + ) + + +def window(y, grid, name, tail, aside, first, *second): + """One tiling window: the grid it leaves, the lanes it runs, and what the pair cost + + Both windows come through here, so the only thing separating the two rows below the + tiling heading is the two numbers each was built from. + """ + text(COL1, y, name, size=11.5, weight="700") + text(TAIL, y, tail, size=10.5, fill=MUTED) + note(COL1, y + 17, aside) + y += 32 + bottom = y + TALL + + tiles(COL1, y, grid) + overlaps(COL1, y, grid) + peak(*box(COL1, y, grid, grid.biggest)) + + flow(COL1 + PANEL + 12, y + TALL / 2, TRACK - 10) + + scale = SPAN / max(grid.load) + start = TRACK + LABEL + last = start + max(grid.load) * scale + + # Idle belongs to a lane and not to the row: three of the four strips run the whole + # length, so one band across every lane would imply that every rank waits. Each lane gets + # a separate tail from the end of its work to the latest rank completion. + for r, ly in enumerate(lanes(y)): + note(TRACK, ly + LANE / 2 + 4, f"rank {r}") + at = start + for n in grid.run[r]: + width = grid.weight[n] * scale + rect(at + 1, ly + 4, width - 2, LANE - 8, *RANK[r], rx=2) + # Checked against the width of this number rather than assumed from the full + # blocks, so a clipped corner tile is either named like the rest or left blank + # instead of overrunning its block. + if width - 2 >= DIGIT * len(str(n)) + 3: + text(at + width / 2, ly + LANE / 2 + 3.5, str(n), anchor="middle", + size=9, font=MONO) + at += width + if last - at >= 1: + rect(at, ly + 4, last - at, LANE - 8, MUTED, "none", rx=1, sw=0, opacity=0.16) + + # Named inside the tail rather than above the row, so the word sits in the lane it is + # true of, and only where the tail is wide enough to hold it. The grid leaves nothing + # to label, which is the comparison: its imbalance is a column, not a picture. + waiting = min(range(RANKS), key=lambda r: grid.load[r]) + soonest = start + grid.load[waiting] * scale + if last - soonest >= 24: + tag((soonest + last) / 2, lanes(y)[waiting] + LANE / 2 + 3, "idle") + + # One tile at a time, and the memory follows the largest of them rather than the first. + heaviest = max(grid.run[0], key=lambda n: grid.weight[n]) + before = sum(grid.weight[n] for n in grid.run[0][:grid.run[0].index(heaviest)]) + peak(start + before * scale - 1.5, lanes(y)[0] + 1.5, + grid.weight[heaviest] * scale + 3, LANE - 3) + + # Nothing crosses between the ranks until every tile is decoded. + collective(last + 10, y, bottom) + collective(last + 24, y, bottom) + tag(last + 26, bottom + 12, "edges, image") + + cap = costs(bottom + 32, grid) + 26 + return max(caption(COL1, cap, first), caption(TRACK, cap, *second)) + + +def tiling(y): + """Tile distribution: a window's worth per call, dealt out, gathered twice""" + y = heading( + y, 2, "Tile distribution", f"the same {word(RANKS)} GPUs, at two windows", + "Peak memory is tile-bound and the decode needs only two collectives, but it " + "repeats more work.", + # Said plainly, because the row otherwise reads as a default. Named in the terms the + # planners take, too: a window and an absolute overlap, not a fraction of a window. + "Window and overlap are yours to set in output pixels, so tune them to your VAE " + "and your GPUs.", + # The figure shows two windows and a reader will take the better-looking one for + # advice, so the disclaimer has to be here rather than left to the docs. + "The two windows below are worked examples, chosen to show the trade rather than " + "to be copied.", + # The honest summary of that trade, and the one thing the columns cannot show: a + # full-width strip is one contiguous span of a row-major tensor, where a grid's + # tile is a stride through every row it touches. + "Full-width strips overlap less and stay contiguous in row-major memory, while a " + "grid holds less at once.", + # Against the sharding row's line in the same place: what the choice costs the + # image. The seam a blend can hide; the norms it cannot, since a tile's are its own + # contents and nothing else, which is why a window can be too small rather than + # merely slow. + "The image is close but not exact: a blend hides seams, but norms over too small a " + "tile can leave the colour blocky.", + ) + + y = window( + y + 4, STRIPS, + "Cut the rows only", + f"{STRIPS.down.window_px} px tall overlapping {STRIPS.down.overlap_px} px, " + "full width", + f"{word(STRIPS.tiles).capitalize()} strips, one per rank, have the same shape as " + "the row-sharded bands, but they overlap and nothing syncs until the end.", + "Cut and overlap the rows", + "One call each, and one rank waits", + "With one strip per rank, there is nothing for the scheduler to decide.", + f"The last strip is {STRIPS.down.extent[-1]} latent rows where the others are " + f"{STRIPS.down.window}.", + # The line that answers the reader who suspects a window was picked to flatter the + # grid below. Not that no split does better, since a thinner blend plainly does: + # that at this depth of blend none does, because the gap is the blend. + "That shortfall is exactly the overlap, so closing it would thin the blend.", + ) + + heavy = max(range(RANKS), key=lambda r: len(TILED.run[r])) + light = min(range(RANKS), key=lambda r: len(TILED.run[r])) + return window( + y + 26, TILED, + "Cut both axes", + f"{TILED.down.window_px} × {TILED.across.window_px} px overlapping " + # One number when one number was asked for, so the row does not imply a per-axis + # decision that was not made. + + (f"{TILED.down.overlap_px} px on both axes" + if TILED.down.overlap_px == TILED.across.overlap_px + else f"{TILED.down.overlap_px} × {TILED.across.overlap_px} px"), + f"{word(TILED.tiles).capitalize()} tiles across {word(RANKS)} ranks let the load " + "be levelled, and a rank now holds a window rather than a strip.", + "Distribute the tiles", + "Each rank decodes its assigned tiles sequentially", + # A run is the cheap shape to blend but a coarse one to balance, so the scheduler + # moves single tiles off it, which is why two lanes hold tiles from either end. + "Each rank starts with a contiguous run, then single tiles move to level it.", + f"Rank {heavy} decodes {word(len(TILED.run[heavy]))} tiles to rank {light}'s " + f"{word(len(TILED.run[light]))}, and they still finish together.", + ) + + +def legend(y): + """What the marks mean, in one row, read before the rows that use them + + Four marks and no swatch for the blend, because the blend is not drawn anywhere. What a + good one leaves is too slight to put on a page at this size without overstating it, so + the tiling heading says it in words instead. + + The four are spaced off the width of the longest label in a fallback font, which is the + widest the row can come out, so the line holds together whichever font renders it. + """ + collective(COL1, y - 9, y + 5) + note(COL1 + 14, y + 2, "collective: every rank waits") + halo(COL1 + 191, [y - 2]) + note(COL1 + 203, y + 2, "halo swap: neighbours only") + peak(COL1 + 380, y - 8, 14, 13) + note(COL1 + 400, y + 2, "held at once") + seam(COL1 + 502, y - 8, 10, 13) + # Two tiles at most joins, four where the corners meet, so the count is left out. + note(COL1 + 518, y + 2, "tile overlap") + return y + 5 + + +def draw(): + text(COL1, 30, "DistVAE parallelism", size=17, weight="700") + # State the shared input once above all comparison rows. + text(COL1, 49, f"The two modes below each decode a {BOUND * SCALE_VAE} × " + f"{BOUND * SCALE_VAE} image from a {BOUND} × {BOUND} latent on " + f"{word(RANKS)} GPUs.", size=11, fill=MUTED) + text(COL1, 65, "The same five metrics compare both alternatives.", + size=11, fill=MUTED) + + # Define figure symbols before the comparison rows. + y = legend(88) + y = sharding(divider(y + 20) + 26) + y = tiling(divider(y + 32) + 26) + height = round(y + 14) + + front = [ + f'', + "" + HATCH + head("fwd", MUTED, width=6) + + head("down", SYNC) + head("up", SYNC, back=True) + "", + f'', + ] + return "\n".join(front + out + [""]) + + +def rasterise(svg): + """Write the PNG beside the SVG, or say why there is no new one""" + try: + import cairosvg + except ImportError: + return "no cairosvg: figure.png left as it was" + cairosvg.svg2png(bytestring=svg.encode(), write_to=PNG, scale=RASTER) + return f"wrote {PNG}" + + +if __name__ == "__main__": + svg = draw() + with open(OUT, "w") as handle: + handle.write(svg) + print(f"wrote {OUT}") + print(rasterise(svg)) diff --git a/docs/strategies.md b/docs/strategies.md new file mode 100644 index 0000000..fcfcd22 --- /dev/null +++ b/docs/strategies.md @@ -0,0 +1,82 @@ +# Choosing a decode path + +DistVAE provides row sharding and whole-tile distribution. The benchmark compares both strategies +with a vanilla unsharded Diffusers decode: + +![Row sharding and two whole-tile distributions for a 1024 by 1024 image on four GPUs, compared by peak activations, work, seams, load imbalance, and synchronization](figure.png) + +[`make_figure.py`](make_figure.py) generates the diagram from the same scheduler used at runtime. +Tile sizes and rank assignments are exact for the example. + +## Comparison + +| | Vanilla unsharded | Row sharding | Whole-tile distribution | +| --------------------------------- | ------------------------- | ------------------------------------------------ | ---------------------------------------------- | +| Work assigned to a rank | Complete decode | A band in every adapted layer | One or more complete tiles | +| Communication during the VAE call | None | Halos, metadata, and normalization statistics | Tile-edge exchange and output assembly | +| Peak activation memory | Full decode on every rank | Usually falls as ranks are added | Usually follows the largest tile | +| Repeated work | None | None | Overlap between tiles | +| Output | Reference | Matches the reference within numerical tolerance | Can differ because normalization sees one tile | + +The benchmark uses vanilla Diffusers decoding as its numerical reference. When that decode fits and +no VAE distribution is needed, DistVAE does not need to replace it. Of DistVAE's two strategies, row +sharding is the default when numerical agreement matters. Whole-tile distribution is useful when the +activation memory of a row-sharded band is still too large. + +Latency and memory depend on the VAE family, input shape, rank count, tile geometry, and +interconnect. Run all available paths on the target system rather than selecting one from rank count +alone. + +## Family guidance + +| Family | Starting point | What to check | +| -------------------------- | ------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------- | +| `AutoencoderKL` and Flux.2 | Whole-tile strips under memory pressure | A full-width or full-height strip can be faster and lighter than row sharding. Compare every planned window. | +| Qwen-Image | Whole-tile strips under memory pressure | Its frame-at-a-time decoder can retain memory outside the spatial tile, so smaller windows do not guarantee a lower peak. | +| Wan | Row sharding for latency | Tiling can reduce peak memory, but row sharding may remain faster. | +| LTX-2 | Row sharding | Tiled plans can use more memory than row sharding. Measure peak memory before enabling them. | +| HunyuanVideo | Tiling when row-sharded memory is too high | Tiling can reduce memory while remaining slower than row sharding. | +| HunyuanVideo 1.5 | Benchmark both distributed paths | A selected tiled plan may reduce both latency and memory. | + +The table lists plausible tradeoffs to measure; none is guaranteed on a given system. The benchmark +uses synthetic weights, so it cannot measure trained-model quality or end-to-end pipeline memory. + +## Communication + +Row sharding communicates inside adapted layers. Convolutions exchange neighboring rows, distributed +group normalization reduces statistics, and uneven outputs require metadata and output gathers. The +number of distributed API calls therefore follows the decoder architecture. + +Whole-tile distribution communicates around independent decoder calls. Ranks exchange tile-edge data +for blending and gather completed pieces for assembly. Most families do this once for the complete +decode. HunyuanVideo repeats the spatial tiling operation for each temporal chunk. + +The unsharded path performs no distributed operation inside the measured VAE call. Benchmark +barriers around timed iterations are excluded from that statement. + +## Whole tiles rather than rows inside them + +DistVAE assigns complete tiles to ranks. Sharding every tile by rows would add patching, halo +exchange, and gathering to each tile. Complete tiles can be decoded independently. + +The scheduler balances tile area, not tile count, because tiles on the last row or column may be +clipped. A tile cannot be divided between ranks, so balance improves when each rank receives several +tiles. If the grid has fewer tiles than ranks, distributed tiling is disabled and every rank decodes +the full grid. + +Latency depends on the VAE, shape, interconnect, and tile geometry. The +[benchmark guide](../bench/README.md) explains how to compare them. + +## Why the window is rectangular + +Height and width affect the grid independently. A rectangular window can reduce clipping or avoid +cutting one axis. The figure's 432 × 296 window produces a 3 × 5 grid with better load balance than +a 384 × 384 window at the same overlap. + +[Choosing a tile window](tiling.md) covers strips, clipping, overlap, and rank count. + +## Video + +Both distributed paths split only spatial axes. Every band or tile keeps its current temporal +extent. [Temporal decoding](tiling.md#temporal-decoding) describes how each VAE iterates over that +extent. diff --git a/docs/tiling.md b/docs/tiling.md new file mode 100644 index 0000000..439f16f --- /dev/null +++ b/docs/tiling.md @@ -0,0 +1,103 @@ +# Choosing a tile window + +The [`Tiling` section of the README](../README.md#tiling) shows the API. This page explains how +window shape, overlap, and rank count affect a tiled decode. Examples use the 128 × 128 latent from +[Choosing a decode path](strategies.md). + +## The two axes cost differently + +`tile_shape_plan` sets height and width separately. Tile area usually determines tile-local +activation memory, while each tiled axis adds overlap. Temporal state and family-specific decoder +behavior can dominate the measured peak. In the figure, full-width strips decode 1.26 times the +latent area; the two-axis grid decodes 1.46 times. + +`tile_overlap_plan` also sets each axis separately. A rectangular window does not require different +overlaps. Use different values only when the two axes have different seam or stride requirements. + +A window at least as wide as the image leaves the width axis untiled and produces full-width strips. +Pass `sample_shape` to `tile_overlap_plan`; an untiled axis must request zero overlap. + +```python +from distvae import vae as vae_api + +height, width = 1024, 1024 + +# 224 rows deep, and wider than the image across, so the across stride clears it in one +# step and the grid comes out as one column of full-width strips. +shape = vae_api.tile_shape_plan(pipe.vae, 224, 1408) +if shape is None: + raise ValueError("this VAE cannot use a 224x1408px tile shape") +vae_api.apply_tile_plan(pipe.vae, shape) +step = vae_api.tile_overlap_plan( + pipe.vae, 56, 0, sample_shape=(height, width) +) +if step is None: + raise ValueError("this VAE cannot use a 56x0px tile overlap") +vae_api.apply_tile_plan(pipe.vae, step) +replacement = vae_api.tiled_decode_for(pipe.vae) +if replacement is not None: + pipe.vae.tiled_decode = replacement +``` + +Strips repeat less work and create fewer seams than a two-axis grid, but retain the full size of the +untiled axis. In the figure, four strips hold 34% of the activations and create three seams. The 3 × +5 grid holds 12% and creates twenty-two seams. + +For a wide image, columns can keep the tiled dimension larger; for a tall image, rows can do the +same. + +Strips can also run faster than a two-axis grid with similar tile area. They decode a few long, +contiguous spans instead of many short tiles. The planner keeps both geometries when neither +dominates on window area, decoded work, load imbalance, and tile columns. Benchmark results decide +which geometry is useful on a device. + +## Clipping unbalances a grid, not the tile count + +The last row and column may contain smaller, clipped tiles. Assigning the same number of tiles to +each rank can therefore assign different amounts of work. + +The figure's four strips cover 43, 43, 43, and 32 latent rows. This leaves the heaviest rank 6.8% +above an even split. The final strip is shorter by exactly the overlap, so the imbalance comes from +blending at the image boundary. + +The 3 × 5 grid is only 0.5% above an even split because each rank receives several tiles. More tiles +give the scheduler more ways to balance clipped edges. + +When each rank receives one tile, full tiles determine peak memory and wall time. Increasing overlap +can sometimes use otherwise idle time without changing either value. This happens only when the +wider overlap does not increase the largest tile or the number of tiles; verify it with the +benchmark. + +## The tile count caps the GPU count + +The window and image shape determine the tile count. Distributed tiling cannot use more ranks than +tiles. Load balance also depends on how full and clipped tiles divide among the ranks. A smaller +window creates more tiles and can improve balance, at the cost of more overlap and seams. + +## Benchmark plan selection + +The DistVAE planners validate an exact request; they do not choose policy for an application. The +benchmark searches grids containing between `max(2, ranks)` and `4 × ranks` tiles. It rejects +unsupported windows and removes candidates that are worse in window area, decoded area, rank +imbalance, and tile columns. It selects up to three distinct windows from the remaining frontier: +`coarse` has the largest window, `fine` has the smallest, and `balanced` minimizes the worst +normalized window-area, decoded-area, and imbalance score among the other candidates. + +The profile names describe geometry. They do not rank latency or peak memory. A coarse plan can be +fastest because it has fewer decoder calls and less overlap. A fine plan has a smaller window than +the coarse plan, but that does not guarantee the lowest measured peak because family state can be +independent of the spatial window. Measure all selected plans. + +The narrow axis must contain at least sixteen latent units and at least one unit per rank. Overlap +on each tiled axis must be at least one quarter of the window. These are conservative limits, not +image-quality measurements. Synthetic weights cannot measure normalization drift or visible seams. +Test the selected window on a trained model before using it in production. + +## Temporal decoding + +DistVAE distributes spatial tiles; the owning diffusers VAE controls temporal iteration. +HunyuanVideo invokes the spatial tiling operation once for each temporal chunk, so tile exchange and +assembly repeat with the chunk count. LTX-2 uses the same spatial loop but normally decodes the full +temporal extent in one call. Enabling its upstream temporal tiling repeats the spatial operation for +each temporal chunk. Wan and Qwen-Image instead step through frames inside each spatial tile while +threading a causal cache. diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..7f10114 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = test +markers = + gloo: spawns multiple ranks over gloo and runs on CPU; no accelerator required diff --git a/setup.py b/setup.py index 2200e31..2b40f1e 100644 --- a/setup.py +++ b/setup.py @@ -11,12 +11,17 @@ name="DistVAE", author="Jinzhe Pan", author_email="eigensystem1318@gmail.com", - packages=find_packages(), - install_requires=["torch>=2.2", "diffusers>=0.35.0", "transformers"], + packages=find_packages(include=["distvae", "distvae.*"]), + # This is the oldest dependency pair covered by compatibility CI. VAE families introduced + # in later diffusers releases are resolved lazily and name the missing class when used. + install_requires=["torch>=2.2", "diffusers>=0.30.3"], extras_require={ + "pipeline": ["transformers"], "dev": [ "pytest", "black", + "mdformat", + "mdformat-gfm", "flake8", "mypy", ], diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 0000000..4a027fb --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,14 @@ +import zlib + +import pytest + + +@pytest.fixture +def master_port(request): + """Return a deterministic port for this test's Gloo rendezvous. + + Ports 20000–29999 are below Linux's default ephemeral range. CRC32 makes the initial port + stable for each test across runs. `run_distributed` retries if two tests still collide. + """ + base = 20000 + return base + zlib.crc32(request.node.nodeid.encode()) % 10000 diff --git a/test/distributed_harness.py b/test/distributed_harness.py new file mode 100644 index 0000000..2deabac --- /dev/null +++ b/test/distributed_harness.py @@ -0,0 +1,177 @@ +"""Scaffolding shared by the multi-rank CPU tests. + +Every adapter test asks the same question: does sharding a module across ranks reproduce what +the unsharded module returns. That means the same preamble (a gloo group over CPU), the same +epilogue (compare on rank 0, then fail everywhere rather than deadlocking the ranks that +passed), and the same spawn call. Only the module under test differs. +""" + +import os +import socket +from datetime import timedelta +from time import monotonic +from typing import Optional + +import torch +import torch.distributed as dist +from torch.multiprocessing import spawn +from torch.multiprocessing.spawn import ProcessRaisedException + +from distvae.utils import ParallelContext + +# How many ports to try before giving up on finding a free one. +_RENDEZVOUS_ATTEMPTS = 4 +_PROCESS_GROUP_TIMEOUT = timedelta(seconds=60) +_DISTRIBUTED_TEST_TIMEOUT_SECONDS = 300 + + +def init_gloo(rank: int, world_size: int, master_port: int) -> torch.device: + """Join this rank to a gloo group over CPU and return the device to build on""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + dist.init_process_group( + backend="gloo", + init_method="env://", + timeout=_PROCESS_GROUP_TIMEOUT, + ) + return torch.device("cpu") + + +def make_parallel_context(patch_dim: int = -2) -> ParallelContext: + """Capture the current test process group, or a one-rank local context.""" + if not dist.is_initialized(): + return ParallelContext(None, rank=0, world_size=1, patch_dim=patch_dim) + group = dist.group.WORLD + world_size = dist.get_world_size(group) + return ParallelContext( + group, + rank=dist.get_rank(group), + world_size=world_size, + patch_dim=patch_dim, + global_ranks=tuple(range(world_size)), + ) + + +def assert_matches_reference( + rank: int, + actual: torch.Tensor, + expected: Optional[torch.Tensor], + what: str, + atol: float = 1e-4, + rtol: float = 1e-3, +) -> None: + """Compare on rank 0, then raise on every rank + + Only rank 0 holds the reference, which is why expected is optional elsewhere. Raising there + alone would leave the other ranks waiting on the next collective, and the test would hang + instead of failing. + """ + detail = "" + ok = torch.ones(1, dtype=torch.int64) + if rank == 0: + if actual.shape != expected.shape: + detail = f"shape {tuple(actual.shape)} != reference {tuple(expected.shape)}" + ok.zero_() + elif not torch.allclose(actual, expected, atol=atol, rtol=rtol): + diff = (actual - expected).abs() + detail = ( + f"max diff {diff.max().item():.3g}, mean diff {diff.mean().item():.3g} " + f"(atol={atol}, rtol={rtol})" + ) + ok.zero_() + dist.broadcast(ok, src=0) + dist.barrier() + if ok.item() == 0: + raise AssertionError(f"{what} did not match the single-rank reference: {detail}") + + +def assert_no_less_precise_than( + rank: int, + actual: torch.Tensor, + stock: Optional[torch.Tensor], + gold: Optional[torch.Tensor], + what: str, + slack: float = 1.5, +) -> None: + """Compare our rounding against the stock operator's, then raise on every rank + + A low-precision dtype cannot reproduce a float32 answer, so asserting equality there is + asserting a failure: every test in this suite runs in float32, where the `.to(x.dtype)` casts + inside the sharded norms are no-ops, and so the precision those casts cost is invisible to all + of them. What can fairly be asked in bf16 is that the replacement rounds no worse than the + operator it replaces. Both are measured against `gold` - the float32 result rounded once, the + best the narrow dtype can hold - and the sharded path is allowed `slack` times the stock op's + own error, floored at one quantum of the dtype so an exact stock answer does not demand one. + """ + detail = "" + ok = torch.ones(1, dtype=torch.int64) + if rank == 0: + scale = gold.float().abs().max().clamp(min=1e-12) + ours = ((actual.float() - gold.float()).abs().max() / scale).item() + theirs = ((stock.float() - gold.float()).abs().max() / scale).item() + allowed = max(theirs * slack, torch.finfo(actual.dtype).eps) + if ours > allowed: + detail = ( + f"{ours * 100:.3f}% of scale, against the stock operator's {theirs * 100:.3f}% " + f"(allowed {allowed * 100:.3f}%, slack x{slack})" + ) + ok.zero_() + dist.broadcast(ok, src=0) + dist.barrier() + if ok.item() == 0: + raise AssertionError(f"{what} rounds worse than the operator it replaces: {detail}") + + +def _free_port() -> int: + """A port nothing is listening on, as of asking""" + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _terminate_processes(context) -> None: + """Stop and reap every rank still owned by a timed-out spawn context.""" + for process in context.processes: + if process.is_alive(): + process.terminate() + for process in context.processes: + process.join(timeout=5) + for process in context.processes: + if process.is_alive(): + process.kill() + process.join(timeout=5) + + +def run_distributed(worker, world_size: int, args: tuple, master_port: int) -> None: + """Spawn world_size ranks running worker(rank, *args); raises if any rank does + + Rank 0 opens the rendezvous socket, so a port taken between the fixture choosing it and rank 0 + binding it fails the test for a reason that has nothing to do with sharding. Retried on a + fresh port, which is the only thing that can be done about it from here: no port can be held + open for the ranks, since rank 0 has to bind it itself. + """ + for attempt in range(_RENDEZVOUS_ATTEMPTS): + context = spawn( + worker, + nprocs=world_size, + args=(world_size, *args, master_port), + join=False, + ) + deadline = monotonic() + _DISTRIBUTED_TEST_TIMEOUT_SECONDS + try: + while not context.join(timeout=1): + if monotonic() >= deadline: + _terminate_processes(context) + raise TimeoutError( + f"{worker.__name__} timed out after " + f"{_DISTRIBUTED_TEST_TIMEOUT_SECONDS}s with {world_size} ranks" + ) + return + except ProcessRaisedException as raised: + _terminate_processes(context) + last = attempt == _RENDEZVOUS_ATTEMPTS - 1 + if last or "EADDRINUSE" not in str(raised): + raise + master_port = _free_port() diff --git a/test/test_ResnetBlock2d.py b/test/test_ResnetBlock2d.py deleted file mode 100644 index baa6507..0000000 --- a/test/test_ResnetBlock2d.py +++ /dev/null @@ -1,83 +0,0 @@ -from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.modules.adapters.resnet_adapters import ResnetBlock2DAdapter -from distvae.utils import DistributedEnv -from torch.nn import GroupNorm - -from diffusers.models.resnet import ResnetBlock2D - -import torch -import random -import argparse -import torch.distributed as dist -from torch import nn -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - DistributedEnv.initialize(None) - - resnet = ResnetBlock2D( - in_channels=64, - out_channels=32, - temb_channels=None, - eps=1e-6, - groups=4, - dropout=0.0, - time_embedding_norm="default", - non_linearity="swish", - output_scale_factor=1.0, - pre_norm=True, - ).to(device) - patch_resnet = ResnetBlock2DAdapter(resnet).to(device) - - hidden_state = torch.randn(1, 64, args.height, args.width, device=device) - - result = resnet(hidden_state, None) - # if rank == 0: - # print("result: ", result) - - patch = Patchify() - depatch = DePatchify() - patch_result = patch_resnet(patch(hidden_state)) - # print("patch_res:", rank, patch_result) - patch_result = depatch(patch_result) - - if dist.get_rank() == 0: - assert torch.allclose(result, patch_result, atol=1e-2), "two hidden states are not equal" - - dist.barrier() - dist.destroy_process_group() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test/test_UpBlock2d.py b/test/test_UpBlock2d.py deleted file mode 100644 index a1a30ae..0000000 --- a/test/test_UpBlock2d.py +++ /dev/null @@ -1,71 +0,0 @@ -from distvae.modules.adapters.unets.unet_2d_blocks_adapters import UpDecoderBlock2DAdapter, UpDecoderBlock2D -from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.utils import DistributedEnv - -import torch -import random -import argparse -import torch.distributed as dist -from torch import nn -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - DistributedEnv.initialize(None) - - up_block = UpDecoderBlock2D(num_layers = 3, in_channels=256, out_channels=128).to(device) - patch_up_block = UpDecoderBlock2DAdapter(up_block).to(device) - - hidden_state = torch.randn(1, 256, args.height, args.width, device=device) - print("hidden state shape: ", hidden_state.shape) - - result = up_block(hidden_state) - # if rank == 0: - # print("result: ", result) - - patch = Patchify() - depatch = DePatchify() - patch_result = patch_up_block(patch(hidden_state)) - # print("patch_res:", rank, patch_result) - patch_result = depatch(patch_result) - print("result shape: ", patch_result.shape) - - if dist.get_rank() == 0: - assert torch.allclose(result, patch_result, atol=1e-3), "two hidden states are not equal" - - dist.barrier() - dist.destroy_process_group() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test/test_adapter_compatibility.py b/test/test_adapter_compatibility.py new file mode 100644 index 0000000..848b77d --- /dev/null +++ b/test/test_adapter_compatibility.py @@ -0,0 +1,35 @@ +import subprocess +import sys +from pathlib import Path + +import pytest + +from distvae.modules.adapters.vae import decoder_adapters + + +ROOT = Path(__file__).parents[1] + + +def test_adapter_packages_do_not_import_implementations_eagerly(): + script = """ +import sys +import distvae.modules.adapters +import distvae.modules.adapters.vae + +loaded = set(sys.modules) +forbidden = { + "distvae.modules.adapters.downsampling_adapters", + "distvae.modules.adapters.upsampling_adapters", + "distvae.modules.adapters.vae.decoder_adapters", + "distvae.modules.adapters.vae.encoder_adapters", +} +assert loaded.isdisjoint(forbidden), sorted(loaded & forbidden) +""" + + subprocess.run([sys.executable, "-c", script], cwd=ROOT, check=True) + + +@pytest.mark.parametrize("option", ["use_profiler", "verbose"]) +def test_decoder_instrumentation_options_point_to_the_benchmark_harness(option): + with pytest.raises(ValueError, match="bench"): + decoder_adapters.DecoderAdapter(object(), **{option: True}) diff --git a/test/test_adapter_parameter_identity.py b/test/test_adapter_parameter_identity.py new file mode 100644 index 0000000..1b24cd2 --- /dev/null +++ b/test/test_adapter_parameter_identity.py @@ -0,0 +1,83 @@ +import pytest +import torch +import torch.nn as nn +from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + +from distvae.modules.adapters.downsampling_adapters import _zero_pad_strided_conv +from distvae.modules.adapters.layers.conv_adapters import ( + Conv2dAdapter, + Conv3dAdapter, + WanCausalConv3dAdapter, +) +from distributed_harness import make_parallel_context + + +def _assert_reuses_parameters_and_gradients( + original, replacement, optimizer, input_shape +): + weight = original.weight + bias = original.bias + + assert replacement.weight is weight + assert replacement.bias is bias + assert optimizer.param_groups[0]["params"][0] is replacement.weight + + replacement(torch.randn(input_shape)).sum().backward() + + assert weight.grad is not None + if bias is not None: + assert bias.grad is not None + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv2d_adapter_reuses_original_parameters(bias): + conv = nn.Conv2d(2, 3, 3, padding=1, bias=bias) + optimizer = torch.optim.SGD(conv.parameters(), lr=0.1) + adapted = Conv2dAdapter( + conv, parallel_context=make_parallel_context() + ).conv2d + + _assert_reuses_parameters_and_gradients( + conv, adapted, optimizer, (1, 2, 5, 5) + ) + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv3d_adapter_reuses_original_parameters(bias): + conv = nn.Conv3d(2, 3, 3, padding=1, bias=bias) + optimizer = torch.optim.SGD(conv.parameters(), lr=0.1) + adapted = Conv3dAdapter( + conv, parallel_context=make_parallel_context() + ).conv3d + + _assert_reuses_parameters_and_gradients( + conv, adapted, optimizer, (1, 2, 4, 5, 5) + ) + + +@pytest.mark.parametrize("bias", [True, False]) +def test_wan_causal_conv3d_adapter_reuses_original_parameters(bias): + conv = WanCausalConv3d(2, 3, 3, padding=1) + if not bias: + conv.register_parameter("bias", None) + optimizer = torch.optim.SGD(conv.parameters(), lr=0.1) + adapted = WanCausalConv3dAdapter( + conv, parallel_context=make_parallel_context() + ).conv3d + + _assert_reuses_parameters_and_gradients( + conv, adapted, optimizer, (1, 2, 4, 5, 5) + ) + + +@pytest.mark.parametrize("bias", [True, False]) +def test_zero_pad_strided_conv_reuses_original_parameters(bias): + conv = nn.Conv2d(2, 3, 3, stride=2, padding=0, bias=bias) + optimizer = torch.optim.SGD(conv.parameters(), lr=0.1) + adapted = _zero_pad_strided_conv( + conv, conv_block_size=0, parallel_context=make_parallel_context() + ) + + _assert_reuses_parameters_and_gradients( + conv, adapted, optimizer, (1, 2, 6, 6) + ) diff --git a/test/test_wanzeropadconv2d.py b/test/test_asymmetric_zero_pad_conv2d.py similarity index 54% rename from test/test_wanzeropadconv2d.py rename to test/test_asymmetric_zero_pad_conv2d.py index 93f9720..7fb2c4f 100644 --- a/test/test_wanzeropadconv2d.py +++ b/test/test_asymmetric_zero_pad_conv2d.py @@ -1,13 +1,13 @@ """ -Multi-rank integration tests for WanZeroPadConv2d (GLOO / CPU). +Multi-rank integration tests for AsymmetricZeroPadConv2d (GLOO / CPU). -Compares merged distributed output (Patchify -> WanZeroPadConv2d -> DePatchify) +Compares merged distributed output (Patchify -> AsymmetricZeroPadConv2d -> DePatchify) to the single-rank reference math (must stay in sync with -distvae.models.layers.wan.zeropadconv2d WanZeroPadConv2d._conv_forward group_world_size==1 branch). +AsymmetricZeroPadConv2d._conv_forward's group_world_size==1 branch). Run from repo root: - pytest test/test_wan_zeropadconv2d_distributed_gloo.py -v -m gloo - python test/test_wan_zeropadconv2d_distributed_gloo.py + pytest test/test_asymmetric_zero_pad_conv2d.py -v -m gloo + python test/test_asymmetric_zero_pad_conv2d.py """ from __future__ import annotations @@ -15,6 +15,7 @@ import argparse import os import sys +import zlib import pytest import torch @@ -22,12 +23,14 @@ import torch.nn.functional as F from torch.multiprocessing import spawn -from distvae.models.layers.wan.zeropadconv2d import WanZeroPadConv2d +from distvae.models.layers.asymmetric_zero_pad_conv2d import AsymmetricZeroPadConv2d from distvae.modules.patch_utils import DePatchify, Patchify -from distvae.utils import DistributedEnv +from distributed_harness import make_parallel_context -def reference_wan_zeropad_conv2d(x: torch.Tensor, module: WanZeroPadConv2d) -> torch.Tensor: +def reference_asymmetric_zero_pad_conv2d( + x: torch.Tensor, module: AsymmetricZeroPadConv2d +) -> torch.Tensor: pad = tuple(module.reversed_zero_padding) x = F.pad(x, pad, mode="constant", value=0) y = F.conv2d( @@ -48,6 +51,9 @@ def worker( world_size: int, patch_dim: int, block_size: int, + height: int, + width: int, + patch_scale_factor: int, seed: int, master_port: int, ) -> None: @@ -57,19 +63,14 @@ def worker( os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) dist.init_process_group(backend="gloo", init_method="env://") - DistributedEnv.initialize(None) torch.manual_seed(seed) in_ch, out_ch = 8, 8 - n, h, w = 1, 16, 16 - if patch_dim == -2: - assert h % world_size == 0, "H must split evenly for Patchify chunk" - else: - assert patch_dim == -1 - assert w % world_size == 0, "W must split evenly for Patchify chunk" + n, h, w = 1, height, width + context = make_parallel_context(patch_dim) x_full = torch.randn(n, in_ch, h, w, device=device, dtype=torch.float32) - layer = WanZeroPadConv2d( + layer = AsymmetricZeroPadConv2d( in_channels=in_ch, out_channels=out_ch, kernel_size=3, @@ -81,24 +82,27 @@ def worker( dtype=torch.float32, reversed_zero_padding=(0, 1, 0, 1), block_size=block_size, - patch_dim=patch_dim, - use_uniform_patch=True, + parallel_context=context, ).eval() - patchify = Patchify(patch_dim=patch_dim, use_uniform_patch=False) - depatchify = DePatchify(patch_dim=patch_dim, use_uniform_patch=False) + patchify = Patchify(context, scale_factor=patch_scale_factor) + depatchify = DePatchify(context) try: with torch.no_grad(): - y_ref = reference_wan_zeropad_conv2d(x_full, layer) + y_ref = reference_asymmetric_zero_pad_conv2d(x_full, layer) x_local = patchify(x_full) y_local = layer(x_local) y_merged = depatchify(y_local) if not torch.allclose(y_ref, y_merged, atol=1e-5, rtol=1e-5): raise AssertionError( - f"WanZeroPadConv2d distributed output mismatch " + f"AsymmetricZeroPadConv2d distributed output mismatch " f"(max diff {(y_ref - y_merged).abs().max().item():.6g})" ) + # Leave together. A rank that tears its Gloo context down while another is still holding + # one exits through std::terminate, which pytest can only report as a spawned process + # dying on SIGABRT - a teardown race wearing the costume of a failed assertion. + dist.barrier() finally: dist.destroy_process_group() @@ -109,11 +113,23 @@ def _run_one( block_size: int, seed: int, master_port: int, + height: int = 16, + width: int = 16, + patch_scale_factor: int = 1, ) -> None: spawn( worker, nprocs=world_size, - args=(world_size, patch_dim, block_size, seed, master_port), + args=( + world_size, + patch_dim, + block_size, + height, + width, + patch_scale_factor, + seed, + master_port, + ), join=True, ) @@ -121,14 +137,16 @@ def _run_one( @pytest.fixture def master_port(request): """Unique port per test to avoid Address already in use when tests run sequentially.""" + # crc32 rather than hash(): the built-in is salted per interpreter, so the port a test binds + # moved every run and a failure could not be reproduced by asking for that test again. base = 29600 nodeid = request.node.nodeid - return base + (hash(nodeid) % 10000) + return base + (zlib.crc32(nodeid.encode()) % 10000) @pytest.mark.gloo @pytest.mark.parametrize("world_size,patch_dim", [(2, -2), (4, -2), (2, -1)]) -def test_wan_zeropadconv2d_gloo_matches_single_rank_reference( +def test_asymmetric_zero_pad_conv2d_gloo_matches_single_rank_reference( world_size, patch_dim, master_port, seed=42 ): """Direct path (block_size=0): merged multi-rank output equals single-rank reference.""" @@ -142,19 +160,42 @@ def test_wan_zeropadconv2d_gloo_matches_single_rank_reference( @pytest.mark.gloo -def test_wan_zeropadconv2d_gloo_chunked_path(master_port, seed=42): - """Chunked path: large H/W and block_size>0 so _use_direct_path is False inside the layer.""" +@pytest.mark.parametrize("block_size", [1, 4]) +def test_asymmetric_zero_pad_conv2d_gloo_chunked_path( + block_size, master_port, seed=42 +): + """Chunked paths clamp every input chunk to at least the kernel size.""" _run_one( world_size=2, patch_dim=-2, - block_size=4, + block_size=block_size, + seed=seed, + master_port=master_port, + ) + + +@pytest.mark.gloo +@pytest.mark.parametrize("patch_dim,block_size", [(-2, 0), (-2, 4), (-1, 0), (-1, 4)]) +def test_asymmetric_zero_pad_conv2d_matches_reference_for_unequal_patch_bands( + patch_dim, block_size, master_port, seed=42 +): + height, width = (40, 16) if patch_dim == -2 else (16, 40) + _run_one( + world_size=3, + patch_dim=patch_dim, + block_size=block_size, + height=height, + width=width, + patch_scale_factor=8, seed=seed, master_port=master_port, ) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="WanZeroPadConv2d GLOO multi-rank tests") + parser = argparse.ArgumentParser( + description="AsymmetricZeroPadConv2d GLOO multi-rank tests" + ) parser.add_argument("--world_size", type=int, default=None) parser.add_argument("--patch_dim", type=int, default=None) parser.add_argument("--seed", type=int, default=42) diff --git a/test/test_cache_cursor.py b/test/test_cache_cursor.py new file mode 100644 index 0000000..6332eb0 --- /dev/null +++ b/test/test_cache_cursor.py @@ -0,0 +1,26 @@ +"""Where the causal decoders are up to in their feature cache, and who owns that position""" + +import unittest + +from distvae.utils import cache_cursor + + +class TestCacheCursor(unittest.TestCase): + + def test_omitting_a_cursor_gets_a_fresh_one_every_time(self): + first, second = cache_cursor(None), cache_cursor(None) + self.assertEqual(first, [0]) + self.assertEqual(second, [0]) + # Identity matters because blocks advance the cursor in place while walking the cache. + # Sharing one list would make the second decode start where the first stopped. + self.assertIsNot(first, second) + + def test_a_cursor_handed_in_is_the_one_used(self): + # A caller walking the cache itself passes its own position in, and gets it back to keep + # advancing rather than a copy that strands its progress here. + mine = [7] + self.assertIs(cache_cursor(mine), mine) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_causal_vae_cache.py b/test/test_causal_vae_cache.py new file mode 100644 index 0000000..7ad6d91 --- /dev/null +++ b/test/test_causal_vae_cache.py @@ -0,0 +1,179 @@ +"""Temporal feature-cache behavior through the public DistVAE VAE API.""" + +import sys + +import pytest +import torch +from distributed_harness import init_gloo, run_distributed + +from distvae.vae.parallel import parallelize_decoder, parallelize_encoder + +diffusers = pytest.importorskip("diffusers") + + +FAMILIES = { + "wan": ( + diffusers.AutoencoderKLWan, + {}, + "WanEncoderAdapter", + "WanDecoderAdapter", + ), + "qwen-image": ( + getattr(diffusers, "AutoencoderKLQwenImage", None), + {"attn_scales": []}, + "QwenImageEncoderAdapter", + "QwenImageDecoderAdapter", + ), +} + + +def _cache_changed(before, after): + for old, new in zip(before, after): + if old is None or new is None: + if old is not new: + return True + elif isinstance(old, torch.Tensor) and isinstance(new, torch.Tensor): + if old.shape != new.shape or not torch.equal(old, new): + return True + elif old != new: + return True + return False + + +def _record_cache_calls(module, records): + original = module.forward + + def recording_forward(*args, **kwargs): + cache = kwargs["feat_cache"] + cursor = kwargs["feat_idx"] + before = [ + value.clone() if isinstance(value, torch.Tensor) else value + for value in cache + ] + record = { + "cache": cache, + "cursor": cursor, + "start": cursor[0], + "nonempty_before": sum(value is not None for value in cache), + "first_chunk": kwargs.get("first_chunk"), + } + result = original(*args, **kwargs) + record.update( + end=cursor[0], + nonempty_after=sum(value is not None for value in cache), + mutated=_cache_changed(before, cache), + ) + records.append(record) + return result + + module.forward = recording_forward + + +def _assert_public_chunks(records, cache_size): + assert len(records) == 2 + assert [record["start"] for record in records] == [0, 0] + ends = [record["end"] for record in records] + assert ends == [ends[0], ends[0]], (ends, cache_size) + assert 0 < ends[0] <= cache_size, (ends, cache_size) + assert records[0]["cache"] is records[1]["cache"] + assert records[0]["cursor"] is not records[1]["cursor"] + assert records[0]["nonempty_before"] == 0 + assert records[0]["nonempty_after"] > 0 + assert records[1]["nonempty_before"] > 0 + assert all(record["mutated"] for record in records) + + +def _assert_omitted_cursor_sessions(adapter, sample, cache_size, **kwargs): + outputs = [] + for _ in range(2): + cache = [None] * cache_size + outputs.append(adapter(sample.clone(), feat_cache=cache, **kwargs)) + assert any(value is not None for value in cache) + torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0) + + +def cache_worker(rank, world_size, family, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + cls, extra, encoder_adapter, decoder_adapter = FAMILIES[family] + + torch.manual_seed(seed) + vae = cls( + base_dim=8, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=1, + **extra, + ).eval() + vae.clear_cache() + encoder_cache_size = vae._enc_conv_num + decoder_cache_size = vae._conv_num + + assert parallelize_encoder(vae, None) == encoder_adapter + assert parallelize_decoder(vae, None) == decoder_adapter + assert len(vae._enc_feat_map) == encoder_cache_size + assert len(vae._feat_map) == decoder_cache_size + + encoder_calls = [] + decoder_calls = [] + _record_cache_calls(vae.encoder.encoder, encoder_calls) + _record_cache_calls(vae.decoder.decoder, decoder_calls) + + pixels = torch.randn(1, 3, 5, 32, 32) + latents = torch.randn(1, 4, 2, 4, 4) + with torch.no_grad(): + encoded = vae.encode(pixels).latent_dist.parameters + decoded = vae.decode(latents).sample + + assert encoded.shape == (1, 8, 2, 4, 4) + assert decoded.shape == (1, 3, 5, 32, 32) + _assert_public_chunks(encoder_calls, encoder_cache_size) + _assert_public_chunks(decoder_calls, decoder_cache_size) + if family == "wan": + assert [record["first_chunk"] for record in decoder_calls] == [True, False] + else: + assert [record["first_chunk"] for record in decoder_calls] == [None, None] + + with torch.no_grad(): + _assert_omitted_cursor_sessions( + vae.encoder, + pixels[:, :, :1], + encoder_cache_size, + ) + decoder_options = {"first_chunk": True} if family == "wan" else {} + _assert_omitted_cursor_sessions( + vae.decoder, + latents[:, :, :1], + decoder_cache_size, + **decoder_options, + ) + finally: + torch.distributed.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("family", FAMILIES, ids=FAMILIES) +def test_public_causal_vae_paths_thread_two_chunks_and_isolate_sessions( + family, master_port, seed=42 +): + if FAMILIES[family][0] is None: + pytest.skip("installed diffusers has no AutoencoderKLQwenImage") + run_distributed(cache_worker, 1, (family, seed), master_port) + + +def test_unavailable_family_skips_before_spawning(monkeypatch): + family = "unavailable" + monkeypatch.setitem(FAMILIES, family, (None, {}, "Encoder", "Decoder")) + spawned = [] + monkeypatch.setattr( + sys.modules[__name__], + "run_distributed", + lambda *args: spawned.append(args), + ) + + with pytest.raises(pytest.skip.Exception): + test_public_causal_vae_paths_thread_two_chunks_and_isolate_sessions( + family, master_port=1 + ) + + assert spawned == [] diff --git a/test/test_conv2d.py b/test/test_conv2d.py index 8269504..cd191f4 100644 --- a/test/test_conv2d.py +++ b/test/test_conv2d.py @@ -1,144 +1,108 @@ -from distvae.models.layers.conv2d import PatchConv2d -from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter -from distvae.utils import DistributedEnv +"""PatchConv2d against nn.Conv2d, over gloo on CPU, at sizes that do not divide evenly. + +This was a torchrun script whose only verdict was a print: it computed the difference, printed +"FAILED" when it was too large, and exited 0 either way, with the one real assertion commented +out at the bottom. Nothing ran it and nothing could have failed it, which is a shame, because +the sizes it swept are the interesting ones - odd extents, and extents that do not divide by the +rank count, at stride 1 and stride 2. Those are what exercise the halo widths and the +global-position cropping, and they are what is kept here. + +Sizes are smaller than the original 1024x1024 at 64 channels, which was sized for a GPU. What +makes a size interesting here is its remainder against the rank count and its parity, not its +magnitude. + +Run from repo root: + pytest test/test_conv2d.py -v +""" -import torch -import random import argparse +import os +import sys + +import pytest +import torch import torch.distributed as dist -from torch import nn -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -class Conv2dModules(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride, padding): - super().__init__() - self.convs = nn.ModuleList([ - # nn.Conv2d(512, 256, kernel_size, stride, padding), - # nn.Conv2d(256, 128, kernel_size, stride, padding), - # nn.Conv2d(128, 64, kernel_size, stride, padding), - nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) - ]) - - def forward(self, x): - for conv in self.convs: - x = conv(x) - return x - - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - DistributedEnv.initialize(None) - in_channels = 64 - out_channels = 3 - - # Test both stride=1 and stride=2 cases - # stride=2 exercises the stride alignment and global-position cropping logic - test_configs = [ - (3, 1, 1), # kernel=3, stride=1, padding=1 (original test) - (3, 2, 1), # kernel=3, stride=2, padding=1 (downsampling with stride alignment) - ] - if args.height != 1024 or args.width != 1024: - test_sizes = [ - (args.height, args.width), - ] - else: - test_sizes = [ - # 1k - (1024, 1024), - (1023, 1025), - (1025, 1023), - # 720p - (720, 1280), - (721, 1281), - (719, 1279), - (1280, 720), - (1281, 721), - (1279, 719), - ] - - for kernel_size, stride, padding in test_configs: - for height, width in test_sizes: - if dist.get_rank() == 0: - print(f"\nTesting kernel={kernel_size}, stride={stride}, padding={padding}, size={height}x{width}", flush=True) - - convs = Conv2dModules(in_channels, out_channels, kernel_size, stride, padding).to(device) - patch_convs = nn.ModuleList() - for conv in convs.convs: - patch_convs.append(Conv2dAdapter(conv)) - patch_convs = patch_convs.to(device) - - hidden_state = torch.randn(1, 64, height, width, device=device) - result = convs(hidden_state) - - - if dist.get_rank() == 0: - print(kernel_size, stride, padding, "start", flush=True) - patch = Patchify() - depatch = DePatchify() - - patch_hidden_state = patch(hidden_state) - for conv in patch_convs: - patch_hidden_state = conv(patch_hidden_state) - ppresult = depatch(patch_hidden_state) - - - - if dist.get_rank() == 0: - print(f"result.shape={result.shape}, ppresult.shape={ppresult.shape}", flush=True) - diff = torch.abs(result - ppresult) - max_diff = diff.max().item() - mean_diff = diff.mean().item() - print(f"Max diff: {max_diff:.2e}, Mean diff: {mean_diff:.2e}", flush=True) - - # Use slightly relaxed tolerance for stride>1 to account for numerical precision - # differences from distributed computation order - tolerance = 1e-5 if stride > 1 else 1e-6 - if not torch.allclose(result, ppresult, atol=tolerance): - print("in kernel size: ", kernel_size, "stride: ", stride, "padding: ", padding, flush=True) - print(f"FAILED with tolerance {tolerance}\n", flush=True) - # Find where the largest differences are - max_diff_idx = torch.argmax(diff) - max_diff_idx = torch.unravel_index(max_diff_idx, diff.shape) - print(f"Largest diff at index {max_diff_idx}: ref={result[max_diff_idx].item():.6f}, patched={ppresult[max_diff_idx].item():.6f}", flush=True) - else: - print(f"{kernel_size} {stride} {padding} end (max_diff={max_diff:.2e}, tol={tolerance:.0e})", flush=True) - - # assert torch.equal(result, ppresult), "two hidden states are not equal" - - dist.barrier() - dist.destroy_process_group() +import torch.nn as nn + +from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter +from distvae.modules.patch_utils import DePatchify, Patchify + +from distributed_harness import ( + assert_matches_reference, + init_gloo, + make_parallel_context, + run_distributed, +) + + +def worker(rank, world_size, size, kernel, stride, padding, patch_dim, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + height, width = size + conv = nn.Conv2d(4, 3, kernel, stride=stride, padding=padding).eval() + x = torch.randn(1, 4, height, width) + + with torch.no_grad(): + expected = conv(x) if rank == 0 else None + context = make_parallel_context(patch_dim) + sharded = Conv2dAdapter(conv, parallel_context=context) + actual = DePatchify(context)(sharded(Patchify(context)(x))) + + assert_matches_reference(rank, actual, expected, "PatchConv2d", atol=1e-5) + finally: + dist.destroy_process_group() + + +# Odd against even, and extents whose remainder against the rank count differs between the two +# axes, so a run cannot pass by having every band the same size. +SIZES = [(32, 32), (33, 31), (31, 33), (45, 28)] + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 4]) +@pytest.mark.parametrize("size", SIZES) +def test_it_matches_conv2d_at_unit_stride(world_size, size, master_port, seed=42): + run_distributed(worker, world_size, (size, 3, 1, 1, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("size", SIZES) +def test_it_matches_conv2d_when_it_halves(world_size, size, master_port, seed=42): + """Stride 2, which is where the crop has to know where its band starts + + At unit stride every output row is an input row and the halo alone lines the bands up. A + A strided convolution uses a global output grid. If a band begins between grid positions, + cropping must start at the first global grid position inside the band rather than at the + band's first row. An even split does not exercise this arithmetic. + """ + run_distributed(worker, world_size, (size, 3, 2, 1, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("size", SIZES) +def test_it_matches_conv2d_when_the_width_is_split(world_size, size, master_port, seed=42): + # The same convolution against the other axis, which the layer supports and nothing above it + # used to check at anything but a square. + run_distributed(worker, world_size, (size, 3, 1, 1, -1, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("kernel,padding", [(1, 0), (5, 2), (7, 3)]) +def test_it_matches_conv2d_across_kernel_widths(world_size, kernel, padding, master_port, seed=42): + # The halo is kernel // 2 rows either side, so a wider kernel asks more of a neighbour than + # the thin bands an uneven split leaves have to spare. + run_distributed(worker, world_size, ((33, 31), kernel, 1, padding, -2, seed), master_port) + if __name__ == "__main__": - main() + parser = argparse.ArgumentParser(description="PatchConv2d GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_conv3d.py b/test/test_conv3d.py index 70da681..2a887de 100644 --- a/test/test_conv3d.py +++ b/test/test_conv3d.py @@ -3,55 +3,88 @@ import pytest import torch import torch.nn as nn -from unittest.mock import patch from distvae.models.layers.conv3d import PatchConv3d +from distributed_harness import make_parallel_context class TestPatchConv3dConstructor: """Tests for PatchConv3d constructor.""" - @pytest.mark.parametrize("patch_dim", [-3, -2, -1, 2, 3, 4]) - def test_valid_patch_dim(self, patch_dim): - module = PatchConv3d(4, 8, 3, patch_dim=patch_dim) - assert module.patch_dim == patch_dim + @pytest.mark.parametrize( + "patch_dim,expected", [(-2, -2), (3, -2), (-1, -1), (4, -1)] + ) + def test_valid_patch_dim(self, patch_dim, expected): + module = PatchConv3d( + 4, 8, 3, parallel_context=make_parallel_context(patch_dim) + ) + assert module.patch_dim == expected assert module.block_size == 0 + @pytest.mark.parametrize("patch_dim", [-3, 2]) + def test_frame_patch_dim_raises(self, patch_dim): + with pytest.raises(ValueError, match="frame axis"): + PatchConv3d( + 4, 8, 3, parallel_context=make_parallel_context(patch_dim) + ) + @pytest.mark.parametrize("patch_dim", [0, 1, 5]) def test_invalid_patch_dim_raises(self, patch_dim): - with pytest.raises(AssertionError) as exc_info: - PatchConv3d(4, 8, 3, patch_dim=patch_dim) - assert "F (-3 or 3) or H (-2 or 2) or W (-1 or 4)" in str(exc_info.value) + with pytest.raises(ValueError): + PatchConv3d( + 4, 8, 3, parallel_context=make_parallel_context(patch_dim) + ) def test_dilation_int_raises(self): with pytest.raises(AssertionError) as exc_info: - PatchConv3d(4, 8, 3, dilation=2) + PatchConv3d( + 4, 8, 3, dilation=2, parallel_context=make_parallel_context() + ) assert "dilation is not supported" in str(exc_info.value) def test_dilation_tuple_raises(self): with pytest.raises(AssertionError) as exc_info: - PatchConv3d(4, 8, 3, dilation=(1, 2, 1)) + PatchConv3d( + 4, + 8, + 3, + dilation=(1, 2, 1), + parallel_context=make_parallel_context(), + ) assert "dilation is not supported" in str(exc_info.value) def test_block_size_int(self): - module = PatchConv3d(4, 8, 3, block_size=0) + module = PatchConv3d( + 4, 8, 3, block_size=0, parallel_context=make_parallel_context() + ) assert module.block_size == 0 def test_block_size_tuple(self): - module = PatchConv3d(4, 8, 3, block_size=(2, 2, 2)) + module = PatchConv3d( + 4, + 8, + 3, + block_size=(2, 2, 2), + parallel_context=make_parallel_context(), + ) assert module.block_size == (2, 2, 2) class TestPatchConv3dSingleRankForward: """Single-rank forward: PatchConv3d matches nn.Conv3d when world size is 1.""" - @patch("distvae.models.layers.conv3d.get_world_size_and_rank") - def test_forward_matches_conv3d(self, mock_get_world_size_and_rank): - mock_get_world_size_and_rank.return_value = (1, 0, 0, 0) + def test_forward_matches_conv3d(self): in_ch, out_ch = 4, 8 k, s, p = 3, 1, 1 ref_conv = nn.Conv3d(in_ch, out_ch, k, stride=s, padding=p) - patch_conv = PatchConv3d(in_ch, out_ch, k, stride=s, padding=p) + patch_conv = PatchConv3d( + in_ch, + out_ch, + k, + stride=s, + padding=p, + parallel_context=make_parallel_context(), + ) with torch.no_grad(): patch_conv.weight.copy_(ref_conv.weight) patch_conv.bias.copy_(ref_conv.bias) @@ -63,11 +96,17 @@ def test_forward_matches_conv3d(self, mock_get_world_size_and_rank): patch_out = patch_conv(x) assert torch.allclose(patch_out, ref_out, atol=1e-5) - @patch("distvae.models.layers.conv3d.get_world_size_and_rank") - def test_forward_padding_mode_zeros(self, mock_get_world_size_and_rank): - mock_get_world_size_and_rank.return_value = (1, 0, 0, 0) + def test_forward_padding_mode_zeros(self): ref_conv = nn.Conv3d(4, 8, 3, stride=1, padding=1, padding_mode="zeros") - patch_conv = PatchConv3d(4, 8, 3, stride=1, padding=1, padding_mode="zeros") + patch_conv = PatchConv3d( + 4, + 8, + 3, + stride=1, + padding=1, + padding_mode="zeros", + parallel_context=make_parallel_context(), + ) with torch.no_grad(): patch_conv.weight.copy_(ref_conv.weight) patch_conv.bias.copy_(ref_conv.bias) @@ -79,20 +118,30 @@ def test_forward_padding_mode_zeros(self, mock_get_world_size_and_rank): class TestPatchConv3dOutputShape: """Single-rank output shape matches standard 3D conv formula.""" - @patch("distvae.models.layers.conv3d.get_world_size_and_rank") - def test_output_shape_k3_s1_p1(self, mock_get_world_size_and_rank): - mock_get_world_size_and_rank.return_value = (1, 0, 0, 0) + def test_output_shape_k3_s1_p1(self): # (F + 2*p - (k-1) - 1) / s + 1 = (4 + 2 - 2) / 1 + 1 = 5 per spatial dim - conv = PatchConv3d(4, 8, kernel_size=3, stride=1, padding=1) + conv = PatchConv3d( + 4, + 8, + kernel_size=3, + stride=1, + padding=1, + parallel_context=make_parallel_context(), + ) x = torch.randn(2, 4, 4, 8, 8) out = conv(x) assert out.shape == (2, 8, 4, 8, 8) - @patch("distvae.models.layers.conv3d.get_world_size_and_rank") - def test_output_shape_k3_s2_p0(self, mock_get_world_size_and_rank): - mock_get_world_size_and_rank.return_value = (1, 0, 0, 0) + def test_output_shape_k3_s2_p0(self): # Standard 3D conv: (L + 2*pad - (k-1) - 1) // stride + 1; L=5,9 k=3 s=2 p=0 -> 2, 4, 4 - conv = PatchConv3d(4, 8, kernel_size=3, stride=2, padding=0) + conv = PatchConv3d( + 4, + 8, + kernel_size=3, + stride=2, + padding=0, + parallel_context=make_parallel_context(), + ) x = torch.randn(1, 4, 5, 9, 9) out = conv(x) assert out.shape == (1, 8, 2, 4, 4) diff --git a/test/test_conv3d_distributed_gloo.py b/test/test_conv3d_distributed_gloo.py index 6fb03fc..d77ab64 100644 --- a/test/test_conv3d_distributed_gloo.py +++ b/test/test_conv3d_distributed_gloo.py @@ -15,12 +15,12 @@ import torch import torch.distributed as dist import torch.nn as nn -from torch.multiprocessing import spawn -from distvae.utils import DistributedEnv from distvae.modules.patch_utils import Patchify, DePatchify from distvae.modules.adapters.layers.conv_adapters import Conv3dAdapter +from distributed_harness import make_parallel_context, run_distributed + def worker( rank: int, @@ -30,6 +30,7 @@ def worker( stride: int, padding: int, block_size: int, + size: tuple, seed: int, master_port: int, ) -> None: @@ -39,7 +40,6 @@ def worker( os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) dist.init_process_group(backend="gloo", init_method="env://") - DistributedEnv.initialize(None) torch.manual_seed(seed) in_ch, out_ch = 4, 8 @@ -47,18 +47,12 @@ def worker( # For stride>1 tests, use sizes that stress-test alignment logic # For stride=1, use sizes divisible by world_size for even splitting n, c, f = 1, in_ch, 4 - if stride > 1: - # Use even sizes for stride>1 tests - # TODO: Add support for odd sizes with stride>1 (currently produces off-by-one errors) - h, w = 8, 8 - else: - # Use sizes divisible by world_size for even splitting - h, w = 8, 8 - if patch_dim == -2: - assert h % world_size == 0 - else: - assert patch_dim == -1 - assert w % world_size == 0 + # Both branches of the conditional this replaces set 8 by 8, so the comment about needing + # even sizes for stride > 1 described the only case there was, and the assertions below it + # enforced an even split that no shipped decode gets. The split axis is now given by the + # caller, so a test can ask for a size that leaves the ranks holding different amounts - + # which is the case the halo widths and the crop are actually difficult for. + h, w = size x_full = torch.randn(n, c, f, h, w, device=device, dtype=torch.float32) ref_conv = nn.Conv3d( @@ -66,9 +60,12 @@ def worker( ).to(device) ref_conv.eval() - patchify = Patchify(patch_dim=patch_dim) - depatchify = DePatchify(patch_dim=patch_dim) - adapter = Conv3dAdapter(ref_conv, block_size=block_size, patch_dim=patch_dim) + context = make_parallel_context(patch_dim) + patchify = Patchify(context) + depatchify = DePatchify(context) + adapter = Conv3dAdapter( + ref_conv, block_size=block_size, parallel_context=context + ) adapter.eval() with torch.no_grad(): @@ -98,31 +95,21 @@ def _run_one( block_size: int, seed: int, master_port: int, + size: tuple = (8, 8), ) -> None: """Spawn processes and run worker; raises on failure.""" - spawn( + # Through the shared harness rather than spawn directly, so a port claimed between being + # found free and being bound is retried rather than failing the test. + run_distributed( worker, - nprocs=world_size, - args=( - world_size, - patch_dim, - kernel_size, - stride, - padding, - block_size, - seed, - master_port, - ), - join=True, + world_size, + (patch_dim, kernel_size, stride, padding, block_size, size, seed), + master_port, ) -@pytest.fixture -def master_port(request): - """Unique port per test to avoid Address already in use when tests run sequentially.""" - base = 29500 - nodeid = request.node.nodeid - return base + (hash(nodeid) % 10000) +# The shared fixture uses CRC32 because Python salts string hashes per process. It also keeps all +# distributed tests in one port range, preventing separate fixtures from selecting the same port. @pytest.mark.gloo @@ -141,6 +128,32 @@ def test_patch_conv3d_gloo_direct(world_size, patch_dim, master_port, seed=42): ) +@pytest.mark.gloo +@pytest.mark.parametrize("world_size,patch_dim", [(4, -2), (4, -1), (3, -2)]) +def test_patch_conv3d_gloo_on_bands_of_different_sizes( + world_size, patch_dim, master_port, seed=42 +): + """The split axis not dividing by the rank count, which is what the 8 by 8 above never gives + + Every band being the same size is the easy case: the halo each rank asks of its neighbour is + the same, and the crop starts at the same offset into each. Nine rows over four ranks gives + 3, 2, 2, 2, and the sizes stop being interchangeable - a rank that assumes its neighbour + matches it reads the wrong rows, and a global quantity derived from a local one is wrong on + every rank but one. + """ + _run_one( + world_size=world_size, + patch_dim=patch_dim, + kernel_size=3, + stride=1, + padding=1, + block_size=0, + seed=seed, + master_port=master_port, + size=(9, 7), + ) + + @pytest.mark.gloo def test_patch_conv3d_gloo_chunked_path(master_port, seed=42): """PatchConv3d with GLOO: chunked path (block_size=4 so _use_direct_path is False, and chunks >= kernel_size=3).""" @@ -159,13 +172,11 @@ def test_patch_conv3d_gloo_chunked_path(master_port, seed=42): @pytest.mark.gloo @pytest.mark.parametrize("world_size,patch_dim", [(4, -2), (2, -1)]) def test_patch_conv3d_stride2_alignment(world_size, patch_dim, master_port, seed=42): - """ - PatchConv3d with stride=2: tests stride alignment and global-position cropping logic. + """PatchConv3d at stride 2, where the crop has to be placed from the global position - This exercises the code path where: - 1. Stride > 1 triggers stride alignment (shift calculation and input trimming) - 2. build_crop_slice uses global_start and global_height for correct output cropping - 3. Ranks would otherwise misalign without this logic + A strided convolution's output grid is set by where a rank's patch begins in the whole + image, not by where it begins in that rank, so build_crop_slice is given the global start + and the ranks would otherwise cut their outputs at offsets that do not join up. """ _run_one( world_size=world_size, @@ -179,6 +190,44 @@ def test_patch_conv3d_stride2_alignment(world_size, patch_dim, master_port, seed ) +@pytest.mark.gloo +@pytest.mark.parametrize("block_size", [0, 2, 4]) +@pytest.mark.parametrize("size", [(9, 7), (15, 11)]) +@pytest.mark.parametrize("world_size,patch_dim", [(4, -2), (3, -2), (2, -1)]) +def test_patch_conv3d_stride2_on_bands_of_different_sizes( + world_size, patch_dim, size, block_size, master_port, seed=42 +): + """Halving an uneven split, which the sizes elsewhere in this file were chosen to avoid + + A TODO used to sit beside those even sizes saying odd extents at stride > 1 produce + off-by-one errors, and the test was shaped around it rather than at it. Written first as a + non-strict xfail so a known bug would not be quietly forgotten, it passed on both of its + cases, so the claim is checked here instead of recorded: three splits, two shapes that + divide by none of the rank counts, and both the direct and the chunked convolution. + + A block of 2 against a kernel of 3 was tried first and cut chunks no convolution can run on. + So did a block of 4, once the frame axis was chunked as well: 4 frames padded to 6, cut in + two at stride 2, ends on a chunk of 2. That is chunk_bounds' to answer rather than each + caller's to avoid, and it now takes no more chunks than leave every one of them a kernel + long, so both blocks work and the small one is kept here. + + If the off-by-one is real it is not this. Should one of these ever fail, it is the arithmetic + that is wrong and not the expectation - a strided convolution over an uneven split has to + match nn.Conv3d, or a decode of any image whose rows do not divide by the rank count is wrong. + """ + _run_one( + world_size=world_size, + patch_dim=patch_dim, + kernel_size=3, + stride=2, + padding=1, + block_size=block_size, + seed=seed, + master_port=master_port, + size=size, + ) + + if __name__ == "__main__": parser = argparse.ArgumentParser(description="PatchConv3d GLOO multi-rank tests") parser.add_argument("--world_size", type=int, default=None) diff --git a/test/test_conv_utils.py b/test/test_conv_utils.py index 14e4595..78c05ed 100644 --- a/test/test_conv_utils.py +++ b/test/test_conv_utils.py @@ -2,13 +2,14 @@ import pytest import torch -from unittest.mock import patch from distvae.models.layers.conv_utils import ( calc_patch_index, calc_top_halo_width, calc_bottom_halo_width, calc_halo_width, + calc_halo_width_unit_stride, + chunk_bounds, correct_end, correct_start, build_crop_slice, @@ -99,33 +100,80 @@ def test_invalid_padding(self): class TestCalcHaloWidth: - """Tests for calc_halo_width.""" - - @patch("distvae.models.layers.conv_utils.DistributedEnv.get_group_world_size") - def test_first_rank_top_zero(self, mock_world_size): - mock_world_size.return_value = 3 - height_index = [0, 8, 16, 24] - top, bottom = calc_halo_width(0, height_index, 3, 0, 1) - assert top == 0 - assert bottom >= 0 - - @patch("distvae.models.layers.conv_utils.DistributedEnv.get_group_world_size") - def test_last_rank_bottom_zero(self, mock_world_size): - mock_world_size.return_value = 3 - height_index = [0, 8, 16, 24] - top, bottom = calc_halo_width(2, height_index, 3, 0, 1) - assert bottom == 0 - assert top >= 0 - - @patch("distvae.models.layers.conv_utils.DistributedEnv.get_group_world_size") - def test_middle_rank_both_nonzero(self, mock_world_size): - mock_world_size.return_value = 3 - height_index = [0, 8, 16, 24] - top, bottom = calc_halo_width(1, height_index, 3, 1, 1) - expected_top = calc_top_halo_width(1, height_index, 3, 1, 1) - expected_bottom = calc_bottom_halo_width(1, height_index, 3, 1, 1) - assert top == expected_top - assert bottom == expected_bottom + """Tests for calc_halo_width. + + Every expectation here is a number worked out by hand from the conv arithmetic. The + halo is how many rows a rank asks its neighbour for, so a wrong-but-non-negative + answer is exactly the bug worth catching: too few rows and the seam is wrong, too + many and the neighbour is asked for rows it does not have. + """ + + def test_first_rank_top_zero(self): + # k=3, p=0, s=1: the rank below reads one row back over the boundary at 8. + assert calc_halo_width(0, [0, 8, 16, 24], 3, 0, 1) == (0, 1) + + def test_last_rank_bottom_zero(self): + assert calc_halo_width(2, [0, 8, 16, 24], 3, 0, 1) == (1, 0) + + def test_middle_rank_both_nonzero(self): + assert calc_halo_width(1, [0, 8, 16, 24], 3, 1, 1) == (1, 1) + + def test_a_strided_middle_rank_reaches_further_one_way_than_the_other(self): + """The case the symmetric ones cannot tell apart + + At stride 1 the two halves of the halo come out equal, so top and bottom can be + swapped, or one computed twice, and every assertion above still holds. Striding + moves the output grid relative to the patch boundary and the two stop matching. + """ + # k=5, p=1, s=2 over even patches: one row above, two below. + assert calc_halo_width(1, [0, 8, 16, 24], 5, 1, 2) == (1, 2) + # k=3, p=0, s=2 over the uneven split: the lower boundary is an output-grid position, so + # a middle rank needs no rows below it. + assert calc_halo_width(1, [0, 9, 17, 24], 3, 0, 2) == (1, 0) + + +class TestCalcHaloWidthUnitStride: + """The stride-1 shortcut has to answer exactly what the gathered boundaries answer. + + It is what every unit-stride convolution uses in place of an all_gather, so if it ever + disagreed with calc_halo_width the ranks would exchange the wrong rows and the seam + between two patches would be quietly wrong rather than loudly broken. + """ + + @pytest.mark.parametrize("kernel_size", [1, 2, 3, 4, 5, 7]) + @pytest.mark.parametrize("padding", [0, 1, 2, 3]) + @pytest.mark.parametrize( + "patch_sizes", + [ + [8, 8], + [8, 8, 8, 8], + [9, 8, 8, 8], # the uneven split Patchify makes when rows do not divide by ranks + [3, 2, 2], # patches barely wider than the kernel + [64, 63, 63, 63], + ], + ) + def test_it_agrees_with_the_gathered_boundaries( + self, patch_sizes, padding, kernel_size + ): + world_size = len(patch_sizes) + if min(patch_sizes) < kernel_size: + # calc_bottom_halo_width asserts its way out of a patch narrower than the kernel + # reaches, so there is no gathered answer to agree with. DistVAE refuses that split + # in Patchify well before a convolution sees it. + pytest.skip("a patch narrower than the kernel is not a split DistVAE makes") + height_index = calc_patch_index([torch.tensor([s]) for s in patch_sizes]) + + for rank in range(world_size): + assert calc_halo_width_unit_stride(rank, world_size, kernel_size) == calc_halo_width( + rank, height_index, kernel_size, padding, 1 + ) + + def test_the_edge_ranks_have_nothing_beyond_them(self): + assert calc_halo_width_unit_stride(0, 4, 3)[0] == 0 + assert calc_halo_width_unit_stride(3, 4, 3)[1] == 0 + + def test_a_lone_rank_needs_no_halo_at_all(self): + assert calc_halo_width_unit_stride(0, 1, 7) == (0, 0) class TestCorrectEnd: @@ -159,6 +207,35 @@ def test_unaligned(self): assert correct_start(3, 2) == 4 +class TestChunkBounds: + """The chunked convolution path cuts every axis with this""" + + @pytest.mark.parametrize("stride", [1, 2]) + @pytest.mark.parametrize("kernel_size", [1, 3, 5]) + @pytest.mark.parametrize("block", [2, 4, 8, 64]) + @pytest.mark.parametrize("extent", [4, 6, 7, 10, 17, 64]) + def test_no_chunk_is_shorter_than_the_kernel(self, extent, block, kernel_size, stride): + """The one property the convolution cannot survive being without + + A chunk shorter than the kernel raises out of torch, so this is not an accuracy question + that a later assertion would catch: it is whether the call can be made at all. Asked over + blocks below the kernel and axes that divide by none of them, which is where the path was + cutting a two-long tail off a six-long frame axis. + """ + if extent < kernel_size: + pytest.skip("an axis shorter than the kernel has no chunking to get right") + for start, end in chunk_bounds(extent, block, kernel_size, stride): + assert end - start >= kernel_size, f"{extent}/{block} k{kernel_size} s{stride}" + + def test_the_chunks_cover_the_axis_and_overlap_by_what_the_kernel_reads(self): + # Eight long, cut in two, kernel 3 at unit stride: the first chunk runs on to the last + # input its final output reads, so the two overlap by the kernel less one. + assert chunk_bounds(8, 4, 3, 1) == [(0, 6), (4, 8)] + + def test_an_axis_that_wants_no_cutting_is_one_chunk(self): + assert chunk_bounds(8, 64, 3, 1) == [(0, 8)] + + class TestBuildCropSlice: """Tests for build_crop_slice (pure).""" diff --git a/test/test_decoderadapter.py b/test/test_decoderadapter.py new file mode 100644 index 0000000..d60a622 --- /dev/null +++ b/test/test_decoderadapter.py @@ -0,0 +1,153 @@ +"""DecoderAdapter against the decoder it shards, over gloo on CPU. + +It is the adapter every AutoencoderKL model decodes through, including xDiT's SD3 and Z-Image. + +Run from repo root: + pytest test/test_decoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") + +CONFIG = dict( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=256, + down_block_types=["DownEncoderBlock2D"] * 4, + up_block_types=["UpDecoderBlock2D"] * 4, +) +LATENT_CHANNELS = 4 + + +def build_decoder(): + return diffusers.AutoencoderKL(**CONFIG).eval().decoder + + +def worker( + rank, + world_size, + height, + width, + conv_block_size, + training, + checkpointing, + seed, + master_port, +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder() + decoder.train(training) + decoder.gradient_checkpointing = checkpointing + runtime_sentinel = object() + decoder.runtime_sentinel = runtime_sentinel + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder() + reference.load_state_dict(weights) + reference.train(training) + reference.gradient_checkpointing = checkpointing + expected = reference(latents) + + adapter = DecoderAdapter( + decoder, vae_group=None, conv_block_size=conv_block_size + ) + assert adapter.decoder is decoder + assert adapter.decoder.runtime_sentinel is runtime_sentinel + assert adapter.training is decoder.training + assert adapter.decoder.training is decoder.training + assert ( + adapter.decoder.gradient_checkpointing + is decoder.gradient_checkpointing + ) + child_contexts = [ + module.parallel_context + for module in adapter.modules() + if hasattr(module, "parallel_context") + ] + assert child_contexts + assert all( + context is adapter.parallel_context for context in child_contexts + ) + actual = adapter(latents) + + # The sharded GroupNorm sums its statistics across ranks in float32 before dividing, so + # its result differs slightly from a single-rank reduction over the same values. + assert_matches_reference(rank, actual, expected, "DecoderAdapter", atol=1e-4) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_sharded_decode_matches_unsharded_decode(world_size, master_port, seed=42): + run_distributed(worker, world_size, (16, 16, 0, False, False, seed), master_port) + + +@pytest.mark.gloo +def test_training_checkpoint_state_survives_adaptation(master_port, seed=42): + run_distributed(worker, 1, (16, 16, 0, True, True, seed), master_port) + + +def grad_enabled_worker(rank, world_size, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder().eval() + decoder.gradient_checkpointing = True + adapter = DecoderAdapter(decoder) + with pytest.raises( + RuntimeError, + match=r"torch\.no_grad.*inference mode", + ): + adapter(torch.randn(1, LATENT_CHANNELS, 16, 16)) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_grad_enabled_forward_directs_callers_to_inference_mode(master_port, seed=42): + run_distributed(grad_enabled_worker, 1, (seed,), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_decodes_the_same(master_port, seed=42): + # A conv_block_size under the feature map size sends PatchConv2d down its chunked path, + # which splits and reassembles each convolution on top of the sharding. + run_distributed(worker, 2, (16, 16, 32, False, False, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # This adapter was never exposed to the pad-and-crop the causal ones used, because + # DecoderAdapter splits after its mid block rather than before. Pinned so it stays that way. + run_distributed(worker, 3, (16, 16, 0, False, False, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="DecoderAdapter GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_distributed_env.py b/test/test_distributed_env.py new file mode 100644 index 0000000..7e66d50 --- /dev/null +++ b/test/test_distributed_env.py @@ -0,0 +1,27 @@ +"""Which distributed backend DistVAE asks for, on machines with and without an accelerator.""" + +import torch + +from distvae.utils import DistributedEnv + + +def test_cpu_only_machines_get_gloo(monkeypatch): + # Sharding is correctness-testable on CPU, so a machine without an accelerator has to be + # offered a backend rather than refused one. + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + if hasattr(torch, "musa"): + monkeypatch.setattr(torch.musa, "is_available", lambda: False) + assert DistributedEnv.get_torch_distributed_backend() == "gloo" + + +def test_cuda_is_still_preferred_where_it_exists(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + assert DistributedEnv.get_torch_distributed_backend() == "nccl" + + +def test_the_device_type_agrees_with_the_backend(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + if hasattr(torch, "musa"): + monkeypatch.setattr(torch.musa, "is_available", lambda: False) + assert DistributedEnv.get_device_type() == "cpu" + assert DistributedEnv.get_device() == torch.device("cpu") diff --git a/test/test_distributed_harness.py b/test/test_distributed_harness.py new file mode 100644 index 0000000..bc460c5 --- /dev/null +++ b/test/test_distributed_harness.py @@ -0,0 +1,66 @@ +"""The comparison the adapter tests rest on, checked against a mismatch it has to catch. + +Every family test reports success by reaching the end of assert_matches_reference, so a bug that +made it accept anything would turn the whole suite green and mean nothing. +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + + +def mismatch_worker(rank, world_size, master_port): + init_gloo(rank, world_size, master_port) + try: + actual = torch.zeros(2, 3) + expected = torch.ones(2, 3) if rank == 0 else None + assert_matches_reference(rank, actual, expected, "a deliberately wrong result") + finally: + dist.destroy_process_group() + + +def shape_mismatch_worker(rank, world_size, master_port): + init_gloo(rank, world_size, master_port) + try: + actual = torch.zeros(2, 3) + expected = torch.zeros(2, 4) if rank == 0 else None + assert_matches_reference(rank, actual, expected, "a deliberately wrong shape") + finally: + dist.destroy_process_group() + + +def agreement_worker(rank, world_size, master_port): + init_gloo(rank, world_size, master_port) + try: + actual = torch.full((2, 3), 0.5) + expected = torch.full((2, 3), 0.5) if rank == 0 else None + assert_matches_reference(rank, actual, expected, "matching results") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("worker", [mismatch_worker, shape_mismatch_worker]) +def test_a_wrong_result_fails_every_rank(worker, master_port): + # Rank 0 is the only one holding a reference, so the failure has to travel: a rank that + # returned instead would sit in the next collective and hang the run rather than fail it. + with pytest.raises(Exception) as caught: + run_distributed(worker, 2, (), master_port) + assert "did not match the single-rank reference" in str(caught.value) + + +@pytest.mark.gloo +def test_a_matching_result_passes(master_port): + run_distributed(agreement_worker, 2, (), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Distributed test harness self-checks") + _, remainder = parser.parse_known_args() + sys.exit(pytest.main([os.path.abspath(__file__), "-v"] + remainder)) diff --git a/test/test_distvae_bench.py b/test/test_distvae_bench.py new file mode 100644 index 0000000..047f154 --- /dev/null +++ b/test/test_distvae_bench.py @@ -0,0 +1,1872 @@ +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from bench.harness import ( + catalog, + cases, + cli, + distributed, + measure, + profile, + report, + shape_costs, +) + + +def test_describe_only_runs_on_cpu_without_distributed_environment( + tmp_path, monkeypatch +): + for name in ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr( + measure.torch.cuda, + "set_device", + lambda *args: pytest.fail("describe-only touched CUDA"), + ) + output = tmp_path / "description.json" + + status = cli.main( + ["--describe-only", "--family", "kl", "--half", "decoder", "--out", str(output)] + ) + + assert status == 0 + record = json.loads(output.read_text()) + assert record["schema_version"] == report.SCHEMA_VERSION + assert record["measurement"]["description"]["adapter"] == "DecoderAdapter" + assert record["composition"]["execution"] == "describe-only" + assert record["runtime"] == {"dtype": "bfloat16", "world_size": 1} + + +def test_catalog_samples_decoder_and_encoder_on_meta(): + spec = catalog.FAMILIES["kl"] + latent = catalog.sample_for(spec, "decoder", 512, 256, "float32", "meta") + image = catalog.sample_for(spec, "encoder", 512, 256, "float32", "meta") + assert tuple(latent.shape) == (1, 16, 64, 32) + assert tuple(image.shape) == (1, 3, 512, 256) + + +def test_exact_cases_do_not_form_a_cartesian_product(): + args = cli.parser().parse_args( + [ + "--case", + "unsharded", + "--case", + "local:256x512@64x32", + "--case", + "tile-runs:384x256@32x16", + ] + ) + + cells = cases.cells_from_args(args) + + assert [cell["name"] for cell in cells] == [ + "unsharded", + "local-256x512-ov64x32", + "tile-runs-384x256-ov32x16", + ] + assert cells[1]["window"] == (256, 512) + assert cells[2]["tile_distribution"] == "runs" + + +def test_default_suite_is_deferred_until_vae_and_world_size_are_known(): + args = cli.parser().parse_args([]) + + assert cases.cells_from_args(args) == [] + + +def test_additional_shapes_are_explicit_and_do_not_mix_with_exact_cases(): + args = cli.parser().parse_args( + ["--shape", "720x1280x81", "--shape", "1080x1920x81"] + ) + + assert cases.shapes_from_args(args) == [ + (720, 1280, 81), + (1080, 1920, 81), + ] + + mixed = cli.parser().parse_args( + ["--shape", "512x512", "--case", "unsharded"] + ) + with pytest.raises(ValueError, match="cannot be combined"): + cases.cells_from_args(mixed) + + +def test_matrix_runs_the_family_shapes_and_yields_to_an_explicit_one(): + """The matrix is a default, not an override: asking for a shape by hand still wins. + + Appending instead would make `--shape` mean "and also", so a one-off check of a single size + would quietly drag the whole family's matrix along with it. + """ + matrix = cli.parser().parse_args(["--family", "wan", "--matrix"]) + assert cases.shapes_from_args(matrix) == [(832, 480, 81), (1280, 720, 81)] + + overridden = cli.parser().parse_args( + ["--family", "wan", "--matrix", "--shape", "512x512x5"] + ) + assert cases.shapes_from_args(overridden) == [(512, 512, 5)] + + single = cli.parser().parse_args(["--family", "wan", "--height", "256"]) + assert cases.shapes_from_args(single) == [(256, 2048, 17)] + + +@pytest.mark.parametrize("family", sorted(catalog.FAMILIES)) +def test_every_catalogued_shape_is_legal_for_its_own_family(family): + """A matrix runs unattended, so an illegal shape has to fail before anything is measured. + + Both bounds come from the family rather than from the shape: an axis has to divide by the + spatial ratio, and a temporal family needs one frame plus a multiple of its ratio. Left to + `sample_for` these surface partway through the third shape, after the first two have been + paid for. + """ + spec = catalog.FAMILIES[family] + if not spec.get("shapes"): + pytest.skip(f"{family} has no canonical shapes") + + for height, width, frames in catalog.matrix_for(family): + assert height % spec["spatial"] == 0 + assert width % spec["spatial"] == 0 + if spec["temporal"]: + assert (frames - 1) % spec["temporal"] == 0 + catalog.sample_for( + spec, "decoder", height, width, "bfloat16", "meta", frames=frames + ) + + +def test_matrix_refuses_a_family_it_has_no_shapes_for(monkeypatch): + # Reached with a family the catalog does not carry shapes for, which is what a newly added + # one looks like before its matrix is chosen. It used to be reached with LTX-2, until LTX-2 + # was given a matrix of its own. + monkeypatch.setitem(catalog.FAMILIES, "shapeless", {"cls": "AutoencoderKL"}) + with pytest.raises(ValueError, match="no canonical shapes"): + catalog.matrix_for("shapeless") + + +def test_every_catalogued_family_carries_a_matrix(): + assert not [ + family for family, spec in catalog.FAMILIES.items() if not spec.get("shapes") + ] + + +@pytest.mark.parametrize( + "value", + ["none", "row:256x256@32x32", "local:256@32x32", "local:256x256"], +) +def test_case_parser_rejects_legacy_or_incomplete_syntax(value): + with pytest.raises(ValueError): + cases.parse_case(value, 512, 256, 1) + + +def test_selector_returns_three_distinct_rectangular_pareto_plans(): + plans = cases.select_plans( + sample_shape=(1024, 2048), + native_overlap=(64, 64), + world_size=4, + normalize=lambda window, overlap: (window, overlap), + ) + + assert [plan["profile"] for plan in plans] == [ + "coarse", + "balanced", + "fine", + ] + assert len({plan["window"] for plan in plans}) == 3 + assert any(height != width for height, width in (p["window"] for p in plans)) + assert all(plan["selection"]["pareto_optimal"] for plan in plans) + assert all(plan["objectives"]["tile_count"] <= 16 for plan in plans) + + +def test_topology_objectives_price_clipped_tiles_and_scheduler_loads(): + objectives = cases.topology_objectives( + window=(72, 72), + overlap=(8, 8), + sample_shape=(128, 128), + world_size=2, + ) + + assert objectives["tile_grid"] == (2, 2) + assert objectives["decoded_area"] == (72 + 64) ** 2 + assert objectives["rank_imbalance"] == pytest.approx(32 / 9248) + + +def test_selector_zeros_overlap_on_inactive_strip_axis(): + plans = cases.select_plans( + sample_shape=(512, 2048), + native_overlap=(64, 96), + world_size=2, + normalize=lambda window, overlap: (window, overlap), + ) + + strips = [ + plan + for plan in plans + if plan["window"][0] >= 512 or plan["window"][1] >= 2048 + ] + assert strips + for plan in strips: + if plan["window"][0] >= 512: + assert plan["overlap"][0] == 0 + if plan["window"][1] >= 2048: + assert plan["overlap"][1] == 0 + + +def test_selector_searches_overlap_and_can_beat_row_sharding(): + """A plan is only a memory win when its window is smaller than a row shard. + + Since `window = pitch + overlap`, fixing overlap at the VAE native value imposes a lower + bound on every window. On this sample the smallest reachable window is 512x512, which equals + the 262144 pixels assigned to one row-sharded rank. Searching smaller overlaps is therefore + required to propose a configuration that reduces memory. + """ + sample_shape, world_size, native = (1024, 1024), 4, (256, 256) + plans = cases.select_plans( + sample_shape=sample_shape, + native_overlap=native, + world_size=world_size, + normalize=lambda window, overlap: (window, overlap), + ) + + row_shard_area = (sample_shape[0] // world_size) * sample_shape[1] + fine = next(plan for plan in plans if plan["profile"] == "fine") + assert fine["objectives"]["window_area"] < row_shard_area + assert fine["objectives"]["beats_row_sharding"] + # The pinned-overlap search could not get below the native value on an active axis. + assert min(fine["overlap"]) < min(native) + + +def test_overlap_ladder_scales_with_pitch_and_keeps_the_native_value(): + # An inactive axis still blends nothing, which the strip cases rely on. + assert cases._overlap_options(1024, 1, 256) == (0,) + + options = cases._overlap_options(1024, 4, 256) + assert 256 in options, "the native overlap must stay reachable for comparability" + assert options == tuple(sorted(options, reverse=True)), "widest first" + assert all(option > 0 for option in options) + # Pitch is 256 here. The ladder stops at a third of the pitch, which is a quarter of the + # window it blends, so halves and thirds survive and the thinner rungs that band are gone. + assert {128, 86} <= set(options) + assert min(options) * 3 >= 256 + + +def test_selector_keeps_every_blend_above_a_quarter_of_its_window(): + """Keep overlap at least one quarter of the normalized window. + + Measured on FLUX.2 at 1024x1024 on four ranks: a 128px window blended 32px is clean, the + same window blended 16px shows banding, and the difference is concentrated at tile + boundaries. Check the bound after window normalization because normalization may enlarge the + window without changing the overlap. + """ + + def grow(window, overlap): + return tuple(-(-axis // 64) * 64 for axis in window), overlap + + plans = cases.select_plans( + sample_shape=(1024, 1024), + native_overlap=(256, 256), + world_size=4, + normalize=grow, + ) + + blends = [ + (blend, size) + for plan in plans + for blend, size in zip(plan["overlap"], plan["window"]) + if blend + ] + assert blends, "an all-strip selection would not exercise the bound" + for blend, size in blends: + assert blend * 4 >= size, f"{blend}px blends a {size}px window" + + +def test_selector_declines_a_fine_profile_that_is_only_a_transpose(): + """Require the fine profile to have a smaller window area than the coarse profile. + + Transposed windows have equal modeled area, work, and imbalance. Measurements showed that a + full-height transpose used 17% more memory than the selected full-width strip. + """ + plans = cases.select_plans( + sample_shape=(1024, 1024), + native_overlap=(256, 256), + world_size=4, + normalize=lambda window, overlap: (window, overlap), + ) + + by_profile = {plan["profile"]: plan for plan in plans} + coarse = by_profile["coarse"] + fine = by_profile.get("fine") + if fine is not None: + assert (fine["objectives"]["window_area"] + < coarse["objectives"]["window_area"]) + assert tuple(reversed(fine["window"])) != coarse["window"] + assert len({plan["window"] for plan in plans}) == len(plans) + + +def test_profiles_bracket_the_tile_axis_rather_than_predicting_a_winner(): + """Coarse is the fewest tiles and fine the most, so the suite spans the axis it is testing. + + On gfx1201, the configurations with the least modeled work were slower and used more memory + than row sharding: 5034 MB versus 3526 MB at 2048x2048 on four ranks. Performance may differ + by device, so profile names describe geometry rather than predicted outcomes. + """ + for sample_shape, world_size in (((1024, 1024), 2), ((2048, 2048), 4)): + plans = cases.select_plans( + sample_shape=sample_shape, + native_overlap=(256, 256), + world_size=world_size, + normalize=lambda window, overlap: (window, overlap), + ) + by_profile = {plan["profile"]: plan for plan in plans} + coarse, fine = by_profile["coarse"], by_profile["fine"] + + assert coarse["objectives"]["tile_count"] == min( + plan["objectives"]["tile_count"] for plan in plans + ), f"{sample_shape} ws={world_size}: coarse must be the fewest tiles" + assert fine["objectives"]["tile_count"] == max( + plan["objectives"]["tile_count"] for plan in plans + ), f"{sample_shape} ws={world_size}: fine must be the most tiles" + + +def test_tile_columns_separate_a_plan_from_its_transpose(): + wide = cases.topology_objectives((128, 1024), (32, 0), (1024, 1024), 4) + tall = cases.topology_objectives((1024, 128), (0, 32), (1024, 1024), 4) + + assert wide["window_area"] == tall["window_area"], "the transpose is the point" + assert wide["tile_columns"] == 1 + assert tall["tile_columns"] > 1 + # Equal on every symmetric objective, so only tile_columns can prefer the cheaper one. + assert cases._dominates(wide, tall) + assert not cases._dominates(tall, wide) + + +def test_row_shard_area_is_recorded_against_every_plan(): + objectives = cases.topology_objectives( + window=(72, 72), + overlap=(8, 8), + sample_shape=(128, 128), + world_size=2, + ) + + assert objectives["row_shard_area"] == 64 * 128 + assert objectives["beats_row_sharding"] is (72 * 72 < 64 * 128) + + +def test_vae_normalizer_rejects_windows_with_too_few_latent_rows(monkeypatch): + vae = object() + monkeypatch.setattr(cases.vae_api, "tile_shape", lambda value: (64, 64)) + monkeypatch.setattr( + cases.vae_api, + "tile_shape_plan", + lambda value, height, width: {"window": (height, width)}, + ) + monkeypatch.setattr(cases, "_latent_shape", lambda value, plan: (3, 3)) + monkeypatch.setattr( + cases.vae_api, + "tile_overlap_plan", + lambda *args, **kwargs: pytest.fail("invalid row window planned overlap"), + ) + + normalize = cases.normalizer_for_vae(vae, (512, 512), world_size=4) + + assert normalize((256, 256), (32, 32)) is None + + +def test_vae_normalizer_rejects_windows_that_band(monkeypatch): + """A tile large enough to shard can still be too small to normalize over. + + Sharding needs one latent row per rank; representative statistics need considerably more. + Searching overlap made small windows reachable for the first time, so this bound is what + stops the memory profile choosing a tile that decodes at a visibly different tone from its + neighbours - a difference the blend smooths into a ramp, which no seam metric detects. + """ + vae = object() + extent = cases.MIN_TILE_LATENT_EXTENT - 1 + assert extent > 4, "the bound must bind harder than the world sizes we run" + monkeypatch.setattr(cases.vae_api, "tile_shape", lambda value: (64, 64)) + monkeypatch.setattr( + cases.vae_api, + "tile_shape_plan", + lambda value, height, width: {"window": (height, width)}, + ) + monkeypatch.setattr(cases, "_latent_shape", lambda value, plan: (extent, extent)) + monkeypatch.setattr( + cases.vae_api, + "tile_overlap_plan", + lambda *args, **kwargs: pytest.fail("a banding window reached overlap planning"), + ) + + normalize = cases.normalizer_for_vae(vae, (512, 512), world_size=4) + + assert normalize((256, 256), (32, 32)) is None + + +def test_vae_normalizer_checks_the_shorter_latent_axis(monkeypatch): + rectangle = { + "tile_sample_min_size": 64, + "tile_sample_min_height": 256, + "tile_sample_min_width": 64, + "tile_latent_min_size": 8, + "tile_latent_min_height": 32, + "tile_latent_min_width": 8, + } + vae = SimpleNamespace(**rectangle) + monkeypatch.setattr(cases.vae_api, "tile_shape", lambda value: (256, 256)) + monkeypatch.setattr( + cases.vae_api, "tile_shape_plan", lambda *args, **kwargs: rectangle + ) + monkeypatch.setattr( + cases.vae_api, + "tile_overlap_plan", + lambda *args, **kwargs: pytest.fail("a narrow window reached overlap planning"), + ) + + normalize = cases.normalizer_for_vae(vae, (512, 512), world_size=4) + + assert normalize((256, 64), (32, 8)) is None + + +def test_tile_latent_area_prefers_keyed_rectangle_over_scalar_threshold(): + vae = SimpleNamespace( + tile_latent_min_size=8, + tile_latent_min_height=32, + tile_latent_min_width=8, + ) + + assert measure._tile_latent_area(vae) == 256 + + +def _bounded_plans(): + return cases.select_plans( + sample_shape=(1024, 2048), + native_overlap=(64, 64), + world_size=4, + normalize=lambda window, overlap: (window, overlap), + ) + + +def test_default_suite_carries_only_selectable_compositions(): + """Local tiling and row-beneath-tiling are not reachable, so they are not the default. + + An orchestrator either marks a VAE for tile parallelism or parallelizes its decoder; it does + not combine those modes. The excluded cases also account for about 60% of the suite's + compute without representing a supported deployment configuration. + """ + plans = _bounded_plans() + + suite = cases.default_suite(plans, 1024, 2048, 1) + + assert [cell["name"] for cell in suite[:2]] == ["unsharded", "row"] + assert len(suite) == 2 + len(plans) + assert sum(cell["tile_distribution"] == "runs" for cell in suite) == len(plans) + assert not [cell for cell in suite if cell["mode"] in ("local", "row-tiled")] + + +def test_diagnostics_restore_the_unreachable_compositions(): + plans = _bounded_plans() + + suite = cases.default_suite(plans, 1024, 2048, 1, diagnostics=True) + + assert len(suite) == 2 + 2 * len(plans) + 1 + assert [cell["name"] for cell in suite[:2]] == ["unsharded", "row"] + assert sum(cell["mode"] == "local" for cell in suite) == len(plans) + lightest = min(plans, key=lambda plan: plan["objectives"]["window_area"]) + assert [ + cell["profile"] + for cell in suite + if cell["sharding"] == "row" and cell["window"] is not None + ] == [lightest["profile"]] + + +def test_encoder_baseline_suite_has_no_decode_only_tiling(): + suite = cases.baseline_suite(720, 1280, 81) + + assert [cell["name"] for cell in suite] == ["unsharded", "row"] + assert all(cell["window"] is None for cell in suite) + + +def test_parser_exposes_tile_shape_cost_controls(): + args = cli.parser().parse_args( + [ + "--tile-shape-costs", + "--tile-shape-batch", + "4", + "--tile-shape-windows", + "8x16,16x32", + ] + ) + + assert args.tile_shape_costs is True + assert args.tile_shape_batch == 4 + assert args.tile_shape_windows == "8x16,16x32" + + +def test_tile_shape_cost_mode_bypasses_ordinary_cell_normalization(monkeypatch): + runtime = SimpleNamespace(rank=0, world_size=1, group=object()) + runtime.close = lambda: None + monkeypatch.setattr( + cases, + "cells_from_args", + lambda args: pytest.fail("shape-cost mode normalized ordinary cells"), + ) + monkeypatch.setattr(cli.Runtime, "start", lambda timeout: runtime) + monkeypatch.setattr(cli, "_measure", lambda *args, **kwargs: []) + monkeypatch.setattr(report, "report_status", lambda records: 0) + monkeypatch.setattr( + cli.dist, + "all_gather_object", + lambda values, value, **kwargs: values.__setitem__(0, value), + ) + + assert cli.main(["--tile-shape-costs"]) == 0 + + +def test_tile_shape_cost_mode_rejects_describe_only(): + with pytest.raises(SystemExit): + cli.main(["--tile-shape-costs", "--describe-only"]) + + +def test_invocation_provenance_is_collected_once_and_reused(monkeypatch): + calls = [] + provenance_data = {"versions": {}, "provenance": {"recorded_at": "once"}} + monkeypatch.setattr( + report, + "provenance", + lambda: calls.append("provenance") or provenance_data, + ) + monkeypatch.setattr( + cli, + "_describe", + lambda args, cells, shared: [ + {"measurement": {}, **shared}, + {"measurement": {}, **shared}, + ], + ) + monkeypatch.setattr(report, "render", lambda *args: None) + + assert ( + cli.main( + ["--describe-only", "--case", "unsharded", "--case", "row"] + ) + == 0 + ) + assert calls == ["provenance"] + + +def test_provenance_records_explicit_hardware_family(monkeypatch): + monkeypatch.setenv("HW_FAMILY", "mi355") + + assert report.provenance()["provenance"]["hardware_family"] == "mi355" + + +def test_provenance_measures_the_device_rather_than_trusting_the_label(monkeypatch): + """HW_FAMILY is whatever the caller typed; the device is what the run actually used. + + For a long time the label was the only hardware field there was, and since nothing set it + every report said null - so two machines' numbers were separable only by hostname. gcnArchName + is the part that distinguishes AMD generations, where the marketing name repeats across them. + """ + monkeypatch.delenv("HW_FAMILY", raising=False) + monkeypatch.setattr(report.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(report.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(report.torch.cuda, "device_count", lambda: 4) + monkeypatch.setattr( + report.torch.cuda, + "get_device_properties", + lambda index: SimpleNamespace( + name="AMD Radeon Graphics", gcnArchName="gfx1201", total_memory=34342961152 + ), + ) + + recorded = report.provenance()["provenance"] + + assert recorded["hardware_family"] is None + assert recorded["device"] == { + "name": "AMD Radeon Graphics", + "arch": "gfx1201", + "total_memory": 34342961152, + "count": 4, + } + + +def test_provenance_survives_a_run_with_no_accelerator(monkeypatch): + monkeypatch.setattr(report.torch.cuda, "is_available", lambda: False) + + assert report.provenance()["provenance"]["device"] is None + + +def test_rank_error_helpers_preserve_original_rank_and_type(monkeypatch): + peer = {"type": "ValueError", "message": "peer", "rank": 1} + runtime = SimpleNamespace(rank=0, world_size=2, group=object()) + monkeypatch.setattr( + distributed.dist, + "all_gather_object", + lambda values, value, **kwargs: values.__setitem__(slice(None), [value, peer]), + ) + + failures = distributed.gather_rank_errors( + distributed.exception_record(RuntimeError("local"), runtime.rank), runtime + ) + aggregate = distributed.aggregate_rank_errors(failures) + + assert aggregate["type"] == "RuntimeError" + assert aggregate["rank"] == 0 + assert aggregate["failed_ranks"] == [0, 1] + assert aggregate["failures"] == [ + {"type": "RuntimeError", "message": "local", "rank": 0}, + peer, + ] + + +def test_parser_exposes_harness_owned_profiler_controls(tmp_path): + args = cli.parser().parse_args( + [ + "--profile", + "--profile-trace", + "--profile-memory", + "--profile-dir", + str(tmp_path), + ] + ) + + assert args.profile is True + assert args.profile_trace is True + assert args.profile_memory is True + assert args.profile_dir == str(tmp_path) + + +def test_disabled_profiler_has_no_runtime_overhead(monkeypatch): + monkeypatch.setattr( + profile.torch.profiler, + "profile", + lambda **kwargs: pytest.fail("disabled profiling touched torch.profiler"), + ) + args = SimpleNamespace( + profile=False, + profile_trace=False, + profile_memory=False, + ) + + assert ( + profile.profile_once(lambda: pytest.fail("disabled profiling ran"), args) + is None + ) + + +def test_profile_without_exports_returns_a_bounded_summary(monkeypatch): + table_calls = [] + + class Averages: + def table(self, **options): + table_calls.append(options) + return "x" * (profile.PROFILE_SUMMARY_LIMIT + 100) + + class FakeProfile: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def key_averages(self): + return Averages() + + monkeypatch.setattr( + profile.torch.profiler, + "ProfilerActivity", + SimpleNamespace(CPU="cpu", CUDA="cuda"), + ) + monkeypatch.setattr( + profile.torch.profiler, "profile", lambda **kwargs: FakeProfile() + ) + args = SimpleNamespace( + profile=True, + profile_trace=False, + profile_memory=False, + profile_dir="unused", + family="kl", + half="encoder", + ) + + result = profile.profile_once( + lambda: object(), + args, + cell={"name": "single", "height": 256, "width": 128, "frames": 1}, + runtime=SimpleNamespace(rank=0, device=SimpleNamespace(type="cuda")), + ) + + assert result["artifacts"] == {} + assert len(result["summary"]) == profile.PROFILE_SUMMARY_LIMIT + assert table_calls == [{"sort_by": "self_cuda_time_total", "row_limit": 20}] + + +def test_profiler_setup_failure_on_a_peer_cleans_up_before_run(monkeypatch): + events = [] + history = [] + peer_failure = { + "type": "RuntimeError", + "message": "profiler enter failed", + "rank": 1, + } + + class FakeProfile: + def __enter__(self): + events.append("enter") + return self + + def __exit__(self, *args): + events.append("exit") + return False + + def gather(values, value, **kwargs): + values[:] = [value, peer_failure] + + monkeypatch.setattr( + profile.torch.profiler, + "ProfilerActivity", + SimpleNamespace(CPU="cpu", CUDA="cuda"), + ) + monkeypatch.setattr( + profile.torch.cuda.memory, + "_record_memory_history", + lambda enabled=None: history.append(enabled), + ) + monkeypatch.setattr(profile.torch.profiler, "profile", lambda **kwargs: FakeProfile()) + monkeypatch.setattr(profile.torch.distributed, "all_gather_object", gather) + args = SimpleNamespace( + profile=True, + profile_trace=False, + profile_memory=True, + profile_dir="unused", + family="kl", + half="decoder", + ) + runtime = SimpleNamespace( + rank=0, + world_size=2, + group=object(), + device=SimpleNamespace(type="cuda"), + ) + + with pytest.raises(distributed.RankError) as caught: + profile.profile_once( + lambda: events.append("run"), + args, + cell={"name": "single", "height": 16, "width": 16, "frames": 1}, + runtime=runtime, + ) + + assert caught.value.rank_error["rank"] == 1 + assert caught.value.rank_error["type"] == "RuntimeError" + assert caught.value.rank_error["failures"] == [peer_failure] + assert events == ["enter", "exit"] + assert history == ["all", None] + + +@pytest.mark.parametrize("failure_step", ["directory", "memory", "enter"]) +def test_profiler_local_setup_failure_is_synchronized_with_peers( + monkeypatch, failure_step +): + history = [] + events = [] + + class FakeProfile: + def __enter__(self): + events.append("enter") + if failure_step == "enter": + raise ValueError("context setup") + return self + + def __exit__(self, *args): + events.append("exit") + return False + + def record_memory(enabled=None): + history.append(enabled) + if failure_step == "memory" and enabled == "all": + raise RuntimeError("memory setup") + + def gather(values, value, **kwargs): + values[:] = [value, None] + + monkeypatch.setattr( + profile.torch.profiler, + "ProfilerActivity", + SimpleNamespace(CPU="cpu", CUDA="cuda"), + ) + monkeypatch.setattr( + profile.torch.cuda.memory, "_record_memory_history", record_memory + ) + monkeypatch.setattr(profile.torch.profiler, "profile", lambda **kwargs: FakeProfile()) + monkeypatch.setattr(profile.torch.distributed, "all_gather_object", gather) + if failure_step == "directory": + monkeypatch.setattr( + profile.Path, + "mkdir", + lambda *args, **kwargs: (_ for _ in ()).throw( + OSError("artifact directory setup") + ), + ) + args = SimpleNamespace( + profile=True, + profile_trace=False, + profile_memory=True, + profile_dir="unused", + family="kl", + half="decoder", + ) + runtime = SimpleNamespace( + rank=0, + world_size=2, + group=object(), + device=SimpleNamespace(type="cuda"), + ) + + with pytest.raises(distributed.RankError) as caught: + profile.profile_once( + lambda: pytest.fail("run began after setup failure"), + args, + cell={"name": "single", "height": 16, "width": 16, "frames": 1}, + runtime=runtime, + ) + + expected_type = { + "directory": "OSError", + "memory": "RuntimeError", + "enter": "ValueError", + }[failure_step] + assert caught.value.rank_error["rank"] == 0 + assert caught.value.rank_error["type"] == expected_type + assert caught.value.rank_error["failed_ranks"] == [0] + if failure_step == "enter": + assert history == ["all", None] + assert events == ["enter"] + + +def test_profiler_exports_harness_named_trace_and_memory_artifacts( + tmp_path, monkeypatch +): + exports = {} + history = [] + + class FakeProfile: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def export_chrome_trace(self, path): + exports["trace"] = Path(path) + + def export_memory_timeline(self, path): + exports["memory"] = Path(path) + + def key_averages(self): + return SimpleNamespace(table=lambda **kwargs: "cuda summary") + + monkeypatch.setattr( + profile.torch.profiler, + "ProfilerActivity", + SimpleNamespace(CPU="cpu", CUDA="cuda"), + ) + monkeypatch.setattr( + profile.torch.cuda.memory, + "_record_memory_history", + lambda enabled=None: history.append(enabled), + ) + monkeypatch.setattr( + profile.importlib, + "import_module", + lambda name: pytest.fail(f"CUDA profiling imported {name}"), + ) + monkeypatch.setattr( + profile.torch.profiler, + "profile", + lambda **kwargs: exports.update(options=kwargs) or FakeProfile(), + ) + args = SimpleNamespace( + profile=True, + profile_trace=True, + profile_memory=True, + profile_dir=str(tmp_path), + family="wan", + half="decoder", + ) + runtime = SimpleNamespace(rank=2, device=SimpleNamespace(type="cuda")) + result = profile.profile_once( + lambda: object(), + args, + cell={"name": "tile-half", "height": 512, "width": 256, "frames": 17}, + runtime=runtime, + ) + + assert result == { + "summary": "cuda summary", + "artifacts": { + "trace": str( + tmp_path / "wan-decoder-tile-half-512x256x17-rank2.trace.json" + ), + "memory": str( + tmp_path / "wan-decoder-tile-half-512x256x17-rank2.memory.html" + ), + }, + } + assert exports["trace"] == Path(result["artifacts"]["trace"]) + assert exports["memory"] == Path(result["artifacts"]["memory"]) + assert exports["options"]["activities"] == ["cpu", "cuda"] + assert exports["options"]["profile_memory"] is True + assert exports["options"]["record_shapes"] is True + assert exports["options"]["with_stack"] is True + assert history == ["all", None] + + +def test_profiler_avoids_overwriting_existing_artifacts(tmp_path): + existing = tmp_path / "kl-decoder-single-16x16x1-rank0.trace.json" + existing.write_text("old") + + stem = profile._artifact_stem( + tmp_path, "kl-decoder-single-16x16x1-rank0", ["trace.json"] + ) + + assert stem == "kl-decoder-single-16x16x1-rank0-2" + + +def test_musa_profiler_is_loaded_lazily_and_uses_musa_memory_history( + tmp_path, monkeypatch +): + exports = {} + history = [] + imported = [] + + class FakeProfile: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def key_averages(self): + return SimpleNamespace(table=lambda **kwargs: "musa summary") + + def export_memory_timeline(self, path): + exports["memory"] = Path(path) + + musa = SimpleNamespace( + memory=SimpleNamespace( + _record_memory_history=lambda enabled=None: history.append(enabled) + ) + ) + monkeypatch.setattr(profile.torch, "musa", musa, raising=False) + monkeypatch.setattr( + profile.importlib, + "import_module", + lambda name: imported.append(name) or object(), + ) + monkeypatch.setattr( + profile.torch.profiler, + "ProfilerActivity", + SimpleNamespace(CPU="cpu", MUSA="musa"), + ) + monkeypatch.setattr( + profile.torch.profiler, + "profile", + lambda **kwargs: exports.update(options=kwargs) or FakeProfile(), + ) + args = SimpleNamespace( + profile=True, + profile_trace=False, + profile_memory=True, + profile_dir=str(tmp_path), + family="wan", + half="decoder", + ) + + result = profile.profile_once( + lambda: object(), + args, + cell={"name": "single", "height": 64, "width": 64, "frames": 5}, + runtime=SimpleNamespace(rank=1, device=SimpleNamespace(type="musa")), + ) + + assert imported == ["torch_musa"] + assert exports["options"]["activities"] == ["cpu", "musa"] + assert exports["memory"] == Path(result["artifacts"]["memory"]) + assert history == ["all", None] + assert result["summary"] == "musa summary" + + +def test_runtime_selects_cuda_without_importing_musa(monkeypatch): + monkeypatch.setattr(distributed.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + distributed.importlib, + "import_module", + lambda name: pytest.fail(f"CUDA runtime imported {name}"), + ) + + name, api, backend = distributed.accelerator_backend() + + assert (name, api, backend) == ("cuda", distributed.torch.cuda, "nccl") + + +def test_runtime_loads_musa_lazily_and_selects_mccl(monkeypatch): + imported = [] + musa = SimpleNamespace(is_available=lambda: True) + monkeypatch.setattr(distributed.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(distributed.torch, "musa", musa, raising=False) + monkeypatch.setattr( + distributed.importlib, + "import_module", + lambda name: imported.append(name) or object(), + ) + + name, api, backend = distributed.accelerator_backend() + + assert imported == ["torch_musa"] + assert (name, api, backend) == ("musa", musa, "mccl") + + +def test_profile_summary_is_embedded_in_measurement(monkeypatch): + sample = measure.torch.zeros(1, 4, 2, 2) + profile = {"summary": "bounded profiler table", "artifacts": {}} + execution_order = [] + monkeypatch.setattr(measure.catalog, "build_vae", lambda *args: object()) + monkeypatch.setattr(measure.catalog, "sample_for", lambda *args: sample) + monkeypatch.setattr( + measure.catalog, "describe_vae", lambda *args: {"adapter": "Adapter"} + ) + monkeypatch.setattr(measure.catalog, "run_half", lambda *args: sample) + monkeypatch.setattr(measure, "configure_sharding", lambda *args: "Adapter") + monkeypatch.setattr(measure, "configure_tiling", lambda *args: {"enabled": False}) + monkeypatch.setattr( + measure.profile, + "profile_once", + lambda *args: execution_order.append("profile") or profile, + ) + monkeypatch.setattr(measure, "across_ranks", lambda *args: {}) + monkeypatch.setattr( + measure, + "timed", + lambda *args: execution_order.append("timed") or {"median_s": 0.0}, + ) + monkeypatch.setattr(measure.torch.cuda, "synchronize", lambda *args: None) + monkeypatch.setattr( + measure.torch.cuda, "reset_peak_memory_stats", lambda *args: None + ) + monkeypatch.setattr(measure.torch.cuda, "max_memory_allocated", lambda *args: 0) + + class Log: + enabled = False + by_call = {} + + def reset(self): + pass + + def report(self): + return {} + + args = SimpleNamespace( + family="kl", + half="decoder", + dtype="float32", + batch=1, + skip_reference=True, + reference_max_latent_elems=0, + phase_timing=False, + profile=True, + profile_trace=False, + profile_memory=False, + warmup=0, + iters=1, + max_rel=None, + ) + cell = { + "name": "single", + "height": 16, + "width": 16, + "frames": 1, + "sharding": "unsharded", + } + runtime = SimpleNamespace( + device=SimpleNamespace(type="cuda"), + device_api=measure.torch.cuda, + rank=0, + world_size=1, + group=object(), + log=Log(), + ) + + _, measurement = measure.measure_cell( + args, + {"spatial": 8, "temporal": None}, + cell, + runtime, + {}, + lambda *args: None, + ) + + assert measurement["profile"] == profile + assert execution_order == ["timed", "profile"] + + +def test_tile_shape_costs_measure_latency_memory_and_batch_scaling(monkeypatch): + vae = object() + calls = [] + monkeypatch.setattr(shape_costs.catalog, "build_vae", lambda *args: vae) + monkeypatch.setattr(shape_costs.vae_api, "tile_shape", lambda value: (64, 64)) + monkeypatch.setattr( + shape_costs.catalog, + "run_half", + lambda value, half, sample: calls.append(tuple(sample.shape)) or sample, + ) + + monkeypatch.setattr(shape_costs.dist, "barrier", lambda *args, **kwargs: None) + monkeypatch.setattr( + shape_costs.dist, + "all_gather_object", + lambda values, value, **kwargs: values.__setitem__(0, value), + ) + monkeypatch.setattr(shape_costs.torch.cuda, "synchronize", lambda *args: None) + monkeypatch.setattr( + shape_costs.torch.cuda, "reset_peak_memory_stats", lambda *args: None + ) + monkeypatch.setattr( + shape_costs.torch.cuda, + "max_memory_allocated", + lambda *args: 10 * 1024 * 1024, + ) + monkeypatch.setattr(shape_costs.torch.cuda, "empty_cache", lambda: None) + args = SimpleNamespace( + family="kl", + dtype="float32", + frames=1, + iters=1, + tile_shape_batch=2, + tile_shape_windows="8x4,4x8", + warmup=0, + ) + spec = {"latent_channels": 16, "spatial": 8, "temporal": None} + + result = shape_costs.tile_shape_costs( + args, + spec, + SimpleNamespace( + device="cpu", + device_api=shape_costs.torch.cuda, + group=object(), + rank=0, + world_size=1, + ), + lambda *parts: None, + ) + + assert [entry["tiles_in_the_call"] for entry in result["shapes"]] == [1, 2, 1, 2] + assert all(entry["median_ms"] >= 0 for entry in result["shapes"]) + assert all(entry["peak_vram_mb"] == 10.0 for entry in result["shapes"]) + assert result["analysis"]["highest_area_cost"]["rows"] in {4, 8} + assert result["analysis"]["worst_batch_scaling"]["tiles_in_the_call"] == 2 + assert result["latent_window"] == (8, 8) + assert result["frames"] is None + assert calls == [(1, 16, 8, 4), (2, 16, 8, 4), (1, 16, 4, 8), (2, 16, 4, 8)] + + +def test_default_shape_costs_reuse_bounded_rectangular_plans(monkeypatch): + vae = object() + monkeypatch.setattr(shape_costs.catalog, "build_vae", lambda *args: vae) + monkeypatch.setattr(shape_costs.vae_api, "tile_shape", lambda value: (64, 32)) + selected = [ + {"window": (64, 48)}, + {"window": (48, 64)}, + {"window": (32, 32)}, + ] + monkeypatch.setattr( + shape_costs.cases, + "plans_for_vae", + lambda value, height, width, world_size: selected, + ) + monkeypatch.setattr( + shape_costs.dist, + "all_gather_object", + lambda values, value, **kwargs: values.__setitem__(0, value), + ) + monkeypatch.setattr(shape_costs.torch, "randn", lambda *args, **kwargs: object()) + monkeypatch.setattr( + shape_costs, "_shape_iterations", lambda *args: ([0.001], None) + ) + device_api = SimpleNamespace( + reset_peak_memory_stats=lambda *args: None, + max_memory_allocated=lambda *args: 0, + empty_cache=lambda: None, + ) + args = SimpleNamespace( + family="kl", + dtype="float32", + frames=1, + iters=1, + tile_shape_batch=1, + tile_shape_windows="", + height=512, + width=1024, + warmup=0, + ) + + result = shape_costs.tile_shape_costs( + args, + {"latent_channels": 16, "spatial": 8, "temporal": None}, + SimpleNamespace( + device="cpu", + device_api=device_api, + group=object(), + rank=0, + world_size=1, + ), + lambda *parts: None, + ) + + assert result["latent_window"] == (8, 4) + assert [(entry["rows"], entry["columns"]) for entry in result["shapes"]] == [ + (8, 6), + (6, 8), + (4, 4), + ] + + +@pytest.mark.parametrize("failure_phase", ["allocation", "decode"]) +def test_tile_shape_oom_is_synchronized_before_the_next_case( + monkeypatch, failure_phase +): + vae = object() + monkeypatch.setattr(shape_costs.catalog, "build_vae", lambda *args: vae) + monkeypatch.setattr(shape_costs.vae_api, "tile_shape", lambda value: (64, 64)) + monkeypatch.setattr(shape_costs.dist, "barrier", lambda *args, **kwargs: None) + gathered = [] + + def gather(values, value, **kwargs): + gathered.append(value) + values[:] = [value, None] + + monkeypatch.setattr(shape_costs.dist, "all_gather_object", gather) + monkeypatch.setattr(shape_costs.torch.cuda, "synchronize", lambda *args: None) + monkeypatch.setattr( + shape_costs.torch.cuda, "reset_peak_memory_stats", lambda *args: None + ) + monkeypatch.setattr(shape_costs.torch.cuda, "empty_cache", lambda: None) + if failure_phase == "allocation": + monkeypatch.setattr( + shape_costs.torch, + "randn", + lambda *args, **kwargs: (_ for _ in ()).throw( + shape_costs.torch.OutOfMemoryError("allocation") + ), + ) + monkeypatch.setattr( + shape_costs.catalog, + "run_half", + lambda *args: pytest.fail("decode ran after allocation failed"), + ) + else: + monkeypatch.setattr( + shape_costs.catalog, + "run_half", + lambda *args: (_ for _ in ()).throw( + shape_costs.torch.OutOfMemoryError("decode") + ), + ) + args = SimpleNamespace( + family="kl", + dtype="float32", + frames=17, + iters=1, + tile_shape_batch=1, + tile_shape_windows="8x8", + warmup=1, + ) + + result = shape_costs.tile_shape_costs( + args, + {"latent_channels": 16, "spatial": 8, "temporal": None}, + SimpleNamespace( + device="cpu", + device_api=shape_costs.torch.cuda, + group=object(), + rank=0, + world_size=2, + ), + lambda *parts: None, + ) + + assert result["shapes"][0]["out_of_memory"] is True + assert result["shapes"][0]["failed_ranks"] == [0] + assert any( + failure and failure["type"] == "OutOfMemoryError" for failure in gathered + ) + + +@pytest.mark.parametrize("failure_phase", ["allocation", "decode"]) +def test_tile_shape_mixed_rank_failure_is_not_treated_as_oom( + monkeypatch, failure_phase +): + monkeypatch.setattr(shape_costs.catalog, "build_vae", lambda *args: object()) + monkeypatch.setattr(shape_costs.vae_api, "tile_shape", lambda value: (64, 64)) + monkeypatch.setattr(shape_costs.dist, "barrier", lambda *args, **kwargs: None) + peer_failure = {"type": "RuntimeError", "message": "fatal peer", "rank": 1} + + def gather(values, value, **kwargs): + values[:] = [value, peer_failure if value is not None else None] + + monkeypatch.setattr(shape_costs.dist, "all_gather_object", gather) + device_api = SimpleNamespace( + synchronize=lambda *args: None, + reset_peak_memory_stats=lambda *args: None, + empty_cache=lambda: None, + ) + if failure_phase == "allocation": + monkeypatch.setattr( + shape_costs.torch, + "randn", + lambda *args, **kwargs: (_ for _ in ()).throw( + shape_costs.torch.OutOfMemoryError("local oom") + ), + ) + else: + monkeypatch.setattr(shape_costs.torch, "randn", lambda *args, **kwargs: object()) + monkeypatch.setattr( + shape_costs.catalog, + "run_half", + lambda *args: (_ for _ in ()).throw( + shape_costs.torch.OutOfMemoryError("local oom") + ), + ) + args = SimpleNamespace( + family="kl", + dtype="float32", + frames=1, + iters=1, + tile_shape_batch=1, + tile_shape_windows="8x8", + warmup=1, + ) + + with pytest.raises(distributed.RankError) as caught: + shape_costs.tile_shape_costs( + args, + {"latent_channels": 16, "spatial": 8, "temporal": None}, + SimpleNamespace( + device="cpu", + device_api=device_api, + group=object(), + rank=0, + world_size=2, + ), + lambda *parts: None, + ) + + assert caught.value.rank_error["failed_ranks"] == [0, 1] + assert [failure["type"] for failure in caught.value.rank_error["failures"]] == [ + "OutOfMemoryError", + "RuntimeError", + ] + + +def test_tile_shape_setup_failure_is_synchronized_before_cases(monkeypatch): + monkeypatch.setattr(shape_costs.catalog, "build_vae", lambda *args: object()) + monkeypatch.setattr(shape_costs.vae_api, "tile_shape", lambda value: (64, 64)) + peer_failure = {"type": "RuntimeError", "message": "setup failed", "rank": 1} + gathered = [] + + def gather(values, value, **kwargs): + gathered.append(value) + values[:] = [value, peer_failure] + + monkeypatch.setattr(shape_costs.dist, "all_gather_object", gather) + monkeypatch.setattr( + shape_costs.torch, + "randn", + lambda *args, **kwargs: pytest.fail("case allocation began before setup vote"), + ) + args = SimpleNamespace( + family="kl", + dtype="float32", + frames=1, + iters=1, + tile_shape_batch=1, + tile_shape_windows="8x8", + warmup=0, + ) + + with pytest.raises(RuntimeError, match="setup failed on rank 1"): + shape_costs.tile_shape_costs( + args, + {"latent_channels": 16, "spatial": 8, "temporal": None}, + SimpleNamespace( + device="cpu", + device_api=shape_costs.torch.cuda, + group=object(), + rank=0, + world_size=2, + ), + lambda *parts: None, + ) + + assert gathered == [None] + + +@pytest.mark.parametrize("is_shape_costs", [False, True]) +def test_measurement_records_effective_dtype_and_world_size( + monkeypatch, is_shape_costs +): + args = SimpleNamespace( + family="kl", + half="decoder", + dtype="float16", + frames=1, + tile_shape_costs=is_shape_costs, + ) + runtime = SimpleNamespace( + rank=0, world_size=3, group=object(), device_api=measure.torch.cuda + ) + cell = { + "name": "single", + "height": 512, + "width": 512, + "frames": 1, + } + monkeypatch.setattr(report, "render", lambda *args: None) + monkeypatch.setattr(cli.dist, "all_gather_object", lambda *args, **kwargs: None) + if is_shape_costs: + monkeypatch.setattr( + shape_costs, + "tile_shape_costs", + lambda *args: { + "latent_window": 64, + "frames": None, + "analysis": {}, + "shapes": [], + }, + ) + else: + monkeypatch.setattr( + measure, + "measure_cell", + lambda *args: ({"sharding": "row"}, {"timing": {}}), + ) + + [record] = cli._measure(args, [cell], runtime) + + assert record["runtime"] == {"dtype": "float16", "world_size": 3} + if is_shape_costs: + assert record["shape"]["frames"] is None + assert record["measurement"]["tile_shape_costs"]["frames"] is None + + +@pytest.mark.parametrize("is_shape_costs", [False, True]) +def test_distributed_cell_errors_preserve_per_rank_details( + monkeypatch, capsys, is_shape_costs +): + args = SimpleNamespace( + family="kl", + half="decoder", + dtype="float32", + frames=1, + tile_shape_costs=is_shape_costs, + ) + runtime = SimpleNamespace( + rank=0, + world_size=2, + group=object(), + device_api=SimpleNamespace(empty_cache=lambda: None), + ) + cell = { + "name": "failing-cell", + "height": 512, + "width": 512, + "frames": 1, + } + peer_error = {"type": "ValueError", "message": "peer failure", "rank": 1} + + def gather(values, value, **kwargs): + values[:] = [value, peer_error] + + monkeypatch.setattr(cli.dist, "all_gather_object", gather) + if is_shape_costs: + monkeypatch.setattr( + shape_costs, + "tile_shape_costs", + lambda *args: (_ for _ in ()).throw(RuntimeError("local failure")), + ) + else: + monkeypatch.setattr( + measure, + "measure_cell", + lambda *args: (_ for _ in ()).throw(RuntimeError("local failure")), + ) + + [record] = cli._measure(args, [cell], runtime) + + assert record["error"]["rank"] == 0 + assert record["error"]["failed_ranks"] == [0, 1] + assert record["error"]["failures"] == [ + {"type": "RuntimeError", "message": "local failure", "rank": 0}, + peer_error, + ] + assert report.report_status([record]) == 1 + assert "RuntimeError: local failure" in capsys.readouterr().out + + +def test_tiled_numerical_disagreement_keeps_raw_verdict_without_enforcement(): + agreement = measure.agreement_with( + measure.torch.tensor([2.0]), + measure.torch.tensor([1.0]), + "float32", + max_rel=0.1, + ) + + assert agreement["disagreement_type"] == "numerical" + assert agreement["ok"] is False + assert "enforced" not in agreement + report.set_agreement_policy(agreement, tiling_enabled=True) + assert agreement["enforced"] is False + assert report.report_status([{"measurement": {"agreement": agreement}}]) == 0 + + +def test_tiled_shape_mismatch_is_enforced(): + agreement = measure.agreement_with( + measure.torch.zeros(1, 2), + measure.torch.zeros(1, 3), + "float32", + max_rel=None, + ) + + assert agreement["disagreement_type"] == "shape" + assert agreement["ok"] is False + assert "enforced" not in agreement + report.set_agreement_policy(agreement, tiling_enabled=True) + assert agreement["enforced"] is True + assert report.report_status([{"measurement": {"agreement": agreement}}]) == 1 + + +def test_enforced_agreement_and_execution_errors_fail(): + mismatch = {"measurement": {"agreement": {"ok": False, "enforced": True}}} + assert report.report_status([mismatch]) == 1 + assert report.report_status([{}, mismatch]) == 1 + assert ( + report.report_status([{"error": {"type": "RuntimeError", "message": "boom"}}]) + == 1 + ) + assert report.report_status([{}, {"error": {"type": "RuntimeError"}}]) == 1 + + +def test_rectangular_tile_shape_and_overlap_use_exact_distvae_plans(monkeypatch): + class Vae: + tile_sample_min_size = 512 + overlap = (128, 128) + + def enable_tiling(self): + pass + + vae = Vae() + calls = [] + + monkeypatch.setattr(measure.vae_api, "require_vae_support", lambda *args: None) + monkeypatch.setattr( + measure.vae_api, + "tile_shape", + lambda value: (value.tile_sample_min_size,) * 2, + ) + monkeypatch.setattr(measure.vae_api, "tile_overlap", lambda value: value.overlap) + monkeypatch.setattr(measure, "latent_rows", lambda value, plan=None: 32) + monkeypatch.setattr( + measure.vae_api, + "tile_shape_plan", + lambda value, height, width: ( + calls.append(("tile_shape_plan", (height, width))) + or {"tile_sample_min_size": height} + ), + ) + monkeypatch.setattr( + measure.vae_api, + "tile_overlap_plan", + lambda value, height, width, sample_shape=None: ( + calls.append(("tile_overlap_plan", (height, width), sample_shape)) + or {"tile_sample_stride_height": 192} + ), + ) + monkeypatch.setattr( + measure.vae_api, + "tiled_decode_for", + lambda value: calls.append(("tiled_decode_for",)) or (lambda sample: sample), + ) + + def apply(value, plan): + calls.append(("apply_tile_plan", dict(plan))) + for name, setting in plan.items(): + setattr(value, name, setting) + if "tile_sample_min_size" in plan: + value.overlap = (64, 64) + if "tile_sample_stride_height" in plan: + value.overlap = (64, 32) + + monkeypatch.setattr(measure.vae_api, "apply_tile_plan", apply) + cell = { + "sharding": "unsharded", + "window": (256, 384), + "height": 2048, + "width": 2048, + "overlap": (64, 32), + "tile_distribution": None, + } + + facts = measure.configure_tiling( + vae, + cell, + SimpleNamespace(world_size=2, group=object()), + "decoder", + lambda *parts: None, + ) + + assert calls == [ + ("tile_shape_plan", (256, 384)), + ("apply_tile_plan", {"tile_sample_min_size": 256}), + ("tile_overlap_plan", (64, 32), (2048, 2048)), + ("apply_tile_plan", {"tile_sample_stride_height": 192}), + ("tiled_decode_for",), + ] + assert facts["native_window_px"] == (512, 512) + assert facts["requested_window_px"] == (256, 384) + assert facts["window_px"] == (256, 384) + assert facts["native_overlap_px"] == (128, 128) + assert facts["overlap"] == (64, 32) + assert "default_overlap" not in facts + assert "narrowest_useful_window_px" not in facts + assert "below_useful_floor" not in facts + + +def test_an_invalid_exact_tile_shape_is_not_silently_snapped(monkeypatch): + class Vae: + def enable_tiling(self): + pass + + monkeypatch.setattr(measure.vae_api, "require_vae_support", lambda *args: None) + monkeypatch.setattr(measure.vae_api, "tile_shape", lambda value: (512, 512)) + monkeypatch.setattr(measure.vae_api, "tile_overlap", lambda value: (128, 128)) + monkeypatch.setattr(measure.vae_api, "tile_shape_plan", lambda *args: None) + + with pytest.raises(ValueError, match=r"tile shape \(255, 257\) is invalid for Vae"): + measure.configure_tiling( + Vae(), + { + "sharding": "unsharded", + "window": (255, 257), + "height": 2048, + "width": 2048, + "overlap": None, + "tile_distribution": None, + }, + SimpleNamespace(world_size=1, group=object()), + "decoder", + lambda *parts: None, + ) + + +def test_custom_overlap_installs_the_per_axis_replacement(monkeypatch): + class Vae: + def enable_tiling(self): + pass + + vae = Vae() + applied = [] + monkeypatch.setattr(measure.vae_api, "require_vae_support", lambda *args: None) + monkeypatch.setattr(measure.vae_api, "tile_shape", lambda value: (512, 512)) + monkeypatch.setattr(measure.vae_api, "tile_overlap", lambda value: (64, 32)) + monkeypatch.setattr(measure, "latent_rows", lambda value, plan=None: 64) + monkeypatch.setattr( + measure.vae_api, + "tile_shape_plan", + lambda value, height, width: {"window": (height, width)}, + ) + monkeypatch.setattr( + measure.vae_api, + "tile_overlap_plan", + lambda value, height, width, sample_shape=None: { + "overlap": (height, width), + "sample_shape": sample_shape, + }, + ) + monkeypatch.setattr( + measure.vae_api, + "apply_tile_plan", + lambda value, plan: applied.append(plan), + ) + replacement = object() + monkeypatch.setattr( + measure.vae_api, "tiled_decode_for", lambda value: replacement + ) + + measure.configure_tiling( + vae, + { + "sharding": "unsharded", + "window": (512, 512), + "height": 2048, + "width": 2048, + "overlap": (64, 32), + "tile_distribution": None, + }, + SimpleNamespace(world_size=1, group=object()), + "decoder", + lambda *parts: None, + ) + + assert applied == [ + {"window": (512, 512)}, + {"overlap": (64, 32), "sample_shape": (2048, 2048)} + ] + assert vae.tiled_decode is replacement + + +def test_report_schema_contains_provenance_and_effective_composition(): + record = report.make_record( + family="kl", + half="decoder", + shape={"height": 512, "width": 512, "frames": 1}, + composition={ + "sharding": "row", + "window": (512, 384), + "overlap": (64, 32), + "tile_distribution": None, + }, + measurement={"timing": {"median_s": 1.0}}, + dtype="float32", + world_size=4, + ) + + # Spelled out rather than derived, so that bumping the schema is a deliberate act with a + # test to edit, instead of something a record can start reporting on its own. + assert report.SCHEMA_VERSION == 7 + assert record["schema_version"] == 7 + assert set(record["versions"]) >= {"torch", "diffusers", "distvae"} + assert "distvae_git_revision" in record["provenance"] + assert record["composition"]["sharding"] == "row" + assert record["measurement"]["timing"]["median_s"] == 1.0 + assert record["runtime"] == {"dtype": "float32", "world_size": 4} + assert any( + path.endswith("bench/harness/measure.py") + for path in record["provenance"]["benchmark"]["implementation"] + ) + + +def test_measured_record_with_description_renders_metrics(capsys): + record = { + "composition": {"execution": "measurement"}, + "measurement": { + "description": {"adapter": "DecoderAdapter"}, + "collectives": { + "by_call": {}, + "by_call_max": {}, + }, + "timing": {"median_s": 0.125}, + "peak_vram_mb": 64, + "agreement": None, + }, + } + + report.render(record, "decoder") + + output = capsys.readouterr().out + assert "median 125.0 ms" in output + assert '"adapter"' not in output + + +def test_a_failure_on_every_rank_leaves_the_group_able_to_continue(): + failures = [ + {"type": "OutOfMemoryError", "message": "no", "rank": rank} for rank in range(4) + ] + + assert not distributed.ranks_diverged(failures) + aggregated = distributed.aggregate_rank_errors(failures) + assert aggregated["failed_ranks"] == [0, 1, 2, 3] + + +def test_a_failure_on_some_ranks_only_is_reported_as_divergence(): + failures = [{"type": "OutOfMemoryError", "message": "no", "rank": 0}, None] + + assert distributed.ranks_diverged(failures) + + +def test_a_crossed_gather_is_recorded_rather_than_raised(): + # When ranks call different collectives, rank 1 is still inside another all_gather_object, so + # its payload arrives instead of a failure record. + failures = [{"type": "OutOfMemoryError", "message": "no", "rank": 0}, [None, None]] + + assert distributed.ranks_diverged(failures) + aggregated = distributed.aggregate_rank_errors(failures) + assert aggregated["failed_ranks"] == [0, 1] + assert any( + failure["type"] == distributed.DESYNCHRONIZED + for failure in aggregated["failures"] + ) + + +def test_no_failure_anywhere_is_not_divergence(): + assert not distributed.ranks_diverged([None, None, None, None]) + assert distributed.aggregate_rank_errors([None, None]) is None + + +def test_overlap_grows_with_a_normalized_window(): + # A quarter of 86 is 22, but 22 is only 17% of the normalized 128px window. + assert cases.blend_for_window((22, 22), (128, 128)) == (32, 32) + # Already a quarter or wider, so left exactly as it is. + assert cases.blend_for_window((32, 64), (128, 128)) == (32, 64) + # An inactive axis blends nothing and stays that way. + assert cases.blend_for_window((0, 22), (128, 128)) == (0, 32) + + +def test_vae_normalizer_returns_overlap_for_the_normalized_window(monkeypatch): + """A coarse window increment requires overlap to grow with the normalized window. + + LTX-2 quantizes windows to 256px. Without overlap adjustment, every candidate at its CI shape + has overlap below one quarter of the normalized window and select_plans rejects all of them. + """ + planned = [] + # This mutable stub allows the normalizer to apply the planned window. + vae = SimpleNamespace() + monkeypatch.setattr(cases.vae_api, "tile_shape", lambda value: (128, 128)) + # Only multiples of 128 are tileable, so a 200px request snaps to 256. + monkeypatch.setattr( + cases.vae_api, + "tile_shape_plan", + lambda value, height, width: ( + {"window": (height, width)} if height % 128 == 0 and width % 128 == 0 else None + ), + ) + monkeypatch.setattr(cases, "_latent_shape", lambda value, plan: (32, 32)) + monkeypatch.setattr( + cases.vae_api, + "tile_overlap_plan", + lambda value, *overlap, **kwargs: planned.append(overlap) or {"overlap": overlap}, + ) + + normalize = cases.normalizer_for_vae(vae, (1024, 1024), world_size=4) + window, overlap = normalize((200, 200), (50, 50)) + + assert window == (256, 256) + assert overlap == (64, 64), "overlap must follow the normalized window" + assert planned == [(64, 64)], "the adjusted overlap must be planned" + + +def test_a_selected_plan_never_blends_thinner_than_a_quarter(): + plans = cases.select_plans( + sample_shape=(1088, 1920), + native_overlap=(64, 64), + world_size=8, + normalize=lambda window, overlap: ( + # A coarse quantum, as LTX-2 has: windows snap up to the next multiple of 256. + (-(-window[0] // 256) * 256, -(-window[1] // 256) * 256), + overlap, + ), + ) + + assert plans + for plan in plans: + for blend, size in zip(plan["overlap"], plan["window"]): + assert blend == 0 or blend * 4 >= size, plan diff --git a/test/test_encoderadapter.py b/test/test_encoderadapter.py new file mode 100644 index 0000000..8b34fba --- /dev/null +++ b/test/test_encoderadapter.py @@ -0,0 +1,137 @@ +"""EncoderAdapter against the encoder it shards, over gloo on CPU. + +This is the encoder every AutoencoderKL model encodes through, and Flux.2's as well, which builds +the same one. It is what an image-to-image or inpainting pipeline runs over a full-sized image, +so it is the encode worth splitting. + +Run from repo root: + pytest test/test_encoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.encoder_adapters import EncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") + +# Four stages, three of which downsample, so the encoder narrows by 8 as the shipped ones do. +CONFIG = dict( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=256, + down_block_types=["DownEncoderBlock2D"] * 4, + up_block_types=["UpDecoderBlock2D"] * 4, +) +SCALE_FACTOR = 8 +IN_CHANNELS = 3 + + +def build_encoder(mid_block_add_attention=True): + vae = diffusers.AutoencoderKL( + **CONFIG, mid_block_add_attention=mid_block_add_attention + ) + return vae.eval().encoder + + +def worker( + rank, world_size, height, width, add_attention, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder(add_attention) + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder(add_attention) + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = EncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + child_contexts = [ + module.parallel_context + for module in adapter.modules() + if hasattr(module, "parallel_context") + ] + assert child_contexts + assert all( + context is adapter.parallel_context for context in child_contexts + ) + actual = adapter(pixels) + + # The sharded GroupNorms inside the down blocks sum their statistics across ranks in + # float32 before dividing, so their results differ slightly from a single-rank reduction. + assert_matches_reference(rank, actual, expected, "EncoderAdapter", atol=1e-4) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_sharded_encode_matches_unsharded_encode(world_size, master_port, seed=42): + run_distributed(worker, world_size, (64, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2]) +def test_a_mid_block_without_attention_encodes_the_same(world_size, master_port, seed=42): + # Flux.2 can be configured either way, and the mid block runs whole on every rank regardless, + # so this checks the split is undone before it either way. + run_distributed(worker, world_size, (64, 64, False, 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_image_taller_than_it_is_wide_still_encodes(master_port, seed=42): + run_distributed(worker, 2, (96, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + # 80 rows at a ratio of 8 is ten bands, which over 3 ranks leaves them different sizes: the + # downsamplers have to read their neighbours' sizes rather than assume they match. + run_distributed(worker, 3, (80, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + # A conv_block_size under the feature map size sends PatchConv2d down its chunked path, which + # splits and reassembles each convolution on top of the sharding. + run_distributed(worker, 2, (64, 64, True, 32, seed), master_port) + + +def test_a_ratio_that_the_down_blocks_do_not_agree_with_is_refused(): + # Sizing bands by a ratio the stages do not actually narrow by would leave a later stage + # halving a band into a row belonging to the next rank, quietly, so it is counted and checked. + encoder = build_encoder() + with pytest.raises(ValueError, match="narrow by 8"): + EncoderAdapter(encoder, vae_group=None, vae_scale_factor=16) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="EncoderAdapter GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_groupnorm.py b/test/test_groupnorm.py deleted file mode 100644 index cfb992c..0000000 --- a/test/test_groupnorm.py +++ /dev/null @@ -1,77 +0,0 @@ -from distvae.modules.patch_utils import Patchify, DePatchify -from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter -from torch.nn import GroupNorm -from distvae.utils import DistributedEnv - -import torch -import random -import argparse -import torch.distributed as dist -from torch import nn -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - parser.add_argument( - "--channels", - type=int, - default=512, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - DistributedEnv.initialize(None) - - norm = GroupNorm(num_groups=32, num_channels=args.channels, eps=1e-6, affine=True).to(device) - patch_norm = GroupNormAdapter(norm).to(device) - - hidden_state = torch.randn(1, args.channels, args.height, args.width, device=device) - - result = norm(hidden_state) - # if rank == 0: - # print("result: ", result) - - patch = Patchify() - depatch = DePatchify() - patch_result = patch_norm(patch(hidden_state)) - # print("patch_res:", rank, patch_result) - patch_result = depatch(patch_result) - - - if dist.get_rank() == 0: - assert torch.allclose(result, patch_result), "two hidden states are not equal" - - dist.barrier() - dist.destroy_process_group() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test/test_hunyuanvideo15decoderadapter.py b/test/test_hunyuanvideo15decoderadapter.py new file mode 100644 index 0000000..653a20d --- /dev/null +++ b/test/test_hunyuanvideo15decoderadapter.py @@ -0,0 +1,104 @@ +"""HunyuanVideo15DecoderAdapter against the decoder it shards, over gloo on CPU. + +HunyuanVideo 1.5 pads by replication like HunyuanVideo does, but normalises with RMS rather than +GroupNorm, and its attention block takes a whole 5D tensor and builds its own mask, so unlike +HunyuanVideo's it can be wrapped on its own and the mid block around it stays sharded. + +Run from repo root: + pytest test/test_hunyuanvideo15decoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import HunyuanVideo15DecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLHunyuanVideo15"): + pytest.skip( + "installed diffusers has no AutoencoderKLHunyuanVideo15", allow_module_level=True + ) + +# Five channel stages exercise every decoder upsampling transition. +CONFIG = dict( + block_out_channels=(8, 8, 16, 16, 16), + layers_per_block=1, + latent_channels=4, +) +LATENT_CHANNELS = 4 + + +def build_decoder(): + return diffusers.AutoencoderKLHunyuanVideo15(**CONFIG).eval().decoder + + +def worker(rank, world_size, frames, height, width, conv_block_size, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder() + # Taken before the adapter runs, which rebuilds the decoder in place. + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder() + reference.load_state_dict(weights) + expected = reference(latents) + + adapter = HunyuanVideo15DecoderAdapter( + decoder, vae_group=None, conv_block_size=conv_block_size + ).eval() + actual = adapter(latents) + + assert_matches_reference(rank, actual, expected, "HunyuanVideo15DecoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_hunyuan15_decode_matches_unsharded_decode(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_latent_taller_than_it_is_wide_still_decodes(master_port, seed=42): + run_distributed(worker, 2, (1, 24, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_more_than_one_frame_still_decodes(master_port, seed=42): + # Its upsampler treats the first frame differently from the rest, so a single frame would + # never reach the branch that shuffles channels into time as well as space. + run_distributed(worker, 2, (5, 16, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_decodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, 32, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # Uneven bands must preserve all 16 rows without padding the decoder input. + run_distributed(worker, 3, (1, 16, 16, 0, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="HunyuanVideo15DecoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_hunyuanvideo15encoderadapter.py b/test/test_hunyuanvideo15encoderadapter.py new file mode 100644 index 0000000..979593c --- /dev/null +++ b/test/test_hunyuanvideo15encoderadapter.py @@ -0,0 +1,103 @@ +"""HunyuanVideo15EncoderAdapter against the encoder it shards, over gloo on CPU. + +HunyuanVideo 1.5 pads by replication like HunyuanVideo does, but normalises with RMS rather than +GroupNorm, and downsamples by folding each pair of rows and columns into channels instead of by +striding. That fold reads one input position per output one, so what it needs is for a rank to +hold whole pairs of rows, which is what cutting bands in multiples of the ratio is for. + +Run from repo root: + pytest test/test_hunyuanvideo15encoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.encoder_adapters import HunyuanVideo15EncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLHunyuanVideo15"): + pytest.skip( + "installed diffusers has no AutoencoderKLHunyuanVideo15", allow_module_level=True + ) + +# Five channel stages exercise every encoder downsampling transition. +CONFIG = dict( + block_out_channels=(8, 8, 16, 16, 16), + layers_per_block=1, + latent_channels=4, +) +# Four stages downsample, so the encoder narrows the rows by this much. +SCALE_FACTOR = 16 +IN_CHANNELS = 3 + + +def build_encoder(): + return diffusers.AutoencoderKLHunyuanVideo15(**CONFIG).eval().encoder + + +def worker(rank, world_size, frames, height, width, conv_block_size, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder() + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder() + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = HunyuanVideo15EncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + actual = adapter(pixels) + + assert_matches_reference(rank, actual, expected, "HunyuanVideo15EncoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_hunyuan15_encode_matches_unsharded_encode(master_port, seed=42): + run_distributed(worker, 2, (5, 64, 64, 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_image_taller_than_it_is_wide_still_encodes(master_port, seed=42): + run_distributed(worker, 2, (5, 96, 64, 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + # 48 rows at a ratio of 16 is 3 bands over 2 ranks, so one rank takes two and the other one. + run_distributed(worker, 2, (5, 48, 64, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (5, 64, 64, 32, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="HunyuanVideo15EncoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_hunyuanvideodecoderadapter.py b/test/test_hunyuanvideodecoderadapter.py new file mode 100644 index 0000000..17c55bb --- /dev/null +++ b/test/test_hunyuanvideodecoderadapter.py @@ -0,0 +1,115 @@ +"""HunyuanVideoDecoderAdapter against the decoder it shards, over gloo on CPU. + +Two things here are unlike the Wan-derived families. The causal convolutions pad by replication, +so a rank left alone would repeat its own edge rows rather than read its neighbour's, and the +norms are GroupNorms, whose statistics span the axis being split. + +Run from repo root: + pytest test/test_hunyuanvideodecoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import HunyuanVideoDecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLHunyuanVideo"): + pytest.skip("installed diffusers has no AutoencoderKLHunyuanVideo", allow_module_level=True) + +# Four channel stages exercise every decoder upsampling transition. +CONFIG = dict( + block_out_channels=(8, 8, 16, 16), + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, +) +LATENT_CHANNELS = 4 + + +def build_decoder(mid_block_add_attention=True): + vae = diffusers.AutoencoderKLHunyuanVideo( + **CONFIG, mid_block_add_attention=mid_block_add_attention + ) + return vae.eval().decoder + + +def worker( + rank, world_size, frames, height, width, add_attention, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder(add_attention) + # Taken before the adapter runs, which rebuilds the decoder in place. + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder(add_attention) + reference.load_state_dict(weights) + expected = reference(latents) + + adapter = HunyuanVideoDecoderAdapter( + decoder, vae_group=None, conv_block_size=conv_block_size + ).eval() + actual = adapter(latents) + + assert_matches_reference(rank, actual, expected, "HunyuanVideoDecoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_hunyuan_decode_matches_unsharded_decode(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_mid_block_without_attention_shards_its_resnets(master_port, seed=42): + # Without attention the mid block is sharded rather than gathered around, which is a + # different path through the adapter and the only one that reaches its resnet adapters. + run_distributed(worker, 2, (1, 16, 16, False, 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_latent_taller_than_it_is_wide_still_decodes(master_port, seed=42): + run_distributed(worker, 2, (1, 24, 16, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_more_than_one_frame_still_decodes(master_port, seed=42): + # The frame axis is not the one being split, and its causal padding stays in the adapter + # rather than moving into PatchConv3d, so this is what checks that split was made correctly. + run_distributed(worker, 2, (5, 16, 16, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_decodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, True, 32, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # Uneven bands must preserve all 16 rows without padding the decoder input. + run_distributed(worker, 3, (1, 16, 16, True, 0, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="HunyuanVideoDecoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_hunyuanvideoencoderadapter.py b/test/test_hunyuanvideoencoderadapter.py new file mode 100644 index 0000000..1e30121 --- /dev/null +++ b/test/test_hunyuanvideoencoderadapter.py @@ -0,0 +1,113 @@ +"""HunyuanVideoEncoderAdapter against the encoder it shards, over gloo on CPU. + +Two things here are unlike the Wan-derived families. The causal convolutions pad by replication, +so a rank left alone would repeat its own edge rows rather than read its neighbour's, and the +encoder ends on a GroupNorm, whose statistics span the axis being split. + +Run from repo root: + pytest test/test_hunyuanvideoencoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.encoder_adapters import HunyuanVideoEncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLHunyuanVideo"): + pytest.skip("installed diffusers has no AutoencoderKLHunyuanVideo", allow_module_level=True) + +# Four channel stages exercise every encoder downsampling transition. +CONFIG = dict( + block_out_channels=(8, 8, 16, 16), + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, +) +# Three stages downsample, so the encoder narrows the rows by this much. +SCALE_FACTOR = 8 +IN_CHANNELS = 3 + + +def build_encoder(mid_block_add_attention=True): + vae = diffusers.AutoencoderKLHunyuanVideo( + **CONFIG, mid_block_add_attention=mid_block_add_attention + ) + return vae.eval().encoder + + +def worker( + rank, world_size, frames, height, width, add_attention, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder(add_attention) + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder(add_attention) + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = HunyuanVideoEncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + actual = adapter(pixels) + + assert_matches_reference(rank, actual, expected, "HunyuanVideoEncoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_hunyuan_encode_matches_unsharded_encode(master_port, seed=42): + run_distributed(worker, 2, (5, 64, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_encoder_whose_mid_block_has_no_attention(master_port, seed=42): + # Without attention the mid block is convolutions alone, so it stays sharded rather than + # being gathered around, which is a different path through the mid block adapter. + run_distributed(worker, 2, (5, 64, 64, False, 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_image_taller_than_it_is_wide_still_encodes(master_port, seed=42): + run_distributed(worker, 2, (5, 96, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + # 64 rows at a ratio of 8 is 8 bands over 3 ranks, so the bands come out 3, 3 and 2 long. + run_distributed(worker, 3, (5, 64, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (5, 64, 64, True, 32, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="HunyuanVideoEncoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_ltx2videodecoderadapter.py b/test/test_ltx2videodecoderadapter.py new file mode 100644 index 0000000..92af9ba --- /dev/null +++ b/test/test_ltx2videodecoderadapter.py @@ -0,0 +1,134 @@ +"""LTX2VideoDecoderAdapter against the decoder it shards, over gloo on CPU. + +LTX-2 is the easiest of these to shard and the hardest to read. Its mid block has no attention +at all, so nothing needs the whole image, and its spatial padding already lives inside the +convolution rather than being applied around it. The shipped LTX-2 pads by reflection, so the +default here is reflect rather than the zeros LTX-2.3 uses. + +Run from repo root: + pytest test/test_ltx2videodecoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import LTX2VideoDecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLLTX2Video"): + pytest.skip("installed diffusers has no AutoencoderKLLTX2Video", allow_module_level=True) + +# Four channel stages preserve the decoder's spatial compression structure. +CONFIG = dict( + block_out_channels=(8, 16, 32, 32), + latent_channels=8, + layers_per_block=(1, 1, 1, 1, 1), + spatial_compression_ratio=32, +) +LATENT_CHANNELS = 8 + + +def build_decoder(spatial_padding_mode="reflect", inject_noise=False): + vae = diffusers.AutoencoderKLLTX2Video( + **CONFIG, + decoder_spatial_padding_mode=spatial_padding_mode, + decoder_inject_noise=inject_noise, + ) + return vae.eval().decoder + + +def worker( + rank, world_size, frames, height, width, padding_mode, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder(padding_mode) + # Taken before the adapter runs, which rebuilds the decoder in place. + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder(padding_mode) + reference.load_state_dict(weights) + expected = reference(latents) + + adapter = LTX2VideoDecoderAdapter( + decoder, vae_group=None, conv_block_size=conv_block_size + ).eval() + actual = adapter(latents) + + assert_matches_reference(rank, actual, expected, "LTX2VideoDecoderAdapter") + finally: + dist.destroy_process_group() + + +def refusal_worker(rank, world_size, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(0) + LTX2VideoDecoderAdapter(build_decoder(inject_noise=True), vae_group=None) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_ltx2_decode_matches_unsharded_decode(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_zeros_padding_ltx23_ships_decodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, "zeros", 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_latent_taller_than_it_is_wide_still_decodes(master_port, seed=42): + run_distributed(worker, 2, (1, 24, 16, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_more_than_one_frame_still_decodes(master_port, seed=42): + # The temporal padding repeats the end frames rather than padding with anything, and it is + # left to run in the wrapped module, so this is what confirms it survived the swap. + run_distributed(worker, 2, (3, 16, 16, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_decodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, "reflect", 32, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # Uneven bands must preserve all 16 rows without padding the decoder input. + run_distributed(worker, 3, (1, 16, 16, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_injected_noise_is_refused_rather_than_drawn_per_rank(master_port): + # Each rank would draw noise for its own rows, and together they would not reconstruct what + # one rank draws, so the decode could not match its reference. No shipped config enables it. + with pytest.raises(Exception) as caught: + run_distributed(refusal_worker, 2, (), master_port) + assert "inject_noise" in str(caught.value) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="LTX2VideoDecoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_ltx2videoencoderadapter.py b/test/test_ltx2videoencoderadapter.py new file mode 100644 index 0000000..136db36 --- /dev/null +++ b/test/test_ltx2videoencoderadapter.py @@ -0,0 +1,122 @@ +"""LTX2VideoEncoderAdapter against the encoder it shards, over gloo on CPU. + +LTX-2 is the easiest of these to shard: its mid block has no attention, so nothing needs the +whole image, and its spatial padding already lives inside the convolution rather than being +applied around it. Its down block is the one place a checkpoint has a real choice, holding either +a space-to-channel downsampler or a plain strided convolution, and both are covered here. + +Run from repo root: + pytest test/test_ltx2videoencoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.encoder_adapters import LTX2VideoEncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLLTX2Video"): + pytest.skip("installed diffusers has no AutoencoderKLLTX2Video", allow_module_level=True) + +# Four channel stages exercise every downsampling mode. The compression ratio matches their +# space-to-channel factors. +CONFIG = dict( + block_out_channels=(8, 16, 32, 32), + latent_channels=8, + layers_per_block=(1, 1, 1, 1, 1), + spatial_compression_ratio=32, +) +SCALE_FACTOR = 32 +IN_CHANNELS = 3 +# The shipped stages, every one of which downsamples by folding space into channels. Only "conv" +# leaves a bare strided convolution behind, and only a stage that does not widen can hold one. +FOLDING = ("spatial", "temporal", "spatiotemporal", "spatiotemporal") +STRIDING = ("spatial", "temporal", "spatiotemporal", "conv") + + +def build_encoder(downsample_type=FOLDING, padding_mode="reflect"): + vae = diffusers.AutoencoderKLLTX2Video( + **CONFIG, + downsample_type=downsample_type, + encoder_spatial_padding_mode=padding_mode, + ) + return vae.eval().encoder + + +def worker( + rank, world_size, frames, height, width, downsample_type, padding_mode, conv_block_size, + seed, master_port, +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder(downsample_type, padding_mode) + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder(downsample_type, padding_mode) + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = LTX2VideoEncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + actual = adapter(pixels) + + assert_matches_reference(rank, actual, expected, "LTX2VideoEncoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_ltx2_encode_matches_unsharded_encode(master_port, seed=42): + run_distributed(worker, 2, (9, 64, 64, FOLDING, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_stage_that_downsamples_by_striding_instead_of_folding(master_port, seed=42): + # A bare strided convolution rather than the space-to-channel downsampler, which reaches the + # sharded convolution by a different route through the down block adapter. + run_distributed(worker, 2, (9, 64, 64, STRIDING, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_zero_padding_ltx23_uses_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (9, 64, 64, FOLDING, "zeros", 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + # 96 rows at a ratio of 32 is 3 bands over 2 ranks, so one takes two and the other one. The + # image is also taller than it is wide, which is where a mis-split shows up as a wrong shape. + run_distributed(worker, 2, (9, 96, 64, FOLDING, "reflect", 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (9, 64, 64, FOLDING, "reflect", 32, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="LTX2VideoEncoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_patch_utils.py b/test/test_patch_utils.py new file mode 100644 index 0000000..a6bbd36 --- /dev/null +++ b/test/test_patch_utils.py @@ -0,0 +1,248 @@ +"""Splitting rows across ranks and gathering them back, over gloo on CPU. + +The pair has to round-trip exactly for row counts that do not divide by the rank count. Padding +changes the computation once convolutions propagate values into retained rows. Uneven bands +preserve the input, so the gather must accept different amounts from each rank. + +Run from repo root: + pytest test/test_patch_utils.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +import torch.nn as nn + +from distvae.models.layers.conv2d import PatchConv2d +from distvae.models.layers.conv3d import PatchConv3d +from distvae.modules.patch_utils import ( + DePatchify, + Patchify, + VAERowSplitError, + gather_patches, + widest_halo, +) +from distvae.utils import ParallelContext, normalize_patch_dim + +from distributed_harness import ( + assert_matches_reference, + init_gloo, + make_parallel_context, + run_distributed, +) + + +def test_patchify_requires_an_explicit_parallel_context(): + with pytest.raises(TypeError, match="parallel_context"): + Patchify() + + +def test_non_integral_vae_rows_raise_a_typed_error(): + context = ParallelContext(group=None, rank=0, world_size=1, patch_dim=-2) + + with pytest.raises(VAERowSplitError) as error: + Patchify(context, scale_factor=2)(torch.randn(1, 2, 45, 4)) + + assert error.value.rows == 45 + assert error.value.factor == 2 + + +def test_the_widest_halo_is_half_the_widest_kernel_on_the_split_axis(): + context = make_parallel_context() + # Only the split axis counts: a kernel is only ever wide across rows a neighbour holds. + stack = nn.Sequential( + PatchConv2d(1, 1, kernel_size=3, parallel_context=context), + PatchConv2d(1, 1, kernel_size=(7, 1), parallel_context=context), + PatchConv2d(1, 1, kernel_size=(1, 9), parallel_context=context), + ) + assert widest_halo(stack) == 3 + + +def test_the_widest_halo_reads_the_axis_the_convolution_was_told_to_split(): + across = nn.Sequential( + PatchConv2d( + 1, 1, kernel_size=(1, 9), parallel_context=make_parallel_context(-1) + ) + ) + assert widest_halo(across) == 4 + + +def test_a_three_dimensional_kernel_is_read_on_its_split_axis_too(): + stack = nn.Sequential( + PatchConv3d( + 1, 1, kernel_size=(9, 5, 9), parallel_context=make_parallel_context() + ) + ) + assert widest_halo(stack) == 2 + + +def test_a_stack_that_shards_nothing_asks_for_no_halo(): + assert widest_halo(nn.Sequential(nn.Conv2d(1, 1, kernel_size=11))) == 0 + + +def round_trip_worker(rank, world_size, rows, scale_factor, patch_dim, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + whole = torch.randn(1, 4, rows, rows) + context = make_parallel_context(patch_dim) + + band = Patchify(context, scale_factor=scale_factor)(whole) + # Every band is a whole number of scale_factor rows, which is what keeps a rank's share + # of a strided convolution on the same grid as the reference's. + assert band.shape[patch_dim] % scale_factor == 0, ( + f"rank {rank} got {band.shape[patch_dim]} rows, not a multiple of {scale_factor}" + ) + rebuilt = DePatchify(context)(band) + + assert_matches_reference(rank, rebuilt, whole if rank == 0 else None, "Patchify round trip") + finally: + dist.destroy_process_group() + + +def gather_worker(rank, world_size, rows, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed + rank) + # Deliberately lopsided: rank r contributes r + 1 rows, so no two ranks agree. + band = torch.full((1, 2, rank + 1, 3), float(rank)) + bands, sizes = gather_patches(band, make_parallel_context()) + + assert sizes == [r + 1 for r in range(world_size)], f"rank {rank} read sizes {sizes}" + for r, gathered in enumerate(bands): + assert gathered.shape[2] == r + 1, f"band {r} has {gathered.shape[2]} rows" + # The pad added for transport must not survive into what callers read back. + assert torch.equal(gathered, torch.full_like(gathered, float(r))), ( + f"band {r} carries transport padding" + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 3, 4]) +def test_rows_that_divide_by_the_rank_count_round_trip(world_size, master_port, seed=42): + run_distributed(round_trip_worker, world_size, (24, 1, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 3, 4]) +def test_rows_that_do_not_divide_by_the_rank_count_round_trip(world_size, master_port, seed=42): + # 25 is prime to every rank count here, so at least one band is short in each case. + run_distributed(round_trip_worker, world_size, (25, 1, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 3]) +def test_bands_stay_whole_multiples_of_the_vae_ratio(world_size, master_port, seed=42): + # 40 rows at a ratio of 8 is 5 bands to share out, which no rank count here divides. + run_distributed(round_trip_worker, world_size, (40, 8, -2, seed), master_port) + + +@pytest.mark.gloo +def test_splitting_along_width_round_trips_too(master_port, seed=42): + run_distributed(round_trip_worker, 3, (25, 1, -1, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 3]) +def test_the_gather_hands_back_every_rank_its_own_rows(world_size, master_port, seed=42): + run_distributed(gather_worker, world_size, (0, seed), master_port) + + +def refusal_worker(rank, world_size, rows, scale_factor, halo, expected, master_port): + init_gloo(rank, world_size, master_port) + try: + with pytest.raises(ValueError, match=expected): + Patchify( + make_parallel_context(), scale_factor=scale_factor, halo=halo + )(torch.randn(1, 2, rows, 4)) + finally: + dist.destroy_process_group() + + +def halo_worker(rank, world_size, rows, halo, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + whole = torch.randn(1, 4, rows, rows) + context = make_parallel_context() + assert torch.equal( + DePatchify(context)(Patchify(context, halo=halo)(whole)), whole + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_rows_that_are_not_a_multiple_of_the_ratio_are_refused(master_port): + # The encoder narrows by 8, so 20 rows cannot be cut into bands whose latent rows line up. + run_distributed(refusal_worker, 2, (20, 8, 0, "multiples of 8"), master_port) + + +@pytest.mark.gloo +def test_more_ranks_than_bands_is_refused(master_port): + # 16 rows at a ratio of 8 leaves two bands, which three ranks cannot share. + run_distributed(refusal_worker, 3, (16, 8, 0, "at most 2 ranks"), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [2, 3, 4]) +def test_a_halo_wider_than_the_thinnest_band_is_refused_by_every_rank(world_size, master_port): + # That every rank refuses is the whole point of asking here. Seven rows over two ranks is a + # band of four and a band of three, so a rank deciding from what it holds would have the wide + # one go on into a halo exchange with a rank that had already stopped - and that is a hang + # rather than a failure, because the rows it waits for are never sent. + run_distributed( + refusal_worker, world_size, (7, 1, 4, "reaching 4 rows past a band"), master_port + ) + + +@pytest.mark.gloo +def test_a_halo_the_thinnest_band_can_just_lend_is_allowed(master_port, seed=42): + # Seven rows over two ranks leaves a band of three, and a halo of three is the last width + # that works rather than the first that does not. The guard has to let it through. + run_distributed(halo_worker, 2, (7, 3, seed), master_port) + + +@pytest.mark.parametrize("patch_dim", [-2, 3]) +def test_video_height_axis_indices_normalize_to_the_same_axis(patch_dim): + assert normalize_patch_dim(patch_dim, ndim=5, spatial_only=True) == -2 + + +@pytest.mark.parametrize("patch_dim", [-1, 4]) +def test_video_width_axis_indices_normalize_to_the_same_axis(patch_dim): + assert normalize_patch_dim(patch_dim, ndim=5, spatial_only=True) == -1 + + +@pytest.mark.parametrize("patch_dim", [-3, 2]) +def test_video_frame_axis_indices_are_rejected(patch_dim): + with pytest.raises(ValueError, match="frame axis"): + normalize_patch_dim(patch_dim, ndim=5, spatial_only=True) + + +def test_patchifiers_keep_their_own_parallel_context(): + first_context = ParallelContext(group=None, rank=0, world_size=2, patch_dim=-2) + first = Patchify(parallel_context=first_context) + second_context = ParallelContext(group=None, rank=1, world_size=2, patch_dim=-1) + second = Patchify(parallel_context=second_context) + whole = torch.arange(24).reshape(1, 1, 4, 6) + + assert torch.equal(first(whole), whole[:, :, :2, :]) + assert torch.equal(second(whole), whole[:, :, :, 3:]) + assert first.parallel_context is first_context + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Patchify and gather GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_patchconv_padding_modes.py b/test/test_patchconv_padding_modes.py new file mode 100644 index 0000000..88f4459 --- /dev/null +++ b/test/test_patchconv_padding_modes.py @@ -0,0 +1,129 @@ +"""PatchConv2d and PatchConv3d under padding modes other than zeros, over gloo on CPU. + +Replicate and reflect padding are what the HunyuanVideo and LTX-2 VAEs convolve with, and they +are the case a sharded convolution can get quietly wrong: left alone, a rank repeats or mirrors +its own edge rows at a boundary that is not an edge of the image at all. + +Run from repo root: + pytest test/test_patchconv_padding_modes.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from distvae.models.layers.conv2d import PatchConv2d +from distvae.models.layers.conv3d import PatchConv3d +from distvae.modules.patch_utils import DePatchify, Patchify + +from distributed_harness import ( + assert_matches_reference, + init_gloo, + make_parallel_context, + run_distributed, +) + + +def worker( + rank, world_size, ndim, padding_mode, kernel_size, padding, block_size, patch_dim, seed, + master_port, +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + in_channels, out_channels = 4, 8 + context = make_parallel_context(patch_dim) + if ndim == 5: + shape = (1, in_channels, 3, 16, 16) + reference = nn.Conv3d( + in_channels, out_channels, kernel_size, padding=padding, + padding_mode=padding_mode, + ).eval() + sharded = PatchConv3d( + in_channels, out_channels, kernel_size, padding=padding, + padding_mode=padding_mode, block_size=block_size, + parallel_context=context, + ).eval() + else: + shape = (1, in_channels, 16, 16) + reference = nn.Conv2d( + in_channels, out_channels, kernel_size, padding=padding, + padding_mode=padding_mode, + ).eval() + sharded = PatchConv2d( + in_channels, out_channels, kernel_size, padding=padding, + padding_mode=padding_mode, block_size=block_size, + parallel_context=context, + ).eval() + sharded.weight.data = reference.weight.data + sharded.bias.data = reference.bias.data + + x = torch.randn(*shape) + patchify = Patchify(context) + depatchify = DePatchify(context) + + with torch.no_grad(): + expected = reference(x) if rank == 0 else None + actual = depatchify(sharded(patchify(x))) + + assert_matches_reference( + rank, actual, expected, f"a {padding_mode}-padded convolution", atol=1e-5 + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("padding_mode", ["replicate", "reflect"]) +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_a_3d_convolution_matches_an_unsharded_one(padding_mode, world_size, master_port, seed=42): + run_distributed(worker, world_size, (5, padding_mode, 3, 1, 0, -2, seed), master_port) + + +@pytest.mark.gloo +def test_circular_padding_is_refused_rather_than_wrapped_within_a_patch(master_port, seed=42): + # Circular padding reads from the far edge of the image, which no neighbour holds. Wrapping + # within the patch instead would be silently wrong, so the convolution has to say so. + with pytest.raises(Exception) as caught: + run_distributed(worker, 2, (5, "circular", 3, 1, 0, -2, seed), master_port) + assert "circular" in str(caught.value) + + +@pytest.mark.gloo +def test_circular_padding_is_still_allowed_on_a_single_rank(master_port, seed=42): + run_distributed(worker, 1, (5, "circular", 3, 1, 0, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("padding_mode", ["replicate", "reflect"]) +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_a_2d_convolution_matches_an_unsharded_one(padding_mode, world_size, master_port, seed=42): + run_distributed(worker, world_size, (4, padding_mode, 3, 1, 0, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("padding_mode", ["replicate", "reflect"]) +def test_the_chunked_path_pads_the_same_way(padding_mode, master_port, seed=42): + run_distributed(worker, 2, (5, padding_mode, 3, 1, 4, -2, seed), master_port) + + +@pytest.mark.gloo +def test_splitting_the_width_instead_pads_the_same_way(master_port, seed=42): + run_distributed(worker, 2, (5, "replicate", 3, 1, 0, -1, seed), master_port) + + +@pytest.mark.gloo +def test_a_wider_kernel_pads_the_same_way(master_port, seed=42): + # Kernel 3 with padding 1 has a fast path of its own; a kernel of 5 goes the other way. + run_distributed(worker, 2, (5, "replicate", 5, 2, 0, -2, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="PatchConv padding mode GLOO tests") + _, remainder = parser.parse_known_args() + sys.exit(pytest.main([os.path.abspath(__file__), "-v"] + remainder)) diff --git a/test/test_patchgroupnorm.py b/test/test_patchgroupnorm.py new file mode 100644 index 0000000..61f242e --- /dev/null +++ b/test/test_patchgroupnorm.py @@ -0,0 +1,240 @@ +"""PatchGroupNorm against nn.GroupNorm over multiple ranks. + +GroupNorm statistics include the split spatial axis, so group sums and variances must be +aggregated across ranks. + +Run from repo root: + pytest test/test_patchgroupnorm.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter +from distvae.modules.patch_utils import DePatchify, Patchify +from distvae.utils import ParallelContext + +from distributed_harness import ( + assert_matches_reference, + assert_no_less_precise_than, + init_gloo, + make_parallel_context, + run_distributed, +) + + +def worker(rank, world_size, shape, num_groups, patch_dim, seed, affine, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + channels = shape[1] + norm = nn.GroupNorm( + num_groups=num_groups, num_channels=channels, eps=1e-6, affine=affine + ).eval() + # Shifted per channel, so the group statistics are not already near zero mean and unit + # variance and an incorrect reduction has somewhere to show up. + x = torch.randn(*shape) * 3.0 + 2.0 + + context = make_parallel_context(patch_dim) + patchify = Patchify(context) + depatchify = DePatchify(context) + sharded = GroupNormAdapter(norm, parallel_context=context) + + with torch.no_grad(): + expected = norm(x) if rank == 0 else None + actual = depatchify(sharded(patchify(x))) + + assert_matches_reference(rank, actual, expected, "PatchGroupNorm", atol=1e-5) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_patch_group_norm_matches_group_norm_on_a_feature_map(world_size, master_port, seed=42): + run_distributed(worker, world_size, ((1, 16, 16, 16), 8, -2, seed, True), master_port) + + +@pytest.mark.gloo +def test_patch_group_norm_matches_group_norm_when_an_odd_height_is_split( + master_port, seed=42 +): + """An uneven, non-square height split detects reduction over the wrong spatial axis. + + An even 16x16 split gives the same element count for height and width, so it cannot detect + use of the wrong axis. Splitting a 15x4 tensor across height produces unequal rank sizes and + exposes that error. + """ + run_distributed( + worker, 2, ((1, 16, 15, 4), 8, -2, seed, True), master_port + ) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2]) +def test_patch_group_norm_matches_group_norm_on_a_video_feature_map( + world_size, master_port, seed=42 +): + # Video GroupNorm reduces over all of (F, H, W), including the axes around the split axis. + run_distributed(worker, world_size, ((1, 16, 3, 8, 8), 4, -2, seed, True), master_port) + + +@pytest.mark.gloo +def test_patch_group_norm_matches_group_norm_on_a_video_map_of_three_different_extents( + master_port, seed=42 +): + # F, H and W all different and the split uneven, so confusing the split axis for either of + # the two it is reduced alongside changes the count rather than cancelling against it. + run_distributed( + worker, 2, ((1, 16, 3, 7, 8), 4, -2, seed, True), master_port + ) + + +@pytest.mark.gloo +def test_patch_group_norm_matches_group_norm_when_the_width_is_split( + master_port, seed=42 +): + run_distributed(worker, 2, ((1, 16, 16, 16), 8, -1, seed, True), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize( + "shape,patch_dim", + [ + pytest.param((1, 16, 10, 8), -2, id="uneven-height"), + pytest.param((1, 16, 8, 10), -1, id="uneven-width"), + ], +) +def test_patch_group_norm_matches_group_norm_on_uneven_spatial_bands_without_affine( + shape, patch_dim, master_port, seed=42 +): + run_distributed( + worker, 3, (shape, 8, patch_dim, seed, False), master_port + ) + + +def test_positive_video_frame_axis_index_is_rejected(): + context = ParallelContext(None, rank=0, world_size=1, patch_dim=2) + norm = GroupNormAdapter(nn.GroupNorm(1, 2), parallel_context=context) + with pytest.raises(ValueError, match="frame axis"): + norm(torch.randn(1, 2, 3, 4, 4)) + + +def test_constructing_a_second_norm_adapter_does_not_reconfigure_the_first(monkeypatch): + first_group, second_group = object(), object() + first_context = ParallelContext(first_group, rank=0, world_size=2, patch_dim=-2) + second_context = ParallelContext(second_group, rank=0, world_size=2, patch_dim=-1) + first = GroupNormAdapter(nn.GroupNorm(1, 2), parallel_context=first_context) + GroupNormAdapter(nn.GroupNorm(1, 2), parallel_context=second_context) + used_groups = [] + + monkeypatch.setattr( + dist, + "all_reduce", + lambda tensor, group=None: used_groups.append(group), + ) + first(torch.randn(1, 2, 2, 2)) + + assert used_groups == [first_group, first_group] + + +@pytest.mark.gloo +def test_patch_group_norm_matches_group_norm_when_an_odd_width_is_split( + master_port, seed=42 +): + """An uneven, non-square width split detects reduction over the wrong spatial axis. + + Splitting a width of 15 across two ranks assigns 8 columns to one rank and 7 to the other. + An even square split gives the same element count for height and width and cannot expose this + error. + """ + run_distributed( + worker, 2, ((1, 16, 4, 15), 8, -1, seed, True), master_port + ) + + +def told_worker(rank, world_size, shape, num_groups, patch_dim, seed, master_port): + """A norm reads its axis from its own context.""" + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + norm = nn.GroupNorm( + num_groups=num_groups, num_channels=shape[1], eps=1e-6, affine=True + ).eval() + x = torch.randn(*shape) * 3.0 + 2.0 + context = make_parallel_context(patch_dim) + + with torch.no_grad(): + expected = norm(x) if rank == 0 else None + sharded = GroupNormAdapter(norm, parallel_context=context) + actual = DePatchify(context)(sharded(Patchify(context)(x))) + + assert_matches_reference(rank, actual, expected, "PatchGroupNorm told", atol=1e-5) + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("patch_dim", [-2, -1]) +def test_the_context_axis_selects_the_uneven_spatial_band(patch_dim, master_port, seed=42): + # Uneven along whichever axis is split, so that being told the wrong one would show. + shape = (1, 16, 15, 4) if patch_dim == -2 else (1, 16, 4, 15) + run_distributed(told_worker, 2, (shape, 8, patch_dim, seed), master_port) + + +def bfloat16_worker(rank, world_size, shape, num_groups, patch_dim, seed, master_port): + """PatchGroupNorm's bf16 rounding against nn.GroupNorm's own, both judged by the fp32 answer""" + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + channels = shape[1] + norm = nn.GroupNorm( + num_groups=num_groups, num_channels=channels, eps=1e-6, affine=True + ).eval() + x = torch.randn(*shape) * 3.0 + 2.0 + + with torch.no_grad(): + # Before the cast: nn.Module.to is in place, and this needs the float32 answer. + gold = norm(x).to(torch.bfloat16) if rank == 0 else None + norm = norm.to(torch.bfloat16) + x = x.to(torch.bfloat16) + stock = norm(x) if rank == 0 else None + + context = make_parallel_context(patch_dim) + patchify = Patchify(context) + depatchify = DePatchify(context) + actual = depatchify( + GroupNormAdapter(norm, parallel_context=context)(patchify(x)) + ) + + assert_no_less_precise_than(rank, actual, stock, gold, "PatchGroupNorm in bfloat16") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2, 4]) +def test_it_rounds_no_worse_than_group_norm_in_bfloat16(world_size, master_port, seed=42): + run_distributed(bfloat16_worker, world_size, ((1, 32, 64, 64), 32, -2, seed), master_port) + + +@pytest.mark.gloo +@pytest.mark.parametrize("world_size", [1, 2]) +def test_it_rounds_no_worse_than_group_norm_in_bfloat16_on_video(world_size, master_port, seed=42): + run_distributed(bfloat16_worker, world_size, ((1, 16, 3, 8, 8), 4, -2, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="PatchGroupNorm GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_public_vae_api.py b/test/test_public_vae_api.py new file mode 100644 index 0000000..1a9804e --- /dev/null +++ b/test/test_public_vae_api.py @@ -0,0 +1,38 @@ +from packaging.version import Version + +from distvae.__version__ import __version__ +from distvae import vae + + +MINIMUM_PUBLIC_VAE_API_VERSION = Version("0.0.0beta9") +PUBLIC_VAE_FUNCTIONS = { + "ParallelContext", + "VAERowSplitError", + "apply_tile_plan", + "context_of", + "decoder_adapter_name", + "encoder_adapter_name", + "encoder_scale_factor", + "is_tile_padding_error", + "latent_rows", + "mark", + "parallelize_decoder", + "parallelize_encoder", + "require_vae_support", + "sharing", + "supports_tile_parallel", + "tile_overlap", + "tile_overlap_plan", + "tile_shape", + "tile_shape_plan", + "tiled_decode_for", +} + + +def test_package_version_meets_the_public_vae_api_minimum(): + assert Version(__version__) >= MINIMUM_PUBLIC_VAE_API_VERSION + + +def test_public_vae_api_exports_xdit_orchestration_functions(): + assert set(vae.__all__) == PUBLIC_VAE_FUNCTIONS + assert all(callable(getattr(vae, name)) for name in PUBLIC_VAE_FUNCTIONS) diff --git a/test/test_qwenimagedecoderadapter.py b/test/test_qwenimagedecoderadapter.py new file mode 100644 index 0000000..5a15478 --- /dev/null +++ b/test/test_qwenimagedecoderadapter.py @@ -0,0 +1,99 @@ +"""QwenImageDecoderAdapter against the decoder it shards, over gloo on CPU. + +Qwen-Image, Qwen-Image-Edit, and Krea-2 share this decoder structure. + +Run from repo root: + pytest test/test_qwenimagedecoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import QwenImageDecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "AutoencoderKLQwenImage"): + pytest.skip("installed diffusers has no AutoencoderKLQwenImage", allow_module_level=True) + +# Four channel stages exercise every decoder upsampling transition. +CONFIG = dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1, attn_scales=[]) +LATENT_CHANNELS = 4 + + +def build_decoder(): + return diffusers.AutoencoderKLQwenImage(**CONFIG).eval().decoder + + +def worker(rank, world_size, frames, height, width, conv_block_size, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder() + # Taken before the adapter runs, which rebuilds the decoder in place. + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder() + reference.load_state_dict(weights) + expected = reference(latents, feat_cache=None, feat_idx=[0]) + + adapter = QwenImageDecoderAdapter( + decoder, vae_group=None, conv_block_size=conv_block_size + ).eval() + actual = adapter(latents) + + assert_matches_reference(rank, actual, expected, "QwenImageDecoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_qwen_decode_matches_unsharded_decode(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_a_latent_taller_than_it_is_wide_still_decodes(master_port, seed=42): + # The patch dimension defaults to H, so a non-square latent is the case where getting the + # split wrong shows up as a wrongly shaped output rather than as wrong values. + run_distributed(worker, 2, (1, 24, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_more_than_one_frame_still_decodes(master_port, seed=42): + # The frame axis is not the one being split, but the causal padding along it is applied by + # the adapter rather than by PatchConv3d, which is where that could go wrong. + run_distributed(worker, 2, (3, 16, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # Padding to an even split changes the decode because convolution and mid-block attention + # propagate padded values into the rows that survive cropping. + run_distributed(worker, 3, (1, 16, 16, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_decodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, 32, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="QwenImageDecoderAdapter GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_qwenimageencoderadapter.py b/test/test_qwenimageencoderadapter.py new file mode 100644 index 0000000..0a1fcef --- /dev/null +++ b/test/test_qwenimageencoderadapter.py @@ -0,0 +1,112 @@ +"""QwenImageEncoderAdapter against the encoder it shards, over gloo on CPU. + +Qwen-Image's encoder is Wan 2.1's laid out flat and renamed, so what this really checks is that +the resample's zero pad and stride-2 convolution survive being split, and that an attention +block in the middle of the down blocks gets gathered rather than left on its own patch. + +Run from repo root: + pytest test/test_qwenimageencoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.encoder_adapters import QwenImageEncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +autoencoder_kl_qwenimage = pytest.importorskip( + "diffusers.models.autoencoders.autoencoder_kl_qwenimage" +) + +CONFIG = dict( + dim=32, + z_dim=16, + dim_mult=[1, 2, 4, 8], + num_res_blocks=1, + temperal_downsample=[False, True, True], + dropout=0.0, +) +SCALE_FACTOR = 8 +IN_CHANNELS = 3 + + +def build_encoder(attn_scales=()): + encoder = autoencoder_kl_qwenimage.QwenImageEncoder3d( + **CONFIG, attn_scales=list(attn_scales) + ) + return encoder.eval() + + +def worker( + rank, world_size, frames, height, width, attn_scales, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder(attn_scales) + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder(attn_scales) + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = QwenImageEncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + actual = adapter(pixels) + + assert_matches_reference(rank, actual, expected, "QwenImageEncoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_qwen_encode_matches_unsharded_encode(master_port, seed=42): + run_distributed(worker, 2, (4, 64, 64, (), 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_attention_block_among_the_down_blocks_is_gathered(master_port, seed=42): + # An attention reduces over every position, so a rank holding one patch of rows cannot do it + # alone. attn_scales=(1.0,) puts one at the first stage, where the feature map is largest. + run_distributed(worker, 2, (4, 64, 64, (1.0,), 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_image_taller_than_it_is_wide_still_encodes(master_port, seed=42): + run_distributed(worker, 2, (4, 96, 64, (), 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + run_distributed(worker, 3, (4, 80, 64, (), 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (4, 64, 64, (), 32, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="QwenImageEncoderAdapter GLOO tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_resnet_adapter_context.py b/test/test_resnet_adapter_context.py new file mode 100644 index 0000000..53a49f8 --- /dev/null +++ b/test/test_resnet_adapter_context.py @@ -0,0 +1,41 @@ + +from diffusers.models.resnet import ResnetBlock2D + +from distvae.modules.adapters.layers.conv_adapters import Conv2dAdapter +from distvae.modules.adapters.layers.norm_adapters import GroupNormAdapter +import distvae.modules.adapters.resnet_adapters as resnet_adapters +from distvae.modules.adapters.resnet_adapters import ResnetBlock2DAdapter +from distvae.utils import ParallelContext + + +def test_resnet_wrappers_receive_the_adapters_parallel_settings(monkeypatch): + context = ParallelContext(group=None, rank=0, world_size=1, patch_dim=-1) + received_norms = [] + received_convs = [] + + def recording_group_norm(norm, parallel_context=None): + received_norms.append(parallel_context) + return GroupNormAdapter(norm, parallel_context=parallel_context) + + def recording_conv(conv, *, block_size=0, parallel_context=None): + received_convs.append(parallel_context) + return Conv2dAdapter( + conv, + block_size=block_size, + parallel_context=parallel_context, + ) + + monkeypatch.setattr(resnet_adapters, "GroupNormAdapter", recording_group_norm) + monkeypatch.setattr(resnet_adapters, "Conv2dAdapter", recording_conv) + source = ResnetBlock2D( + in_channels=4, + out_channels=8, + temb_channels=None, + groups=1, + dropout=0.0, + ) + + ResnetBlock2DAdapter(source, parallel_context=context) + + assert received_norms == [context, context] + assert received_convs == [context, context, context] diff --git a/test/test_smoke_families.py b/test/test_smoke_families.py new file mode 100644 index 0000000..21af11f --- /dev/null +++ b/test/test_smoke_families.py @@ -0,0 +1,29 @@ +"""Every benchmark family builds on the meta device with representative samples.""" + +import pytest +import torch + +from bench.harness.catalog import FAMILIES, sample_for + +diffusers = pytest.importorskip("diffusers") + + +@pytest.mark.parametrize("family", sorted(FAMILIES)) +def test_benchmark_family_builds_without_weights_or_an_accelerator(family): + spec = FAMILIES[family] + cls = getattr(diffusers, spec["cls"], None) + if cls is None: + pytest.skip(f"{spec['cls']} is not in this diffusers") + + with torch.device("meta"): + vae = cls(**spec["config"]).eval() + latent = sample_for( + spec, "decoder", 512, 512, torch.bfloat16, "meta", frames=17 + ) + pixels = sample_for( + spec, "encoder", 512, 512, torch.bfloat16, "meta", frames=17 + ) + + assert sum(parameter.numel() for parameter in vae.parameters()) > 0 + assert latent.device.type == "meta" + assert pixels.device.type == "meta" diff --git a/test/test_tile_overlap_absolute.py b/test/test_tile_overlap_absolute.py new file mode 100644 index 0000000..7ffe78b --- /dev/null +++ b/test/test_tile_overlap_absolute.py @@ -0,0 +1,115 @@ +from types import SimpleNamespace + +import torch +import torch.nn.functional as functional + +from distvae.vae import tiling + + +class StubVAE: + def __init__(self, **attrs): + for name, value in attrs.items(): + setattr(self, name, value) + + +def overlap_vae(height=256, width=256, latent_height=32, latent_width=32): + return StubVAE( + tile_sample_min_height=height, + tile_sample_min_width=width, + tile_latent_min_height=latent_height, + tile_latent_min_width=latent_width, + tile_overlap_factor=0.25, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + ) + + +def stride_vae(): + cls = type("AutoencoderKLQwenImage", (StubVAE,), {}) + return cls( + tile_sample_min_height=256, + tile_sample_min_width=384, + tile_sample_stride_height=192, + tile_sample_stride_width=288, + spatial_compression_ratio=8, + config=SimpleNamespace(), + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + decoder=lambda tile: tile, + post_quant_conv=lambda tile: tile, + clear_cache=lambda: None, + ) + + +def test_tile_overlap_reports_absolute_pixels_for_both_storage_families(): + assert tiling.tile_overlap(overlap_vae(height=240, width=320)) == (60, 80) + assert tiling.tile_overlap(stride_vae()) == (64, 96) + + +def test_exact_per_axis_overlap_plans_keyed_factors_and_scalar_only_when_equal(): + vae = overlap_vae(height=240, width=320, latent_height=30, latent_width=40) + + plan = tiling.tile_overlap_plan(vae, 40, 64) + + assert plan == { + "tile_overlap_factor_height": 1 / 6, + "tile_overlap_factor_width": 0.2, + } + tiling.apply_tile_plan(vae, plan) + assert tiling.tile_overlap(vae) == (40, 64) + assert vae.tile_overlap_factor == 0.25 + + square = overlap_vae() + equal = tiling.tile_overlap_plan(square, 64, 64) + assert equal["tile_overlap_factor"] == 0.25 + + +def test_overlap_plan_is_exact_and_rejects_unrepresentable_requests(): + vae = overlap_vae(height=240, width=320, latent_height=30, latent_width=40) + + assert tiling.tile_overlap_plan(vae, 41, 64) is None + assert tiling.tile_overlap_plan(vae, 240, 64) is None + assert tiling.tile_overlap_plan(vae, True, 64) is None + assert tiling.tile_overlap_plan(vae, -1, 64) is None + + +def test_sample_shape_requires_zero_on_inactive_axes_and_sets_both_factors(): + vae = overlap_vae(height=240, width=320, latent_height=30, latent_width=40) + + assert tiling.tile_overlap_plan(vae, 0, 64, sample_shape=(240, 640)) == { + "tile_overlap_factor_height": 0.0, + "tile_overlap_factor_width": 0.2, + } + assert tiling.tile_overlap_plan(vae, 40, 64, sample_shape=(240, 640)) is None + assert tiling.tile_overlap_plan(vae, 0, 0, sample_shape=(240, 320)) == { + "tile_overlap_factor": 0.0, + "tile_overlap_factor_height": 0.0, + "tile_overlap_factor_width": 0.0, + } + + +def test_stored_stride_plan_sets_both_axes_and_requires_exact_granularity(): + vae = stride_vae() + + assert tiling.tile_overlap_plan(vae, 64, 128) == { + "tile_sample_stride_height": 192, + "tile_sample_stride_width": 256, + } + assert tiling.tile_overlap_plan(vae, 63, 128) is None + assert tiling.tile_overlap_plan(vae, 0, 128, sample_shape=(256, 768)) == { + "tile_sample_stride_height": 256, + "tile_sample_stride_width": 256, + } + + +def test_replacement_decode_uses_rectangular_keyed_overlap_factors(): + vae = overlap_vae(height=16, width=24, latent_height=2, latent_width=3) + vae.decoder = lambda tile: functional.interpolate(tile, scale_factor=8, mode="nearest") + plan = tiling.tile_overlap_plan(vae, 8, 8) + assert plan is not None + tiling.apply_tile_plan(vae, plan) + + decode = tiling.tiled_decode_for(vae) + sample = decode(torch.randn(1, 4, 4, 6)).sample + + assert sample.shape == (1, 4, 32, 48) diff --git a/test/test_upsample2D.py b/test/test_upsample2D.py deleted file mode 100644 index 82449a1..0000000 --- a/test/test_upsample2D.py +++ /dev/null @@ -1,73 +0,0 @@ -from distvae.modules.adapters.upsampling_adapters import Upsample2DAdapter -from distvae.modules.patch_utils import Patchify, DePatchify -from diffusers.models.upsampling import Upsample2D -from distvae.utils import DistributedEnv - -import torch -import random -import argparse -import torch.distributed as dist -from torch import nn -import os -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - DistributedEnv.initialize(None) - - upsampler = Upsample2D(64, use_conv=True, out_channels=64).to(device) - patch_upsampler = Upsample2DAdapter(upsampler).to(device) - - hidden_state = torch.randn(1, 64, args.height, args.width, device=device) - print("hidden state shape: ", hidden_state.shape) - - result = upsampler(hidden_state) - # if rank == 0: - # print("result: ", result) - - patch = Patchify() - depatch = DePatchify() - patch_result = patch_upsampler(patch(hidden_state)) - # print("patch_res:", rank, patch_result) - patch_result = depatch(patch_result) - print("result shape: ", patch_result.shape) - - if dist.get_rank() == 0: - assert torch.allclose(result, patch_result), "two hidden states are not equal" - - dist.barrier() - dist.destroy_process_group() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test/test_vae_decoder.py b/test/test_vae_decoder.py deleted file mode 100644 index bb70bb4..0000000 --- a/test/test_vae_decoder.py +++ /dev/null @@ -1,77 +0,0 @@ -from diffusers.models.autoencoders.vae import Decoder -from distvae.modules.adapters.vae.decoder_adapters import DecoderAdapter -from distvae.utils import DistributedEnv - -import time -import torch -import random -import argparse -import torch.distributed as dist -import os -from torch.cuda import set_device, device_count -from torch.cuda import manual_seed as device_manual_seed -try: - import torch_musa - from torch_musa.core.device import set_device, device_count - from torch_musa.core.random import manual_seed as device_manual_seed -except ModuleNotFoundError: - pass - -def set_seed(seed: int = 42): - random.seed(seed) - torch.manual_seed(seed) - device_manual_seed(seed) - -@torch.no_grad() -def main(): - set_seed() - torch.backends.cudnn.deterministic = True - parser = argparse.ArgumentParser() - parser.add_argument( - "--height", - type=int, - default=1024, - help="The height of image", - ) - parser.add_argument( - "--width", - type=int, - default=1024, - help="The width of image", - ) - args = parser.parse_args() - backend = DistributedEnv.get_torch_distributed_backend() - dist.init_process_group(backend=backend) - device = torch.distributed.get_rank() % device_count() - set_device(device) - # input - # create vae.decoder instance - decoder = Decoder( - in_channels=4, - out_channels=3, - up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], - block_out_channels=(128, 256, 512, 512), - layers_per_block=2, - norm_num_groups=32, - act_fn="silu", - ).to(device) - # transform vae.decoder to distvae.decoder - patch_decoder = DecoderAdapter(decoder, conv_block_size=1024).to(device) - # forward - hidden_state = torch.randn(1, 4, args.height // 8, args.width // 8, device=device) - result = decoder(hidden_state) - - DistributedEnv.record_memory_history() - start_time = time.time() - patch_result = patch_decoder(hidden_state) - end_time = time.time() - peak_memory = DistributedEnv.get_peak_memory(device) - if dist.get_rank() == 0: - assert torch.allclose(result, patch_result, atol=1e-2), "two hidden states are not equal" - print(f"VAE: resolution: {args.height}x{args.width}, time: {end_time - start_time} sec, peak memory: {peak_memory / 2 ** 30} GB") - - dist.barrier() - dist.destroy_process_group() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test/test_vae_parallel.py b/test/test_vae_parallel.py new file mode 100644 index 0000000..91c3fd9 --- /dev/null +++ b/test/test_vae_parallel.py @@ -0,0 +1,363 @@ +"""Which VAE classes DistVAE can shard, checked against real VAEs rather than a list. + +The adapters assert their block types from inside a half-built replacement, so a VAE they cannot +take has to be recognised before wrapping. These build each VAE class an integration loads and +demand the answer for both halves of it, so a model declaring use_parallel_vae or +use_parallel_vae_encoder cannot quietly become unshardable when diffusers reworks a block. +""" + +import os +import unittest +from types import SimpleNamespace +from unittest import mock + +import diffusers +import torch.distributed as dist +import torch.nn as nn + +from distvae.vae import parallel as vae_parallel + +import test_vae_tiling + +# The same tiny builds the tiling tests decode through, taken from there rather than repeated, so +# one VAE class is described in one place. Imported as a module, since binding its TestCase here +# would have unittest collect and run those decodes a second time. Only the config is needed: +# picking an adapter reads the decoder's block types and never runs it. +CONFIGS = { + name: config + for name, (config, _, _) in test_vae_tiling.TestEverySupportedVAE.VAES.items() +} + + +def _vae_class(name): + """Return an installed optional VAE class, or skip only the current test/subtest.""" + cls = getattr(diffusers, name, None) + if cls is None: + raise unittest.SkipTest(f"{name} is not in diffusers {diffusers.__version__}") + return cls + + +# The adapter each class needs, or None where DistVAE has nothing for its decoder. Every class an +# integration loads is shardable as of DistVAE's +# QwenImage, HunyuanVideo and LTX-2 adapters, so a None appearing here again would mean a newly +# supported model arrived ahead of the adapter for its VAE. +EXPECTED = { + "AutoencoderKL": vae_parallel.TWO_D, + "AutoencoderKLFlux2": vae_parallel.TWO_D, + "AutoencoderKLWan": vae_parallel.WAN, + "AutoencoderKLQwenImage": vae_parallel.QWEN_IMAGE, + "AutoencoderKLHunyuanVideo": vae_parallel.HUNYUAN_VIDEO, + "AutoencoderKLHunyuanVideo15": vae_parallel.HUNYUAN_VIDEO_15, + "AutoencoderKLLTX2Video": vae_parallel.LTX2_VIDEO, +} + +# The encoder adapter each class needs. DistVAE reached full encoder coverage alongside its +# decoders, so a None here would mean an encoder adapter was lost rather than never written. +EXPECTED_ENCODERS = { + "AutoencoderKL": vae_parallel.TWO_D_ENCODER, + "AutoencoderKLFlux2": vae_parallel.TWO_D_ENCODER, + "AutoencoderKLWan": "WanEncoderAdapter", + "AutoencoderKLQwenImage": "QwenImageEncoderAdapter", + "AutoencoderKLHunyuanVideo": "HunyuanVideoEncoderAdapter", + "AutoencoderKLHunyuanVideo15": "HunyuanVideo15EncoderAdapter", + "AutoencoderKLLTX2Video": "LTX2VideoEncoderAdapter", +} + + +class TestDiffusersCompatibility(unittest.TestCase): + def test_an_unavailable_optional_vae_is_skipped(self): + with mock.patch.object(diffusers, "AutoencoderKLFlux2", None, create=True): + with self.assertRaises(unittest.SkipTest): + _vae_class("AutoencoderKLFlux2") + + def test_an_installed_vae_class_is_returned(self): + self.assertIs(_vae_class("AutoencoderKL"), diffusers.AutoencoderKL) + + +class TestDecoderAdapterChoice(unittest.TestCase): + + def test_every_vae_class_gets_the_adapter_it_needs(self): + for name, expected in EXPECTED.items(): + with self.subTest(vae=name): + vae = _vae_class(name)(**CONFIGS[name]) + self.assertEqual(vae_parallel.decoder_adapter_name(vae), expected) + + def test_a_vae_with_no_decoder_is_not_shardable(self): + class Bare: + pass + + self.assertIsNone(vae_parallel.decoder_adapter_name(Bare())) + + def test_an_ltx2_decoder_that_injects_noise_is_not_shardable(self): + # Every rank would draw noise for its own rows, and together they would not reconstruct + # what one rank draws, so the decode could not match an unsharded one. No released LTX-2 + # checkpoint enables this, which is why the shardable config above is the shipped shape. + vae = _vae_class("AutoencoderKLLTX2Video")( + **CONFIGS["AutoencoderKLLTX2Video"], decoder_inject_noise=True + ) + self.assertIsNone(vae_parallel.decoder_adapter_name(vae)) + + def test_a_two_d_decoder_without_group_norm_is_not_shardable(self): + # DecoderAdapter replaces conv_norm_out with a sharded GroupNorm and asserts it found one, + # so a decoder normalising some other way is out even with the right up blocks. + vae = diffusers.AutoencoderKL(**CONFIGS["AutoencoderKL"]) + vae.decoder.conv_norm_out = nn.Identity() + self.assertIsNone(vae_parallel.decoder_adapter_name(vae)) + + +class TestEncoderAdapterChoice(unittest.TestCase): + """Which adapter shards each class's encoder, read off the family its decoder names""" + + def test_every_vae_class_gets_the_encoder_adapter_it_needs(self): + for name, expected in EXPECTED_ENCODERS.items(): + with self.subTest(vae=name): + vae = _vae_class(name)(**CONFIGS[name]) + self.assertEqual(vae_parallel.encoder_adapter_name(vae), expected) + + def test_an_encoder_with_no_down_blocks_is_not_shardable(self): + class Bare: + pass + + self.assertIsNone(vae_parallel.encoder_adapter_name(Bare())) + + def test_the_two_halves_are_recognised_independently(self): + # Sharding either half replaces its blocks with adapters, so an encoder read off the + # decoder would come back unrecognised once the decoder had been done first, which is the + # order the supported integrations shard them in. + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + expected = vae_parallel.encoder_adapter_name(vae) + vae.decoder.conv_norm_out = nn.Identity() + vae.decoder.up_blocks = nn.ModuleList() + self.assertIsNone(vae_parallel.decoder_adapter_name(vae)) + self.assertEqual(vae_parallel.encoder_adapter_name(vae), expected) + + +class TestEncoderScaleFactor(unittest.TestCase): + """The spatial factor passed from a VAE's encoder structure into its adapter.""" + + def test_a_vae_that_does_not_patch_uses_its_spatial_ratio(self): + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + self.assertEqual(vae_parallel.encoder_scale_factor(vae), 8) + + def test_patching_is_divided_out(self): + # Cosmos 3's 16 is 8 from the encoder's convolutions and 2 from patching, and the adapter + # shards the convolutions. + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + vae.register_to_config(scale_factor_spatial=16, patch_size=2) + self.assertEqual(vae_parallel.encoder_scale_factor(vae), 8) + + def test_a_vae_with_no_spatial_ratio_falls_back(self): + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + vae.register_to_config(scale_factor_spatial=None) + self.assertEqual(vae_parallel.encoder_scale_factor(vae), 8) + + def test_flux2s_pair_of_patch_sizes_is_not_mistaken_for_a_ratio(self): + # Flux.2 names patch_size for how its latents are packed for the transformer, which its + # convolutions know nothing about, and names it as a pair. Dividing by that both divides + # by the wrong thing and cannot be compared against a number in the first place. + vae = _vae_class("AutoencoderKLFlux2")(**CONFIGS["AutoencoderKLFlux2"]) + self.assertEqual(vae.config.patch_size, (2, 2)) + self.assertEqual(vae_parallel.encoder_scale_factor(vae), 8) + + def test_a_two_d_encoder_is_counted_off_its_stages(self): + # These VAEs record no ratio, so the four-stage 8 has to be counted rather than assumed. + # A three-stage one narrows by 4, and sizing its bands by 8 would leave its last stage + # halving a band into a row belonging to the next rank. + config = dict(CONFIGS["AutoencoderKL"]) + config["block_out_channels"] = [8, 8, 16] + config["down_block_types"] = ["DownEncoderBlock2D"] * 3 + config["up_block_types"] = ["UpDecoderBlock2D"] * 3 + self.assertEqual( + vae_parallel.encoder_scale_factor(diffusers.AutoencoderKL(**config)), 4 + ) + + +class _StubAdapter(nn.Module): + """Stands in for a DistVAE adapter, which needs a process group to build""" + + def __init__(self, decoder, vae_group=None, **kwargs): + super().__init__() + self.wrapped = decoder + self.kwargs = kwargs + # WanDecoderAdapter crops by the factor it upsamples; the 2D one has no such step. + self.patchify = SimpleNamespace(scale_factor=1) + + +class TestWrappingReadsEveryVAEConfig(unittest.TestCase): + """Wrapping supports the configuration attribute layouts used by each VAE class.""" + + def _parallelize(self, vae): + with mock.patch.object(vae_parallel, "_adapter", return_value=_StubAdapter): + vae_parallel.parallelize_decoder(vae, vae_group=None) + return vae.decoder + + def test_every_shardable_vae_class_can_be_wrapped(self): + for name, adapter in EXPECTED.items(): + if adapter is None: + continue + with self.subTest(vae=name): + vae = _vae_class(name)(**CONFIGS[name]) + self.assertIsInstance(self._parallelize(vae), _StubAdapter) + + def test_a_patching_vae_tells_the_adapter_its_factor(self): + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + vae.register_to_config(patch_size=2) + self.assertEqual(self._parallelize(vae).patchify.scale_factor, 2) + + def test_flux_2s_patching_is_not_read_as_a_factor(self): + # Flux 2 declares patch_size (2, 2) for the pixel unshuffle at its boundary, which is not + # the single factor an adapter's patchify takes. + vae = _vae_class("AutoencoderKLFlux2")(**CONFIGS["AutoencoderKLFlux2"]) + self.assertEqual(self._parallelize(vae).patchify.scale_factor, 1) + + +class TestUnshardableIsRefused(unittest.TestCase): + + def test_it_names_the_contract_and_points_at_the_alternative(self): + # DistVAE now fits every VAE class an integration loads, so the refusal is provoked with + # a decoder taken out of the shape its adapter needs rather than with a real VAE. + vae = diffusers.AutoencoderKL(**CONFIGS["AutoencoderKL"]) + vae.decoder.conv_norm_out = nn.Identity() + decoder = vae.decoder + with self.assertRaises(ValueError) as caught: + vae_parallel.parallelize_decoder(vae, vae_group=None) + message = str(caught.exception) + self.assertIn("DistVAE cannot shard this VAE decoder", message) + self.assertIn("Diffusers VAE tiling", message) + # Refused before touching anything, so the caller is left a working decode to fall back on. + self.assertIs(vae.decoder, decoder) + + def test_encoding_is_refused_for_a_vae_no_encoder_adapter_fits(self): + # Provoked with an encoder taken out of the shape its adapter needs, since DistVAE fits + # every VAE class an integration loads. + vae = diffusers.AutoencoderKL(**CONFIGS["AutoencoderKL"]) + vae.encoder.down_blocks = nn.ModuleList() + encoder = vae.encoder + with self.assertRaises(ValueError) as caught: + vae_parallel.parallelize_encoder(vae, vae_group=None) + self.assertIn("Parallel VAE encoding is not available", str(caught.exception)) + # Refused before touching anything, so the caller is left a working encode. + self.assertIs(vae.encoder, encoder) + + +class TestBothHalvesShardTogether(unittest.TestCase): + """Every VAE class an integration loads has both halves replaced, in the caller's order + + Naming an adapter and installing it are different things: the adapters rebuild a half in + place, so the half done first no longer answers to the blocks it was recognised by. These + tests choose both names from intact blocks before constructing either adapter. + """ + + @classmethod + def setUpClass(cls): + cls.owns_group = not dist.is_initialized() + if cls.owns_group: + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "24118") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + dist.init_process_group(backend="gloo", init_method="env://") + + @classmethod + def tearDownClass(cls): + if cls.owns_group: + dist.destroy_process_group() + + def test_every_vae_class_shards_both_halves(self): + for name, config in CONFIGS.items(): + with self.subTest(vae=name): + vae = _vae_class(name)(**config).eval() + expected = vae_parallel.encoder_adapter_name(vae) + encoder, decoder = vae.encoder, vae.decoder + # Integrations shard the decoder first, which is what makes the order matter. + self.assertEqual( + vae_parallel.parallelize_decoder(vae, vae_group=None), + EXPECTED[name], + ) + self.assertEqual(vae_parallel.encoder_adapter_name(vae), expected) + self.assertEqual( + vae_parallel.parallelize_encoder(vae, vae_group=None), + EXPECTED_ENCODERS[name], + ) + self.assertIsNot(vae.decoder, decoder) + self.assertIsNot(vae.encoder, encoder) + + +class _ConvLosingAdapter(_StubAdapter): + """A stand-in that keeps none of the convolutions it replaced, as the real adapters do + + The count a recounting VAE makes is over the convolutions still answering to its own class, + and after sharding there are none: that is the whole of what has to be reproduced here. + """ + + def __init__(self, half, vae_group=None, **kwargs): + super().__init__(half, vae_group=vae_group, **kwargs) + self.wrapped = nn.Identity() + + +class TestCausalCacheLength(unittest.TestCase): + """A causal VAE's feature cache has to keep its length once its convolutions are replaced + + Qwen-Image sizes that cache by counting its causal convolutions on every call, so sharding + them away leaves an empty list and the first convolution to want its entry indexes off the + end. The failure is an IndexError from inside diffusers on the first decode, well after the + point where anything says which VAE stopped being decodable. + """ + + def _parallelize(self, vae, half="decoder"): + with mock.patch.object( + vae_parallel, "_adapter", return_value=_ConvLosingAdapter + ): + if half == "decoder": + vae_parallel.parallelize_decoder(vae, vae_group=None) + else: + vae_parallel.parallelize_encoder(vae, vae_group=None) + + def test_a_recounting_vae_still_caches_one_entry_per_convolution(self): + vae = _vae_class("AutoencoderKLQwenImage")(**CONFIGS["AutoencoderKLQwenImage"]) + vae.clear_cache() + expected = vae._conv_num + self.assertGreater(expected, 0) + + self._parallelize(vae) + vae.clear_cache() + + self.assertEqual(vae._conv_num, expected) + self.assertEqual(len(vae._feat_map), expected) + + def test_the_second_half_is_counted_before_the_first_is_replaced(self): + # Both halves are counted on the first call, because by the time the encoder is sharded + # the decoder has been, and a count taken then would already be short. + vae = _vae_class("AutoencoderKLQwenImage")(**CONFIGS["AutoencoderKLQwenImage"]) + vae.clear_cache() + expected = vae._enc_conv_num + self.assertGreater(expected, 0) + + self._parallelize(vae, "decoder") + self._parallelize(vae, "encoder") + vae.clear_cache() + + self.assertEqual(vae._enc_conv_num, expected) + self.assertEqual(len(vae._enc_feat_map), expected) + + def test_a_vae_that_counts_once_when_built_is_left_as_it_is(self): + # Wan caches its counts at construction, so nothing here has to hold them; this is that + # the holding does not disturb the ones that never needed it. + vae = _vae_class("AutoencoderKLWan")(**CONFIGS["AutoencoderKLWan"]) + vae.clear_cache() + expected = (vae._conv_num, vae._enc_conv_num) + + self._parallelize(vae) + vae.clear_cache() + + self.assertEqual((vae._conv_num, vae._enc_conv_num), expected) + + def test_a_vae_with_no_cache_at_all_is_wrapped_anyway(self): + vae = diffusers.AutoencoderKL(**CONFIGS["AutoencoderKL"]) + self.assertFalse(hasattr(vae, "clear_cache")) + self._parallelize(vae) + self.assertIsInstance(vae.decoder, _ConvLosingAdapter) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vae_tile_parallel.py b/test/test_vae_tile_parallel.py new file mode 100644 index 0000000..cc6e27a --- /dev/null +++ b/test/test_vae_tile_parallel.py @@ -0,0 +1,595 @@ +"""Tile-decode call distribution and assembly over a process group.""" + +import itertools +import os +import random +import socket +import unittest +import warnings +from datetime import timedelta +from types import SimpleNamespace +from typing import List, Optional, Tuple +from unittest import mock + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from distvae.vae import tile_parallel as vae_tile_parallel +from distvae.utils import ParallelContext + +# Shapes differing by rank share and by size, the way a tile grid's edges do: the last two are +# clipped, and no rank holds both of them. +SHAPES = ((2, 3), (2, 3), (2, 3), (2, 3), (1, 3), (2, 1)) + + +def _result(index: int, device=None) -> torch.Tensor: + """What call `index` returns: its own number, in a shape only it has""" + return torch.full(SHAPES[index], float(index), dtype=torch.float32, device=device) + + +def _load(weights, owner, world_size: int) -> List[int]: + """What each rank carries under an assignment""" + load = [0] * world_size + for n, weight in enumerate(weights): + load[owner[n]] += weight + return load + + +def _by_runs(weights, world_size: int) -> List[int]: + """The assignment before any tile is moved to level it""" + return [ + rank + for rank, (start, stop) in enumerate( + vae_tile_parallel.runs(weights, world_size) + ) + for _ in range(start, stop) + ] + + +def _every_split(tiles: int, world_size: int): + """Every way to cut `tiles` into `world_size` contiguous non-empty runs""" + for cuts in itertools.combinations(range(1, tiles), world_size - 1): + edges = (0,) + cuts + (tiles,) + yield list(zip(edges, edges[1:])) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _cores_allowed() -> int: + """The process CPU quota expressed as a core count. + + A cgroup quota can be lower than the affinity-visible count returned by `os.cpu_count()`. + Worker thread pools must fit the quota shared by all spawned ranks. + """ + try: + quota, period = open("/sys/fs/cgroup/cpu.max").read().split() + if quota != "max": + return max(1, int(quota) // int(period)) + except (OSError, ValueError): + pass + return os.cpu_count() or 1 + + +def _share_the_cores(world_size: int) -> None: + """Divide the available process threads among ranks sharing the host.""" + torch.set_num_threads(max(1, _cores_allowed() // world_size)) + + +def _backend_for(world_size: int) -> Tuple[str, Optional[str]]: + """Select a device collective when every rank has a device, otherwise a host collective.""" + if torch.cuda.is_available() and torch.cuda.device_count() >= world_size: + return dist.Backend.NCCL, "cuda" + return dist.Backend.GLOO, None + + +def _assert_tiled_decode_matches(got, expected, device) -> None: + if device is None: + torch.testing.assert_close(got, expected, rtol=0, atol=0) + return + # Device collectives may change floating-point accumulation order across ranks. + torch.testing.assert_close(got, expected, rtol=1e-5, atol=3e-6) + + +def _dispatch_in_a_group( + rank: int, world_size: int, port: int, calls_made: int +) -> None: + """One rank of the group, asserting for itself; mp.spawn re-raises what it fails on""" + _share_the_cores(world_size) + backend, device = _backend_for(world_size) + if device is not None: + torch.cuda.set_device(rank) + dist.init_process_group( + backend, + rank=rank, + world_size=world_size, + init_method=f"tcp://127.0.0.1:{port}", + timeout=timedelta(seconds=120), + ) + try: + made = [] + + def call(index): + made.append(index) + return _result(index, device) + + indices = list(range(calls_made)) + results = vae_tile_parallel.dispatch_over(dist.group.WORLD)( + [lambda index=index: call(index) for index in indices] + ) + + assert len(results) == len( + indices + ), f"{len(results)} results for {len(indices)} calls" + for index, result in zip(indices, results): + # Every rank ends holding every result, whoever computed it, in the order the calls + # were given rather than the order they were made. + torch.testing.assert_close(result, _result(index, device), rtol=0, atol=0) + + if len(indices) < world_size: + # Too few to divide, so each rank makes them all rather than leaving a rank with + # nothing to send. + assert made == indices, f"rank {rank} made {made}, not all of {indices}" + else: + assert made == indices[rank::world_size], f"rank {rank} made {made}" + finally: + dist.destroy_process_group() + + +class TestDispatchOverAGroup(unittest.TestCase): + """A rank makes its share of the calls and comes away with what every other rank made""" + + def _spawn(self, world_size: int, calls: int) -> None: + mp.spawn( + _dispatch_in_a_group, + args=(world_size, _free_port(), calls), + nprocs=world_size, + join=True, + ) + + def test_the_calls_are_divided_and_the_results_shared(self): + # Six calls over two ranks divides evenly, over four it does not, and the odd rank out + # sends one call's worth less than the others. + for world_size in (2, 4): + with self.subTest(world_size=world_size): + self._spawn(world_size, len(SHAPES)) + + def test_a_group_of_one_makes_every_call_itself(self): + self._spawn(1, len(SHAPES)) + + def test_fewer_calls_than_ranks_leaves_every_rank_making_them_all(self): + self._spawn(4, 3) + + +class TestParallelContext(unittest.TestCase): + def test_metadata_is_distvae_owned_and_immutable(self): + vae = SimpleNamespace() + context = ParallelContext(group=None, rank=0, world_size=1, patch_dim=-2) + + vae_tile_parallel.mark(vae, context) + + self.assertIs(vae_tile_parallel.context_of(vae), context) + self.assertIsNone(vae_tile_parallel.group_of(vae)) + self.assertEqual(vars(vae), {"_distvae_tile_parallel_context": context}) + + def test_dispatch_uses_a_captured_context_without_global_lookup(self): + context = ParallelContext(group=None, rank=0, world_size=1, patch_dim=-2) + calls = [lambda: torch.tensor(1), lambda: torch.tensor(2)] + + results = vae_tile_parallel.dispatch_over(context)(calls) + + self.assertEqual([result.item() for result in results], [1, 2]) + + +class TestDiffusersCompatibility(unittest.TestCase): + def test_an_unavailable_run_vae_is_skipped_before_spawning(self): + import diffusers + + with mock.patch.object(diffusers, "AutoencoderKLQwenImage", None, create=True): + with self.assertRaises(unittest.SkipTest): + _require_run_vae(self, "AutoencoderKLQwenImage") + + def test_an_installed_tiling_api_is_accepted(self): + _require_run_vae(self, "AutoencoderKL") + + +RUN_VAES = { + "AutoencoderKL": ( + dict( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=256, + down_block_types=["DownEncoderBlock2D"] * 4, + up_block_types=["UpDecoderBlock2D"] * 4, + ), + False, + ), + "AutoencoderKLWan": ( + dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1), + True, + ), + "AutoencoderKLQwenImage": ( + dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1), + True, + ), +} + + +def _require_run_vae(testcase, name: str) -> None: + """Skip one parent subtest before spawning when its Diffusers VAE cannot tile.""" + import diffusers + + cls = getattr(diffusers, name, None) + if cls is None: + testcase.skipTest(f"{name} is not in diffusers {diffusers.__version__}") + vae = cls(**RUN_VAES[name][0]) + if not hasattr(vae, "use_tiling"): + testcase.skipTest(f"diffusers {diffusers.__version__} cannot tile {name}") + + +# A quarter-overlap grid spanning two windows down and three across has 3x4 tiles. Distributing +# those tiles over four ranks makes each contiguous run start or end mid-row, exercising run +# boundaries that cannot be represented by whole-row assignment. +WINDOWS_DOWN, WINDOWS_ACROSS = 2, 3 + + +def _blend(deep_down: int, deep_across: int): + """A Blend carrying nothing but its depths, which is all `_wanted` reads""" + return vae_tile_parallel.Blend( + down=None, + across=None, + deep_down=deep_down, + deep_across=deep_across, + crop=None, + tile_down=64, + tile_across=64, + ) + + +def _tiled_vae(name: str, device=None, overlap: Optional[Tuple[int, int]] = None): + """The same small VAE and latents on every rank, at a window several tiles across""" + import diffusers + + from distvae.vae import tiling as vae_tiling + + kwargs, video = RUN_VAES[name] + # Seeded because every rank builds its own and they have to agree to the bit. + torch.manual_seed(0) + cls = getattr(diffusers, name, None) + if cls is None: + raise RuntimeError( + f"{name} disappeared after the parent process availability check" + ) + vae = cls(**kwargs).eval() + vae.enable_tiling() + native = vae_tiling.tile_shape(vae) + shape = tuple(axis // 4 for axis in native) + plan = vae_tiling.tile_shape_plan(vae, *shape) + assert plan is not None, f"{name} cannot use exact tile shape {shape}" + vae_tiling.apply_tile_plan(vae, plan) + if overlap is not None: + step = vae_tiling.tile_overlap_plan(vae, *overlap) + assert step is not None, f"{name} cannot step its tiles at {overlap}" + vae_tiling.apply_tile_plan(vae, step) + + window = vae_tiling.latent_rows(vae, plan) + down, across = window * WINDOWS_DOWN, window * WINDOWS_ACROSS + shape = (1, 4, 2, down, across) if video else (1, 4, down, across) + torch.manual_seed(1) + latents = torch.randn(*shape) + if device is not None: + vae, latents = vae.to(device), latents.to(device) + return vae, latents + + +def _runs_in_a_group( + rank: int, + world_size: int, + port: int, + name: str, + overlap: Optional[Tuple[int, int]] = None, +) -> None: + """One rank blending its own run, checked against the whole grid blended by one rank""" + from distvae.vae import tile_parallel as vae_tile_parallel + from distvae.vae import tiling as vae_tiling + + _share_the_cores(world_size) + backend, device = _backend_for(world_size) + if device is not None: + torch.cuda.set_device(rank) + device = f"cuda:{rank}" + dist.init_process_group( + backend, + rank=rank, + world_size=world_size, + init_method=f"tcp://127.0.0.1:{port}", + timeout=timedelta(seconds=300), + ) + try: + vae, latents = _tiled_vae(name, device, overlap) + with torch.no_grad(): + expected = vae.tiled_decode(latents).sample + + dispatch, assemble = vae_tile_parallel.sharing(dist.group.WORLD) + decode = vae_tiling.tiled_decode_for(vae, dispatch, assemble) + assert decode is not None, f"no reimplemented loop for {name}" + got = decode(latents).sample + + assert got.shape == expected.shape, f"{got.shape} != {expected.shape}" + _assert_tiled_decode_matches(got, expected, device) + finally: + dist.destroy_process_group() + + +class TestBackendAgreement(unittest.TestCase): + def test_host_collective_requires_bit_exact_output(self): + with self.assertRaises(AssertionError): + _assert_tiled_decode_matches( + torch.tensor([1.0]), torch.tensor([1.0 + 1e-6]), device=None + ) + + def test_device_collective_allows_accumulation_order_rounding(self): + _assert_tiled_decode_matches( + torch.tensor([1.0]), torch.tensor([1.0 + 2.1e-6]), device="cuda" + ) + + +class TestRuns(unittest.TestCase): + """Tiles split into a contiguous run per rank, blended locally, gathered back whole""" + + def test_fewer_tiles_than_ranks_warns_that_every_rank_repeats_the_decode(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with mock.patch.object( + vae_tile_parallel, "_distributed", return_value=(None, 0, 8) + ): + result = vae_tile_parallel.assemble_in_runs( + None, + rows=2, + columns=2, + decode=mock.Mock(), + blend=_blend(1, 1), + weights=[1, 1, 1, 1], + ) + + self.assertIsNone(result) + self.assertEqual(len(caught), 1) + self.assertIs(caught[0].category, RuntimeWarning) + self.assertRegex( + str(caught[0].message), + "4 tiles for 8 ranks.*every rank will decode all 4 tiles", + ) + + def test_tiles_of_equal_weight_are_split_as_evenly_as_they_divide(self): + # Evenly means no run heavier than it has to be, which for equal weights is the share + # rounded up. It does not mean the runs are the same length: ten tiles over three ranks + # divides 4, 4, 2, and the 2 costs nothing because the 4s are what the decode waits for. + for tiles, world_size in ((4, 2), (5, 4), (3, 1), (9, 4), (10, 3), (12, 5)): + split = vae_tile_parallel.runs([1] * tiles, world_size) + self.assertEqual(len(split), world_size, split) + self.assertEqual(split[0][0], 0, split) + self.assertEqual(split[-1][1], tiles, split) + for (_, stop), (start, _) in zip(split, split[1:]): + self.assertEqual(stop, start, split) + longest = max(stop - start for start, stop in split) + self.assertEqual(longest, -(-tiles // world_size), split) + + def test_the_heaviest_run_is_the_lightest_it_can_be(self): + # Against every contiguous split there is, at sizes small enough to enumerate them all. + random.seed(7) + for tiles in range(2, 9): + for world_size in range(2, min(tiles, 5) + 1): + for _ in range(20): + weights = [random.randint(1, 9) for _ in range(tiles)] + mine = max( + sum(weights[start:stop]) + for start, stop in vae_tile_parallel.runs(weights, world_size) + ) + best = min( + max(sum(weights[start:stop]) for start, stop in split) + for split in _every_split(tiles, world_size) + ) + self.assertEqual(mine, best, f"{weights} over {world_size}") + + def test_the_lighter_tiles_at_the_end_do_not_all_land_on_one_rank(self): + # The case that a run split by count got wrong: the latent bounds clip the last row and + # the last column, so an equal count of tiles is an unequal amount of work. Weighing the + # split is not enough on its own here - the best contiguous cut of these nine is 31 + # against 23 - so it takes the levelling to move one tile across and even them up. + weights = [9, 9, 4, 9, 9, 4, 4, 4, 2] + owner = vae_tile_parallel.shares(weights, 2) + held = [ + sum(weight for n, weight in enumerate(weights) if owner[n] == rank) + for rank in range(2) + ] + self.assertLessEqual(max(held) / min(held), 1.1, held) + + def test_levelling_never_leaves_a_rank_worse_off_than_the_runs_it_started_from( + self, + ): + random.seed(13) + for tiles in range(2, 20): + for world_size in range(2, min(tiles, 6) + 1): + for _ in range(10): + weights = [random.randint(1, 9) for _ in range(tiles)] + owner = vae_tile_parallel.shares(weights, world_size) + self.assertEqual(set(owner), set(range(world_size)), weights) + self.assertEqual(len(owner), tiles) + self.assertLessEqual( + max(_load(weights, owner, world_size)), + max(_load(weights, _by_runs(weights, world_size), world_size)), + f"{weights} over {world_size}", + ) + + def test_large_grids_keep_the_weighted_runs_when_one_round_exceeds_the_budget( + self, + ): + world_size = 4 + weights = [128 * 128] * 483 + [64 * 64] + candidates = ( + len(weights) * (world_size - 1) + + len(weights) * (len(weights) - 1) // 2 + ) + self.assertGreater(candidates, vae_tile_parallel.MAX_LEVEL_CANDIDATES) + + vae_tile_parallel._shares.cache_clear() + try: + self.assertEqual( + vae_tile_parallel.shares(weights, world_size), + _by_runs(weights, world_size), + ) + finally: + vae_tile_parallel._shares.cache_clear() + + def test_a_tile_moves_across_where_a_run_cannot_be_levelled(self): + # The grid measured on four ranks: nine tiles, the last row and column clipped. Contiguity + # alone leaves the heaviest rank a quarter above the lightest possible; a tile moving + # across takes that back. + weights = [16384, 16384, 8192, 16384, 16384, 8192, 8192, 8192, 4096] + by_runs = max(_load(weights, _by_runs(weights, 4), 4)) + levelled = max(_load(weights, vae_tile_parallel.shares(weights, 4), 4)) + self.assertEqual(by_runs, 32768) + self.assertEqual(levelled, 28672) + + def test_a_swap_escapes_a_move_only_local_optimum(self): + weights = [3136, 2464, 2464, 1936] + + owner = vae_tile_parallel.shares(weights, 2) + + self.assertEqual(max(_load(weights, owner, 2)), 5072) + self.assertEqual(set(owner), {0, 1}) + + def test_levelling_can_cross_a_global_makespan_plateau(self): + extents = [28, 28, 28, 22] + weights = [height * width for height in extents for width in extents] + + owner = vae_tile_parallel.shares(weights, 4) + + self.assertEqual(max(_load(weights, owner, 4)), 2836) + self.assertEqual(set(owner), {0, 1, 2, 3}) + + def test_levelling_is_deterministic(self): + extents = [28, 28, 28, 22] + weights = [height * width for height in extents for width in extents] + expected = vae_tile_parallel.shares(weights, 4) + + for _ in range(10): + self.assertEqual(vae_tile_parallel.shares(weights, 4), expected) + + def test_equal_weights_reuse_the_cached_assignment(self): + weights = [101, 103, 107, 109, 113, 127, 131] + + with mock.patch.object( + vae_tile_parallel, "runs", wraps=vae_tile_parallel.runs + ) as split: + first = vae_tile_parallel.shares(weights, 3) + second = vae_tile_parallel.shares(list(weights), 3) + + split.assert_called_once() + self.assertEqual(first, second) + self.assertIsNot(first, second) + + def test_equal_balance_prefers_fewer_tiles_displaced_from_the_runs(self): + # Moving tile 1 and swapping tiles 0 and 2 both produce loads [2, 2, 4]. The move leaves + # only one tile outside its original run, while the swap leaves two. + self.assertEqual( + vae_tile_parallel.shares([1, 1, 2, 4], 3), + [0, 0, 1, 2], + ) + + def test_equal_moves_prefer_a_tile_beside_its_receiving_rank(self): + # Any of rank 0's three unit tiles gives loads [3, 2, 2] on rank 1. Tile 2 touches rank + # 1's run already, so moving it adds fewer remote boundaries than moving tile 0 or 1. + self.assertEqual( + vae_tile_parallel.shares([1, 1, 1, 1, 3], 3), + [0, 0, 1, 1, 2], + ) + + def test_the_edges_asked_for_are_the_edges_the_blending_reaches_for(self): + random.seed(17) + for rows in range(1, 6): + for columns in range(1, 6): + for world_size in range(2, min(rows * columns, 5) + 1): + weights = [random.randint(1, 9) for _ in range(rows * columns)] + owner = vae_tile_parallel.shares(weights, world_size) + blend = _blend(1, 1) + wanted = vae_tile_parallel._wanted(owner, columns, blend) + for n in range(rows * columns): + row, column = divmod(n, columns) + reaches = [] + if row > 0 and owner[n - columns] != owner[n]: + reaches += [n - columns] + ( + [n - columns - 1] if column else [] + ) + if column > 0 and owner[n - 1] != owner[n]: + reaches += [n - 1] + ([n - columns - 1] if row else []) + for at in reaches: + self.assertIn( + at, wanted, f"{rows}x{columns}/{world_size} at {n}" + ) + + def test_a_blend_no_rows_deep_asks_for_no_edges_on_that_axis(self): + # A wide enough stride leaves the tiles touching rather than overlapping, and then there + # is nothing to blend and nothing to send. Worth its own case because a depth of zero is + # not inert where the edges are sliced: `tile[..., -0:, :]` is the whole tile, so a loop + # that only skipped the blending would still put the entire grid on the wire. + weights = [1] * 12 + owner = vae_tile_parallel.shares(weights, 4) + self.assertEqual(vae_tile_parallel._wanted(owner, 4, _blend(0, 0)), set()) + # One axis at a time, since the strides are set per axis and only one of them can run + # out. Each asks for less than both do - the corner tile a blended edge is rebuilt from + # is only reached when both blends happen - and neither asks for anything the pair does + # not. + both = vae_tile_parallel._wanted(owner, 4, _blend(1, 1)) + for depths in ((1, 0), (0, 1)): + wanted = vae_tile_parallel._wanted(owner, 4, _blend(*depths)) + self.assertTrue(wanted, depths) + self.assertTrue(wanted < both, depths) + + def test_every_rank_holds_a_tile_and_every_tile_is_held_once(self): + random.seed(11) + for tiles in range(1, 12): + for world_size in range(1, min(tiles, 6) + 1): + weights = [random.randint(1, 9) for _ in range(tiles)] + split = vae_tile_parallel.runs(weights, world_size) + self.assertEqual(len(split), world_size, f"{weights}/{world_size}") + self.assertTrue(all(stop > start for start, stop in split), split) + covered = [n for start, stop in split for n in range(start, stop)] + self.assertEqual(covered, list(range(tiles)), f"{tiles}/{world_size}") + + def test_a_run_decode_is_what_one_rank_blending_everything_gives(self): + for name in RUN_VAES: + for world_size in (2, 4): + with self.subTest(vae=name, world_size=world_size): + _require_run_vae(self, name) + mp.spawn( + _runs_in_a_group, + args=(world_size, _free_port(), name, None), + nprocs=world_size, + join=True, + ) + + def test_tiles_that_do_not_overlap_at_all_still_assemble(self): + # A caller can widen the stride until the tiles touch rather than overlap, and + # then the blends are no rows deep. Checked against the same VAE's own loop at the same + # stride, so what this proves is that dividing the work changes nothing: a depth of zero + # has to mean no blending and no edges, and not the whole tile taken as its own edge. + for name in RUN_VAES: + with self.subTest(vae=name): + _require_run_vae(self, name) + mp.spawn( + _runs_in_a_group, + args=(4, _free_port(), name, (0, 0)), + nprocs=4, + join=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vae_tiling.py b/test/test_vae_tiling.py new file mode 100644 index 0000000..0cfe7fa --- /dev/null +++ b/test/test_vae_tiling.py @@ -0,0 +1,1048 @@ +import unittest +from unittest import mock + +from distvae.vae import tiling as vae_tiling + + +class StubVAE: + """Minimal VAE stub that stores the supplied tiling attributes.""" + + def __init__(self, **attrs): + for name, value in attrs.items(): + setattr(self, name, value) + + +def _diffusers_vae(testcase, name, kwargs, *, require_tiling=False): + """Build an installed VAE, skipping only this test/subtest when its API is unavailable.""" + import diffusers + + cls = getattr(diffusers, name, None) + if cls is None: + testcase.skipTest(f"{name} is not in diffusers {diffusers.__version__}") + vae = cls(**kwargs).eval() + if require_tiling and not hasattr(vae, "use_tiling"): + testcase.skipTest(f"diffusers {diffusers.__version__} cannot tile {name}") + return vae + + +def legacy_pair_vae(): + """Stub with shared square sample and latent windows and one overlap factor.""" + return StubVAE( + tile_sample_min_size=256, tile_latent_min_size=32, tile_overlap_factor=0.25 + ) + + +def stride_vae(): + """Stub with per-axis sample windows and explicit pixel strides.""" + return StubVAE( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_sample_stride_height=192, + tile_sample_stride_width=192, + spatial_compression_ratio=8, + ) + + +def overlap_hw_vae(): + """Stub with per-axis sample windows, latent windows, and overlap factors.""" + return StubVAE( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_latent_min_height=32, + tile_latent_min_width=32, + tile_overlap_factor_height=0.25, + tile_overlap_factor_width=0.25, + ) + + +def asymmetric_vae(): + """Stub with rectangular per-axis sample and latent windows.""" + return StubVAE( + tile_sample_min_height=240, + tile_sample_min_width=360, + tile_latent_min_height=30, + tile_latent_min_width=45, + tile_overlap_factor_height=1 / 6, + tile_overlap_factor_width=0.2, + ) + + +def overlap_factor_vae(sample=256): + """Stub with shared square tiling attributes and blend methods.""" + return StubVAE( + tile_sample_min_size=sample, + tile_latent_min_size=sample // 8, + tile_overlap_factor=0.25, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + ) + + +def overlap_keyed_vae(): + """Stub with per-axis windows, one shared overlap factor, and blend methods.""" + return StubVAE( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_latent_min_height=16, + tile_latent_min_width=16, + tile_overlap_factor=0.25, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + ) + + +def per_axis_overlap_vae(): + """Stub with unequal per-axis overlap factors and latent windows.""" + return StubVAE( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_latent_min_height=32, + tile_latent_min_width=40, + tile_overlap_factor_height=0.25, + tile_overlap_factor_width=0.2, + ) + + +class TestDiffusersCompatibility(unittest.TestCase): + def test_an_unavailable_optional_vae_is_skipped(self): + import diffusers + + with mock.patch.object(diffusers, "AutoencoderKLFlux2", None, create=True): + with self.assertRaises(unittest.SkipTest): + _diffusers_vae(self, "AutoencoderKLFlux2", {}, require_tiling=True) + + def test_a_class_without_a_tiling_api_is_skipped(self): + import diffusers + + vae = StubVAE() + vae.eval = lambda: vae + with mock.patch.object(diffusers, "AutoencoderKLFlux2", return_value=vae): + with self.assertRaises(unittest.SkipTest): + _diffusers_vae(self, "AutoencoderKLFlux2", {}, require_tiling=True) + + +class TestSupportProbe(unittest.TestCase): + + def test_the_method_alone_does_not_count_as_support(self): + # Diffusers provides enable_tiling through a mixin whether or not the class implements it, + # so a VAE can carry the method and still raise NotImplementedError when called. + unsupported = StubVAE(enable_tiling=lambda: None) + with self.assertRaises(ValueError): + vae_tiling.require_vae_support(unsupported, "tiling", "--enable_tiling") + + def test_the_state_flag_counts_as_support(self): + vae_tiling.require_vae_support( + StubVAE(use_tiling=False), "tiling", "--enable_tiling" + ) + vae_tiling.require_vae_support( + StubVAE(use_slicing=False), "slicing", "--enable_slicing" + ) + + +class TestTilePaddingError(unittest.TestCase): + """The padding failure a too-thin tile causes, told apart from failures with other causes""" + + # Verbatim from AutoencoderKLLTX2Video decoding a 16x16 latent at a 128px window. + REAL = ( + "Argument #4: Padding size should be less than the corresponding input dimension, " + "but got: padding (1, 1) at dimension 4 of input [1, 8, 3, 4, 1]" + ) + + def test_the_padding_failure_is_recognised(self): + self.assertTrue(vae_tiling.is_tile_padding_error(RuntimeError(self.REAL))) + + def test_other_decode_failures_are_not(self): + for message in ( + "expected scalar type BFloat16 but found Float", + "Expected all tensors to be on the same device", + "shape '[1, 8, 16, 16]' is invalid for input of size 1024", + "CUDA error: an illegal memory access was encountered", + ): + with self.subTest(error=message): + self.assertFalse( + vae_tiling.is_tile_padding_error(RuntimeError(message)) + ) + + +class TestTileShape(unittest.TestCase): + + def test_reads_the_pixel_shape_of_each_family(self): + self.assertEqual(vae_tiling.tile_shape(legacy_pair_vae()), (256, 256)) + self.assertEqual(vae_tiling.tile_shape(stride_vae()), (256, 256)) + self.assertEqual(vae_tiling.tile_shape(overlap_hw_vae()), (256, 256)) + + def test_a_vae_without_a_shape_reports_none(self): + vae = StubVAE(tile_overlap_h=0.25) + self.assertIsNone(vae_tiling.tile_shape(vae)) + self.assertIsNone(vae_tiling.tile_shape_plan(vae, 128, 128)) + + def test_a_native_rectangle_is_preserved(self): + self.assertEqual(vae_tiling.tile_shape(asymmetric_vae()), (240, 360)) + + def test_spatial_ratio_falls_back_from_config_to_the_module(self): + self.assertEqual(vae_tiling.spatial_ratio(stride_vae()), 8) + self.assertIsNone(vae_tiling.spatial_ratio(legacy_pair_vae())) + + +class TestSquareTileShapePlan(unittest.TestCase): + + def test_every_attribute_is_rescaled_by_the_same_factor(self): + plan = vae_tiling.tile_shape_plan(stride_vae(), 128, 128) + self.assertEqual( + plan, + { + "tile_sample_min_height": 128, + "tile_sample_min_width": 128, + "tile_sample_stride_height": 96, + "tile_sample_stride_width": 96, + }, + ) + + def test_a_window_that_does_not_divide_whole_is_refused(self): + # 100px would put the latent window at 12.5, which no VAE can hold. + self.assertIsNone(vae_tiling.tile_shape_plan(legacy_pair_vae(), 100, 100)) + + def test_an_overlap_that_does_not_land_whole_is_refused(self): + # 200px gives a latent window of 25, and 25 x 0.75 truncates to a stride the pixel crop + # does not agree with, which assembles an image of the wrong size. + self.assertIsNone(vae_tiling.tile_shape_plan(legacy_pair_vae(), 200, 200)) + self.assertIsNone(vae_tiling.tile_shape_plan(overlap_hw_vae(), 200, 200)) + self.assertIsNotNone(vae_tiling.tile_shape_plan(legacy_pair_vae(), 192, 192)) + + def test_each_axis_must_keep_its_latent_step_and_pixel_crop_consistent(self): + # The width maps 256 pixels to 40 latents, a non-integral 6.4x ratio. Its latent stride + # and pixel crop cannot describe the same distance, even though 40 x 0.8 is whole. + self.assertIsNone( + vae_tiling.tile_shape_plan(per_axis_overlap_vae(), 256, 256) + ) + self.assertIsNone( + vae_tiling.tile_shape_plan(per_axis_overlap_vae(), 224, 224) + ) + + def test_an_unkeyed_overlap_fraction_covers_both_axes(self): + vae = StubVAE( + tile_sample_min_height=256, + tile_sample_min_width=256, + tile_latent_min_height=16, + tile_latent_min_width=16, + tile_overlap_factor=0.25, + ) + self.assertIsNotNone(vae_tiling.tile_shape_plan(vae, 64, 64)) + # 32px puts each latent window at 2, and 2 x 0.75 truncates to a stride of 1. + self.assertIsNone(vae_tiling.tile_shape_plan(vae, 32, 32)) + + def test_a_stride_below_one_latent_pixel_is_refused(self): + # 8px would leave a 6px stride, under this VAE's 8px latent pixel, and diffusers steps + # through the latents in a range() that would then be empty. + self.assertIsNone(vae_tiling.tile_shape_plan(stride_vae(), 8, 8)) + + def test_a_window_above_the_default_still_plans(self): + # _apply_vae_tile_size declines these itself, having the config to say why. + plan = vae_tiling.tile_shape_plan(stride_vae(), 512, 512) + self.assertEqual(plan["tile_sample_stride_height"], 384) + + +class TestTileShapePlan(unittest.TestCase): + + def test_a_legacy_square_window_can_be_planned_rectangularly(self): + vae = legacy_pair_vae() + + plan = vae_tiling.tile_shape_plan(vae, 128, 192) + + self.assertEqual( + plan, + { + "tile_sample_min_size": 128, + "tile_sample_min_height": 128, + "tile_sample_min_width": 192, + "tile_latent_min_size": 16, + "tile_latent_min_height": 16, + "tile_latent_min_width": 24, + }, + ) + vae_tiling.apply_tile_plan(vae, plan) + self.assertEqual( + vae_tiling.overlap_windows(vae), ((16, 24), (128, 192)) + ) + + def test_a_stored_stride_is_rescaled_independently_on_each_axis(self): + self.assertEqual( + vae_tiling.tile_shape_plan(stride_vae(), 128, 192), + { + "tile_sample_min_height": 128, + "tile_sample_min_width": 192, + "tile_sample_stride_height": 96, + "tile_sample_stride_width": 144, + }, + ) + + def test_either_non_integral_axis_rejects_the_rectangle(self): + self.assertIsNone( + vae_tiling.tile_shape_plan(legacy_pair_vae(), 128, 100) + ) + # The scaled width stride is 99 pixels, which cannot step an 8-pixel + # latent grid without truncating. + self.assertIsNone(vae_tiling.tile_shape_plan(stride_vae(), 128, 132)) + + def test_the_shape_reader_never_squares_a_native_rectangle(self): + self.assertEqual(vae_tiling.tile_shape(legacy_pair_vae()), (256, 256)) + self.assertEqual(vae_tiling.tile_shape(asymmetric_vae()), (240, 360)) + + def test_square_planning_uses_the_rectangular_mechanics(self): + self.assertEqual( + vae_tiling.tile_shape_plan(legacy_pair_vae(), 128, 128), + { + "tile_sample_min_size": 128, + "tile_sample_min_height": 128, + "tile_sample_min_width": 128, + "tile_latent_min_size": 16, + "tile_latent_min_height": 16, + "tile_latent_min_width": 16, + }, + ) + + def test_rectangular_legacy_windows_install_a_local_replacement(self): + import torch + import torch.nn.functional as functional + + vae = overlap_factor_vae() + vae.decoder = lambda tile: functional.interpolate( + tile, scale_factor=8, mode="nearest" + ) + plan = vae_tiling.tile_shape_plan(vae, 128, 192) + vae_tiling.apply_tile_plan(vae, plan) + + decode = vae_tiling.tiled_decode_for(vae) + + self.assertIsNotNone(decode) + sample = decode(torch.randn(1, 4, 24, 32)).sample + self.assertEqual(sample.shape, (1, 4, 192, 256)) + + def test_legacy_threshold_enters_tiling_when_the_smaller_axis_is_exceeded(self): + import torch + from diffusers.models.autoencoders.vae import DecoderOutput + + kwargs, _, _ = TestEverySupportedVAE.VAES["AutoencoderKL"] + vae = _diffusers_vae(self, "AutoencoderKL", kwargs, require_tiling=True) + vae.enable_tiling() + plan = vae_tiling.tile_shape_plan(vae, 128, 384) + vae_tiling.apply_tile_plan(vae, plan) + vae.tiled_decode = mock.Mock( + return_value=DecoderOutput(sample=torch.empty(1, 4, 136, 192)) + ) + + vae._decode(torch.randn(1, 4, 17, 24)) + + self.assertEqual(vae.tile_latent_min_size, 16) + vae.tiled_decode.assert_called_once() + + def test_native_keyed_rectangles_keep_the_upstream_local_loop(self): + vae = overlap_keyed_vae() + plan = vae_tiling.tile_shape_plan(vae, 128, 384) + vae_tiling.apply_tile_plan(vae, plan) + + self.assertIsNotNone(vae_tiling.tiled_decode_for(vae)) + + +class TestLatentRows(unittest.TestCase): + """How many rows a planned tile leaves available for spatial sharding""" + + def test_rectangular_window_uses_height_instead_of_the_smaller_axis(self): + vae = legacy_pair_vae() + self.assertEqual( + vae_tiling.latent_rows( + vae, vae_tiling.tile_shape_plan(vae, 256, 64) + ), + 32, + ) + + def test_rows_come_from_the_latent_window_where_the_vae_carries_one(self): + vae = legacy_pair_vae() + self.assertEqual( + vae_tiling.latent_rows( + vae, vae_tiling.tile_shape_plan(vae, 128, 128) + ), + 16, + ) + + def test_rows_come_from_the_compression_ratio_otherwise(self): + vae = stride_vae() + self.assertEqual( + vae_tiling.latent_rows( + vae, vae_tiling.tile_shape_plan(vae, 128, 128) + ), + 16, + ) + + def test_a_vae_that_says_neither_reports_none(self): + vae = StubVAE(tile_sample_min_height=256, tile_sample_min_width=256) + self.assertIsNone( + vae_tiling.latent_rows( + vae, vae_tiling.tile_shape_plan(vae, 128, 128) + ) + ) + + def test_with_no_plan_the_vae_s_own_window_is_the_plan(self): + # DistVAE must validate the VAE's default window when tiling was enabled before the + # integration applied an explicit plan; every tile is subsequently split across ranks. + self.assertEqual(vae_tiling.latent_rows(legacy_pair_vae()), 32) + self.assertEqual(vae_tiling.latent_rows(stride_vae()), 32) + self.assertIsNone(vae_tiling.latent_rows(StubVAE(tile_overlap_factor=0.25))) + + def test_a_plan_is_read_ahead_of_what_the_vae_still_holds(self): + # The plan describes what is about to be set, so a caller weighing one against the ranks + # has to be answered about the plan and not about the window it is replacing. + vae = legacy_pair_vae() + self.assertEqual( + vae_tiling.latent_rows( + vae, vae_tiling.tile_shape_plan(vae, 128, 128) + ), + 16, + ) + self.assertEqual(vae_tiling.latent_rows(vae), 32) + + +class TestEverySupportedVAE(unittest.TestCase): + """Every supported VAE accepts a resized tile window without changing output size""" + + # Minimal configs preserve each class's decoder topology. LTX2 pins its compression ratio + # because the default describes more encoder stages than its decoder upsamples. + VAES = { + "AutoencoderKL": ( + dict( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=256, + down_block_types=["DownEncoderBlock2D"] * 4, + up_block_types=["UpDecoderBlock2D"] * 4, + ), + False, + 4, + ), + "AutoencoderKLFlux2": ( + dict( + block_out_channels=[8, 8, 16, 16], + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=256, + ), + False, + 4, + ), + "AutoencoderKLWan": ( + dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1), + True, + 4, + ), + "AutoencoderKLQwenImage": ( + dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1), + True, + 4, + ), + "AutoencoderKLHunyuanVideo": ( + dict( + block_out_channels=(8, 8, 16, 16), + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + ), + True, + 4, + ), + "AutoencoderKLHunyuanVideo15": ( + dict( + block_out_channels=(8, 8, 16, 16, 16), + layers_per_block=1, + latent_channels=4, + ), + True, + 4, + ), + "AutoencoderKLLTX2Video": ( + dict( + block_out_channels=(8, 16, 32, 32), + latent_channels=8, + layers_per_block=(1, 1, 1, 1, 1), + spatial_compression_ratio=32, + ), + True, + 8, + ), + } + # Large enough that the output is several tiles across once the window is halved, since a + # decode that fits in one tile would pass without tiling anything. + LATENT_GRID = 16 + + def test_a_halved_window_decodes_to_the_same_size(self): + import torch + + for name, (kwargs, video, channels) in self.VAES.items(): + with self.subTest(vae=name): + vae = _diffusers_vae(self, name, kwargs, require_tiling=True) + + grid = self.LATENT_GRID + shape = ( + (1, channels, 1, grid, grid) if video else (1, channels, grid, grid) + ) + torch.manual_seed(0) + latents = torch.randn(*shape) + with torch.no_grad(): + vae.disable_tiling() + expected = vae.decode(latents).sample.shape[-2:] + + # Same order as the caller: turn tiling on, then size its window. + vae.enable_tiling() + window = vae_tiling.tile_shape(vae) + self.assertIsNotNone( + window, f"{name} tiles but exposes no window this can read" + ) + shape = tuple(axis // 2 for axis in window) + plan = vae_tiling.tile_shape_plan(vae, *shape) + self.assertIsNotNone( + plan, f"{name} refused exact tile shape {shape}" + ) + for attr, value in plan.items(): + setattr(vae, attr, value) + with torch.no_grad(): + got = vae.decode(latents).sample.shape[-2:] + self.assertEqual( + got, expected, f"{name} decoded at tile shape {shape}" + ) + + +class TestTileOverlap(unittest.TestCase): + """The exact output-pixel overlap between neighbouring tiles.""" + + def test_stride_and_overlap_factor_layouts_report_absolute_pixels(self): + self.assertEqual(vae_tiling.tile_overlap(legacy_pair_vae()), (64, 64)) + self.assertEqual(vae_tiling.tile_overlap(stride_vae()), (64, 64)) + self.assertIsNone(vae_tiling.tile_overlap(StubVAE(tile_sample_min_size=256))) + + def test_reporting_a_step_is_not_knowing_what_moving_it_does(self): + self.assertIsNotNone(vae_tiling.tile_overlap(stride_vae())) + self.assertIsNone(vae_tiling.tile_overlap_plan(stride_vae(), 32, 32)) + self.assertIsNone(vae_tiling.tile_overlap_plan(overlap_hw_vae(), 32, 32)) + + def test_a_column_strip_requires_zero_overlap_on_its_inactive_axis(self): + vae = StubVAE( + tile_sample_min_height=120, + tile_sample_min_width=128, + tile_latent_min_height=15, + tile_latent_min_width=16, + tile_overlap_factor=0.25, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + ) + self.assertEqual( + vae_tiling.tile_overlap_plan(vae, 0, 16, sample_shape=(120, 512)), + { + "tile_overlap_factor_height": 0.0, + "tile_overlap_factor_width": 0.125, + }, + ) + self.assertIsNone( + vae_tiling.tile_overlap_plan(vae, 8, 16, sample_shape=(120, 512)) + ) + + def test_a_row_strip_accepts_a_distinct_height_overlap(self): + vae = StubVAE( + tile_sample_min_height=120, + tile_sample_min_width=128, + tile_latent_min_height=15, + tile_latent_min_width=16, + tile_overlap_factor=0.25, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + ) + self.assertEqual( + vae_tiling.tile_overlap_plan(vae, 16, 0, sample_shape=(480, 128)), + { + "tile_overlap_factor_height": 2 / 15, + "tile_overlap_factor_width": 0.0, + }, + ) + + def test_a_stride_walked_strip_sets_both_strides(self): + cls = type("AutoencoderKLQwenImage", (StubVAE,), {}) + vae = cls( + tile_sample_min_height=120, + tile_sample_min_width=128, + tile_sample_stride_height=96, + tile_sample_stride_width=96, + spatial_compression_ratio=8, + blend_v=lambda above, tile, extent: tile, + blend_h=lambda left, tile, extent: tile, + decoder=lambda tile: tile, + post_quant_conv=lambda tile: tile, + clear_cache=lambda: None, + ) + self.assertEqual( + vae_tiling.tile_overlap_plan(vae, 0, 16, sample_shape=(120, 512)), + { + "tile_sample_stride_height": 120, + "tile_sample_stride_width": 112, + }, + ) + + def test_a_single_tile_sets_zero_on_both_axes(self): + self.assertEqual( + vae_tiling.tile_overlap_plan( + overlap_factor_vae(), 0, 0, sample_shape=(256, 256) + ), + { + "tile_overlap_factor": 0.0, + "tile_overlap_factor_height": 0.0, + "tile_overlap_factor_width": 0.0, + }, + ) + + def test_exact_pixel_requests_keep_both_loop_truncations_agreeing(self): + for build in (overlap_factor_vae, overlap_keyed_vae): + for asked in (0, 16, 32, 64): + with self.subTest(vae=build.__name__, overlap=asked): + vae = build() + plan = vae_tiling.tile_overlap_plan(vae, asked, asked) + self.assertIsNotNone(plan) + vae_tiling.apply_tile_plan(vae, plan) + factors = ( + vae.tile_overlap_factor_height, + vae.tile_overlap_factor_width, + ) + (down, across), (deep, wide) = vae_tiling.overlap_windows(vae) + for latent, pixel, factor in zip( + (down, across), (deep, wide), factors + ): + stride = int(latent * (1.0 - factor)) + self.assertGreaterEqual(stride, 1) + self.assertEqual( + pixel - int(pixel * factor), stride * (pixel // latent) + ) + + def test_unrepresentable_overlap_is_refused_without_rounding(self): + self.assertIsNone( + vae_tiling.tile_overlap_plan(overlap_factor_vae(), 1, 64) + ) + + def test_zero_overlap_is_a_step_of_the_whole_window(self): + vae = overlap_factor_vae() + vae_tiling.apply_tile_plan(vae, vae_tiling.tile_overlap_plan(vae, 0, 0)) + self.assertEqual(vae.tile_overlap_factor, 0.0) + self.assertEqual(vae_tiling.tile_overlap(vae), (0, 0)) + + def test_an_overlap_as_wide_as_the_window_is_refused(self): + vae = overlap_factor_vae() + self.assertIsNone(vae_tiling.tile_overlap_plan(vae, 256, 64)) + + +class TestTiledDecode(unittest.TestCase): + """Tests DistVAE's replacement for overlap-factor tiled-decode loops.""" + + # These classes derive tile strides from overlap factors. HunyuanVideo 1.5 uses per-axis + # window attributes and one shared overlap factor. + FAMILY = ("AutoencoderKL", "AutoencoderKLFlux2", "AutoencoderKLHunyuanVideo15") + # Three windows per axis exercise both full and clipped boundary tiles. + WINDOWS_ACROSS = 3 + # Two frames exercise the video path without treating the frame axis as a singleton. + FRAMES = 2 + + def test_overlap_factor_detection_requires_a_supported_vae_class(self): + # Stored-stride VAEs use a different replacement loop. CogVideoX is excluded because its + # spatial loop also performs temporal tiling. + self.assertTrue(vae_tiling.tiles_by_overlap_factor(overlap_factor_vae())) + self.assertTrue(vae_tiling.tiles_by_overlap_factor(overlap_keyed_vae())) + self.assertFalse(vae_tiling.tiles_by_overlap_factor(stride_vae())) + self.assertFalse(vae_tiling.tiles_by_overlap_factor(overlap_hw_vae())) + self.assertIsNone(vae_tiling.overlap_tiled_decode(stride_vae())) + self.assertIsNotNone(vae_tiling.overlap_tiled_decode(overlap_factor_vae())) + self.assertIsNotNone(vae_tiling.overlap_tiled_decode(overlap_keyed_vae())) + + def test_shared_and_per_axis_window_attributes_produce_the_same_pair(self): + self.assertEqual( + vae_tiling.overlap_windows(overlap_factor_vae()), ((32, 32), (256, 256)) + ) + self.assertEqual( + vae_tiling.overlap_windows(overlap_keyed_vae()), ((16, 16), (256, 256)) + ) + self.assertIsNone(vae_tiling.overlap_windows(stride_vae())) + + def _sample(self, decoded): + """Return the sample tensor from either supported tiled-decode return type.""" + return getattr(decoded, "sample", decoded) + + def _tiled_vae(self, name, batch=1): + """Build a small tiled VAE and an input spanning several tiles.""" + import torch + + kwargs, video, channels = TestEverySupportedVAE.VAES[name] + vae = _diffusers_vae(self, name, kwargs, require_tiling=True) + vae.enable_tiling() + + window = vae_tiling.tile_shape(vae) + shape = tuple(axis // 4 for axis in window) + plan = vae_tiling.tile_shape_plan(vae, *shape) + self.assertIsNotNone( + plan, f"{name} refused exact tile shape {shape}" + ) + vae_tiling.apply_tile_plan(vae, plan) + self.assertTrue( + vae_tiling.tiles_by_overlap_factor(vae), + f"{name} was expected to tile by overlap fraction", + ) + + (latent_down, _), _ = vae_tiling.overlap_windows(vae) + grid = latent_down * self.WINDOWS_ACROSS + torch.manual_seed(0) + shape = ( + (batch, channels, self.FRAMES, grid, grid) + if video + else (batch, channels, grid, grid) + ) + return vae, torch.randn(*shape) + + def _counted(self, vae): + """Replace the decoder with a wrapper that records each input shape.""" + import torch.nn as nn + + class CountingDecoder(nn.Module): + def __init__(self, decoder): + super().__init__() + self.decoder = decoder + self.shapes = [] + + def forward(self, x): + self.shapes.append(tuple(x.shape)) + return self.decoder(x) + + @property + def rows(self): + """Return the leading dimension of every decoder input.""" + return [shape[0] for shape in self.shapes] + + counted = CountingDecoder(vae.decoder) + vae.decoder = counted + return counted + + def test_it_decodes_a_tile_at_a_time_exactly_as_upstream_does(self): + import torch + + # With no dispatcher, this loop must preserve the VAE's call boundaries and produce a + # bit-identical sample. Each tile remains a separate decoder call because convolution + # arithmetic depends on the rows grouped into that call. + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + with torch.no_grad(): + expected = self._sample(vae.tiled_decode(latents)) + counted = self._counted(vae) + got = self._sample(vae_tiling.overlap_tiled_decode(vae)(latents)) + self.assertEqual(got.shape, expected.shape) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + self.assertEqual(set(counted.rows), {1}) + + def test_a_wider_step_decodes_fewer_tiles_to_the_same_image_size(self): + import torch + + # Zero overlap increases the stride and reduces the tile count. Compare with the + # upstream loop at the same settings to verify both output size and values. + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + counted = self._counted(vae) + with torch.no_grad(): + before = self._sample(vae.tiled_decode(latents)) + at_own = len(counted.shapes) + + plan = vae_tiling.tile_overlap_plan(vae, 0, 0) + self.assertIsNotNone(plan, f"{name} refused a step of its whole window") + vae_tiling.apply_tile_plan(vae, plan) + counted.shapes.clear() + with torch.no_grad(): + expected = self._sample(vae.tiled_decode(latents)) + at_zero = len(counted.shapes) + counted.shapes.clear() + with torch.no_grad(): + got = self._sample(vae_tiling.overlap_tiled_decode(vae)(latents)) + + self.assertGreater( + at_own, 0, f"{name} decoded nothing through its decoder" + ) + self.assertLess(at_zero, at_own) + self.assertEqual(len(counted.shapes), at_zero) + self.assertEqual(got.shape, before.shape) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_the_replacement_preserves_the_upstream_return_type(self): + import torch + + # Most tiled-decode methods return DecoderOutput when requested. HunyuanVideo 1.5 returns + # a tensor directly. The replacement must preserve each class's return convention. + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + wraps = vae_tiling._returns_decoder_output(vae) + with torch.no_grad(): + upstream = vae.tiled_decode(latents) + ours = vae_tiling.overlap_tiled_decode(vae)(latents) + self.assertEqual(wraps, not isinstance(upstream, torch.Tensor)) + self.assertIs(type(ours), type(upstream)) + + def test_a_latent_batch_is_decoded_as_it_stands(self): + import torch + + # Each tile carries every sample in the batch, so a call already decodes as many rows as + # there are samples and the decoder is handed the tensor upstream would have given it. + vae, latents = self._tiled_vae("AutoencoderKL", batch=2) + with torch.no_grad(): + expected = vae.tiled_decode(latents).sample + counted = self._counted(vae) + got = vae_tiling.overlap_tiled_decode(vae)(latents).sample + self.assertEqual(set(counted.rows), {2}) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_tile_parallel_support_requires_a_distvae_owned_loop(self): + self.assertTrue(vae_tiling.supports_tile_parallel(overlap_factor_vae())) + self.assertTrue(vae_tiling.supports_tile_parallel(overlap_keyed_vae())) + self.assertFalse(vae_tiling.supports_tile_parallel(stride_vae())) + self.assertFalse(vae_tiling.supports_tile_parallel(overlap_hw_vae())) + + def test_the_dispatcher_is_given_every_call_and_the_image_is_unchanged(self): + import torch + + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + seen = [] + + def dispatch(calls): + seen.append(len(calls)) + return [call() for call in calls] + + with torch.no_grad(): + expected = self._sample(vae.tiled_decode(latents)) + counted = self._counted(vae) + got = self._sample( + vae_tiling.overlap_tiled_decode(vae, dispatch)(latents) + ) + # The replacement submits every tile in one dispatch call. + self.assertEqual(seen, [len(counted.shapes)]) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_the_calls_can_be_made_in_any_order(self): + import torch + + # Independent tiles may execute in any order; assembly restores grid order. + def backwards(calls): + return list(reversed([call() for call in reversed(calls)])) + + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + with torch.no_grad(): + expected = self._sample(vae.tiled_decode(latents)) + got = self._sample( + vae_tiling.overlap_tiled_decode(vae, backwards)(latents) + ) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_a_tiled_decode_that_fits_in_one_tile_still_works(self): + import torch + + # A single tile means one call and no blending pass at all. + vae, _ = self._tiled_vae("AutoencoderKL") + latents = torch.randn(1, vae.config.latent_channels, 4, 4) + with torch.no_grad(): + expected = vae.tiled_decode(latents).sample + got = vae_tiling.overlap_tiled_decode(vae)(latents).sample + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +class TestStrideTiledDecode(unittest.TestCase): + """Tests DistVAE's replacement for stored-stride tiled-decode loops.""" + + # Wan and Qwen-Image decode each tile frame by frame with a tile-local feature cache. + # HunyuanVideo and LTX-2 decode each spatial tile in one call. + FAMILY = ( + "AutoencoderKLWan", + "AutoencoderKLQwenImage", + "AutoencoderKLHunyuanVideo", + "AutoencoderKLLTX2Video", + ) + # LTX-2 passes a positional timestep embedding through tiled_decode. + CONDITIONED = ("AutoencoderKLLTX2Video",) + # The grid spans several tiles after halving the window. Two frames exercise cache reuse. + LATENT_GRID = 16 + FRAMES = 2 + + def _tiled_vae(self, name, **extra): + """Build a small video VAE and an input spanning several tiles.""" + import torch + + kwargs, _, channels = TestEverySupportedVAE.VAES[name] + vae = _diffusers_vae(self, name, {**kwargs, **extra}, require_tiling=True) + vae.enable_tiling() + + window = vae_tiling.tile_shape(vae) + shape = tuple(axis // 2 for axis in window) + plan = vae_tiling.tile_shape_plan(vae, *shape) + self.assertIsNotNone( + plan, f"{name} refused exact tile shape {shape}" + ) + vae_tiling.apply_tile_plan(vae, plan) + self.assertTrue( + vae_tiling.tiles_by_stored_stride(vae), + f"{name} was expected to tile by a stride it stores", + ) + + torch.manual_seed(0) + grid = self.LATENT_GRID + return vae, torch.randn(1, channels, self.FRAMES, grid, grid) + + def _conditioning(self, vae): + """Return positional conditioning arguments required by tiled_decode.""" + return (None,) if type(vae).__name__ in self.CONDITIONED else () + + def test_stored_stride_detection_requires_a_supported_vae_class(self): + # Matching stride attributes is insufficient; support is limited to known loop + # implementations. + self.assertFalse(vae_tiling.tiles_by_stored_stride(stride_vae())) + self.assertFalse(vae_tiling.tiles_by_stored_stride(overlap_factor_vae())) + # HunyuanVideo 1.5 derives its stride from an overlap factor. + self.assertFalse(vae_tiling.tiles_by_stored_stride(overlap_keyed_vae())) + + def test_reimplemented_stride_tiling_matches_native_tiled_decode(self): + import torch + + for name, extra in ( + ("AutoencoderKLWan", {}), + # Wan 2.2 uses pixel unshuffle during decode. The channel count and spatial + # compression ratio must include its patch size. + ( + "AutoencoderKLWan", + { + "patch_size": 2, + "in_channels": 12, + "out_channels": 12, + "scale_factor_spatial": 16, + }, + ), + ("AutoencoderKLQwenImage", {}), + ("AutoencoderKLHunyuanVideo", {}), + ("AutoencoderKLLTX2Video", {}), + ): + with self.subTest(vae=name, **extra): + vae, latents = self._tiled_vae(name, **extra) + args = self._conditioning(vae) + with torch.no_grad(): + expected = vae.tiled_decode(latents, *args).sample + got = vae_tiling.strided_tiled_decode(vae)(latents, *args).sample + # Local execution must match the upstream loop exactly. + self.assertEqual(got.shape, expected.shape) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_a_tile_is_a_call_and_they_can_be_made_in_any_order(self): + import torch + + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + args = self._conditioning(vae) + seen = [] + + def backwards(calls): + seen.append(len(calls)) + return list(reversed([call() for call in reversed(calls)])) + + with torch.no_grad(): + expected = vae.tiled_decode(latents, *args).sample + got = vae_tiling.strided_tiled_decode(vae, backwards)( + latents, *args + ).sample + # State may be shared between frames within one tile, but not between tiles. + # Tile execution order therefore cannot affect the output. + stride_height = ( + vae.tile_sample_stride_height // vae.spatial_compression_ratio + ) + stride_width = ( + vae.tile_sample_stride_width // vae.spatial_compression_ratio + ) + rows = len(range(0, latents.shape[-2], stride_height)) + columns = len(range(0, latents.shape[-1], stride_width)) + self.assertEqual(seen, [rows * columns]) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_a_wider_step_decodes_fewer_tiles_to_the_same_image_size(self): + import torch + + # Stored pixel strides are converted to latent-grid steps by the compression ratio and, + # for pixel-unshuffle decoders, to crop steps by the patch size. Compare against the + # upstream loop to catch inconsistent integer conversion. + for name in self.FAMILY: + with self.subTest(vae=name): + vae, latents = self._tiled_vae(name) + args = self._conditioning(vae) + with torch.no_grad(): + before = vae.tiled_decode(latents, *args).sample + at_own = self._tiles_across(vae, latents) + + plan = vae_tiling.tile_overlap_plan(vae, 0, 0) + self.assertIsNotNone(plan, f"{name} refused a step of its whole window") + vae_tiling.apply_tile_plan(vae, plan) + self.assertLess(self._tiles_across(vae, latents), at_own) + + with torch.no_grad(): + expected = vae.tiled_decode(latents, *args).sample + got = vae_tiling.strided_tiled_decode(vae)(latents, *args).sample + self.assertEqual(got.shape, before.shape) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def _tiles_across(self, vae, latents): + """Return the number of tile columns at the VAE's current stride.""" + # Count grid positions because cached VAEs make several decoder calls per tile. + stride = vae.tile_sample_stride_width // vae.spatial_compression_ratio + return len(range(0, latents.shape[-1], stride)) + + def test_temporal_chunking_reaches_the_installed_spatial_loop(self): + import torch + + # HunyuanVideo's temporal loop calls the installed spatial loop once per frame chunk. + vae, _ = self._tiled_vae("AutoencoderKLHunyuanVideo") + ratio = vae.spatial_compression_ratio + latent_stride = vae.tile_sample_stride_width // ratio + chunk = vae.tile_sample_stride_num_frames // vae.temporal_compression_ratio + # Exceed the spatial window by one latent pixel and provide two temporal chunks. + grid = vae.tile_sample_min_width // ratio + 1 + torch.manual_seed(0) + latents = torch.randn(1, vae.config.latent_channels, 2 * chunk, grid, grid) + + seen = [] + + def counted(calls): + seen.append(len(calls)) + return [call() for call in calls] + + with torch.no_grad(): + expected = vae.decode(latents).sample + vae.tiled_decode = vae_tiling.strided_tiled_decode(vae, counted) + got = vae.decode(latents).sample + + across = len(range(0, grid, latent_stride)) + self.assertGreater(across, 1, "the grid was too narrow to tile") + self.assertEqual( + seen, + [across * across] * 2, + "the temporal loop did not reach the installed loop", + ) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_the_native_stride_loop_is_kept_without_a_dispatcher(self): + vae, _ = self._tiled_vae("AutoencoderKLWan") + self.assertTrue(vae_tiling.supports_tile_parallel(vae)) + self.assertIsNone(vae_tiling.tiled_decode_for(vae)) + self.assertIsNotNone(vae_tiling.tiled_decode_for(vae, lambda calls: [])) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_wandecoderadapter.py b/test/test_wandecoderadapter.py new file mode 100644 index 0000000..be09654 --- /dev/null +++ b/test/test_wandecoderadapter.py @@ -0,0 +1,81 @@ +"""WanDecoderAdapter against the decoder it shards, over gloo on CPU. + +Run from repo root: + pytest test/test_wandecoderadapter.py -v +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.distributed as dist + +from distvae.modules.adapters.vae.decoder_adapters import WanDecoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed + +diffusers = pytest.importorskip("diffusers") + +# Four channel stages exercise every decoder upsampling transition. +CONFIG = dict(base_dim=8, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=1) +LATENT_CHANNELS = 4 + + +def build_decoder(): + return diffusers.AutoencoderKLWan(**CONFIG).eval().decoder + + +def worker(rank, world_size, frames, height, width, seed, master_port): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + decoder = build_decoder() + # Taken before the adapter runs, which rebuilds the decoder in place. + weights = decoder.state_dict() + + latents = torch.randn(1, LATENT_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_decoder() + reference.load_state_dict(weights) + expected = reference(latents, feat_cache=None, feat_idx=[0]) + + adapter = WanDecoderAdapter(decoder, vae_group=None).eval() + actual = adapter(latents) + + assert_matches_reference(rank, actual, expected, "WanDecoderAdapter") + finally: + dist.destroy_process_group() + + +@pytest.mark.gloo +def test_sharded_wan_decode_matches_unsharded_decode(master_port, seed=42): + run_distributed(worker, 2, (1, 16, 16, seed), master_port) + + +@pytest.mark.gloo +def test_a_latent_taller_than_it_is_wide_still_decodes(master_port, seed=42): + # The patch dimension defaults to H, so a non-square latent is the case where getting the + # split wrong shows up as a wrongly shaped output rather than as wrong values. + run_distributed(worker, 2, (1, 24, 16, seed), master_port) + + +@pytest.mark.gloo +def test_latent_rows_that_do_not_divide_by_the_rank_count(master_port, seed=42): + # Padding to an even split changes the decode because convolution and mid-block attention + # propagate padded values into the rows that survive cropping. + run_distributed(worker, 3, (1, 16, 16, seed), master_port) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="WanDecoderAdapter GLOO multi-rank tests") + parser.add_argument("--world_size", type=int, default=None) + args, remainder = parser.parse_known_args() + pytest_args = [os.path.abspath(__file__), "-v"] + remainder + if args.world_size is not None: + pytest_args.extend(["-k", f"[{args.world_size}]"]) + sys.exit(pytest.main(pytest_args)) diff --git a/test/test_wanencoderadapter.py b/test/test_wanencoderadapter.py index cd32705..e3cea36 100644 --- a/test/test_wanencoderadapter.py +++ b/test/test_wanencoderadapter.py @@ -1,8 +1,11 @@ -""" -Test WanEncoderAdapter with real WanEncoder3d from diffusers. +"""WanEncoderAdapter against the encoder it shards, over gloo on CPU. + +Wan ships two encoder shapes: 2.2 groups each stage into a WanResidualDownBlock, 2.1 lays the +same residual blocks, attentions and resamples out flat. Both are covered here, since the +encoder class alone does not say which one a checkpoint carries. -This test uses actual Wan encoder classes instead of mocks to demonstrate -the exact output matching between distributed and single-rank encoders. +Run from repo root: + pytest test/test_wanencoderadapter.py -v """ import argparse @@ -12,201 +15,106 @@ import pytest import torch import torch.distributed as dist -import torch.nn as nn -from torch.multiprocessing import spawn -from distvae.utils import DistributedEnv +from distvae.modules.adapters.vae.encoder_adapters import WanEncoderAdapter + +from distributed_harness import assert_matches_reference, init_gloo, run_distributed autoencoder_kl_wan = pytest.importorskip( "diffusers.models.autoencoders.autoencoder_kl_wan" ) -def worker( - rank: int, - world_size: int, - height: int, - width: int, - seed: int, - master_port: int, -) -> None: - device = torch.device("cpu") - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - os.environ["RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - dist.init_process_group(backend="gloo", init_method="env://") - DistributedEnv.initialize(None) - - torch.manual_seed(seed) - - # Create a small encoder: 3 spatial downsamples (8x reduction) - WanEncoder3d = autoencoder_kl_wan.WanEncoder3d - encoder = WanEncoder3d( - in_channels=3, - dim=32, # Small for testing - z_dim=16, - dim_mult=[1, 2, 4, 8], # 32 -> 32 -> 64 -> 128 -> 256 - num_res_blocks=1, # Minimal for speed - attn_scales=[], # No attention - temperal_downsample=[False, True, True, False], # 3 spatial downsamples, no temporal - dropout=0.0, - is_residual=False, - ) - encoder = encoder.to(device) - encoder.eval() - - # Save state dict before patching - encoder_state_dict = encoder.state_dict() - - # Create distributed adapter - from distvae.modules.adapters.vae.encoder_adapters import WanEncoderAdapter - - encoder_adapter = WanEncoderAdapter( - encoder, - vae_group=None, - conv_block_size=0, - patch_dim=-2, - vae_scale_factor=8, - use_uniform_patch=True, - ) - encoder_adapter.eval() - - # Input: (B, C, F, H, W) - n, c, f = 1, 3, 4 - x_full = torch.randn(n, c, f, height, width, device=device, dtype=torch.float32) - - with torch.no_grad(): - # Reference: create unwrapped encoder on rank 0 only - if rank == 0: - encoder_ref = WanEncoder3d( - in_channels=3, - dim=32, - z_dim=16, - dim_mult=[1, 2, 4, 8], - num_res_blocks=1, - attn_scales=[], - temperal_downsample=[False, True, True, False], - dropout=0.0, - is_residual=False, - ) - encoder_ref.load_state_dict(encoder_state_dict) - encoder_ref = encoder_ref.to(device) - encoder_ref.eval() - y_ref = encoder_ref(x_full) - else: - y_ref = None - - # Distributed: run through adapter on all ranks - y_dist = encoder_adapter(x_full) - - success = torch.ones(1, dtype=torch.int64, device=device) - if rank == 0: - # Always write debug info - with open("/tmp/real_wan_encoder_test_debug.txt", "a") as f: - f.write(f"\nTest executed: world_size={world_size}, height={height}, width={width}\n") - f.write(f"y_ref shape: {y_ref.shape}, y_dist shape: {y_dist.shape}\n") - - # Check shapes match - if y_ref.shape != y_dist.shape: - print(f"Shape mismatch: ref={y_ref.shape}, dist={y_dist.shape}", flush=True) - success.zero_() - # Check values are close (use realistic tolerance for real encoder with padding/cropping) - elif not torch.allclose(y_ref, y_dist, atol=1e-4, rtol=1e-3): - diff = torch.abs(y_ref - y_dist) - max_diff = diff.max().item() - mean_diff = diff.mean().item() - rel_diff = (diff / (torch.abs(y_ref) + 1e-8)).mean().item() - print(f"Values mismatch: max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, rel_diff={rel_diff:.2e}", flush=True) - with open("/tmp/real_wan_encoder_test_debug.txt", "a") as f: - f.write(f"FAILED tolerance check: max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, rel_diff={rel_diff:.2e}\n") - # Sample some values for debugging - f.write(f"y_ref sample: {y_ref.flatten()[:10].tolist()}\n") - f.write(f"y_dist sample: {y_dist.flatten()[:10].tolist()}\n") - success.zero_() - else: - max_diff = torch.abs(y_ref - y_dist).max().item() - mean_diff = torch.abs(y_ref - y_dist).mean().item() - msg = f"\n{'='*60}\n" - msg += f"SUCCESS: WanEncoderAdapter output matches reference encoder\n" - msg += f"Input shape: {x_full.shape}\n" - msg += f"Output shape: {y_dist.shape}\n" - msg += f"Max absolute difference: {max_diff:.2e}\n" - msg += f"Mean absolute difference: {mean_diff:.2e}\n" - msg += f"{'='*60}\n" - print(msg, flush=True) - # Write to file - with open("/tmp/real_wan_encoder_test_results.txt", "a") as f: - f.write(f"\nTest: world_size={world_size}, height={height}, width={width}\n") - f.write(msg) - - dist.broadcast(success, src=0) - dist.barrier() - dist.destroy_process_group() - - if success.item() == 0: - raise AssertionError("WanEncoderAdapter output did not match reference encoder") - - -def _run_one( - world_size: int, - height: int, - width: int, - seed: int, - master_port: int, -) -> None: - """Spawn processes and run worker; raises on failure.""" - spawn( - worker, - nprocs=world_size, - args=(world_size, height, width, seed, master_port), - join=True, +# Three spatial downsamples, so the encoder narrows by 8 the way the shipped ones do. +CONFIG = dict( + in_channels=3, + dim=32, + z_dim=16, + # The shipped ratios, kept because Wan 2.2's shortcut averages space into channels and + # asserts the two divide; a last stage that widens would fail to build at all. + dim_mult=[1, 2, 4, 4], + num_res_blocks=1, + attn_scales=[], + dropout=0.0, +) +SCALE_FACTOR = 8 +IN_CHANNELS = 3 + + +def build_encoder(is_residual=False): + # Wan 2.2's grouped down block averages frames in its shortcut whenever a stage downsamples + # time, while the stage itself only does so through the feature cache. Called without one, + # as here, the two disagree on the frame count and diffusers' own encoder raises. So the + # grouped shape is exercised without temporal downsampling; splitting is spatial regardless. + temporal = [False] * 4 if is_residual else [False, True, True, False] + encoder = autoencoder_kl_wan.WanEncoder3d( + **CONFIG, temperal_downsample=temporal, is_residual=is_residual ) + return encoder.eval() -@pytest.fixture -def master_port(request): - """Unique port per test to avoid conflicts.""" - base = 29700 - nodeid = request.node.nodeid - return base + (hash(nodeid) % 10000) +def worker( + rank, world_size, frames, height, width, is_residual, conv_block_size, seed, master_port +): + init_gloo(rank, world_size, master_port) + try: + torch.manual_seed(seed) + encoder = build_encoder(is_residual) + # Taken before the adapter runs, which rebuilds the encoder in place. + weights = encoder.state_dict() + + pixels = torch.randn(1, IN_CHANNELS, frames, height, width) + + with torch.no_grad(): + expected = None + if rank == 0: + reference = build_encoder(is_residual) + reference.load_state_dict(weights) + expected = reference(pixels) + + adapter = WanEncoderAdapter( + encoder, + vae_group=None, + vae_scale_factor=SCALE_FACTOR, + conv_block_size=conv_block_size, + ).eval() + actual = adapter(pixels) + + assert_matches_reference(rank, actual, expected, "WanEncoderAdapter") + finally: + dist.destroy_process_group() @pytest.mark.gloo -@pytest.mark.parametrize("world_size", [1, 2, 4]) -def test_real_wan_encoder_even_sizes(world_size, master_port, seed=42): - """Real WanEncoder3d with even sizes divisible by downsampling factor (8).""" - height, width = 64, 64 # After 3x stride=2: 64 -> 32 -> 16 -> 8 - _run_one( - world_size=world_size, - height=height, - width=width, - seed=seed, - master_port=master_port, - ) +def test_sharded_wan_encode_matches_unsharded_encode(master_port, seed=42): + run_distributed(worker, 2, (4, 64, 64, False, 0, seed), master_port) @pytest.mark.gloo -def test_real_wan_encoder_larger_input(master_port, seed=42): - """Real WanEncoder3d with larger input.""" - height, width = 128, 128 # After 3x stride=2: 128 -> 64 -> 32 -> 16 - _run_one( - world_size=2, - height=height, - width=width, - seed=seed, - master_port=master_port, - ) +def test_the_grouped_wan22_down_blocks_encode_the_same(master_port, seed=42): + run_distributed(worker, 2, (4, 64, 64, True, 0, seed), master_port) + + +@pytest.mark.gloo +def test_an_image_taller_than_it_is_wide_still_encodes(master_port, seed=42): + run_distributed(worker, 2, (4, 96, 64, False, 0, seed), master_port) + + +@pytest.mark.gloo +def test_rows_that_do_not_divide_by_the_rank_count_still_encode(master_port, seed=42): + # 80 rows at a ratio of 8 is ten bands, which over 3 ranks leaves them different sizes. + run_distributed(worker, 3, (4, 80, 64, False, 0, seed), master_port) + + +@pytest.mark.gloo +def test_the_chunked_convolution_path_encodes_the_same(master_port, seed=42): + run_distributed(worker, 2, (4, 64, 64, False, 32, seed), master_port) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Real WanEncoder3d GLOO multi-rank tests") + parser = argparse.ArgumentParser(description="WanEncoderAdapter GLOO tests") parser.add_argument("--world_size", type=int, default=None) - parser.add_argument("--seed", type=int, default=42) args, remainder = parser.parse_known_args() - - # Pass through any remaining args to pytest pytest_args = [os.path.abspath(__file__), "-v"] + remainder if args.world_size is not None: pytest_args.extend(["-k", f"[{args.world_size}]"]) - sys.exit(pytest.main(pytest_args))