lite_boost.BoostManager

class lite_boost.BoostManager(target=None, config=None)

Modify a supported model in-place for distributed parallel inference.

BoostManager wraps a supported model or pipeline and patches it in-place for multi-NPU parallel inference. Two parallelism strategies are applied automatically based on the detected model components:

• Ulysses Sequence Parallel (USP) for DiT models — patches the forward method and attention layers to enable sequence-dimension parallelism via all_to_all communication. Each device holds full model weights and operates on a slice of the sequence.

• Data Parallel (DP) temporal tiling for VAE models — replaces vae.encode and vae.decode with DP temporal slicing versions that split the video along the temporal dimension into overlapping chunks, distribute them across devices, and gather results.

When a pipeline object (e.g., WanT2V) is passed, both strategies are applied: USP for the DiT model and DP for the VAE. When a raw WanModel is passed, only USP is applied.

The model is modified in-place and returned as-is, so all existing attributes and methods (.to, .cpu, .eval, etc.) continue to work normally.

Parameters:

• target (object, optional) – A supported pipeline object to be parallelized. Supported classes include WanT2V and WanTI2V. When omitted, a BoostManager instance is returned for later use.

• config (str, optional) – Path to a YAML file that lets the user select which module optimizations to enable, parsed into a dict and forwarded to the model's boost function.

Returns:

object, the same instance modified in-place with USP-patched forward and attention methods (for DiT) and DP-patched encode/decode methods (for VAE).

Raises:

RuntimeError – If the model type is not supported by lite_boost.

Examples

>>> import os
>>> os.environ["RANK"] = "0"
>>> os.environ["WORLD_SIZE"] = "1"
>>> from lite_boost import BoostManager
>>> from lite_boost.parallel import initialize_usp
>>> from wan.textimage2video import WanTI2V
>>> initialize_usp()
>>> pipe = WanTI2V(config=cfg, checkpoint_dir=ckpt_dir, ...)
>>> boost_manager = BoostManager()
>>> pipe = boost_manager(pipe)
>>> # or with a per-module optimization config:
>>> pipe = boost_manager(pipe, config="boost.yaml")