LiteGen

A light-weight and high-efficient training framework for accelerating diffusion tasks.

📜 About

LiteGen is a lightweight and high efficient training acceleration framework specifically designed for diffusion tasks, and has been applied and validated on video generation project Vchitech-2.0. This framework integrates multiple training optimization techniques and offers a user-friendly interface, allowing researchers and developers to easily scale from single-GPU setups to multi-node, multi-GPU environments.

🚀 Installation

LiteGen uses uv to manage Python and project dependencies. The locked core environment targets Python 3.12 on Linux x86_64, with PyTorch 2.13.0, CUDA 13.2, and Triton 3.7.1. FlashAttention and Apex are optional CUDA extensions and are not installed by the core environment.

LiteGen itself declares version ranges (PyTorch >=2.8,<3) and contains no CUDA-version-specific code, so a project that embeds LiteGen can resolve it against its own PyTorch and CUDA build instead of the versions locked here. Only the combination locked in uv.lock has been tested.

Create the project virtual environment and install the locked dependencies:

uv sync --locked

The optional extensions require a separate GPU-node build after the core environment is ready. They are not built or tested by the standard upgrade checks:

# Optional, CUDA-node-only installs (not part of core validation):
uv sync --locked --extra flash-attn
uv sync --locked --extra apex

FlashAttention needs a CUDA toolkit matching torch.version.cuda, ninja, packaging, and psutil; use a non-isolated source build if the selected release does not provide a matching wheel. Apex needs matching nvcc and compiler/C++ ABI settings, and should be built with APEX_CPP_EXT=1 APEX_CUDA_EXT=1 and build isolation disabled. Neither extension is imported, compiled, or tested by the core environment checks.

Run commands inside the environment with uv run:

uv run python your_script.py

Alternatively, activate the generated environment directly:

source .venv/bin/activate

Code quality

GitHub Actions runs Ruff lint and formatting checks on every pull request and push to any branch. Both checks report failures independently. To require them before merging, select Ruff lint and Ruff format as required status checks in the branch rules for the branches you want to protect.

Ruff's rules live in pyproject.toml; its version is locked in uv.lock. The dev group includes the lightweight lint group. With the development environment prepared, run the same read-only checks manually:

uv run --offline --no-sync ruff check --no-fix .
uv run --offline --no-sync ruff format --check .

CI installs only the lint group, without LiteGen or CUDA dependencies. For local lint-only work on Linux x86_64, use a separate environment so the training environment is preserved:

UV_PROJECT_ENVIRONMENT=.venv-lint uv sync --locked --only-group lint
UV_PROJECT_ENVIRONMENT=.venv-lint uv run --offline --no-sync ruff check --no-fix .
UV_PROJECT_ENVIRONMENT=.venv-lint uv run --offline --no-sync ruff format --check .

Formatting and lint fixes are explicit edits to review before committing:

uv run --offline --no-sync ruff check --fix .
uv run --offline --no-sync ruff format .

No Git commit or push hooks are installed or required, including when using Paseo's commit and push buttons. If an older checkout installed pre-commit, uninstall those hooks using that installation's pre-commit uninstall and pre-commit uninstall --hook-type pre-push commands.

✨ Support Features

• vae:
  • sliced VAE encoding helper (litegen.utils.encode.vae_encode)
• ema model:
  • sharded EMA (Exponential Moving Average)
• distributed optimization
  • DDP
  • ZeRO-1 (via zero1)
  • FSDP1 (via fsdp1)
  • FSDP2 and HSDP (via fsdp2 / hsdp)
  • ZeRO-2 configuration placeholder
• precision
  • optional whole-model dtype cast and FSDP2 mixed precision
  • FP32 master weights for AdamW under ZeRO-1, FSDP2 and HSDP
• memory optimization
  • Grad activation checkpointing
  • selective checkpointing
  • activation offload
• tensor artifact cache (litegen.cache)

We also provide easy-to-use interfaces for common operations such as model loading and saving, etc. LiteGen allows users to focus on generative algorithm development and training without getting bogged down in implementation details.

🔨 Usage

Implementing LiteGen's optimizations involves two straightforward steps:

