Want to Train a Sora-Class Model? Luchen Tech's OpenDiT Delivers 80% Speedup | BlueRun Ventures Headlines

蓝驰创投蓝驰创投··0·0

OpenDiT, an easy-to-use, fast, and memory-efficient system

How to reproduce Sora, the first blockbuster of 2024? BlueRun Ventures recommends checking out this approach — recently, Luchen Tech open-sourced Open-Sora, a complete architectural blueprint for replicating Sora that cuts reproduction costs by 46% and extends model training input sequence length to 819K patches. BlueRun Ventures was Luchen Tech's angel-round investor. We look forward to continued iteration and optimization of Open-Sora.

As the explosive opening to 2024, Sora set an entirely new benchmark to chase. Every text-to-video researcher wants to reproduce its results in the shortest possible time.

According to the technical report disclosed by OpenAI, one of Sora's core innovations is converting visual data into a unified patch representation, then combining Transformer and diffusion models to demonstrate remarkable scaling properties. After the report's release, the paper Scalable Diffusion Models with Transformers — co-authored by Sora core researcher William Peebles and New York University assistant professor of computer science Saining Xie — became a focal point for researchers. Many hoped to use the DiT architecture proposed in the paper as a breakthrough point for exploring viable paths to replicate Sora.

Recently, a project called OpenDiT, open-sourced by Yang You's team at National University of Singapore, has opened new possibilities for training and deploying DiT models.

OpenDiT is an easy-to-use, fast, and memory-efficient system specifically designed to improve training and inference efficiency for DiT applications, including text-to-video and text-to-image generation.

Project repository: https://github.com/NUS-HPC-AI-Lab/OpenDiT

OpenDiT Methodology

OpenDiT provides a high-performance implementation of Diffusion Transformer (DiT) powered by Colossal-AI. During training, video and conditioning information are fed into their respective encoders as inputs to the DiT model. The model then undergoes training and parameter updates through diffusion methods, with updated parameters ultimately synchronized to an EMA (Exponential Moving Average) model. During inference, the EMA model is used directly, taking conditioning information as input to generate corresponding outputs.

Image source: https://www.zhihu.com/people/berkeley-you-yang

OpenDiT leverages ZeRO parallelism to distribute DiT model parameters across multiple machines, initially reducing GPU memory pressure. To achieve a better balance between performance and precision, OpenDiT also employs mixed-precision training. Specifically, model parameters and optimizers are stored in float32 to ensure update accuracy. During computation, the research team designed a float16 and float32 mixed-precision approach tailored for DiT models, accelerating calculations while maintaining model precision.

The EMA method used in DiT models is a strategy for smoothing parameter updates, effectively improving model stability and generalization. However, it creates an additional copy of parameters, increasing memory burden. To further reduce this memory overhead, the team shards the EMA model and distributes it across different GPUs. During training, each GPU only needs to compute and store its portion of EMA parameters, synchronizing updates after each step once ZeRO completes its updates.

FastSeq

In visual generation models like DiT, sequence parallelism is essential for effective long-sequence training and low-latency inference.

However, existing approaches such as DeepSpeed-Ulysses and Megatron-LM Sequence Parallelism face limitations when applied to these tasks — either introducing excessive sequence communication or lacking efficiency when handling small-scale sequence parallelism.

To address this, the team proposed FastSeq, a novel sequence parallelism suitable for both large sequences and small-scale parallelism. FastSeq minimizes sequence communication by using only two communication operators per transformer layer, leverages AllGather to improve communication efficiency, and strategically employs asynchronous ring communication to overlap AllGather with qkv computation, further optimizing performance.

Operator Optimization

Introducing the adaLN module in DiT models to incorporate conditioning information into visual content, while critical for performance, brings numerous element-wise operations that are frequently invoked throughout the model, reducing overall computational efficiency. To solve this, the team proposed an efficient Fused adaLN Kernel that merges multiple operations into one, increasing computational efficiency and reducing I/O overhead for visual information.

Image source: https://www.zhihu.com/people/berkeley-you-yang

In summary, OpenDiT offers the following performance advantages:

1. Up to 80% GPU acceleration, 50% memory savings

  • Efficient operators designed specifically for DiT, including Fused AdaLN, as well as FlashAttention, Fused Layernorm, and HybridAdam.
  • Hybrid parallelism methods including ZeRO, Gemini, and DDP. Sharding the EMA model further reduces memory costs.

2. FastSeq: A novel sequence parallelism method

  • Purpose-built for DiT-like workloads where sequences are typically long but parameters are smaller compared to LLMs.
  • Intra-node sequence parallelism saves up to 48% in communication volume.
  • Breaks through single-GPU memory constraints, reducing overall training and inference time.

3. Ease of use

  • Massive performance gains with just a few lines of code changes.
  • No need for users to understand distributed training implementations.

