Files

T

Bubbliiiing 403f1f7b78 Fix Bug in Self Forcing in Multi-Gpus Infernece && Update Training Code and Docs && Update Reward Models (#498 )

2026-07-14 10:30:54 +08:00

21 KiB

Raw Blame History

Z-Image Full Parameter Training Guide

This document provides a complete workflow for full parameter training of Z-Image Diffusion Transformer, including environment configuration, data preparation, distributed training, and inference testing.

Note

: Z-Image has two model variants: Z-Image (standard) and Z-Image-Turbo (fast inference). This guide uses Z-Image as the default. For Z-Image-Turbo, simply replace the model path accordingly.

1. Environment Configuration
2. Data Preparation
3. Full Parameter Training
4. Inference Testing
5. Additional Resources

1. Environment Configuration

Method 1: Using requirements.txt

pip install -r requirements.txt

Method 2: Manual Dependency Installation

pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
pip install yunchang xfuser modelscope openpyxl
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
pip install opencv-python-headless
pip install deepspeed==0.17.0 numpy==1.26.4

Method 3: Using Docker

When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands:

# pull image
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun

# enter image
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun

2. Data Preparation

2.1 Quick Test Dataset

We provide a test dataset containing several training samples.

# Download official demo dataset
modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun-Images-Demo

2.2 Dataset Structure

📦 datasets/
├── 📂 my_dataset/
│   ├── 📂 train/
│   │   ├── 📄 image001.jpg
│   │   ├── 📄 image002.png
│   │   └── 📄 ...
│   └── 📄 metadata.json

2.3 metadata.json Format

Relative Path Format (example):

[
  {
    "file_path": "train/image001.jpg",
    "text": "A beautiful sunset over the ocean, golden hour lighting",
    "width": 1024,
    "height": 1024
  },
  {
    "file_path": "train/image002.png",
    "text": "Portrait of a young woman, studio lighting, high quality",
    "width": 1328,
    "height": 1328
  }
]

Absolute Path Format:

[
  {
    "file_path": "/mnt/data/images/sunset.jpg",
    "text": "A beautiful sunset over the ocean",
    "width": 1024,
    "height": 1024
  }
]

Key Fields Description:

file_path: Image path (relative or absolute)
text: Image description (English prompt)
width / height: Image dimensions (recommended to provide for bucket training; if not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS)
- You can use scripts/process_json_add_width_and_height.py to add width and height fields to JSON files without these fields, supporting both images and videos
- Usage: python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Images-Demo/metadata.json --output_file datasets/X-Fun-Images-Demo/metadata_add_width_height.json

2.4 Relative vs Absolute Path Usage

Relative Paths:

If your data uses relative paths, configure the training script as follows:

export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"

Absolute Paths:

If your data uses absolute paths, configure the training script as follows:

export DATASET_NAME=""
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"

💡 Recommendation: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths.

3. Full Parameter Training

3.1 Download Pretrained Model

# Create model directory
mkdir -p models/Diffusion_Transformer

# Download Z-Image official weights
modelscope download --model PAI/Z-Image --local_dir models/Diffusion_Transformer/Z-Image

# (Optional) Download Z-Image-Turbo for fast inference
modelscope download --model PAI/Z-Image-Turbo --local_dir models/Diffusion_Transformer/Z-Image-Turbo

3.2 Quick Start (DeepSpeed-Zero-2)

If you have downloaded the data as per 2.1 Quick Test Dataset and the weights as per 3.1 Download Pretrained Model, you can directly copy and run the quick start command.

DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example.

The difference between DeepSpeed-Zero-2 and FSDP lies in whether the model weights are sharded. If VRAM is insufficient when using multiple GPUs with DeepSpeed-Zero-2, you can switch to FSDP.

export MODEL_NAME="models/Diffusion_Transformer/Z-Image"
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO

accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/z_image/train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --train_batch_size=1 \
  --image_sample_size=1328 \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-05 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_z_image" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
  --max_grad_norm=0.05 \
  --enable_bucket \
  --uniform_sampling \
  --trainable_modules "."

3.3 Common Training Parameters

Key Parameter Descriptions:

Parameter	Description	Example Value
`--pretrained_model_name_or_path`	Path to pretrained model	`models/Diffusion_Transformer/Z-Image`
`--train_data_dir`	Training data directory	`datasets/X-Fun-Images-Demo/`
`--train_data_meta`	Training data metadata file	`datasets/X-Fun-Images-Demo/metadata_add_width_height.json`
`--train_batch_size`	Samples per batch	1
`--image_sample_size`	Maximum training resolution, auto bucketing	1328
`--gradient_accumulation_steps`	Gradient accumulation steps (equivalent to larger batch)	1
`--dataloader_num_workers`	DataLoader subprocesses	8
`--num_train_epochs`	Number of training epochs	100
`--checkpointing_steps`	Save checkpoint every N steps	50
`--learning_rate`	Initial learning rate	2e-05
`--lr_scheduler`	Learning rate scheduler	`constant_with_warmup`
`--lr_warmup_steps`	Learning rate warmup steps	100
`--seed`	Random seed	42
`--output_dir`	Output directory	`output_dir_z_image`
`--gradient_checkpointing`	Enable activation checkpointing	-
`--mixed_precision`	Mixed precision: `fp16/bf16`	`bf16`
`--adam_weight_decay`	AdamW weight decay	3e-2
`--adam_epsilon`	AdamW epsilon value	1e-10
`--vae_mini_batch`	Mini-batch size for VAE encoding	1
`--max_grad_norm`	Gradient clipping threshold	0.05
`--enable_bucket`	Enable bucket training: trains entire images grouped by resolution without center cropping	-
`--random_hw_adapt`	Auto-scale images to random size in range `[512, image_sample_size]`	-
`--resume_from_checkpoint`	Resume training from checkpoint path, use `"latest"` to auto-select latest	None
`--uniform_sampling`	Uniform timestep sampling	-
`--trainable_modules`	Trainable modules (`"."` means all modules)	`"."`
`--validation_steps`	Execute validation every N steps	100
`--validation_epochs`	Execute validation every N epochs	100
`--validation_prompts`	Prompts used during validation	`"a young girl..."`

3.4 Training Validation

You can configure validation parameters to periodically generate test images during training, allowing you to monitor training progress and model quality.

Validation Parameters:

Parameter	Description	Recommended Value
`--validation_steps`	Execute validation every N steps	100
`--validation_epochs`	Execute validation every N epochs	100
`--validation_prompts`	Prompt for validation image generation. Use multiple space-separated prompt strings	Space-separated prompt strings

Example:

  --validation_steps=100 \
  --validation_epochs=100 \
  --validation_prompts="a young girl with flowing long hair, wearing a white halter dress"

Notes:

Validation images will be saved to the output_dir directory
For multi-prompt validation, use: --validation_prompts "prompt1" "prompt2" "prompt3"

3.5 Training with FSDP

If VRAM is insufficient when using multiple GPUs with DeepSpeed-Zero-2, you can switch to FSDP.

export MODEL_NAME="models/Diffusion_Transformer/Z-Image"
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO

accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap ZImageTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/z_image/train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --train_batch_size=1 \
  --image_sample_size=1328 \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-05 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_z_image" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
  --max_grad_norm=0.05 \
  --enable_bucket \
  --uniform_sampling \
  --trainable_modules "."

3.6 Training Without DeepSpeed or FSDP

This approach is not recommended as it lacks VRAM-saving backends and may easily cause out-of-memory errors. This is provided for reference only.

export MODEL_NAME="models/Diffusion_Transformer/Z-Image"
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO

accelerate launch --mixed_precision="bf16" scripts/z_image/train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --train_batch_size=1 \
  --image_sample_size=1328 \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-05 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_z_image" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
  --max_grad_norm=0.05 \
  --enable_bucket \
  --uniform_sampling \
  --trainable_modules "."

3.7 Multi-Machine Distributed Training

Suitable for: Ultra-large-scale datasets, faster training speed

3.7.1 Environment Configuration

Assuming 2 machines with 8 GPUs each:

Machine 0 (Master):

export MODEL_NAME="models/Diffusion_Transformer/Z-Image"
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100"  # Master machine IP
export MASTER_PORT=10086
export WORLD_SIZE=2                  # Total number of machines
export NUM_PROCESS=16                # Total processes = machines × 8
export RANK=0                        # Current machine rank (0 or 1)
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO

accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/z_image/train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --train_batch_size=1 \
  --image_sample_size=1328 \
  --gradient_accumulation_steps=1 \
  --dataloader_num_workers=8 \
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-05 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_z_image" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
  --max_grad_norm=0.05 \
  --enable_bucket \
  --uniform_sampling \
  --trainable_modules "."

Machine 1 (Worker):

export MODEL_NAME="models/Diffusion_Transformer/Z-Image"
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
export MASTER_ADDR="192.168.1.100"  # Same as Master
export MASTER_PORT=10086
export WORLD_SIZE=2
export NUM_PROCESS=16
export RANK=1  # Note this is 1
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO

# Use the same accelerate launch command as Machine 0

3.7.2 Multi-Machine Training Notes

Network Requirements:
- RDMA/InfiniBand recommended (high performance)
- Without RDMA, add environment variables:
```
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1
```
Data Synchronization: All machines must be able to access the same data paths (NFS/shared storage)

4. Inference Testing

4.1 Inference Parameters

Key Parameter Descriptions:

Parameter	Description	Example Value
`GPU_memory_mode`	GPU memory mode, see table below for options	`model_cpu_offload`
`ulysses_degree`	Head dimension parallelization degree, 1 for single GPU	1
`ring_degree`	Sequence dimension parallelization degree, 1 for single GPU	1
`fsdp_dit`	Use FSDP for Transformer in multi-GPU inference to save VRAM	`False`
`fsdp_text_encoder`	Use FSDP for text encoder in multi-GPU inference	`False`
`compile_dit`	Compile Transformer to accelerate inference (effective at fixed resolution)	`False`
`model_name`	Model path	`models/Diffusion_Transformer/Z-Image`
`sampler_name`	Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++`	`Flow`
`transformer_path`	Path to trained Transformer weights	`None`
`vae_path`	Path to trained VAE weights	`None`
`lora_path`	LoRA weights path	`None`
`sample_size`	Generated image resolution `[height, width]`	`[1728, 992]`
`weight_dtype`	Model weight precision, use `torch.float16` for GPUs without bf16 support	`torch.bfloat16`
`prompt`	Positive prompt describing the content to generate	`"a young girl..."`
`negative_prompt`	Negative prompt for content to avoid	`"low resolution, low quality..."`
`guidance_scale`	Guidance strength, recommended 0.0 for Turbo model	4.0 / 0.0
`seed`	Random seed for reproducibility	43
`num_inference_steps`	Inference steps, can be greatly reduced for Turbo model	25 / 9
`lora_weight`	LoRA weight strength	0.55
`save_path`	Generated image save path	`samples/z-image-t2i`

GPU Memory Mode Description:

Mode	Description	VRAM Usage
`model_full_load`	Load entire model to GPU	Highest
`model_full_load_and_qfloat8`	Full load + FP8 quantization	High
`model_cpu_offload`	Offload model to CPU after use	Medium
`model_cpu_offload_and_qfloat8`	CPU offload + FP8 quantization	Medium-Low
`model_group_offload`	Layer group offload between CPU/CUDA	Low
`sequential_cpu_offload`	Offload each layer individually (slowest)	Lowest

4.2 Single GPU Inference

Z-Image (Standard)

Run single GPU inference with:

python examples/z_image/predict_t2i.py

Edit examples/z_image/predict_t2i.py according to your needs. For first-time inference, focus on these parameters. For other parameters, see the Inference Parameters section above.

# Choose based on your GPU VRAM
GPU_memory_mode = "model_cpu_offload"
# Your actual model path
model_name = "models/Diffusion_Transformer/Z-Image"  
# Trained weights path, e.g. "output_dir_z_image/checkpoint-xxx/diffusion_pytorch_model.safetensors"
transformer_path = None  
# Write based on content to generate
prompt = "a young girl with flowing long hair, wearing a white halter dress"  
# ...

Z-Image-Turbo (Fast)