Spaces:
Running on Zero
Running on Zero
| # -*- coding: utf-8 -*- | |
| """ | |
| Context Parallel (Ulysses-style) Implementation for Long Video Training | |
| Ulysses-style context parallelism: | |
| - partition the input tensor along the temporal (sequence) dimension across GPUs | |
| - redistribute activations with all-to-all communication inside attention | |
| - each device computes attention on its own sequence shard | |
| Reference: https://arxiv.org/abs/2601.20540 Section 3.3.3 | |
| """ | |
| import datetime | |
| import torch | |
| import torch.distributed as dist | |
| from typing import Optional, Tuple | |
| from dataclasses import dataclass | |
| from fastvideo.utils.parallel_states import nccl_info | |
| _CP_TIMEOUT = datetime.timedelta(hours=2) | |
| class ContextParallelConfig: | |
| """Context-parallel configuration.""" | |
| enabled: bool = False | |
| cp_size: int = 1 # number of GPUs in the context-parallel group | |
| cp_rank: int = 0 # rank of this GPU inside the group | |
| cp_group: Optional[dist.ProcessGroup] = None | |
| _cp_config = ContextParallelConfig() | |
| def initialize_context_parallel(cp_size: int = 1): | |
| """ | |
| Initialize the context-parallel group. | |
| Args: | |
| cp_size: number of GPUs; 1 disables context parallelism. | |
| """ | |
| global _cp_config | |
| if cp_size <= 1: | |
| _cp_config = ContextParallelConfig(enabled=False, cp_size=1, cp_rank=0, cp_group=None) | |
| return | |
| world_size = dist.get_world_size() | |
| rank = dist.get_rank() | |
| assert world_size % cp_size == 0, f"world_size ({world_size}) must be divisible by cp_size ({cp_size})" | |
| num_cp_groups = world_size // cp_size | |
| for i in range(num_cp_groups): | |
| ranks = list(range(i * cp_size, (i + 1) * cp_size)) | |
| group = dist.new_group(ranks, timeout=_CP_TIMEOUT) | |
| if rank in ranks: | |
| _cp_config = ContextParallelConfig( | |
| enabled=True, | |
| cp_size=cp_size, | |
| cp_rank=rank - i * cp_size, | |
| cp_group=group, | |
| ) | |
| if rank == 0: | |
| print(f"[Context Parallel] Initialized with cp_size={cp_size}, num_groups={num_cp_groups}") | |
| def warmup_context_parallel(hidden_dim: int = 4096, num_heads: int = 32, | |
| seq_len: int = 1024, device: torch.device = None): | |
| """ | |
| Warm up NCCL communication. | |
| Runs a dummy all-to-all before the first forward so NCCL finishes communicator | |
| initialization and topology setup; otherwise the lazy init during the first forward | |
| can time out on some ranks. | |
| Args: | |
| hidden_dim: model hidden dimension | |
| num_heads: number of attention heads | |
| seq_len: dummy sequence length | |
| device: GPU device | |
| """ | |
| if not _cp_config.enabled: | |
| return | |
| if device is None: | |
| device = torch.cuda.current_device() | |
| rank = dist.get_rank() | |
| head_dim = hidden_dim // num_heads | |
| cp_size = _cp_config.cp_size | |
| group = _cp_config.cp_group | |
| if rank == 0: | |
| print(f"[CP Warmup] starting NCCL warmup: cp_size={cp_size}, " | |
| f"hidden={hidden_dim}, heads={num_heads}, seq={seq_len}") | |
| dummy = torch.zeros(1, seq_len // cp_size, num_heads, head_dim, device=device, dtype=torch.bfloat16) | |
| _ = _all_to_all_single(dummy, scatter_dim=2, gather_dim=1, group=group) | |
| dummy2 = torch.zeros(1, seq_len, num_heads // cp_size, head_dim, device=device, dtype=torch.bfloat16) | |
| _ = _all_to_all_single(dummy2, scatter_dim=1, gather_dim=2, group=group) | |
| dummy3 = torch.zeros(seq_len // cp_size, device=device, dtype=torch.bfloat16) | |
| gathered = [torch.zeros_like(dummy3) for _ in range(cp_size)] | |
| dist.all_gather(gathered, dummy3, group=group) | |
| dist.barrier() | |
| del dummy, dummy2, dummy3, gathered | |
| torch.cuda.empty_cache() | |
| if rank == 0: | |
| print(f"[CP Warmup] NCCL warmup done") | |
| def get_cp_config() -> ContextParallelConfig: | |
| """Return the context-parallel configuration.""" | |
| return _cp_config | |
| def destroy_context_parallel(): | |
| """Destroy the context-parallel group.""" | |
| global _cp_config | |
| _cp_config = ContextParallelConfig() | |
| def _all_to_all_single( | |
| input_: torch.Tensor, | |
| scatter_dim: int, | |
| gather_dim: int, | |
| group: dist.ProcessGroup, | |
| ) -> torch.Tensor: | |
| """ | |
| All-to-all communication used by Ulysses-style context parallelism. | |
| Scatters along scatter_dim and gathers along gather_dim. | |
| Args: | |
| input_: input tensor | |
| scatter_dim: dimension to scatter | |
| gather_dim: dimension to gather | |
| group: process group | |
| Returns: | |
| the redistributed tensor | |
| """ | |
| world_size = dist.get_world_size(group) | |
| if world_size == 1: | |
| return input_ | |
| input_list = [chunk.contiguous() for chunk in torch.chunk(input_, world_size, dim=scatter_dim)] | |
| output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] | |
| dist.all_to_all(output_list, input_list, group=group) | |
| return torch.cat(output_list, dim=gather_dim) | |
| class _SeqAllToAll(torch.autograd.Function): | |
| """ | |
| Sequence-dimension all-to-all with autograd support. | |
| Used by Ulysses-style context parallelism. | |
| """ | |
| def forward( | |
| ctx, | |
| input_: torch.Tensor, | |
| scatter_dim: int, | |
| gather_dim: int, | |
| group: dist.ProcessGroup, | |
| ) -> torch.Tensor: | |
| ctx.scatter_dim = scatter_dim | |
| ctx.gather_dim = gather_dim | |
| ctx.group = group | |
| return _all_to_all_single(input_, scatter_dim, gather_dim, group) | |
| def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None, None]: | |
| grad_input = _all_to_all_single( | |
| grad_output, | |
| ctx.gather_dim, | |
| ctx.scatter_dim, | |
| ctx.group | |
| ) | |
| return grad_input, None, None, None | |
| def seq_all_to_all( | |
| input_: torch.Tensor, | |
| scatter_dim: int, | |
| gather_dim: int, | |
| ) -> torch.Tensor: | |
| """ | |
| Sequence-dimension all-to-all. | |
| Args: | |
| input_: input tensor | |
| scatter_dim: dimension to scatter | |
| gather_dim: dimension to gather | |
| Returns: | |
| the redistributed tensor | |
| """ | |
| if not _cp_config.enabled: | |
| return input_ | |
| return _SeqAllToAll.apply(input_, scatter_dim, gather_dim, _cp_config.cp_group) | |
| def scatter_sequence(input_: torch.Tensor, dim: int = 1) -> torch.Tensor: | |
| """ | |
| Scatter the sequence across all context-parallel GPUs. | |
| Args: | |
| input_: [B, S, ...] or [B, ..., S, ...] | |
| dim: sequence dimension | |
| Returns: | |
| the scattered tensor [B, S//cp_size, ...] | |
| """ | |
| if not _cp_config.enabled: | |
| return input_ | |
| seq_len = input_.shape[dim] | |
| assert seq_len % _cp_config.cp_size == 0, \ | |
| f"Sequence length ({seq_len}) must be divisible by cp_size ({_cp_config.cp_size})" | |
| chunks = torch.chunk(input_, _cp_config.cp_size, dim=dim) | |
| return chunks[_cp_config.cp_rank].contiguous() | |
| def gather_sequence(input_: torch.Tensor, dim: int = 1) -> torch.Tensor: | |
| """ | |
| Gather the sequence from all context-parallel GPUs. | |
| Args: | |
| input_: this GPU's sequence shard | |
| dim: sequence dimension | |
| Returns: | |
| the gathered full sequence | |
| """ | |
| if not _cp_config.enabled: | |
| return input_ | |
| # All-gather | |
| world_size = _cp_config.cp_size | |
| gathered = [torch.empty_like(input_) for _ in range(world_size)] | |
| dist.all_gather(gathered, input_.contiguous(), group=_cp_config.cp_group) | |
| return torch.cat(gathered, dim=dim) | |
| class _GatherForward(torch.autograd.Function): | |
| """Gather on forward, scatter on backward (for loss computation).""" | |
| def forward(ctx, input_: torch.Tensor, dim: int) -> torch.Tensor: | |
| ctx.dim = dim | |
| return gather_sequence(input_, dim) | |
| def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: | |
| return scatter_sequence(grad_output, ctx.dim), None | |
| def gather_for_loss(input_: torch.Tensor, dim: int = 1) -> torch.Tensor: | |
| """Gather for loss computation; gradients are scattered automatically on backward.""" | |
| if not _cp_config.enabled: | |
| return input_ | |
| return _GatherForward.apply(input_, dim) | |
| def apply_ulysses_attention( | |
| q: torch.Tensor, | |
| k: torch.Tensor, | |
| v: torch.Tensor, | |
| heads: int, | |
| attention_fn, | |
| mask: Optional[torch.Tensor] = None, | |
| ) -> torch.Tensor: | |
| """ | |
| Attention with Ulysses-style context parallelism. | |
| 1. All-to-all: scatter sequence dim, gather head dim | |
| 2. run standard attention (each GPU handles a subset of heads over the full sequence) | |
| 3. All-to-all: scatter head dim, gather sequence dim | |
| Args: | |
| q: Query [B, S, H*D] | |
| k: Key [B, S, H*D] | |
| v: Value [B, S, H*D] | |
| heads: number of attention heads | |
| attention_fn: attention function | |
| mask: optional attention mask | |
| Returns: | |
| attention output [B, S, H*D] | |
| """ | |
| if not _cp_config.enabled: | |
| return attention_fn(q, k, v, heads, mask) | |
| B, S, HD = q.shape | |
| H = heads | |
| D = HD // H | |
| cp_size = _cp_config.cp_size | |
| assert H % cp_size == 0, f"Number of heads ({H}) must be divisible by cp_size ({cp_size})" | |
| # 1. Reshape to expose head dimension: [B, S, H, D] | |
| q = q.view(B, S, H, D) | |
| k = k.view(B, S, H, D) | |
| v = v.view(B, S, H, D) | |
| # 2. All-to-all: scatter heads (dim=2), gather sequence (dim=1) | |
| # [B, S_local, H, D] -> [B, S_full, H//cp_size, D] | |
| q = seq_all_to_all(q, scatter_dim=2, gather_dim=1) | |
| k = seq_all_to_all(k, scatter_dim=2, gather_dim=1) | |
| v = seq_all_to_all(v, scatter_dim=2, gather_dim=1) | |
| # 3. Reshape back for attention: [B, S*cp_size, (H//cp_size)*D] | |
| _, S_full, H_local, _ = q.shape | |
| q = q.view(B, S_full, H_local * D) | |
| k = k.view(B, S_full, H_local * D) | |
| v = v.view(B, S_full, H_local * D) | |
| # 4. Execute attention with local heads | |
| if mask is not None: | |
| mask = gather_sequence(mask, dim=-1) if mask.shape[-1] == S else mask | |
| out = attention_fn(q, k, v, H_local, mask) | |
| # 5. Reshape for all-to-all back: [B, S_full, H_local, D] | |
| out = out.view(B, S_full, H_local, D) | |
| # 6. All-to-all: scatter sequence (dim=1), gather heads (dim=2) | |
| # [B, S_full, H//cp_size, D] -> [B, S_local, H, D] | |
| out = seq_all_to_all(out, scatter_dim=1, gather_dim=2) | |
| # 7. Reshape to original format: [B, S, H*D] | |
| out = out.view(B, S, HD) | |
| return out | |
| def pad_to_cp_divisible(tensor: torch.Tensor, dim: int = 1) -> Tuple[torch.Tensor, int]: | |
| """ | |
| Pad a tensor so the dimension is divisible by cp_size. | |
| Args: | |
| tensor: input tensor | |
| dim: dimension to pad | |
| Returns: | |
| (padded_tensor, original_length) | |
| """ | |
| if not _cp_config.enabled: | |
| return tensor, tensor.shape[dim] | |
| original_length = tensor.shape[dim] | |
| cp_size = _cp_config.cp_size | |
| if original_length % cp_size == 0: | |
| return tensor, original_length | |
| pad_length = cp_size - (original_length % cp_size) | |
| pad_spec = [0] * (2 * tensor.ndim) | |
| pad_idx = 2 * (tensor.ndim - 1 - dim) | |
| pad_spec[pad_idx + 1] = pad_length # pad at the end of this dimension | |
| padded = torch.nn.functional.pad(tensor, pad_spec, mode='constant', value=0) | |
| return padded, original_length | |
| def unpad_from_cp(tensor: torch.Tensor, original_length: int, dim: int = 1) -> torch.Tensor: | |
| """ | |
| Remove context-parallel padding. | |
| Args: | |
| tensor: padded tensor | |
| original_length: original length | |
| dim: padded dimension | |
| Returns: | |
| the unpadded tensor | |
| """ | |
| if not _cp_config.enabled or tensor.shape[dim] == original_length: | |
| return tensor | |
| return tensor.narrow(dim, 0, original_length) | |
| # ============================================================================= | |
| # ============================================================================= | |
| def prepare_cp_inputs( | |
| video_latent: torch.Tensor, | |
| audio_latent: Optional[torch.Tensor] = None, | |
| video_positions: Optional[torch.Tensor] = None, | |
| audio_positions: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], dict]: | |
| """ | |
| Prepare inputs for context parallelism. | |
| Shards the latents and positional encodings across GPUs. | |
| Args: | |
| video_latent: [B, T, ...] patchified video latent | |
| audio_latent: [B, T, ...] patchified audio latent | |
| video_positions: [B, 1, T, D] video positional encoding | |
| audio_positions: [B, 1, T, D] audio positional encoding | |
| Returns: | |
| the sharded (video_latent, audio_latent, video_positions, audio_positions, metadata) | |
| """ | |
| metadata = { | |
| 'video_original_length': video_latent.shape[1] if video_latent is not None else 0, | |
| 'audio_original_length': audio_latent.shape[1] if audio_latent is not None else 0, | |
| } | |
| if not _cp_config.enabled: | |
| return video_latent, audio_latent, video_positions, audio_positions, metadata | |
| if video_latent is not None: | |
| video_latent = scatter_sequence(video_latent, dim=1) | |
| if audio_latent is not None: | |
| audio_latent = scatter_sequence(audio_latent, dim=1) | |
| if video_positions is not None: | |
| video_positions = scatter_sequence(video_positions, dim=2) | |
| if audio_positions is not None: | |
| audio_positions = scatter_sequence(audio_positions, dim=2) | |
| return video_latent, audio_latent, video_positions, audio_positions, metadata | |
| def gather_cp_outputs( | |
| video_output: torch.Tensor, | |
| audio_output: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: | |
| """ | |
| Gather context-parallel outputs. | |
| Args: | |
| video_output: video output shard | |
| audio_output: audio output shard | |
| Returns: | |
| the gathered (video_output, audio_output) | |
| """ | |
| if not _cp_config.enabled: | |
| return video_output, audio_output | |
| video_output = gather_sequence(video_output, dim=1) | |
| if audio_output is not None: | |
| audio_output = gather_sequence(audio_output, dim=1) | |
| return video_output, audio_output | |
| def get_cp_world_size() -> int: | |
| """Return the context-parallel world size.""" | |
| return _cp_config.cp_size if _cp_config.enabled else 1 | |
| def get_cp_rank() -> int: | |
| """Return the context-parallel rank.""" | |
| return _cp_config.cp_rank if _cp_config.enabled else 0 | |
| def is_cp_enabled() -> bool: | |
| """Return True when context parallelism is enabled.""" | |
| return _cp_config.enabled | |
| def compute_cp_divisible_frames(frames: int, cp_size: int, temporal_stride: int = 8) -> int: | |
| """ | |
| Compute a frame count divisible by both cp_size and the temporal stride. | |
| The VAE requires (frames - 1) % 8 == 0 and context parallelism requires | |
| latent_frames % cp_size == 0. | |
| Args: | |
| frames: target frame count | |
| cp_size: CP size | |
| temporal_stride: VAE temporal stride | |
| Returns: | |
| the adjusted frame count | |
| """ | |
| # LTX2 VAE: latent_frames = (frames - 1) // 8 + 1 | |
| frames = ((frames - 1) // temporal_stride) * temporal_stride + 1 | |
| latent_frames = (frames - 1) // temporal_stride + 1 | |
| if latent_frames % cp_size != 0: | |
| latent_frames = ((latent_frames // cp_size) + 1) * cp_size | |
| frames = (latent_frames - 1) * temporal_stride + 1 | |
| return frames | |