1. Configuration: Adjust the relevant fields in your config file to enable desired optimizations.
2. Integration: Utilize the api from the LiteGen instance in your codebase. These simple steps allow you seamlessly integrating optimizations into your existing workflow.

Quick Start Guide

Follow these steps to integrate LiteGen into your project:

1. Create LiteGen

Load your YAML configuration into a validated MainConfig and create a LiteGen instance from it (see Config file for the available fields and loading options):

import litegen
from litegen.utils.config import MainConfig

config = MainConfig.load_yaml("configs/config.yaml")
gen = litegen.initialize(config)

2. Initialize Components

Use prepare to hand native PyTorch objects to the runtime. It accepts models, optimizers, datasets, and existing data loaders, and returns prepared objects in the same order as the input arguments. Other objects, including functions such as vae.encode, are rejected. A dataset is wrapped in a DataLoader whose sampler is built from global_batch_size, max_steps and global_seed; unless the dataset defines collate_fn, each sample is expected to be a (data, caption) pair.

model, optimizer, text_encoder, dataloader = gen.prepare(
    model,  # A model (multiple models are supported)
    optimizer,  # Optimizer for the model
    text_encoder,  # Untrainable models (e.g., encoders in diffusion tasks)
    dataset,  # Your dataset
)

The two steps described above constitute the minimal code changes required to implement LiteGen's optimizations. This approach allows for quick integration while leveraging LiteGen's performance enhancements.

In the following sections, we provide a detailed explanation of the specific optimizations LiteGen offers and how to configure the corresponding key-value pairs in the config file.

Optimizations

DDP, ZeRO-1, FSDP1, FSDP2, and HSDP

Select the trainable model's parallel strategy under optimization.parallel:

optimization:
  parallel:
    strategy: ddp    # ddp, zero1, fsdp1, fsdp2, or hsdp

zero2 is a reserved name and currently raises NotImplementedError when configured. zero1, fsdp2, and hsdp require a CUDA distributed runtime with a data parallel size of at least 2, and HSDP requires at least two nodes; HSDP has not yet been validated on real multi-node runs. FSDP2 and HSDP do not support activation offload or optimization.compile.enabled yet, and HSDP does not support EMA. Strategy-specific options live in parallel.zero1 and parallel.fsdp2 (see the field reference below). prepare() accepts multiple models and processes them before dependent optimizers.

Selective Activation Checkpointing

LiteGen incorporates activation checkpointing, a common optimization technique for reducing memory usage, and simplifies its usage. Furthermore, when sufficient memory is available, we allow for selective application of activation checkpointing to specific modules, thereby reducing performance overhead.

Example configuration:

optimization:
  memory:
    activation_checkpoint:
      enabled: true
      selective_ratio: 0.2  # 0: all detected blocks, 1: no blocks

Note:

1. Activation checkpointing applies to every model prepared with an optimization profile that enables it; use optimization.profiles and prepare(..., profile=...) to enable it for some models only.
2. Implement get_fsdp_wrap_module_list() in your model class to specify modules for checkpointing.
3. If not implemented, LiteGen automatically detects and applies checkpointing to repetitive module structures in the model (e.g., repeated transformer blocks in DiT models).
4. With fsdp2 or hsdp, the candidates are the planned sharding units instead, or the modules listed in activation_checkpoint.target_paths.

Activation Offload

To further conserve GPU memory, we have implemented CPU offloading for activations in our system. This technique effectively overlaps computation and communication, significantly reducing memory usage during training with minimal additional performance overhead.

You can enable this feature using the following configuration:

optimization:
  memory:
    activation_offload: true

Note: Activation offload applies to the blocks selected by activation checkpointing, so it only takes effect when activation checkpointing is enabled with selective_ratio below 1. It is not supported with fsdp2 or hsdp.

Sequence Parallel

Sequence Parallel is represented by optimization.parallel.sp_size and is reserved for a future implementation. Values greater than one currently raise NotImplementedError when LiteGen builds its parallel plan on CUDA.

You can configure this feature as follows:

optimization:
  parallel:
    sp_size: 8  # Sequence parallel degree (reserved; >1 is not implemented)

Note:

When this capability is implemented, it will require scatter/gather operations and an attention processor mapping for supported models.

EMA Model