4. Complete text-to-image and text-to-video generation pipelines

  • Researchers and engineers can easily use OpenDiT pipelines and apply them to real-world applications without modifying parallel components.
  • The team validated OpenDiT's accuracy through text-to-image training on ImageNet and released checkpoints.

Installation and Usage

To use OpenDiT, first install the prerequisites:

  • Python >= 3.10
  • PyTorch >= 1.13 (version >2.0 recommended)
  • CUDA >= 11.6

We recommend using Anaconda to create a new environment (Python >= 3.10) to run the examples:

conda create -n opendit python=3.10 -y
conda activate opendit

Install ColossalAI:

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
git checkout adae123df3badfb15d044bd416f0cf29f250bc86
pip install -e .

Install OpenDiT:

git clone https://github.com/oahzxl/OpenDiT
cd OpenDiT
pip install -e .

(Optional but recommended) Install libraries to speed up training and inference:

pip install triton
pip install flash-attn
git clone https://github.com/NVIDIA/apex.git
cd apex
git checkout 741bdf50825a97664db08574981962d66436d16a
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ --global-option="--cuda_ext" --global-option="--cpp_ext"

Image Generation

You can train the DiT model by executing the following command:

bash train_img.sh
# Use command line
torchrun --standalone --nproc_per_node=2 train.py \
    --model DiT-XL/2 \
    --batch_size 2

All acceleration methods are disabled by default. Here are details on key elements during training:

  • plugin: Supports booster plugins used by ColossalAI, zero2, and ddp. Default is zero2; enabling zero2 is recommended.
  • mixed_precision: Data type for mixed precision training, default is fp16.
  • grad_checkpoint: Whether to enable gradient checkpointing. This saves memory during training. Default is False. Recommended to disable when memory is sufficient.
  • enable_modulate_kernel: Whether to enable modulate kernel optimization to speed up training. Default is False; recommended to enable when GPU < H100.
  • enable_layernorm_kernel: Whether to enable layernorm kernel optimization to speed up training. Default is False; recommended to enable.
  • enable_flashattn: Whether to enable FlashAttention to speed up training. Default is False; recommended to enable.
  • sequence_parallel_size: Sequence parallelism degree. Sequence parallelism is enabled when set > 1. Default is 1; recommended to disable when memory is sufficient.

If you want to use the DiT model for inference, you can run the following code, replacing the checkpoint path with your own trained model.

# Use script
bash sample_img.sh
# Use command line
python sample.py --model DiT-XL/2 --image_size 256 --ckpt ./model.pt

Video Generation

You can train the video DiT model by executing the following command:

# train with command line
torchrun --standalone --nproc_per_node=2 train.py \
    --model vDiT-XL/222 \
    --use_video \
    --data_path ./videos/demo.csv \
    --batch_size 1 \
    --num_frames 16 \
    --image_size 256 \
    --frame_interval 3

python preprocess.py

Code for performing video inference with the DiT model is shown below:

# Use command line
python sample.py \
    --model vDiT-XL/222 \
    --use_video \
    --ckpt ckpt_path \
    --num_frames 16 \
    --image_size 256 \
    --frame_interval 3

DiT Reproduction Results

To validate OpenDiT's accuracy, the team trained DiT using OpenDiT's original method, training the model from scratch on ImageNet for 80k steps on 8xA100. Here are some results from the trained DiT:

Loss also aligns with results listed in the DiT paper:

To reproduce the above results, change the dataset in train_img.py and execute the following command:

torchrun --standalone --nproc_per_node=8 train.py \
    --model DiT-XL/2 \
    --batch_size 180 \
    --enable_layernorm_kernel \
    --enable_flashattn \
    --mixed_precision fp16

Originating in Silicon Valley, BlueRun Ventures was established in 2005 as a venture capital firm focused on early-stage startups.

Currently, BlueRun Ventures manages multiple USD and RMB dual-currency funds in China, with assets under management exceeding RMB 15 billion, making it one of the largest early-stage funds domestically. Its investment stage concentrates on Pre-A and Series A rounds, covering technology, consumer, and healthcare sectors. It has invested in nearly 200 startups, including Li Auto, Waterdrop, QingCloud, Guazi.com, Qudian, Ganji.com, Energy Monster, Gaussian Robotics, Songguo Chuxing, Yuntu Semiconductor, Machenike, Clouds Intelligence, Anxin Wangdun, and BioMap.

BlueRun Ventures has been ranked #1 on Zero2IPO's "China Top 30 Early-Stage Investment Institutions" and ChinaVenture's "China Best Early-Stage Venture Capital Firms TOP30," and was named among Preqin's Top 10 VC fund managers globally for sustained high returns.

Additionally, BlueRun Ventures has received consecutive honors from Forbes China, 36Kr, Cyzone, Caixin Media, CBNweekly, Jiemian, and other media organizations, including "China's Best Early-Stage Firm of the Year," "China's Top Venture Capital Firm," "Most Entrepreneur-Friendly Early-Stage Firm of the Year," and "Most Influential Early-Stage Firm of the Year."