indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-07-07

Browse Source
This commit is contained in:
github-actions[bot]
2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/docs/multigpu/README.md
+78
View File
@@ -0,0 +1,78 @@
# Multi-GPU Inference
Run LTX-2 pipelines across several GPUs on a single machine.
> ## ⚠️ Important
>
> **Multi-GPU (MGPU) is a latency tool, not a memory tool.** It is designed to reduce
> the latency of a single generation on multi-GPU servers (H100, B200) by splitting
> each denoising step and the VAE decode across GPUs.
>
> **MGPU is not a way to fit a bigger model.** The mutable **working copy** of the
> transformer is a **full replica on every GPU** (each rank builds the whole model;
> LoRAs are fused into it in place). MGPU therefore cannot make a checkpoint that
> doesn't fit on one GPU suddenly fit — for that use FP8 quantization and weight
> offloading (see [Optimization Tips](../optimization.md)).
>
> Each rank *also* holds a second, immutable copy of the **clean (pre-LoRA)
> weights** — kept for LoRA hot-swap (reset + broadcast) — but that copy is
> **sharded** across GPUs (`ShardedSD`, ~1/world_size per rank), not replicated.
> Sequence parallelism additionally splits **activation** memory across ranks. See
> the [weight tracker](pipeline-setup.md#transformerweighttracker--working-copy--sharded-clean-weights)
> for the exact layout. The headline purpose is **latency**, not memory.
>
> **Single machine only.** One process per GPU, `MASTER_ADDR=localhost`, one rank
> per GPU. No multi-node.
## Requirements
- **Linux** -- NCCL and CUDA-IPC peer buffers are Linux-only (no macOS/Windows).
- **>=2 CUDA GPUs on a single node** with P2P access (NVLink/PCIe). No multi-node.
- **PyTorch with CUDA.**
- **`ltx-kernels` built** -- the SP all2all kernel is mandatory. Build with
`uv sync --group kernels` (needs a CUDA toolkit / nvcc and a C++ compiler, gcc or
clang). See the root README.
## Capabilities
| Technique | Purpose |
| --------- | ------- |
| [Sequence parallelism (SP)](sequence-parallel.md) | Split the token sequence across GPUs; faithful — numerically equivalent to single-GPU |
| [Tiled data parallelism (TDP)](tiled-data-parallel.md) | One spatial (height x width) tile per GPU; for resolutions outside the training distribution. **Upscale only** |
| [Distributed decoder](distributed-decoder.md) | Decode latent tiles in parallel, assemble on the driver |
| [Distributed Gemma](gemma.md) | Shard Gemma across GPUs via Accelerate `device_map`, or replicate + split prompts |
| [MGPU controller](controller.md) | Persistent worker fleet; dispatch a job, stream results |
| [Pipeline setup](pipeline-setup.md) | Swap single-GPU builders for MGPU builders; share one weights registry |
## Architecture overview
The [`MGPUController`](controller.md) spawns one worker process per GPU and runs a
user-defined **runner** (a subclass of `MGPURunner`) in SPMD lockstep. A runner's
`setup()` builds a standard pipeline, then **swaps** each block's builder for an MGPU
builder (SP / TDP / distributed decoder / distributed Gemma). All builders share one
`StateDictRegistry` so the checkpoint loads from disk once per process.
Two runners are provided, each with a CLI:
- [`ltx_pipelines.ti2vid_two_stages_mgpu`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py) — SP stage 1 + TDP stage 2 + Accelerate Gemma + distributed VAE.
- [`ltx_pipelines.distilled_mgpu`](../../src/ltx_pipelines/distilled_mgpu.py) — SP (shared stage) + Accelerate Gemma + distributed VAE.
```bash
# Two-stage on all visible GPUs
python -m ltx_pipelines.ti2vid_two_stages_mgpu \
--checkpoint-path path/to/checkpoint.safetensors \
--distilled-lora path/to/distilled_lora.safetensors 1.0 \
--spatial-upsampler-path path/to/upsampler.safetensors \
--gemma-root path/to/gemma \
--prompt "A beautiful sunset over the ocean" \
--output-path output.mp4
```
## Pages
- **[Controller](controller.md)** — `MGPUController` / `MGPURunner` / `Stream`, lifecycle, one-job-at-a-time contract, threading, error handling.
- **[Pipeline setup](pipeline-setup.md)** — swapping builders, the shared weights registry, the LoRA-hot-swap weight tracker.
- **[Sequence parallelism](sequence-parallel.md)** — faithful token-dim split, the all2all kernels, `AttentionManager`, `SequenceParallelBuilder`.
- **[Tiled data parallelism](tiled-data-parallel.md)** — out-of-distribution resolutions, position normalization, shared negative (reference) positions, `TiledDataParallelBuilder`.
- **[Distributed decoder](distributed-decoder.md)** — inter-GPU vs intra-GPU tiling, `DistributedDecoderBuilder`.
- **[Gemma](gemma.md)** — `AccelerateGemmaBuilder` (Accelerate `device_map` sharding) and `BatchParallelGemmaBuilder` (replicated; not for the distilled pipeline).
packages/ltx-pipelines/docs/multigpu/controller.md
+115
View File
@@ -0,0 +1,115 @@
# MGPU Controller
**Source**: [`multigpu/controller.py`](../../src/ltx_pipelines/multigpu/controller.py), [`multigpu/runner.py`](../../src/ltx_pipelines/multigpu/runner.py)
The controller is a persistent, one-job-at-a-time GPU fleet. It spawns one worker
process per GPU, runs a user-defined **runner** in SPMD lockstep, and streams
results back.
## Public classes
### `MGPUController`
```python
MGPUController(
runner_cls: type[MGPURunner],
*,
num_gpus: int | None = None, # GPUs 0..num_gpus-1 (default: all visible)
devices: Sequence[int] | None = None, # place on specific physical GPUs, e.g. [2, 3]
logs_specs: LogsSpecs | None = None,
)
```
`num_gpus` and `devices` are mutually exclusive. `devices=[2, 3]` puts rank r on
`cuda:devices[r]`, so two controllers can share one machine on disjoint GPU sets.
Lifecycle:
| Method | What it does |
| ------ | ------------ |
| `start(*, timeout=30min, **setup_kwargs)` | Spawn the fleet, run `setup(**setup_kwargs)` on every rank, block until all report ready. `timeout` bounds NCCL init + CUDA init + `setup()`; it must exceed the slowest model load. |
| `stream(*, timeout=None, **kwargs) -> Stream` | Dispatch one job and return **immediately**. Iterate the returned `Stream` to collect. |
| `shutdown(*, graceful_timeout=60.0)` | Tear down the fleet; also force-terminates it — safe to call from another thread to recover a job that cannot be drained. |
| `is_alive` (property) | True while the fleet is up and unpoisoned. |
### `MGPURunner`
`MGPURunner` is the abstract base class implemented per pipeline. The controller
ships the subclass to every worker by value (a runner defined in `__main__` or a test
module is supported), builds one instance per worker, injects the NCCL groups, calls
`setup()` once, then invokes the instance per job.
```python
class MyRunner(MGPURunner):
@torch.inference_mode()
def setup(self, *, checkpoint_path: str, ...) -> None:
# build the pipeline + swap in MGPU builders (see pipeline-setup.md)
...
@torch.inference_mode()
def __call__(self, *, prompt: str, ...) -> Iterator[...]:
video, audio = self._pipeline(...)
yield output_path # __call__ MUST be a generator (use `yield`, even once)
```
- `setup()` and `__call__()` run on **every** rank. `self.groups` gives the
per-component `NCCLGroups` (`gemma_group`, `transformer_group`, `vae_group`).
- The framework does **not** apply inference mode — decorate `setup`/`__call__` explicitly.
## Usage
```python
from ltx_pipelines.multigpu import MGPUController
controller = MGPUController(MyRunner, num_gpus=8)
controller.start(checkpoint_path="...", gemma_root="...") # setup kwargs
stream = controller.stream(prompt="a cat", seed=42)
try:
for item in stream: # one element per yield, as it arrives (NOT gathered across ranks)
show(item)
finally:
stream.drain() # free the controller even on early exit
controller.shutdown()
```
### Passing tensors
Tensors are transparent:
- **Inputs.** Pass them as **top-level** kwargs (`stream(latent=t, steps=30)`) and
the relay (rank 0) broadcasts them to every rank over NCCL — `__call__` receives
them already on the local GPU. An input tensor nested inside a list/dict kwarg is
**not** broadcast; it falls back to the (slower) pickle path.
- **Outputs.** Yield tensors back (including nested inside a dict) and they return
via the result queue by shared memory / CUDA IPC — no pickling, regardless of
nesting.
Everything else must be picklable and small.
## Contract and limitations
- **Single machine only.** `MASTER_ADDR=localhost`, `RANK == LOCAL_RANK`, one rank per GPU.
- **One job at a time.** No job queue, no pipelining. Consume the `Stream` to the
end before the next `stream()`. Abandoning it is **not** cleaned up: the next
`stream()` raises `ControllerBusyError` until `stream.drain()` or `shutdown()` is
called. The recommended pattern is `try: ... finally: stream.drain()`.
- **SPMD lockstep.** Yields are forwarded individually (in result-queue order), not
gathered. Only per-rank terminals are collected to end the stream.
- **Thread ownership (baton-lock).** Any thread may call `stream()`. Each job
belongs to its **dispatching** thread — only that
thread may iterate or `drain()` its `Stream` (enforced in `Stream.__next__`). A
single lock guards only the in-flight check-and-set: among concurrent `stream()`
callers one proceeds and the rest raise `ControllerBusyError`.
## Error handling
| Situation | Outcome |
| --------- | ------- |
| Runner raises an **unexpected** exception | **Fatal** — a desynced NCCL collective cannot be unwound. The controller is poisoned and a new one must be constructed. |
| Runner raises `RunnerError` (or `ValueError`, auto-converted) **identically on every rank** | Recoverable. Iterating the `Stream` re-raises `SymmetricRunnerError`; the fleet survives — fix the input and retry. Raise it outside any collective (e.g. validating broadcast kwargs before the first one). |
| Some ranks raise `RunnerError`, others finish clean | `AsymmetricRunnerError` — surfaced prominently (a latent hang risk), but does not terminate the fleet. |
| Worker death / exceeded per-job `timeout` | Surfaced when the `Stream` is next iterated; the controller is poisoned. |
The public API (`from ltx_pipelines.multigpu import ...`) exports `MGPUController`,
`MGPURunner`, `Stream`, `RunnerError`, `SymmetricRunnerError`,
`AsymmetricRunnerError`, `ControllerBusyError`, and `NCCLGroups`.
packages/ltx-pipelines/docs/multigpu/distributed-decoder.md
+72
View File
@@ -0,0 +1,72 @@
# Distributed VAE Decoder
**Source**: [`ltx_core/multigpu/vae/distributed_decoder.py`](../../../ltx-core/src/ltx_core/multigpu/vae/distributed_decoder.py), [`multigpu/vae_builders.py`](../../src/ltx_pipelines/multigpu/vae_builders.py)
## What it is
VAE decode is expensive and embarrassingly parallel over space/time tiles. The
distributed decoder splits the latent into tiles, assigns them **round-robin** to
ranks (the tile count may exceed the GPU count), and every rank decodes its tiles in
parallel.
Workers ship their decoded tiles to the **driver rank** over an `mp.Queue` (CUDA
IPC — zero-copy handle sharing); the driver blends overlaps and yields the assembled
frames as temporal batches spread across the GPUs.
## Inter-GPU tiling vs intra-GPU tiling
Two independent tilings, commonly conflated:
| | Controls | Config | Set by |
| --- | --- | --- | --- |
| **Inter-GPU** (MGPU) | Which rank decodes which tile — **parallelism** | `vae_tiling: TileCountConfig` (at build time) | `DistributedDecoderBuilder` |
| **Intra-GPU** (SGPU) | Chunking *within* a rank's tile to bound **VRAM** | `tiling_config: TilingConfig` (per call) | the pipeline's usual tiling kwarg |
They compose — a rank can further chunk its assigned tile for VRAM — with one
guard: **multi-GPU temporal tiling and single-GPU temporal tiling cannot both be
on.** If `vae_tiling.frames.num_tiles > 1` and `tiling_config.temporal_config` is
set, `decode_video` raises, because two causal temporal splits would conflict.
## API
### `DistributedDecoderBuilder`
```python
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig
DistributedDecoderBuilder(
inner: BuilderProtocol, # the block's single-GPU decoder builder
queue: Queue, # spawn-context mp.Queue, shared across ranks (CUDA IPC)
vae_group: dist.ProcessGroup, # self.groups.vae_group
vae_tiling: TileCountConfig, # inter-GPU split
driver_rank: int, # rank that collects + assembles (usually 0)
registry: Registry,
)
```
`build()` returns a `DistributedVideoDecoder`; its `decode_video(latent,
tiling_config=None, ...)` returns an iterator of temporal batches on the driver, and
an empty iterator on workers (they only `put` their tiles).
## Usage
```python
# The queue is created once, in the CLI __main__, and passed to controller.start(...)
# as a setup kwarg so every worker shares it:
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
# inside runner.setup():
vae_tiling = TileCountConfig(height=DimensionTilingConfig(num_tiles=8, overlap=4))
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(
inner=pipeline.video_decoder._decoder_builder,
queue=vae_queue,
vae_group=self.groups.vae_group,
vae_tiling=vae_tiling,
driver_rank=0,
registry=registry,
)
```
The shipped runners tile the VAE across **height** (8 tiles, overlap 4). Only the
driver ends up with the assembled video, which is why the runner's `__call__`
encodes the file on `driver_rank` and yields `None` on the others.
packages/ltx-pipelines/docs/multigpu/gemma.md
+80
View File
@@ -0,0 +1,80 @@
# Gemma Text Encoder (Multi-GPU)
**Source**: [`multigpu/gemma_builders.py`](../../src/ltx_pipelines/multigpu/gemma_builders.py), [`multigpu/bp_gemma_builder.py`](../../src/ltx_pipelines/multigpu/bp_gemma_builder.py)
Two ways to run the Gemma text encoder across the fleet. Both swap in for
`pipeline.prompt_encoder._text_encoder_builder` and broadcast the resulting
embeddings to every rank (so the transformer ranks all have them).
## `AccelerateGemmaBuilder` (Accelerate `device_map` — the default)
Loads Gemma **once, on the source rank**, with Accelerate `device_map="auto"`, which
**shards Gemma's layers across the available GPUs**. Non-source ranks receive a
lightweight `AccelerateGemmaWrapper` stub that receives the encoded embeddings over
NCCL. The source rank fuses all prompts into one Gemma call, then broadcasts each
output.
The first `build()` loads via HuggingFace `from_pretrained` and caches the full
state dict (including non-persistent buffers) in the registry; later builds recreate
the model from cache and reinstall the dispatch hooks — no disk I/O.
```python
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
AccelerateGemmaBuilder(
gemma_root_path: str,
gemma_group: dist.ProcessGroup | None, # self.groups.gemma_group
broadcast_group: dist.ProcessGroup | None, # self.groups.transformer_group
registry: Registry,
*,
src_rank: int, # rank that loads + encodes (usually 0)
dtype: torch.dtype = torch.bfloat16,
)
```
Usage (in `runner.setup()`):
```python
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
gemma_root_path=gemma_root,
gemma_group=self.groups.gemma_group,
broadcast_group=self.groups.transformer_group,
registry=registry,
src_rank=0,
dtype=pipeline.dtype,
)
```
The shipped runners (`ti2vid_two_stages_mgpu`, `ti2vid_two_stages_hq_mgpu`, `distilled_mgpu`) use this builder.
## `BatchParallelGemmaBuilder` (replicated — data-parallel over prompts)
Every rank materialises a **full** `GemmaTextEncoder` on its own GPU via the standard
`SingleGPUModelBuilder` path (no Accelerate, no `device_map`, no per-layer dispatch
hooks). The wrapper (`BatchParallelGemmaWrapper`) then **partitions the prompt list
across ranks** in `encode` and broadcasts each prompt's output, so the forwards run
concurrently on different GPUs. Non-deterministic prompt enhancement
(`enhance_t2v` / `enhance_i2v`) is routed through a single `src_rank`.
```python
from ltx_pipelines.multigpu.bp_gemma_builder import BatchParallelGemmaBuilder
BatchParallelGemmaBuilder(
gemma_root_path: str,
broadcast_group: dist.ProcessGroup | None,
registry: Registry,
*,
src_rank: int,
dtype: torch.dtype = torch.bfloat16,
)
```
### Not for the distilled pipeline
Batch-parallel is beneficial only when there is **more than one prompt to encode** —
the typical CFG case, positive + negative (B=2 on 2 ranks = one prompt per rank, both
forwards concurrent). The **distilled** pipeline runs **without CFG**: its `__call__`
accepts a single `prompt` and no `negative_prompt`, so there is only one prompt to
encode and no work to partition; batch-parallel provides no speedup in that case. Use
`AccelerateGemmaBuilder` for the distilled pipeline (as the shipped `distilled` runner
does).
packages/ltx-pipelines/docs/multigpu/pipeline-setup.md
+111
View File
@@ -0,0 +1,111 @@
# Setting Up an MGPU Pipeline
**Source**: [`ti2vid_two_stages_mgpu.py`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py), [`multigpu/weight_tracker.py`](../../src/ltx_pipelines/multigpu/weight_tracker.py)
An MGPU pipeline **is** a single-GPU pipeline with its per-block builders swapped
for MGPU builders. Build the standard pipeline, then replace each block's
`_transformer_builder` / `_text_encoder_builder` / `_decoder_builder`.
## The pattern
This is performed inside a runner's `setup()` (which runs on every rank — see
[Controller](controller.md)).
```python
from ltx_core.loader.registry import StateDictRegistry
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
from ltx_core.multigpu.transformer.attention import AttentionManager
# 1. ONE shared registry for every builder in this process.
registry = StateDictRegistry()
# 2. Build the normal pipeline, handing it the registry.
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=..., distilled_lora=..., spatial_upsampler_path=...,
gemma_root=..., loras=[], registry=registry, quantization=...,
)
# 3. One weight tracker per transformer process group (shared by the stages).
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
# 4. Swap each block's builder.
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder, attn_mgr=attn_mgr,
registry=registry, tracker=tracker,
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder, group=self.groups.transformer_group,
tiling=tdp_tiling, registry=registry, tracker=tracker,
)
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(...)
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(...)
```
Each MGPU builder **wraps** the block's existing single-GPU builder (`inner=...`),
so it inherits the checkpoint path, quantization, compilation, and LoRA config —
only the parallelism is added. See the per-technique pages for each builder's
constructor.
> **`with_builder` vs direct assignment.** `DiffusionStage.with_builder(builder)`
> returns a *new* stage with the builder swapped (functional, never mutates). The
> runners assign `stage._transformer_builder = ...` directly because they mutate the
> pipeline once, in place, during `setup()`. Both reach the same builder slot.
## The shared weights registry
`StateDictRegistry` is an in-process cache of loaded state dicts, keyed by
`(resolved paths, sd_ops name)`. Passing **one** registry to every builder means:
- The transformer checkpoint is read from disk **once per process**, even though
stage 1 (SP) and stage 2 (TDP) are separate builders on the same file.
- Gemma and the VAE cache their weights the same way (rebuild the module tree from
the cached tensors, skip disk I/O).
The registry is **per process** — it is not shared across ranks. Each worker loads
its own copy, so the full checkpoint is resident on every GPU (see the
[memory disclaimer](README.md)).
## `TransformerWeightTracker` — working copy + sharded clean weights
```python
TransformerWeightTracker(group: dist.ProcessGroup, bucket_mb=256, no_lora_swap=False)
```
The tracker is shared by the transformer stage builders that operate on the same
checkpoint. It does **not** own weights — it references the tensors in the registry
and receives a builder at `build()` time. Two copies of the weights exist per rank,
and they are **not** the same shape of memory:
- **Working copy** — the model the builder returns, backed by the registry's
tensors. This is a **full replica on every GPU**. LoRAs are fused into it
**in place**; `broadcast_sd` (a zero-copy `ShardedSD` view over these tensors)
broadcasts each owner rank's freshly fused shards — bucketed, `bucket_mb` at a
time — so all ranks converge on identical working weights.
- **Clean weights** (`stored_sd`) — an immutable, cloned backup of the original
(pre-LoRA) weights, held so the working copy can be reset before a different LoRA
set is applied. This copy is **sharded** across ranks (deterministic
`md5(key) % world_size` ownership): each rank stores only its ~1/world_size slice,
not a full clone.
So per-GPU transformer memory is one full working model **plus** a ~1/world_size
clean-weights shard — the clean backup is distributed, the working copy is not.
This allows a two-stage pipeline to apply the distilled LoRA to stage 2 and reset it
for stage 1 without reloading the checkpoint. Pass `no_lora_swap=True` when the
LoRA set is fixed (none, or one set for the whole run): the clean-weights clone is
skipped (`stored_sd` becomes a zero-copy view) and any swap/reset raises — saves the
~1/N shard, and guards against accidental swaps.
## Full example
The shipped runners are the reference: read
[`ti2vid_two_stages_mgpu.py`](../../src/ltx_pipelines/ti2vid_two_stages_mgpu.py)
(`setup()` lines ~54–132) and
[`distilled_mgpu.py`](../../src/ltx_pipelines/distilled_mgpu.py). Each ends with a
`__main__` block wiring the runner into an `MGPUController` behind the standard
two-stage CLI parser.
packages/ltx-pipelines/docs/multigpu/sequence-parallel.md
+100
View File
@@ -0,0 +1,100 @@
# Sequence Parallelism (SP)
**Source**: [`ltx_core/multigpu/transformer/sequence_parallel.py`](../../../ltx-core/src/ltx_core/multigpu/transformer/sequence_parallel.py), [`multigpu/sp_builder.py`](../../src/ltx_pipelines/multigpu/sp_builder.py)
## What it is
SP splits the **token (sequence) dimension** of the video across GPUs. Each rank
holds a slice of the tokens, runs the transformer forward on its slice, and the
outputs are gathered back to all ranks. Self-attention still needs every token to
see every other token, so the Q/K/V heads are exchanged across ranks with a custom
**all2all** kernel: each rank ends up with all tokens for a subset of heads, does
local attention, then the results are shuffled back.
**SP is faithful — numerically equivalent to single-GPU inference.** Attention stays
global (all2all preserves the full token interaction); only the floating-point
reduction order changes. The all2all kernels move bytes only — the round-trip
`gather(send(x)) == x` is byte-exact. SP is the appropriate choice whenever the
single-GPU result is required at lower latency.
This is the default for **stage 1** (`ti2vid_two_stages_mgpu`) and the **shared stage**
(`distilled_mgpu`, where one SP wrapping covers both the half-res and full-res calls).
## How the forward pass works
Per denoising step, [`SequenceParallelModelWrapper`](../../../ltx-core/src/ltx_core/multigpu/transformer/sequence_parallel.py):
1. Pads the video seq dim up to a multiple of `world_size` (padded keys are masked
out; padded rows sliced off after the gather) so every rank gets an equal shard.
2. Tiles latent/timesteps/positions to this rank's slice.
3. Runs the model — video self-attention (`attn1`) and video→audio cross-attention
are patched to route Q/K/V through the all2all kernel.
4. `all_gather`s the output tokens back to full length on every rank and unpads.
## The all2all kernels (`ltx-kernels`)
The custom op is `ltx_kernels.All2All` (from the `ltx-kernels` package); the CUDA
kernels use CUDA-IPC peer buffers to exchange tokens directly between ranks' GPUs.
**`ltx-kernels` must be installed** — the SP builder imports it.
## API
### `AttentionManager`
```python
from ltx_core.multigpu.transformer.attention import AttentionManager
attn_mgr = AttentionManager(
max_tokens: int, # upper bound on total video tokens (raises above it)
num_heads: int, # transformer.num_attention_heads
head_dim: int, # transformer.attention_head_dim
tensor_dtype: torch.dtype,
group: dist.ProcessGroup, # self.groups.transformer_group
copy_out_: bool = False,
)
```
Owns the all2all buffers (sized `ceil(max_tokens / world_size)` tokens per rank) and, per step,
`set_seqlen_all2all(...)` updates the per-rank token counts. `num_heads` must be
divisible by `world_size`.
### `SequenceParallelBuilder`
```python
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
SequenceParallelBuilder(
inner: ModelBuilderProtocol, # the stage's single-GPU transformer builder
attn_mgr: AttentionManager,
registry: Registry,
tracker: TransformerWeightTracker,
)
```
Wraps a `SingleGPUModelBuilder` (raises otherwise), injects the all2all attention
module-ops, and `build()` returns a `SequenceParallelModelWrapper`.
## Usage
```python
# inside runner.setup(), per stage:
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
attn_mgr = AttentionManager(
max_tokens=32768,
num_heads=model_cfg["num_attention_heads"],
head_dim=model_cfg["attention_head_dim"],
tensor_dtype=pipeline.dtype,
group=self.groups.transformer_group,
)
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
inner=pipeline.stage_1._transformer_builder,
attn_mgr=attn_mgr,
registry=registry,
tracker=tracker,
)
```
`max_tokens` must cover the largest step. Reference: stage 1 at 512x768x121 is
~6k video tokens; the distilled shared stage's full-res call (1024x1536x121) is
~24k — both ship with `sp_max_tokens=32768`. Exceeding it raises with a clear
"use a smaller resolution or fewer frames" message.
packages/ltx-pipelines/docs/multigpu/tiled-data-parallel.md
+130
View File
@@ -0,0 +1,130 @@
# Tiled Data Parallelism (TDP)
**Source**: [`ltx_core/multigpu/transformer/tiled_data_parallel.py`](../../../ltx-core/src/ltx_core/multigpu/transformer/tiled_data_parallel.py), [`multigpu/tdp_builder.py`](../../src/ltx_pipelines/multigpu/tdp_builder.py)
## What it is
TDP splits the patchified `(frames, height, width)` latent into **tiles** and gives
each tile to a GPU. Every rank runs the full transformer on its own tile(s),
overlapping regions are blended with trapezoidal masks, and a single `all_reduce`
sums the blended tiles into the final result (masks sum to 1 globally). Tiles are
assigned **round-robin**, so the tile count may exceed the GPU count (16 tiles on
4 GPUs = 4 tiles/rank). Audio is processed untiled on every tile forward and averaged
across tiles.
Unlike [sequence parallelism](sequence-parallel.md), TDP is **not** bit-faithful to
single-GPU: each tile is denoised with only local context and blended, so it is an
approximation.
> ## ⚠️ Do not use the TDP stage's audio output
>
> Audio is **not** tiled. It is denoised on **every** tile's forward pass — each with
> a different, partial video context — and those results are **averaged across all
> tiles**. That average is not a meaningful audio latent. Take the final audio from
> the first (SP) stage and keep it frozen through the TDP upscale; treat the TDP
> stage's audio only as the video-conditioning context it needs internally, never as
> output.
## When to use it
TDP is an **upscaler**. It produces video at resolutions the model never saw during
training by running each tile at a resolution the model handles well and blending the
results. This is why the shipped two-stage runner uses TDP for **stage 2** (the
high-resolution upscale).
TDP can also be **faster** than running the whole frame on one GPU: self-attention is
quadratic in the token count, so splitting `N` tokens into `T` tiles drops per-tile
attention cost from `O(N^2)` to `O((N/T)^2)`.
> **Do not run TDP as the first stage.** Starting from pure noise (a high first
> sigma), each tile denoises independently and produces **unrelated content** — the
> tiles never converge on a single coherent video. Generate the first stage with
> [SP](sequence-parallel.md) (faithful, full-frame), then **upscale** that result
> with TDP.
>
> Even as the upscale stage, tiles can **drift** apart, and the drift grows with the
> **first sigma** of the TDP stage (more noise re-injected means more freedom per
> tile). For consistency, either condition on the stage-1 result with
> **negative-index image conditioning** (for i2v), or use a **smaller first sigma**.
## Position normalization
A tile's tokens must carry positions in the range the model was trained on — not the
global positions of a tile in the corner of a large frame, which the model was never
trained to handle. With `normalize_positions=True` (default), each tile's positions
are shifted so the tile's **generated** tokens start at zero in every dimension:
```python
offset = gen_pos[..., 0].amin(dim=2, keepdim=True)... # min start per (batch, dim)
positions = positions - offset # shift generated + conditioning tokens
```
Interval widths are preserved (only the origin moves), so RoPE sees a valid,
in-distribution position grid per tile.
## Shared negative (reference) positions
Conditioning tokens are appended after the generated tokens. A tile keeps a
conditioning token when its `[start, end)` interval overlaps the tile in all three
dimensions — **or** when it has a **negative time coordinate**. Negative-time tokens
are **reference tokens** (e.g. reference frames / audio references): they are kept by
**every** tile so all tiles share the same reference context. To avoid
double-counting a token kept by several tiles, its blend weight is `1 / (number of
tiles that kept it)`.
## API
### Tiling config (`ltx_core.tiling`)
```python
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig
TileCountConfig(
frames: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
height: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
width: DimensionTilingConfig = DimensionTilingConfig(num_tiles=1, overlap=0),
)
DimensionTilingConfig(num_tiles: int, overlap: int = 0) # counts, not sizes; overlap in latent grid units
```
`TileCountConfig` specifies tile **counts** per dimension (contrast the single-GPU
VAE `TilingConfig`, which specifies tile **sizes**).
### `TiledDataParallelBuilder`
```python
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
TiledDataParallelBuilder(
inner: ModelBuilderProtocol, # the stage's single-GPU transformer builder
group: dist.ProcessGroup, # self.groups.transformer_group
tiling: TileCountConfig,
registry: Registry,
tracker: TransformerWeightTracker,
normalize_positions: bool = True,
)
```
Wraps a `SingleGPUModelBuilder`. Its `build()` requires a `video_tools` kwarg (the
`VideoLatentTools` for the target shape) so the wrapper can compute tiles — the
pipeline passes this through automatically.
## Usage
```python
# inside runner.setup(), stage 2 -- balanced 2D spatial (height x width) grid over the group:
from ltx_core.tiling import TileCountConfig, DimensionTilingConfig, balanced_tile_split
h_tiles, w_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
tdp_tiling = TileCountConfig(
height=DimensionTilingConfig(num_tiles=h_tiles, overlap=5),
width=DimensionTilingConfig(num_tiles=w_tiles, overlap=5),
)
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
inner=pipeline.stage_2._transformer_builder,
group=self.groups.transformer_group,
tiling=tdp_tiling,
registry=registry,
tracker=tracker,
)
```