The Exponential Moving Average (EMA) model is a common technique used to smooth parameter updates and achieve better training results. LiteGen integrates EMA Model functionality with support for parameter sharding to conserve GPU memory, while providing an easy-to-use interface.

Example configuration:

ema:
  enable: True   # Enable EMA
  decay: 0.9999  # Default decay used by update_ema()

When EMA is enabled, prepare creates the EMA state for the first trainable model and initializes it as a copy of that model. Under ddp, zero1, and fsdp2, the EMA state is sharded across data-parallel ranks; under fsdp1, the EMA copy is wrapped with FSDP. HSDP does not support EMA yet. The ema.sharded field is accepted but not currently read.

User Interface: LiteGen provides a simple method to update the EMA model:

gen.update_ema()              # uses ema.decay
gen.update_ema(decay=0.9999)  # explicit decay

Checkpoint saving and loading

LiteGen provides interfaces for saving and loading checkpoints without dealing with the intricacies of distributed state.

Saving Checkpoints

We offer three separate interfaces for saving model, optimizer, and EMA model states. Call them on every rank:

gen.save_model(output_folder=None, filename=None, step=None)
gen.save_optimizer(output_folder=None, filename=None, step=None, state_format="auto")
gen.save_ema(output_folder=None, filename=None, step=None)

1. Specify output_folder and filename to determine the checkpoint file location.
2. If output_folder is unspecified, the files are written to checkpoint_dir from the config (default: results_dir, or results/<exp_name>).
3. Without a specified filename, the system uses the exp_name from the config as the checkpoint prefix:
   • Model: [exp_name].pth
   • Optimizer: [exp_name].optim_state.pth
   • EMA model: [exp_name].ema.pth
4. If step is provided without a filename, the system appends the step information to the exp_name prefix: [exp_name]_step[StepNum]. The step is only used in the file name and is not stored in the model file.
5. save_model gathers the full model state and rank 0 writes one file; a filename ending in .safetensors selects the safetensors format.
6. Under ZeRO-1, save_optimizer writes rank-local shards plus a small manifest by default; state_format="full" writes one consolidated file instead.

FSDP2 and DDP models can also be saved and loaded with PyTorch Distributed Checkpoint through gen.save_model_state(path, model=..., optimizer=..., ema=...) and gen.load_model_state(...).

Loading Checkpoints

prepare applies the following config fields:

1. init_from:

   • Loads full model weights (.pth or .safetensors) into the first model passed to prepare, before it is wrapped by the parallel strategy.
   • Loads only model weights, not optimizer state or EMA weights.

   Example:

   init_from: 'path_to_the_init_model/model.pth'

2. resume_from:

   • After prepare, loads the model weights into the trainable model and the optimizer state from the file with the same prefix, e.g. model_10.optim_state.pth for model_10.pth.
   • The training step, data loader position, RNG state, and EMA state are not restored.

   Example:

   resume_from: 'path_to_the_resumed_model/model_10.pth'

   Note: resume_from is applied after init_from, so its weights take precedence if both are specified.

gen.load(model, path) loads model weights into a prepared model at any time.

auto_resume and ema.resume_from are accepted by the config but are not currently applied by prepare. Do not set auto_resume together with resume_from: loading the optimizer state then raises an error.

Config file

LiteGen ships a typed configuration module, litegen.utils.config, built on Pydantic:

• Config is the base class. It keeps unknown fields and converts nested dictionaries (also inside lists) into Config objects recursively, so every value can be accessed with the config.key syntax (config["key"], config.get("key", default) and "key" in config work as well; a missing key raises KeyError). Dictionaries with non-string keys, such as an epoch schedule {1: 0.1, 2: 0.01}, are kept as plain dictionaries so their keys are preserved.
• MainConfig is the entry class. It declares the runtime, optimization, algorithm, and training fields read by LiteGen, validates types and ranges when the config is loaded, and types the nested runtime, optimization, and ema sections.
• User-defined fields for your model, dataset, or training script are kept as extra fields on the same object. They belong at the top level (or in sections of your own); the runtime, optimization and ema sections and everything nested in them are defined by LiteGen and reject unknown keys.
• Named optimization profiles under optimization.profiles inherit omitted fields from the general settings; select one with gen.prepare(..., profile="<name>").

1. Define a YAML file. Example:

exp_name: 'video_generation_exp1'
results_dir: 'path_to_the_results_dir'

# Checkpoint loading
init_from: 'path_to_the_init_model/model.pth'
auto_resume: False

# Runtime and performance optimization
runtime:
  device: cuda
  distributed:
    backend: nccl
    dist_url: env://
    dist_timeout: 1800
optimization:
  parallel:
    strategy: fsdp1
  memory:
    activation_checkpoint:
      enabled: true
      selective_ratio: 0
    activation_offload: true

... # Other arguments

2. Load and validate the config file in Python:

import argparse
from litegen.utils.config import MainConfig

parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="configs/config.yaml", help="config file")
args = parser.parse_args()
config = MainConfig.load_yaml(args.config)

Other constructors are available for configs that do not come from a YAML file:

config = MainConfig.from_dict({"optimization": {"precision": {"cast_model": "bf16"}}})  # plain dictionary
config = MainConfig.from_omegaconf(cfg)  # Hydra / OmegaConf DictConfig, requires the optional `omegaconf` package
config = MainConfig(optimization={"parallel": {"strategy": "ddp"}})  # keyword arguments

Invalid values raise a ValueError that lists every offending field:

Error to load config, found the following problems:
  - field 'optimization.precision.cast_model': Input should be 'fp32', 'bf16' or 'fp16'
  - field 'ema.decay': Input should be less than or equal to 1
Please check your configuration.

config.to_dict() and config.to_yaml() export the validated config, for example to log it or to save it next to a checkpoint.

3. Initialize the LiteGen instance using the config:

gen = litegen.initialize(config)

initialize requires a validated MainConfig; load YAML with MainConfig.load_yaml(...) (or construct a MainConfig explicitly) before passing it in. Plain dictionaries and attribute-style objects are not accepted by the runtime entry point. The validated config is available as gen.config.

Here we outline the configuration fields defined by LiteGen together with their default values. Explicit configuration is recommended to prevent errors and ambiguity.

# experiment and filepath
exp_name: experiment name, used as the checkpoint filename prefix (default: default_exp)
results_dir: experiment output directory, used as the default checkpoint_dir (default: None)
checkpoint_dir: default directory for saved checkpoints (default: results_dir, or results/<exp_name>)
init_from: full model weights loaded into the first model passed to prepare (default: None)
resume_from: model checkpoint to resume model weights and optimizer state from (default: None)
auto_resume: accepted but not currently applied by prepare (default: False)

