Open-Sora Plan

Open-Sora Plan v1.5.0 introduces several architectural and scaling improvements over v1.3.0:

• SUV (Sparse DiT with U-shaped structure): An evolution of the skiparse attention mechanism. It extends sparse DiT to a U-shaped sparse structure, achieving performance comparable to dense DiT while maintaining speed advantages.
• High-Compression WFVAE: Introduces a WFVAE with an 8x8x8 downsampling rate. This achieves performance comparable to the common 4x8x8 downsampling rate while halving the latent shape, which significantly reduces attention sequence length.
• Data and Model Scaling: The release utilizes 1.1B high-quality images and 40m high-quality videos, with model parameters scaled up to 8.5B.
• Adaptive Grad Clipping: Replaces the complex 'tainted batch discarding' strategy from v1.3.0 with a simpler adaptive gradient norm thresholding and clipping mechanism, making it more compatible with various parallel training strategies.

The CausalVideoVAE is designed for efficient spatial-temporal compression (4×8×8) and supports joint image-video training. It inherits from the Stable-Diffusion Image VAE structure with two key modifications:

1. CausalConv3D: Replaces Conv2D with CausalConv3D to enable joint training. It applies special treatment to the first frame (which lacks preceding frames) to allow seamless encoding of both single images and video sequences.
2. Tail Initialization: Instead of standard average or center initialization, the model uses a specific 'tail initialization' method. This ensures the model can reconstruct images and videos effectively even before training begins.

Inference Tip: To handle high-resolution or long-duration videos without exceeding GPU memory (e.g., on an 80GB GPU), use tile convolution. This allows for inference with nearly constant memory usage regardless of resolution or duration.

Skiparse Attention is a method to accelerate training of Full 3D Attention models by reducing the sequence length in the attention operation by a sparse ratio $k$, while maintaining global spatiotemporal modeling capabilities. It modifies only the Attention component within the Transformer Block.

Single Skip Mode

Organizes elements into $k$ sub-sequences where each token performs attention with tokens spaced $k-1$ apart.

Group Skip Mode

Groups adjacent tokens in segments of length $k$, then bundles these groups with other groups spaced $k-1$ groups apart into a single scope. This introduces a degree of locality while preserving a large receptive field.

Implementation Note: Always apply RoPE (Rotary Positional Embedding) before the Skiparse operation, as the rearrangement causes the sequence to lose its original spatial positions.

# Single Skip Implementation
# x.shape: (B,N,C)
def single_skip_rearrange(x, sparse_k):
	return rearrange(x, 'b (g k) d -> (k b) g d', k=sparse_k)

def reverse_sparse(x, sparse_k):
	return rearrange(x, '(k b) g d -> b (g k) d', k=sparse_k)

q, k, v = Q(x), K(x), V(x)
q = add_rope(q)
k = add_rope(k)
q = single_skip_rearrange(q)
k = single_skip_rearrange(k)
v = single_skip_rearrange(v)
hidden_states = F.scaled_dot_product_attention(q=q,k=k,v=v)
output = reverse_sparse(hidden_states)

# Group Skip Implementation
# x.shape: (B,N,C)
def group_skip_rearrange(x, sparse_k):
	return rearrange(x, ' b (n m k) d -> (m b) (n k) d', m=sparse_k, k=sparse_k)

def reverse_sparse(x, sparse_k):
	return rearrange(x, '(m b) (n k) d -> b (n m k) d', m=sparse_k, k=sparse_k)

q, k, v = Q(x), K(x), V(x)
q = add_rope(q)
k = add_rope(k)
q = group_skip_rearrange(q)
k = group_skip_rearrange(k)
v = group_skip_rearrange(v)
hidden_states = F.scaled_dot_product_attention(q=q,k=k,v=v)
output = reverse_sparse(hidden_states)

Open-Sora Plan uses a Bucket strategy for dynamic training to support arbitrary video lengths and resolutions without the inefficiencies of padding or the complexity of Patch n' Pack. This strategy is decoupled from the model code and acts as a plug-and-play video sampling strategy.

How the Bucket Strategy Works

1. Sort by frame: Count frames in all video data and sort them to group similar data.
2. Group megabatch: Divide sorted data into groups called 'megabatches'. Most videos in a megabatch will have similar frame counts.
3. Re-organize megabatch: For boundary cases (e.g., a 1-frame video in a 61-frame megabatch), randomly replace the minority data with majority data to ensure uniform frame counts within the megabatch.
4. Shuffle megabatch: Shuffle data both within and between megabatches to maintain randomness.

By replacing each sample's frame sequence with (frame × height × width), the data dimension processed by each GPU remains identical every step, preventing synchronization delays where one GPU waits for another processing a longer video.