Two-Stage Split Training

This document introduces split training, which can automatically divide the training process into two stages, reducing VRAM usage while accelerating training speed.

(Split training is an experimental feature that has not yet undergone large-scale validation. If you encounter any issues while using it, please submit an issue on GitHub.)

Split Training

In the training process of most models, a large amount of computation occurs in “preprocessing,” i.e., “computations unrelated to the denoising model,” including VAE encoding, text encoding, etc. When the corresponding model parameters are fixed, the results of these computations are repetitive. For each data sample, the computational results are identical across multiple epochs. Therefore, we provide a “split training” feature that can automatically analyze and split the training process.

For standard supervised training of ordinary text-to-image models, the splitting process is straightforward. It only requires splitting the computation of all Pipeline Units into the first stage, storing the computational results to disk, and then reading these results from disk in the second stage for subsequent computations. However, if gradient backpropagation is required during preprocessing, the situation becomes extremely complex. To address this, we introduced a computational graph splitting algorithm to analyze how to split the computation.

Enabling Split Training

Split training already supports Standard Supervised Training and Direct Distillation Training. The --task parameter in the training command controls this. Taking LoRA training of the Qwen-Image model as an example, the pre-split training command is:

modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image/*" --local_dir ./data/diffsynth_example_dataset

accelerate launch examples/qwen_image/model_training/train.py \
  --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image \
  --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image/metadata.csv \
  --max_pixels 1048576 \
  --dataset_repeat 50 \
  --model_id_with_origin_paths "Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors,Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
  --learning_rate 1e-4 \
  --num_epochs 5 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Qwen-Image_lora" \
  --lora_base_model "dit" \
  --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1" \
  --lora_rank 32 \
  --use_gradient_checkpointing \
  --dataset_num_workers 8 \
  --find_unused_parameters

After splitting, in the first stage, make the following modifications:

Change --dataset_repeat to 1 to avoid redundant computation
Change --output_path to the path where the first-stage computation results are saved
Add the additional parameter --task "sft:data_process"
Fill in offload_models with the models that do not require forward computation, in the same format as model_id_with_origin_paths
- Alternatively, you can directly remove from --model_id_with_origin_paths the models that do not require forward computation. However, you must ensure that the corresponding models are not indirectly invoked in the pipeline, which means you need to understand the internal details of the Pipeline.

accelerate launch examples/qwen_image/model_training/train.py \
  --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image \
  --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image/metadata.csv \
  --max_pixels 1048576 \
  --dataset_repeat 1 \
  --model_id_with_origin_paths "Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors,Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors" \
  --offload_models "Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors" \
  --learning_rate 1e-4 \
  --num_epochs 5 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Qwen-Image-LoRA-splited-cache" \
  --lora_base_model "dit" \
  --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1" \
  --lora_rank 32 \
  --use_gradient_checkpointing \
  --dataset_num_workers 8 \
  --find_unused_parameters \
  --task "sft:data_process"

In the second stage, make the following modifications:

Change --dataset_base_path to the --output_path of the first stage
Remove --dataset_metadata_path
Add the additional parameter --task "sft:train"
Fill in offload_models with the models that do not require forward computation, in the same format as model_id_with_origin_paths
- Alternatively, you can directly remove from --model_id_with_origin_paths the models that do not require forward computation. However, you must ensure that the corresponding models are not indirectly invoked in the pipeline, which means you need to understand the internal details of the Pipeline.

accelerate launch examples/qwen_image/model_training/train.py \
  --dataset_base_path "./models/train/Qwen-Image-LoRA-splited-cache" \
  --max_pixels 1048576 \
  --dataset_repeat 50 \
  --model_id_with_origin_paths "Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors,Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors" \
  --offload_models "Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
  --learning_rate 1e-4 \
  --num_epochs 5 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Qwen-Image-LoRA-splited" \
  --lora_base_model "dit" \
  --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1" \
  --lora_rank 32 \
  --use_gradient_checkpointing \
  --dataset_num_workers 8 \
  --find_unused_parameters \
  --task "sft:train"

We provide sample training scripts and validation scripts located at examples/qwen_image/model_training/special/split_training.

Principles of the Computational Graph Splitting Algorithm

The training framework splits the computational units in the Pipeline through the split_pipeline_units method of DiffusionTrainingModule. The following describes the detailed principles of the computational graph splitting algorithm.

Problem Definition

To precisely characterize the splitting process, we first formalize the computation pipeline. Suppose the pipeline consists of \(n\) computational units (Pipeline Unit), and let the set of units be \(V=\{u_1,u_2,\dots,u_n\}\). Each unit \(u\in V\) has the following properties:

Input parameter set \(\operatorname{in}(u)\): declared by input_params, input_params_posi and input_params_nega, representing the data items that must be read before the computation of \(u\);
Output parameter set \(\operatorname{out}(u)\): declared by output_params, representing the data items produced and written into the data cache after the computation of \(u\);
Associated model set \(\mathcal{M}(u)\): declared by onload_model_names, representing the models that the computation of \(u\) depends on.

All parameters constitute the parameter space \(\mathcal{P}=\bigcup_{u\in V}\left(\operatorname{in}(u)\cup\operatorname{out}(u)\right)\).

Definition 1 (Data Dependency Edge) Let \(p\in\mathcal{P}\) be a parameter. If there exist units \(u_i,u_j\in V\) such that \(p\in\operatorname{out}(u_i)\cap\operatorname{in}(u_j)\), and \(u_i\) is the most recent producer of \(p\) (i.e., the unit with the latest execution order among all units that produce \(p\)), then there exists a data dependency edge \((u_i,u_j)\) between \(u_i\) and \(u_j\), whose semantics is that the computation of \(u_j\) must occur after the computation of \(u_i\) completes.

Accordingly, the computation pipeline is abstracted as a directed acyclic graph \(G=(V,E)\), where \(E\) is the set of all data dependency edges.

Definition 2 (Directly Related Unit) Given a set of models \(\mathcal{W}\) that require gradient backpropagation (specified by trainable_models and lora_base_model, which are respectively the model components being trained and the model components being trained with LoRA). If a unit \(u\in V\) satisfies \(\mathcal{M}(u)\cap\mathcal{W}\neq\varnothing\), then \(u\) is called a directly related unit, whose computation involves the invocation of trainable models.

Definition 3 (Computational Graph Splitting Problem) Given a graph \(G=(V,E)\) and a model set \(\mathcal{W}\), find a bipartition \((V_1,V_2)\) of \(V\) such that \(V_1\) is the minimal set that contains all directly related units and satisfies the following closure conditions, with \(V_2=V\setminus V_1\):

(C1) Forward closure: if \(u\in V_1\) and \((u,v)\in E\), then \(v\in V_1\); that is, \(V_2\) contains no unit that depends on the outputs of \(V_1\);

(C2) Updating-chain closure: for any parameter \(p\in\mathcal{P}\), let its updating chain \(\mathbf{c}(p)=(u^{(1)},u^{(2)},\dots,u^{(k)})\) be the sequence of all units that produce \(p\) in execution order. If \(p\) is first consumed at \(u^{(i)}\) within \(V_1\) and \(i<k\), then \(u^{(i+1)},\dots,u^{(k)}\in V_1\); that is, the order in which parameters are consumed and updated inside \(V_1\) is consistent with the execution order of the whole graph.

Purpose of splitting: the computational results of the units in \(V_2\) are independent of the model parameters and can be reused across multiple epochs. Therefore, they are executed in the data preprocessing stage (the first stage) and cached to disk; the computation of the units in \(V_1\) requires gradient backpropagation and is executed in the training stage (the second stage).

Algorithm Design

Consider the set operator \(T:2^V\to 2^V\):

\[ T(X)=X\cup F(X)\cup U(X), \]

where \(F\) is the forward reachability operator, and \(F(X)\) is the transitive closure of all successor units reachable from \(X\) along data dependency edges (in the direction producer \(\to\) consumer); \(U\) is the updating-chain backtracking operator, defined as follows: for any parameter \(p\), if \(p\) is first consumed at \(u^{(i)}\) within \(X\) and \(i<k\) (\(k\) being the length of the updating chain \(\mathbf{c}(p)\)), then \(U(X)\) contains those units among the subsequent updating units \(u^{(i+1)},\dots,u^{(k)}\) of \(\mathbf{c}(p)\) that are not in \(X\).

Proposition 1 (Monotonicity and Termination) \(T\) is a monotone operator, i.e., \(X\subseteq Y\Rightarrow T(X)\subseteq T(Y)\). Starting from the set of directly related units \(X_0=\{u\in V\mid\mathcal{M}(u)\cap\mathcal{W}\neq\varnothing\}\), iterate \(X_{k+1}=T(X_k)\). Since \(X_k\) is monotonically non-decreasing and \(X_k\subseteq V\) (\(|V|=n\)), the iteration terminates within at most \(n\) steps at the least fixed point \(X^*\), satisfying \(T(X^*)=X^*\).

Proposition 2 (Consistency) Let \(V_1=X^*\) and \(V_2=V\setminus V_1\). Then \((V_1,V_2)\) satisfies conditions (C1) and (C2). From the fixed-point property of \(X^*\), \(F(X^*)\subseteq X^*\) implies (C1), and \(U(X^*)\subseteq X^*\) implies (C2).

Algorithm Pseudocode

The implementation of the above algorithm is as follows:

def split_pipeline_units(units, model_names):
    # Step 1: initialize X_0, the set of directly related units
    related = {id for id, unit in enumerate(units)
               if unit.onload_model_names is not None
               and any(m in unit.onload_model_names for m in model_names)}

    # Step 2: build the data dependency edges E and the updating chain c(p) of each parameter
    edges = build_edges(units)
    chains = build_chains(units)

    # Step 3: fixed-point iteration X_{k+1} = X_k ∪ F(X_k) ∪ U(X_k)
    while True:
        before = len(related)
        related = forward_reachable(edges, related)       # F(X_k)
        related = updating_units(units, chains, related)  # U(X_k)
        if len(related) == before:
            break

    # Step 4: output the bipartition (V1, V2)
    related_units   = [units[i] for i in sorted(related)]
    unrelated_units = [units[i] for i in range(len(units)) if i not in related]
    return related_units, unrelated_units

The construction details of each sub-function are as follows.

build_edges scans the unit sequence and maintains a mapping from “parameter \(\to\) most recent producer”: when the input parameter \(p\) of unit \(u_j\) already has a producer \(u_i\), an edge \((u_i,u_j)\) is recorded; subsequently, the mapping is updated with each output parameter \(p\) of \(u_j\).

def build_edges(units):
    last_producer = {}
    edges = []
    for id, unit in enumerate(units):
        for param in unit.fetch_input_params():
            if param in last_producer:
                edges.append((last_producer[param], id))
        for param in unit.fetch_output_params():
            last_producer[param] = id
    return edges

build_chains records the updating chain \(\mathbf{c}(p)\) of each parameter in execution order.

def build_chains(units):
    params = sorted(set(sum([unit.fetch_input_params() + unit.fetch_output_params()
                             for unit in units], [])))
    chains = {param: [] for param in params}
    for id, unit in enumerate(units):
        for param in unit.fetch_output_params():
            chains[param].append(id)
    return chains

forward_reachable implements the forward reachability operator \(F\), iteratively computing the closure along the producer\(\to\)consumer direction until no new units are added.

def forward_reachable(edges, related):
    while True:
        neighbors = {target for source, target in edges
                     if source in related and target not in related}
        if not neighbors:
            break
        related |= neighbors
    return related

updating_units implements the updating-chain backtracking operator \(U\). For each input parameter \(p\) of every unit in the related set \(X\), it determines the unit \(u^{(i)}\) where \(p\) is first consumed within \(X\); if \(i<k\), all units after \(u^{(i)}\) in the updating chain \(\mathbf{c}(p)\) are merged into \(X\), thereby guaranteeing condition (C2).

def updating_units(units, chains, related):
    # Determine, in execution order (ascending unit id), the unit where p is first consumed in X
    first_consumer = {}
    for unit_id in sorted(related):
        for param in units[unit_id].fetch_input_params():
            if param not in first_consumer:
                first_consumer[param] = unit_id
    for param, first in first_consumer.items():
        chain = chains[param]
        if first in chain and chain.index(first) != len(chain) - 1:
            for later in chain[chain.index(first) + 1:]:
                related.add(later)
    return related

Two-Stage Splitting Workflow

After obtaining the bipartition \((V_1,V_2)\), the training framework configures the Pipeline according to the task type and delegates execution to launch_data_process_task and launch_training_task, respectively:

def split_pipeline_units(task, pipe, trainable_models, lora_base_model, ...):
    models_require_backward = []
    if trainable_models is not None:
        models_require_backward += trainable_models.split(",")
    if lora_base_model is not None:
        models_require_backward.append(lora_base_model)

    if task.endswith(":data_process"):           # Stage 1 (data preprocessing)
        other_units, pipe.units = pipe.split_pipeline_units(models_require_backward)
        # pipe.units = V2, executes only computations unrelated to the models (e.g., VAE encoding, text encoding)
        # Optional: append GeneralUnit_RemoveCache to drop redundant cached items and reduce cache size
    elif task.endswith(":train"):                # Stage 2 (training)
        pipe.units, _ = pipe.split_pipeline_units(models_require_backward)
        # pipe.units = V1, executes only computations related to the models

Stage 1 (data preprocessing) In torch.no_grad() mode, iterate over the dataset, execute the units in \(V_2\) for each sample, and serialize the intermediate results to disk sharded by process:

for data in dataloader:
    with torch.no_grad():
        cache = model(data)                      # executes only the units in V2
        torch.save(cache, cache_path)            # saves each sample as a .pth file

Stage 2 (training) The data loader takes the cache produced in the first stage as input, executes the units in \(V_1\), and performs forward propagation, backpropagation, and parameter updates:

for data in dataloader:                          # reads the first-stage cache
    with accelerator.accumulate(model):
        loss = model({}, inputs=data)            # executes only the units in V1
        accelerator.backward(loss)               # gradients flow only through trainable models
        optimizer.step()

Proposition 3 (Cache Reusability) By Definition 3, the units in \(V_2\) do not depend on any trainable model, and their computational results are independent of the model parameters \(\theta\). Therefore, the cache produced in the first stage can be directly reused in any subsequent epoch without redundant computation. Based on this property, --dataset_repeat should be set to 1 in the first stage, while the second stage can keep its original value.