# runtime and distributed environment
runtime:
  device: runtime device, cuda or cpu; npu is reserved and raises NotImplementedError (default: cuda when available, otherwise cpu)
  distributed:
    backend: torch.distributed backend (default: nccl on CUDA, gloo on CPU when device is also omitted)
    dist_url: init method of the process group (default: env://)
    dist_timeout: process group initialization timeout in seconds (default: 1800)

# performance optimization
optimization:
  parallel:
    strategy: one of ['ddp', 'zero1', 'zero2', 'hsdp', 'fsdp1', 'fsdp2']; zero2 is not implemented (default: ddp)
    zero1:  # only with strategy zero1
      bucket_size_mb: communication bucket size in MB (default: 64)
      parameter_sync: one of ['bucketed', 'native', 'native_bucket_view'] (default: bucketed)
      overlap_comm: overlap the final backward of an accumulation window with communication (default: True)
      release_grad_buffer: release the model-dtype gradient partition after casting to FP32 (default: True)
      normalize_parameter_layout: make trainable parameter storage contiguous before sharding (default: False)
    fsdp2:  # only with strategy fsdp2 or hsdp
      wrap_strategy: one of ['manual', 'repeated', 'size_based', 'class_based'] (default: repeated)
      wrap_paths: module paths to shard, for wrap_strategy manual (default: [])
      wrap_class_names: module class names to shard, for wrap_strategy class_based (default: [])
      min_num_params: parameter threshold, for wrap_strategy size_based (default: 1)
      exclude_paths: frozen modules kept out of every sharding unit (default: [])
      reshard_after_forward: FSDP2 reshard_after_forward (default: None, torch default)
      root_reshard_after_forward: override for the root unit (default: None)
      last_unit_reshard_after_forward: override for the last unit in traversal order (default: None)
      allow_empty_param_unit_sharding: also shard planned units without parameters (default: False)
    tp_size: tensor parallel degree; values > 1 are not implemented (default: 1)
    sp_size: sequence parallel degree; values > 1 are not implemented (default: 1)
    cp_size: context parallel degree; values > 1 are not implemented (default: 1)
    pp_size: pipeline parallel degree; values > 1 are not implemented (default: 1)
  precision:  # prepare keeps the model's own dtype unless cast_model is set
    cast_model: convert every floating parameter and buffer to one of ['fp32', 'bf16', 'fp16'] (default: None, no conversion)
    cast_exclude_paths: module paths the conversion skips (default: [])
    fsdp2_mixed:  # torch MixedPrecisionPolicy, applied by fsdp2/hsdp only; defaults match torch
      param_dtype: dtype parameters are all-gathered and computed in (default: None, storage dtype)
      reduce_dtype: gradient reduce-scatter dtype (default: None, gradient dtype)
      output_dtype: dtype shard-unit outputs are cast to (default: None, no cast)
      cast_forward_inputs: cast shard-unit inputs to param_dtype when it is set (default: True)
    allow_tf32: set TF32 matmul process-wide, read from the general profile (default: True)
  optimizer:
    fp32_master_weights: maintain FP32 master weights for AdamW under zero1/fsdp2/hsdp when a bf16/fp16 parameter is present (default: False)
    backend: one of ['single_tensor', 'foreach', 'fused']; fsdp2/hsdp master weights require single_tensor (default: single_tensor)
  compile:  # reserved: not applied to models yet; fsdp2/hsdp reject enabled: true
    enabled: (default: False)
    mode: (default: default)
    dynamic: (default: False)
  memory:
    activation_checkpoint.enabled: whether to enable activation checkpointing (default: False)
    activation_checkpoint.selective_ratio: ratio without checkpointing, in [0, 1] (default: 1.0)
    activation_checkpoint.target_paths: modules to checkpoint, fsdp2/hsdp only (default: [])
    activation_offload: whether to offload checkpointed activations; not supported by fsdp2/hsdp (default: False)
  kernels:
    fused_layernorm: replace LayerNorm with Apex FusedLayerNorm when Apex is installed (default: False)
  profiles: named partial optimization configs selected by prepare(profile=...) (default: {})

# algorithm settings, stored for training scripts; LiteGen does not create optimizers
lr: learning rate (default: 1.0e-4)
weight_decay: weight decay (default: 0.0)

# ema
ema:
  enable: whether to enable ema, True or False (default: False)
  sharded: accepted but not currently read (default: True)
  decay: default decay used by update_ema() (default: 0.9999)
  resume_from: accepted but not currently applied by prepare (default: None)

# training settings
global_seed: global random seed (default: 0)
max_steps: max steps number (default: 1)
num_workers: number of workers of dataloader (default: 0)
pin_memory: whether to enable pin_memory for dataloader (default: False)
global_batch_size: total samples used across all ranks in one optimizer step; must be a multiple of the data parallel size (default: 1)

Additionally, users can define custom configuration fields to meet specific requirements for algorithm construction and training script needs. They are kept on the same MainConfig object; nested dictionaries (including those inside lists) become Config objects, while dictionaries with non-string keys stay plain dictionaries. To validate your own fields as well, subclass MainConfig and declare them:

from litegen.utils.config import MainConfig


class LatteConfig(MainConfig):
    num_frames: int = 16
    image_size: int = 256


config = LatteConfig.load_yaml(args.config)
gen = litegen.initialize(config)

🚀 Performance

LiteGen includes activation checkpointing and activation offload support. The configuration keeps placeholders for future tensor/sequence/context/pipeline parallelism; parallel dimensions greater than one currently raise NotImplementedError. The chart below is a historical benchmark from the earlier implementation and is not a claim about the current runtime, which implements data-parallel strategies only.

(AO: Activation Offload, SP: Sequence Parallel)

Sequence parallelism is not available in the current runtime; the benchmark is retained for context only.

🔑 License

This code is licensed under Apache-2.0. The framework is fully open for academic research and also allows free commercial usage. To apply for a commercial license or for other questions or collaborations, please contact yangzhenyu@pjlab.org.cn.