Ode training && Update Lens model && Update LTX2 upsampler (#497)
This commit is contained in:
Bubbliiiing
@@ -0,0 +1,525 @@
|
||||
# Lens Full Parameter Training Guide
|
||||
|
||||
This document provides a complete workflow for full parameter training of Lens Diffusion Transformer, including environment configuration, data preparation, distributed training, and inference testing.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
- [1. Environment Configuration](#1-environment-configuration)
|
||||
- [2. Data Preparation](#2-data-preparation)
|
||||
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
|
||||
- [2.2 Dataset Structure](#22-dataset-structure)
|
||||
- [2.3 metadata.json Format](#23-metadatajson-format)
|
||||
- [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
|
||||
- [3. Full Parameter Training](#3-full-parameter-training)
|
||||
- [3.1 Download Pretrained Model](#31-download-pretrained-model)
|
||||
- [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
|
||||
- [3.3 Common Training Parameters](#33-common-training-parameters)
|
||||
- [3.4 Training Validation](#34-training-validation)
|
||||
- [3.5 Training with FSDP](#35-training-with-fsdp)
|
||||
- [3.6 Other Backends](#36-other-backends)
|
||||
- [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
|
||||
- [4. Inference Testing](#4-inference-testing)
|
||||
- [4.1 Inference Parameters](#41-inference-parameters)
|
||||
- [4.2 Single GPU Inference](#42-single-gpu-inference)
|
||||
- [4.3 Multi-GPU Parallel Inference](#43-multi-gpu-parallel-inference)
|
||||
- [5. Additional Resources](#5-additional-resources)
|
||||
|
||||
---
|
||||
|
||||
## 1. Environment Configuration
|
||||
|
||||
**Method 1: Using requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Method 2: Manual Dependency Installation**
|
||||
|
||||
```bash
|
||||
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 deepspeed==0.17.0 numpy==1.26.4
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
```
|
||||
|
||||
**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.
|
||||
|
||||
```bash
|
||||
# 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):
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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**:
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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:
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
|
||||
```
|
||||
|
||||
**Absolute Paths**:
|
||||
|
||||
If your data uses absolute paths, configure the training script as follows:
|
||||
|
||||
```bash
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata.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
|
||||
|
||||
```bash
|
||||
# Create model directory
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# Download Lens official weights
|
||||
modelscope download --model microsoft/Lens --local_dir models/Diffusion_Transformer/Lens
|
||||
```
|
||||
|
||||
### 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.
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/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_lens" \
|
||||
--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/Lens` |
|
||||
| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` |
|
||||
| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.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_lens` |
|
||||
| `--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**:
|
||||
|
||||
```bash
|
||||
--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.
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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 LensTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/lens/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_lens" \
|
||||
--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.
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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/lens/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_lens" \
|
||||
--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)**:
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/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_lens" \
|
||||
--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)**:
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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:
|
||||
```bash
|
||||
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/Lens` |
|
||||
| `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 | `"1girl, black_hair..."` |
|
||||
| `negative_prompt` | Negative prompt for content to avoid | `"低分辨率，低画质..."` |
|
||||
| `guidance_scale` | Guidance strength | 4.5 |
|
||||
| `seed` | Random seed for reproducibility | 43 |
|
||||
| `num_inference_steps` | Inference steps | 40 |
|
||||
| `lora_weight` | LoRA weight strength | 0.55 |
|
||||
| `save_path` | Generated image save path | `samples/lens-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
|
||||
|
||||
Run single GPU inference with:
|
||||
|
||||
```bash
|
||||
python examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
Edit `examples/ernie_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.
|
||||
|
||||
```python
|
||||
# Choose based on your GPU VRAM
|
||||
GPU_memory_mode = "model_cpu_offload"
|
||||
# Your actual model path
|
||||
model_name = "models/Diffusion_Transformer/Lens"
|
||||
# Trained weights path, e.g. "output_dir_lens/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
transformer_path = None
|
||||
# Write based on content to generate
|
||||
prompt = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
# ...
|
||||
```
|
||||
|
||||
### 4.3 Multi-GPU Parallel Inference
|
||||
|
||||
**Suitable for**: High-resolution generation, accelerated inference
|
||||
|
||||
#### Install Parallel Inference Dependencies
|
||||
|
||||
```bash
|
||||
pip install xfuser==0.4.2 yunchang==0.6.2
|
||||
```
|
||||
|
||||
#### Configure Parallel Strategy
|
||||
|
||||
Edit `examples/ernie_image/predict_t2i.py`:
|
||||
|
||||
```python
|
||||
# Ensure ulysses_degree × ring_degree = number of GPUs
|
||||
# For example, using 2 GPUs:
|
||||
ulysses_degree = 2 # Head dimension parallelization
|
||||
ring_degree = 1 # Sequence dimension parallelization
|
||||
```
|
||||
|
||||
**Configuration Principles**:
|
||||
- `ulysses_degree` must evenly divide the model's number of heads
|
||||
- `ring_degree` splits on sequence dimension, affecting communication overhead; avoid using it when heads can be divided
|
||||
|
||||
**Example Configurations**:
|
||||
|
||||
| GPU Count | ulysses_degree | ring_degree | Description |
|
||||
|---------|---------------|-------------|------|
|
||||
| 1 | 1 | 1 | Single GPU |
|
||||
| 4 | 4 | 1 | Head parallelization |
|
||||
| 8 | 8 | 1 | Head parallelization |
|
||||
| 8 | 4 | 2 | Hybrid parallelization |
|
||||
|
||||
#### Run Multi-GPU Inference
|
||||
|
||||
```bash
|
||||
torchrun --nproc-per-node=2 examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
## 5. Additional Resources
|
||||
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
Executable
+536
@@ -0,0 +1,536 @@
|
||||
# Lens LoRA Fine-Tuning Training Guide
|
||||
|
||||
This document provides a complete workflow for Lens LoRA fine-tuning training, including environment configuration, data preparation, multiple distributed training strategies, and inference testing.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
- [1. Environment Configuration](#1-environment-configuration)
|
||||
- [2. Data Preparation](#2-data-preparation)
|
||||
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
|
||||
- [2.2 Dataset Structure](#22-dataset-structure)
|
||||
- [2.3 metadata.json Format](#23-metadatajson-format)
|
||||
- [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
|
||||
- [3. LoRA Training](#3-lora-training)
|
||||
- [3.1 Download Pretrained Model](#31-download-pretrained-model)
|
||||
- [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
|
||||
- [3.3 LoRA-Specific Parameters](#33-lora-specific-parameters)
|
||||
- [3.4 Training Validation](#34-training-validation)
|
||||
- [3.5 Training with FSDP](#35-training-with-fsdp)
|
||||
- [3.6 Training Without DeepSpeed or FSDP](#36-training-without-deepspeed-or-fsdp)
|
||||
- [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
|
||||
- [4. Inference Testing](#4-inference-testing)
|
||||
- [4.1 Inference Parameter Parsing](#41-inference-parameter-parsing)
|
||||
- [4.2 Single GPU Inference](#42-single-gpu-inference)
|
||||
- [4.3 Multi-GPU Parallel Inference](#43-multi-gpu-parallel-inference)
|
||||
- [5. Additional Resources](#5-additional-resources)
|
||||
|
||||
---
|
||||
|
||||
## 1. Environment Configuration
|
||||
|
||||
**Method 1: Using requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**Method 2: Manual Dependency Installation**
|
||||
|
||||
```bash
|
||||
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 deepspeed==0.17.0 numpy==1.26.4
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
```
|
||||
|
||||
**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.
|
||||
|
||||
```bash
|
||||
# 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):
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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**:
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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:
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
|
||||
```
|
||||
|
||||
**Absolute Paths**:
|
||||
|
||||
If your data uses absolute paths, configure the training script as follows:
|
||||
|
||||
```bash
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata.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. LoRA Training
|
||||
|
||||
### 3.1 Download Pretrained Model
|
||||
|
||||
```bash
|
||||
# Create model directory
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# Download Lens official weights
|
||||
modelscope download --model microsoft/Lens --local_dir models/Diffusion_Transformer/Lens
|
||||
```
|
||||
|
||||
### 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.
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 3.3 LoRA-Specific Parameters
|
||||
|
||||
**LoRA Key Parameter Descriptions**:
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|-----|------|-------|
|
||||
| `--pretrained_model_name_or_path` | Path to pretrained model | `models/Diffusion_Transformer/Lens` |
|
||||
| `--train_data_dir` | Training data directory | `datasets/internal_datasets/` |
|
||||
| `--train_data_meta` | Training data metadata file | `datasets/internal_datasets/metadata.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 | 100 |
|
||||
| `--learning_rate` | Initial learning rate (recommended for LoRA) | 1e-04 |
|
||||
| `--lr_warmup_steps` | Learning rate warmup steps | 100 |
|
||||
| `--seed` | Random seed (for reproducible training) | 42 |
|
||||
| `--output_dir` | Output directory | `output_dir_lens_lora` |
|
||||
| `--gradient_checkpointing` | Enable activation checkpointing | - |
|
||||
| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
|
||||
| `--enable_bucket` | Enable bucket training: trains entire images grouped by resolution without center cropping | - |
|
||||
| `--uniform_sampling` | Uniform timestep sampling (recommended) | - |
|
||||
| `--resume_from_checkpoint` | Resume training from checkpoint path, use `"latest"` to auto-select latest | None |
|
||||
| `--rank` | Dimension of LoRA update matrices (higher rank = stronger expressiveness but more VRAM usage) | 64 |
|
||||
| `--network_alpha` | Scaling factor of LoRA update matrices (typically set to half of rank) | 32 |
|
||||
| `--target_name` | Components/modules to apply LoRA, separated by commas | `img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp` |
|
||||
| `--low_vram` | Low VRAM mode, offloads text encoder and VAE to CPU | - |
|
||||
| `--validation_steps` | Execute validation every N steps | 100 |
|
||||
| `--validation_epochs` | Execute validation every N epochs | 100 |
|
||||
| `--validation_prompts` | Prompts used during validation | `"1girl, black_hair, ..."` |
|
||||
|
||||
### 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**:
|
||||
|
||||
```bash
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/lens/train_lora.py \
|
||||
# ... (other training parameters)
|
||||
--validation_steps=100 \
|
||||
--validation_epochs=100 \
|
||||
--validation_prompts="1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
```
|
||||
|
||||
**Parameter Descriptions**:
|
||||
|
||||
| Parameter | Description | Recommended Value |
|
||||
|-----------|-------------|-------------------|
|
||||
| `--validation_steps` | Execute validation every N steps. If your dataset is large and you want to save validation time, you can set a larger value (e.g., 100 or 500) | 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 |
|
||||
|
||||
**Notes**:
|
||||
- Validation images will be saved to the `output_dir` directory
|
||||
- Setting `--validation_steps=1` means validation is performed every step, which may slow down training. Adjust according to your needs
|
||||
- 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.
|
||||
|
||||
> ✅ **Recommended**: FSDP has been thoroughly tested in this repository, with fewer errors and greater stability.
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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=LensTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 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.
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 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)**:
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
**Machine 1 (Worker)**:
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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:
|
||||
```bash
|
||||
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 Parameter Parsing
|
||||
|
||||
**Key Parameter Descriptions**:
|
||||
|
||||
| Parameter | Description | Example Value |
|
||||
|------|------|-------|
|
||||
| `GPU_memory_mode` | VRAM management mode, see table below for options | `model_cpu_offload` |
|
||||
| `ulysses_degree` | Head dimension parallelism degree, set to 1 for single GPU | 1 |
|
||||
| `ring_degree` | Sequence dimension parallelism degree, set to 1 for single GPU | 1 |
|
||||
| `fsdp_dit` | Use FSDP for Transformer during multi-GPU inference to save VRAM | `False` |
|
||||
| `fsdp_text_encoder` | Use FSDP for text encoder during multi-GPU inference | `False` |
|
||||
| `compile_dit` | Compile Transformer for faster inference (effective at fixed resolution) | `False` |
|
||||
| `model_name` | Model path | `models/Diffusion_Transformer/Lens` |
|
||||
| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow` |
|
||||
| `transformer_path` | Path to load trained Transformer weights | `None` |
|
||||
| `vae_path` | Path to load 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 generation content | `"1girl, black_hair..."` |
|
||||
| `negative_prompt` | Negative prompt for content to avoid | `" "` |
|
||||
| `guidance_scale` | Guidance strength | 4.5 |
|
||||
| `seed` | Random seed for reproducible results | 43 |
|
||||
| `num_inference_steps` | Number of inference steps | 40 |
|
||||
| `lora_weight` | LoRA weight strength | 0.55 |
|
||||
| `save_path` | Path to save generated images | `samples/lens-t2i` |
|
||||
|
||||
**VRAM Management 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 groups switch between CPU/CUDA | Low |
|
||||
| `sequential_cpu_offload` | Sequential layer offload (slowest) | Lowest |
|
||||
|
||||
### 4.2 Single GPU Inference
|
||||
|
||||
#### Quick Start
|
||||
|
||||
Run the following command for single GPU inference:
|
||||
|
||||
```bash
|
||||
python examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
Edit `examples/lens/predict_t2i.py` according to your needs. For first-time inference, focus on these parameters. For other parameters, refer to the inference parameter parsing above.
|
||||
|
||||
```python
|
||||
# Choose based on GPU VRAM
|
||||
GPU_memory_mode = "model_cpu_offload"
|
||||
# Based on actual model path
|
||||
model_name = "models/Diffusion_Transformer/Lens"
|
||||
# LoRA weights path, e.g., "output_dir_lens_lora/checkpoint-xxx/lora_weights.safetensors"
|
||||
lora_path = None
|
||||
# LoRA weight strength
|
||||
lora_weight = 0.55
|
||||
# Write based on generation content
|
||||
prompt = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
# ...
|
||||
```
|
||||
|
||||
### 4.3 Multi-GPU Parallel Inference
|
||||
|
||||
**Suitable for**: High-resolution generation, accelerated inference
|
||||
|
||||
#### Install Parallel Inference Dependencies
|
||||
|
||||
```bash
|
||||
pip install xfuser==0.4.2 yunchang==0.6.2
|
||||
```
|
||||
|
||||
#### Configure Parallel Strategy
|
||||
|
||||
Edit `examples/lens/predict_t2i.py`:
|
||||
|
||||
```python
|
||||
# Ensure ulysses_degree × ring_degree = number of GPUs
|
||||
# For example, using 2 GPUs:
|
||||
ulysses_degree = 2 # Head dimension parallelization
|
||||
ring_degree = 1 # Sequence dimension parallelization
|
||||
```
|
||||
|
||||
**Configuration Principles**:
|
||||
- `ulysses_degree` must evenly divide the model's number of heads
|
||||
- `ring_degree` splits on sequence dimension, affecting communication overhead; avoid using it when heads can be divided
|
||||
|
||||
**Example Configurations**:
|
||||
|
||||
| GPU Count | ulysses_degree | ring_degree | Description |
|
||||
|---------|---------------|-------------|------|
|
||||
| 1 | 1 | 1 | Single GPU |
|
||||
| 4 | 4 | 1 | Head parallelization |
|
||||
| 8 | 8 | 1 | Head parallelization |
|
||||
| 8 | 4 | 2 | Hybrid parallelization |
|
||||
|
||||
#### Run Multi-GPU Inference
|
||||
|
||||
```bash
|
||||
torchrun --nproc-per-node=2 examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
## 5. Additional Resources
|
||||
|
||||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|
||||
@@ -0,0 +1,537 @@
|
||||
# Lens LoRA 微调训练指南
|
||||
|
||||
本文档提供 Lens LoRA 微调训练的完整流程，包括环境配置、数据准备、多种分布式训练策略和推理测试。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [2.1 快速测试数据集](#21-快速测试数据集)
|
||||
- [2.2 数据集结构](#22-数据集结构)
|
||||
- [2.3 metadata.json 格式](#23-metadatajson-格式)
|
||||
- [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
|
||||
- [三、LoRA 训练](#三lora-训练)
|
||||
- [3.1 下载预训练模型](#31-下载预训练模型)
|
||||
- [3.2 快速开始（DeepSpeed-Zero-2）](#32-快速开始deepspeed-zero-2)
|
||||
- [3.3 LoRA 专用参数解析](#33-lora-专用参数解析)
|
||||
- [3.4 训练验证](#34-训练验证)
|
||||
- [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
|
||||
- [3.6 不使用 DeepSpeed 与 FSDP 训练](#36-不使用-deepspeed-与-fsdp-训练)
|
||||
- [3.7 多机分布式训练](#37-多机分布式训练)
|
||||
- [四、推理测试](#四推理测试)
|
||||
- [4.1 推理参数解析](#41-推理参数解析)
|
||||
- [4.2 单卡推理](#42-单卡推理)
|
||||
- [4.3 多卡并行推理](#43-多卡并行推理)
|
||||
- [五、更多资源](#五更多资源)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式 1：使用requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式 2：手动安装依赖**
|
||||
|
||||
```bash
|
||||
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 deepspeed==0.17.0 numpy==1.26.4
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
```
|
||||
|
||||
**方式 3：使用docker**
|
||||
|
||||
使用docker的情况下，请保证机器中已经正确安装显卡驱动与CUDA环境，然后以此执行以下命令：
|
||||
|
||||
```
|
||||
# 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.1 快速测试数据集
|
||||
|
||||
我们提供了一个测试的数据集，其中包含若干训练数据。
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun-Images-Demo
|
||||
```
|
||||
|
||||
### 2.2 数据集结构
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 image001.jpg
|
||||
│ │ ├── 📄 image002.png
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
### 2.3 metadata.json 格式
|
||||
|
||||
**相对路径格式**（示例格式）：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**绝对路径格式**：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "/mnt/data/images/sunset.jpg",
|
||||
"text": "A beautiful sunset over the ocean",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**关键字段说明**：
|
||||
- `file_path`：图片路径（相对或绝对路径）
|
||||
- `text`：图片描述（英文提示词）
|
||||
- `width` / `height`：图片宽高（**最好提供**，用于分桶训练，如果不提供则自动在训练时读取，当数据存储在如oss这样的速度较慢的系统上时，可能会影响训练速度）。
|
||||
- 可以使用`scripts/process_json_add_width_and_height.py`文件对无width与height字段的json进行提取，支持处理图片与视频。
|
||||
- 使用方案为`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 相对路径与绝对路径使用方案
|
||||
|
||||
**相对路径**：
|
||||
|
||||
如果数据的路径为相对路径，则在训练脚本中设置：
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
|
||||
```
|
||||
|
||||
**绝对路径**：
|
||||
|
||||
如果数据的路径为绝对路径，则在训练脚本中设置：
|
||||
|
||||
```bash
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata.json"
|
||||
```
|
||||
|
||||
> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。
|
||||
|
||||
---
|
||||
|
||||
## 三、LoRA 训练
|
||||
|
||||
### 3.1 下载预训练模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Lens 官方权重
|
||||
modelscope download --model microsoft/Lens --local_dir models/Diffusion_Transformer/Lens
|
||||
```
|
||||
|
||||
### 3.2 快速开始（DeepSpeed-Zero-2）
|
||||
|
||||
如果按照 **2.1 快速测试数据集下载数据** 与 **3.1 下载预训练模型下载权重**后，直接复制快速开始的启动指令进行启动。
|
||||
|
||||
推荐使用 DeepSpeed-Zero-2 与 FSDP 方案进行训练。这里使用 DeepSpeed-Zero-2 为例配置 shell 文件。
|
||||
|
||||
本文中 DeepSpeed-Zero-2 与 FSDP 的差别在于是否对模型权重进行分片，**如果使用多卡且使用 DeepSpeed-Zero-2 的情况下显存不足**，可以切换使用 FSDP 进行训练。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 3.3 LoRA 专用参数解析
|
||||
|
||||
**LoRA 关键参数说明**：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----|------|-------|
|
||||
| `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Lens` |
|
||||
| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` |
|
||||
| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` |
|
||||
| `--train_batch_size` | 每批次样本数 | 1 |
|
||||
| `--image_sample_size` | 最大训练分辨率，代码会自动分桶 | 1328 |
|
||||
| `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 |
|
||||
| `--dataloader_num_workers` | DataLoader 子进程数 | 8 |
|
||||
| `--num_train_epochs` | 训练 epoch 数 | 100 |
|
||||
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
|
||||
| `--learning_rate` | 初始学习率（LoRA 推荐值） | 1e-04 |
|
||||
| `--lr_warmup_steps` | 学习率预热步数 | 100 |
|
||||
| `--seed` | 随机种子（可复现训练） | 42 |
|
||||
| `--output_dir` | 输出目录 | `output_dir_lens_lora` |
|
||||
| `--gradient_checkpointing` | 激活重计算 | - |
|
||||
| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
|
||||
| `--enable_bucket` | 启用分桶训练，不裁剪图片，按分辨率分组训练整个图像 | - |
|
||||
| `--uniform_sampling` | 均匀采样 timestep（推荐启用） | - |
|
||||
| `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | None |
|
||||
| `--rank` | LoRA 更新矩阵的维度（rank 越大表达能力越强，但显存占用越高） | 64 |
|
||||
| `--network_alpha` | LoRA 更新矩阵的缩放系数（通常设置为 rank 的一半） | 32 |
|
||||
| `--target_name` | 应用 LoRA 的组件/模块，用逗号分隔 | `img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp` |
|
||||
| `--low_vram` | 低显存模式，对文本编码器和 VAE 进行 CPU offload | - |
|
||||
| `--validation_steps` | 每 N 步执行一次验证 | 100 |
|
||||
| `--validation_epochs` | 每 N 个epoch执行一次验证 | 100 |
|
||||
| `--validation_prompts` | 验证时使用的提示词 | `"1girl, black_hair, ..."` |
|
||||
|
||||
|
||||
### 3.4 训练验证
|
||||
|
||||
你可以配置验证参数，在训练过程中定期生成测试图像，以便监控训练进度和模型质量。
|
||||
|
||||
**验证参数配置**：
|
||||
|
||||
```bash
|
||||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/lens/train_lora.py \
|
||||
# ... (其他训练参数)
|
||||
--validation_steps=100 \
|
||||
--validation_epochs=100 \
|
||||
--validation_prompts="1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
```
|
||||
|
||||
**参数说明**：
|
||||
|
||||
| 参数 | 说明 | 推荐值 |
|
||||
|------|------|--------|
|
||||
| `--validation_steps` | 每 N 步执行一次验证。如果数据集较大，想节省验证时间，可以设置更大的值（如100或500） | 100 |
|
||||
| `--validation_epochs` | 每 N 个epoch执行一次验证 | 100 |
|
||||
| `--validation_prompts` | 验证图像生成的提示词。可以设置多个提示词，用空格分隔 | 多个空格分隔的提示词 |
|
||||
|
||||
**注意事项**：
|
||||
- 验证图像会保存到 `output_dir` 目录中
|
||||
- 设置 `--validation_steps=1` 表示每一步都进行验证，可能会拖慢训练速度，可根据实际需求调整
|
||||
- 多提示词验证格式：`--validation_prompts "prompt1" "prompt2" "prompt3"`
|
||||
|
||||
### 3.5 使用 FSDP 训练
|
||||
|
||||
**如果使用多卡且使用 DeepSpeed-Zero-2 的情况下显存不足**，可以切换使用 FSDP 进行训练。
|
||||
|
||||
> ✅ **推荐**：FSDP 在当前仓库中经过充分测试，错误更少、更稳定。
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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=LensTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 3.6 不使用 DeepSpeed 与 FSDP 训练
|
||||
|
||||
**该方案并不被推荐，因为没有显存节约后端，容易造成显存不足**。这里仅提供训练 Shell 用于参考训练。
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
### 3.7 多机分布式训练
|
||||
|
||||
**适合场景**：超大规模数据集、需要更快的训练速度
|
||||
|
||||
#### 3.7.1 环境配置
|
||||
|
||||
假设有 2 台机器，每台 8 张 GPU：
|
||||
|
||||
**机器 0（Master）**：
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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 机器 IP
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2 # 机器总数
|
||||
export NUM_PROCESS=16 # 总进程数 = 机器数 × 8
|
||||
export RANK=0 # 当前机器 rank（0 或 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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp" \
|
||||
--low_vram \
|
||||
--uniform_sampling
|
||||
```
|
||||
|
||||
**机器 1（Worker）**：
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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 相同
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2
|
||||
export NUM_PROCESS=16
|
||||
export RANK=1 # 注意这里是 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
|
||||
|
||||
# 使用与机器 0 相同的 accelerate launch 命令
|
||||
```
|
||||
|
||||
#### 3.7.2 多机训练注意事项
|
||||
|
||||
- **网络要求**：
|
||||
- 推荐 RDMA/InfiniBand（高性能）
|
||||
- 无 RDMA 时添加环境变量：
|
||||
```bash
|
||||
export NCCL_IB_DISABLE=1
|
||||
export NCCL_P2P_DISABLE=1
|
||||
```
|
||||
|
||||
- **数据同步**：所有机器必须能够访问相同的数据路径（NFS/共享存储）
|
||||
|
||||
---
|
||||
|
||||
## 四、推理测试
|
||||
|
||||
### 4.1 推理参数解析
|
||||
|
||||
**关键参数说明**：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|------|------|-------|
|
||||
| `GPU_memory_mode` | 显存管理模式，可选值见下表 | `model_cpu_offload` |
|
||||
| `ulysses_degree` | Head 维度并行度，单卡时为 1 | 1 |
|
||||
| `ring_degree` | Sequence 维度并行度，单卡时为 1 | 1 |
|
||||
| `fsdp_dit` | 多卡推理时对 Transformer 使用 FSDP 节省显存 | `False` |
|
||||
| `fsdp_text_encoder` | 多卡推理时对文本编码器使用 FSDP | `False` |
|
||||
| `compile_dit` | 编译 Transformer 加速推理（固定分辨率下有效） | `False` |
|
||||
| `model_name` | 模型路径 | `models/Diffusion_Transformer/Lens` |
|
||||
| `sampler_name` | 采样器类型：`Flow`、`Flow_Unipc`、`Flow_DPM++` | `Flow` |
|
||||
| `transformer_path` | 加载训练好的 Transformer 权重路径 | `None` |
|
||||
| `vae_path` | 加载训练好的 VAE 权重路径 | `None` |
|
||||
| `lora_path` | LoRA 权重路径 | `None` |
|
||||
| `sample_size` | 生成图像分辨率 `[高度, 宽度]` | `[1728, 992]` |
|
||||
| `weight_dtype` | 模型权重精度，不支持 bf16 的显卡使用 `torch.float16` | `torch.bfloat16` |
|
||||
| `prompt` | 正向提示词，描述生成内容 | `"1girl, black_hair..."` |
|
||||
| `negative_prompt` | 负向提示词，避免生成的内容 | `" "` |
|
||||
| `guidance_scale` | 引导强度 | 4.5 |
|
||||
| `seed` | 随机种子，用于复现结果 | 43 |
|
||||
| `num_inference_steps` | 推理步数 | 40 |
|
||||
| `lora_weight` | LoRA 权重强度 | 0.55 |
|
||||
| `save_path` | 生成图像保存路径 | `samples/lens-t2i` |
|
||||
|
||||
**显存管理模式说明**：
|
||||
|
||||
| 模式 | 说明 | 显存占用 |
|
||||
|------|------|---------|
|
||||
| `model_full_load` | 整个模型加载到 GPU | 最高 |
|
||||
| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
|
||||
| `model_cpu_offload` | 使用后将模型卸载到 CPU | 中等 |
|
||||
| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
|
||||
| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
|
||||
| `sequential_cpu_offload` | 逐层卸载（速度最慢） | 最低 |
|
||||
|
||||
### 4.2 单卡推理
|
||||
|
||||
#### 快速开始
|
||||
|
||||
单卡推理运行如下命令：
|
||||
|
||||
```bash
|
||||
python examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
根据需求修改编辑 `examples/lens/predict_t2i.py`，初次推理重点关注如下参数，如果对其他参数感兴趣，请查看上方的推理参数解析。
|
||||
|
||||
```python
|
||||
# 根据显卡显存选择
|
||||
GPU_memory_mode = "model_cpu_offload"
|
||||
# 根据实际模型路径
|
||||
model_name = "models/Diffusion_Transformer/Lens"
|
||||
# LoRA 权重路径，如 "output_dir_lens_lora/checkpoint-xxx/lora_weights.safetensors"
|
||||
lora_path = None
|
||||
# LoRA 权重强度
|
||||
lora_weight = 0.55
|
||||
# 根据生成内容编写
|
||||
prompt = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
# ...
|
||||
```
|
||||
|
||||
### 4.3 多卡并行推理
|
||||
|
||||
**适合场景**：高分辨率生成、加速推理
|
||||
|
||||
#### 安装并行推理依赖
|
||||
|
||||
```bash
|
||||
pip install xfuser==0.4.2 yunchang==0.6.2
|
||||
```
|
||||
|
||||
#### 配置并行策略
|
||||
|
||||
编辑 `examples/lens/predict_t2i.py`：
|
||||
|
||||
```python
|
||||
# 确保 ulysses_degree × ring_degree = GPU 数量
|
||||
# 例如使用 2 张 GPU：
|
||||
ulysses_degree = 2 # Head 维度并行
|
||||
ring_degree = 1 # Sequence 维度并行
|
||||
```
|
||||
|
||||
**配置原则**：
|
||||
- `ulysses_degree` 必须能整除模型的head数。
|
||||
- `ring_degree` 会在sequence上切分，影响通信开销，在head数能切分的时候尽量不用。
|
||||
|
||||
**示例配置**：
|
||||
|
||||
| GPU 数量 | ulysses_degree | ring_degree | 说明 |
|
||||
|---------|---------------|-------------|------|
|
||||
| 1 | 1 | 1 | 单卡 |
|
||||
| 4 | 4 | 1 | Head 并行 |
|
||||
| 8 | 8 | 1 | Head 并行 |
|
||||
| 8 | 4 | 2 | 混合并行 |
|
||||
|
||||
#### 运行多卡推理
|
||||
|
||||
```bash
|
||||
torchrun --nproc-per-node=2 examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
## 五、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
@@ -0,0 +1,525 @@
|
||||
# Lens 全量参数训练指南
|
||||
|
||||
本文档提供 Lens Diffusion Transformer 全量参数训练的完整流程,包括环境配置、数据准备、分布式训练和推理测试。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
- [一、环境配置](#一环境配置)
|
||||
- [二、数据准备](#二数据准备)
|
||||
- [2.1 快速测试数据集](#21-快速测试数据集)
|
||||
- [2.2 数据集结构](#22-数据集结构)
|
||||
- [2.3 metadata.json 格式](#23-metadatajson-格式)
|
||||
- [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
|
||||
- [三、全量参数训练](#三全量参数训练)
|
||||
- [3.1 下载预训练模型](#31-下载预训练模型)
|
||||
- [3.2 快速开始（DeepSpeed-Zero-2）](#32-快速开始deepspeed-zero-2)
|
||||
- [3.3 训练常用参数解析](#33-训练常用参数解析)
|
||||
- [3.4 训练验证](#34-训练验证)
|
||||
- [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
|
||||
- [3.6 其他后端](#36-其他后端)
|
||||
- [3.7 多机分布式训练](#37-多机分布式训练)
|
||||
- [四、推理测试](#四推理测试)
|
||||
- [4.1 推理参数解析](#41-推理参数解析)
|
||||
- [4.2 单卡推理](#42-单卡推理)
|
||||
- [4.3 多卡并行推理](#43-多卡并行推理)
|
||||
- [五、更多资源](#五更多资源)
|
||||
|
||||
---
|
||||
|
||||
## 一、环境配置
|
||||
|
||||
**方式 1：使用requirements.txt**
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**方式 2：手动安装依赖**
|
||||
|
||||
```bash
|
||||
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 deepspeed==0.17.0 numpy==1.26.4
|
||||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||||
pip install opencv-python-headless
|
||||
```
|
||||
|
||||
**方式 3：使用docker**
|
||||
|
||||
使用docker的情况下，请保证机器中已经正确安装显卡驱动与CUDA环境，然后以此执行以下命令：
|
||||
|
||||
```
|
||||
# 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.1 快速测试数据集
|
||||
|
||||
我们提供了一个测试的数据集，其中包含若干训练数据。
|
||||
|
||||
```bash
|
||||
# 下载官方示例数据集
|
||||
modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun-Images-Demo
|
||||
```
|
||||
|
||||
### 2.2 数据集结构
|
||||
|
||||
```
|
||||
📦 datasets/
|
||||
├── 📂 my_dataset/
|
||||
│ ├── 📂 train/
|
||||
│ │ ├── 📄 image001.jpg
|
||||
│ │ ├── 📄 image002.png
|
||||
│ │ └── 📄 ...
|
||||
│ └── 📄 metadata.json
|
||||
```
|
||||
|
||||
### 2.3 metadata.json 格式
|
||||
|
||||
**相对路径格式**（示例格式）：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"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
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**绝对路径格式**：
|
||||
```json
|
||||
[
|
||||
{
|
||||
"file_path": "/mnt/data/images/sunset.jpg",
|
||||
"text": "A beautiful sunset over the ocean",
|
||||
"width": 1024,
|
||||
"height": 1024
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
**关键字段说明**：
|
||||
- `file_path`：图片路径（相对或绝对路径）
|
||||
- `text`：图片描述（英文提示词）
|
||||
- `width` / `height`：图片宽高（**最好提供**，用于分桶训练，如果不提供则自动在训练时读取，当数据存储在如oss这样的速度较慢的系统上时，可能会影响训练速度）。
|
||||
- 可以使用`scripts/process_json_add_width_and_height.py`文件对无width与height字段的json进行提取，支持处理图片与视频。
|
||||
- 使用方案为`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 相对路径与绝对路径使用方案
|
||||
|
||||
**相对路径**：
|
||||
|
||||
如果数据的路径为相对路径，则在训练脚本中设置：
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
|
||||
```
|
||||
|
||||
**绝对路径**：
|
||||
|
||||
如果数据的路径为绝对路径，则在训练脚本中设置：
|
||||
|
||||
```bash
|
||||
export DATASET_NAME=""
|
||||
export DATASET_META_NAME="/mnt/data/metadata.json"
|
||||
```
|
||||
|
||||
> 💡 **建议**：如果数据集较小且存储在本地，推荐使用相对路径；如果数据集存储在外部存储（如 NAS、OSS）或多个机器共享存储，推荐使用绝对路径。
|
||||
|
||||
---
|
||||
|
||||
## 三、全量参数训练
|
||||
|
||||
### 3.1 下载预训练模型
|
||||
|
||||
```bash
|
||||
# 创建模型目录
|
||||
mkdir -p models/Diffusion_Transformer
|
||||
|
||||
# 下载 Lens 官方权重
|
||||
modelscope download --model microsoft/Lens --local_dir models/Diffusion_Transformer/Lens
|
||||
```
|
||||
|
||||
### 3.2 快速开始（DeepSpeed-Zero-2）
|
||||
|
||||
如果按照 **2.1 快速测试数据集下载数据** 与 **3.1 下载预训练模型下载权重**后，直接复制快速开始的启动指令进行启动。
|
||||
|
||||
推荐使用DeepSpeed-Zero-2与FSDP方案进行训练。这里使用DeepSpeed-Zero-2为例配置shell文件。
|
||||
|
||||
本文中DeepSpeed-Zero-2与FSDP的差别在于是否对模型权重进行分片，**如果使用多卡且使用DeepSpeed-Zero-2的情况下显存不足**，可以切换使用FSDP进行训练。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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/lens/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_lens" \
|
||||
--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 训练常用参数解析
|
||||
|
||||
**关键参数说明**：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|-----|------|-------|
|
||||
| `--pretrained_model_name_or_path` | 预训练模型路径 | `models/Diffusion_Transformer/Lens` |
|
||||
| `--train_data_dir` | 训练数据目录 | `datasets/internal_datasets/` |
|
||||
| `--train_data_meta` | 训练数据元文件 | `datasets/internal_datasets/metadata.json` |
|
||||
| `--train_batch_size` | 每批次样本数 | 1 |
|
||||
| `--image_sample_size` | 最大训练分辨率,代码会自动分桶 | 1328 |
|
||||
| `--gradient_accumulation_steps` | 梯度累积步数（等效增大 batch） | 1 |
|
||||
| `--dataloader_num_workers` | DataLoader 子进程数 | 8 |
|
||||
| `--num_train_epochs` | 训练 epoch 数 | 100 |
|
||||
| `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
|
||||
| `--learning_rate` | 初始学习率 | 2e-05 |
|
||||
| `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
|
||||
| `--lr_warmup_steps` | 学习率预热步数 | 100 |
|
||||
| `--seed` | 随机种子 | 42 |
|
||||
| `--output_dir` | 输出目录 | `output_dir_lens` |
|
||||
| `--gradient_checkpointing` | 激活重计算 | - |
|
||||
| `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
|
||||
| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
|
||||
| `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
|
||||
| `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
|
||||
| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
|
||||
| `--enable_bucket` | 启用分桶训练，不裁剪图片，按分辨率分组训练整个图像 | - |
|
||||
| `--random_hw_adapt` | 自动缩放图片到 `[512, image_sample_size]` 范围内的随机尺寸 | - |
|
||||
| `--resume_from_checkpoint` | 恢复训练路径，使用 `"latest"` 自动选择最新 checkpoint | None |
|
||||
| `--uniform_sampling` | 均匀采样 timestep | - |
|
||||
| `--trainable_modules` | 可训练模块（`"."` 表示所有模块） | `"."` |
|
||||
| `--validation_steps` | 每 N 步执行一次验证 | 100 |
|
||||
| `--validation_epochs` | 每 N 个epoch执行一次验证 | 100 |
|
||||
| `--validation_prompts` | 验证图像生成的提示词 | `"一位年轻女子..."` |
|
||||
|
||||
|
||||
### 3.4 训练验证
|
||||
|
||||
你可以配置验证参数，在训练过程中定期生成测试图像，以便监控训练进度和模型质量。
|
||||
|
||||
**验证参数说明**：
|
||||
|
||||
| 参数 | 说明 | 推荐值 |
|
||||
|------|------|--------|
|
||||
| `--validation_steps` | 每 N 步执行一次验证 | 100 |
|
||||
| `--validation_epochs` | 每 N 个epoch执行一次验证 | 100 |
|
||||
| `--validation_prompts` | 验证图像生成的提示词,可用空格分隔多个提示词 | 多个空格分隔的提示词 |
|
||||
|
||||
**示例**：
|
||||
|
||||
```bash
|
||||
--validation_steps=100 \
|
||||
--validation_epochs=100 \
|
||||
--validation_prompts="一位年轻女子站在阳光明媚的海岸线上，白裙在轻拂的海风中微微飘动。"
|
||||
```
|
||||
|
||||
**注意事项**：
|
||||
- 验证图像会保存到 `output_dir` 目录中
|
||||
- 多提示词验证格式：`--validation_prompts "prompt1" "prompt2" "prompt3"`
|
||||
|
||||
### 3.5 使用 FSDP 训练
|
||||
|
||||
**如果使用多卡且使用DeepSpeed-Zero-2的情况下显存不足**，可以切换使用FSDP进行训练。
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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 LensTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/lens/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_lens" \
|
||||
--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 不使用 DeepSpeed 与 FSDP 训练
|
||||
|
||||
**该方案并不被推荐，因为没有显存节约后端，容易造成显存不足**。这里仅提供训练Shell用于参考训练。
|
||||
|
||||
```sh
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata.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/lens/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_lens" \
|
||||
--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 多机分布式训练
|
||||
|
||||
**适合场景**：超大规模数据集、需要更快的训练速度
|
||||
|
||||
#### 3.7.1 环境配置
|
||||
|
||||
假设有 2 台机器，每台 8 张 GPU：
|
||||
|
||||
**机器 0（Master）**：
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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 机器 IP
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2 # 机器总数
|
||||
export NUM_PROCESS=16 # 总进程数 = 机器数 × 8
|
||||
export RANK=0 # 当前机器 rank（0 或 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/lens/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_lens" \
|
||||
--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 "."
|
||||
```
|
||||
|
||||
**机器 1（Worker）**：
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Lens"
|
||||
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 相同
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2
|
||||
export NUM_PROCESS=16
|
||||
export RANK=1 # 注意这里是 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
|
||||
|
||||
# 使用与机器 0 相同的 accelerate launch 命令
|
||||
```
|
||||
|
||||
#### 3.7.2 多机训练注意事项
|
||||
|
||||
- **网络要求**：
|
||||
- 推荐 RDMA/InfiniBand（高性能）
|
||||
- 无 RDMA 时添加环境变量：
|
||||
```bash
|
||||
export NCCL_IB_DISABLE=1
|
||||
export NCCL_P2P_DISABLE=1
|
||||
```
|
||||
|
||||
- **数据同步**：所有机器必须能够访问相同的数据路径（NFS/共享存储）
|
||||
|
||||
## 四、推理测试
|
||||
|
||||
### 4.1 推理参数解析
|
||||
|
||||
**关键参数说明**：
|
||||
|
||||
| 参数 | 说明 | 示例值 |
|
||||
|------|------|-------|
|
||||
| `GPU_memory_mode` | 显存管理模式，可选值见下表 | `model_cpu_offload` |
|
||||
| `ulysses_degree` | Head 维度并行度，单卡时为 1 | 1 |
|
||||
| `ring_degree` | Sequence 维度并行度，单卡时为 1 | 1 |
|
||||
| `fsdp_dit` | 多卡推理时对 Transformer 使用 FSDP 节省显存 | `False` |
|
||||
| `fsdp_text_encoder` | 多卡推理时对文本编码器使用 FSDP | `False` |
|
||||
| `compile_dit` | 编译 Transformer 加速推理（固定分辨率下有效） | `False` |
|
||||
| `model_name` | 模型路径 | `models/Diffusion_Transformer/Lens` |
|
||||
| `sampler_name` | 采样器类型：`Flow`、`Flow_Unipc`、`Flow_DPM++` | `Flow` |
|
||||
| `transformer_path` | 加载训练好的 Transformer 权重路径 | `None` |
|
||||
| `vae_path` | 加载训练好的 VAE 权重路径 | `None` |
|
||||
| `lora_path` | LoRA 权重路径 | `None` |
|
||||
| `sample_size` | 生成图像分辨率 `[高度, 宽度]` | `[1728, 992]` |
|
||||
| `weight_dtype` | 模型权重精度，不支持 bf16 的显卡使用 `torch.float16` | `torch.bfloat16` |
|
||||
| `prompt` | 正向提示词,描述生成内容 | `"1girl, black_hair..."` |
|
||||
| `negative_prompt` | 负向提示词,避免生成的内容 | `"低分辨率,低画质..."` |
|
||||
| `guidance_scale` | 引导强度 | 4.5 |
|
||||
| `seed` | 随机种子,用于复现结果 | 43 |
|
||||
| `num_inference_steps` | 推理步数 | 40 |
|
||||
| `lora_weight` | LoRA 权重强度 | 0.55 |
|
||||
| `save_path` | 生成图像保存路径 | `samples/lens-t2i` |
|
||||
|
||||
**显存管理模式说明**：
|
||||
|
||||
| 模式 | 说明 | 显存占用 |
|
||||
|------|------|---------|
|
||||
| `model_full_load` | 整个模型加载到 GPU | 最高 |
|
||||
| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
|
||||
| `model_cpu_offload` | 使用后将模型卸载到 CPU | 中等 |
|
||||
| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
|
||||
| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
|
||||
| `sequential_cpu_offload` | 逐层卸载（速度最慢） | 最低 |
|
||||
|
||||
### 4.2 单卡推理
|
||||
|
||||
单卡推理运行如下命令：
|
||||
|
||||
```bash
|
||||
python examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
根据需求修改编辑 `examples/ernie_image/predict_t2i.py`,初次推理重点关注如下参数,如果对其他参数感兴趣,请查看上方的推理参数解析。
|
||||
|
||||
```python
|
||||
# 根据显卡显存选择
|
||||
GPU_memory_mode = "model_cpu_offload"
|
||||
# 根据实际模型路径
|
||||
model_name = "models/Diffusion_Transformer/Lens"
|
||||
# 训练好的权重路径,如 "output_dir_lens/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||||
transformer_path = None
|
||||
# 根据生成内容编写
|
||||
prompt = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||||
# ...
|
||||
```
|
||||
|
||||
### 4.3 多卡并行推理
|
||||
|
||||
**适合场景**：高分辨率生成、加速推理
|
||||
|
||||
#### 安装并行推理依赖
|
||||
|
||||
```bash
|
||||
pip install xfuser==0.4.2 yunchang==0.6.2
|
||||
```
|
||||
|
||||
#### 配置并行策略
|
||||
|
||||
编辑 `examples/ernie_image/predict_t2i.py`:
|
||||
|
||||
```python
|
||||
# 确保 ulysses_degree × ring_degree = GPU 数量
|
||||
# 例如使用 2 张 GPU：
|
||||
ulysses_degree = 2 # Head 维度并行
|
||||
ring_degree = 1 # Sequence 维度并行
|
||||
```
|
||||
|
||||
**配置原则**：
|
||||
- `ulysses_degree` 必须能整除模型的head数。
|
||||
- `ring_degree` 会在sequence上切分，影响通信开销，在head数能切分的时候尽量不用。
|
||||
|
||||
**示例配置**：
|
||||
|
||||
| GPU 数量 | ulysses_degree | ring_degree | 说明 |
|
||||
|---------|---------------|-------------|------|
|
||||
| 1 | 1 | 1 | 单卡 |
|
||||
| 4 | 4 | 1 | Head 并行 |
|
||||
| 8 | 8 | 1 | Head 并行 |
|
||||
| 8 | 4 | 2 | 混合并行 |
|
||||
|
||||
#### 运行多卡推理
|
||||
|
||||
```bash
|
||||
torchrun --nproc-per-node=2 examples/lens/predict_t2i.py
|
||||
```
|
||||
|
||||
## 五、更多资源
|
||||
|
||||
- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,33 @@
|
||||
export MODEL_NAME="../CogVideoX-Fun-Github/models/Diffusion_Transformer/Lens"
|
||||
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/lens/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=100 \
|
||||
--learning_rate=2e-05 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens" \
|
||||
--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 \
|
||||
--low_vram \
|
||||
--uniform_sampling \
|
||||
--trainable_modules "."
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,33 @@
|
||||
export MODEL_NAME="../CogVideoX-Fun-Github/models/Diffusion_Transformer/Lens"
|
||||
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/lens/train_lora.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=100 \
|
||||
--learning_rate=1e-04 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_lens_lora" \
|
||||
--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 \
|
||||
--low_vram \
|
||||
--uniform_sampling \
|
||||
--rank=64 \
|
||||
--network_alpha=32 \
|
||||
--target_name="img_qkv,txt_qkv,to_out.0,to_add_out,img_mod.1,txt_mod.1,img_mlp,txt_mlp"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,36 @@
|
||||
export MODEL_NAME="models/Diffusion_Transformer/LTX-2"
|
||||
export DATASET_NAME="datasets/internal_datasets/"
|
||||
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.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/ltx2/train_upsampler.py \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$DATASET_META_NAME \
|
||||
--video_sample_size=1280 \
|
||||
--video_sample_stride=1 \
|
||||
--video_sample_n_frames=121 \
|
||||
--train_batch_size=1 \
|
||||
--video_repeat=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=50 \
|
||||
--learning_rate=5e-05 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_ltx2_upsampler" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--vae_mini_batch=1 \
|
||||
--max_grad_norm=1.0 \
|
||||
--random_hw_adapt \
|
||||
--enable_bucket \
|
||||
--spatial_scale=2.0 \
|
||||
--trainable_modules "."
|
||||
@@ -36,6 +36,16 @@ def process_media_sample(sample, base_dir=None):
|
||||
if not file_path_str:
|
||||
return sample
|
||||
|
||||
# --- MODIFICATION START ---
|
||||
# If file_path is a list, take the first element
|
||||
if isinstance(file_path_str, list):
|
||||
if len(file_path_str) > 0:
|
||||
file_path_str = file_path_str[0]
|
||||
else:
|
||||
# Empty list, cannot process
|
||||
return sample
|
||||
# --- MODIFICATION END ---
|
||||
|
||||
# Handle path resolution
|
||||
file_path_obj = Path(file_path_str)
|
||||
|
||||
|
||||
@@ -225,7 +225,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_ode
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
@@ -288,6 +288,17 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
|
||||
| `--independent_first_frame` | First frame is independent (`[1, N, N, ...]` block pattern) | - |
|
||||
| `--context_noise` | Context noise level (matches downstream Self-Forcing distillation config) | 0 |
|
||||
|
||||
**Validation Parameters (Optional)**:
|
||||
|
||||
| Parameter | Description | Example |
|
||||
|-----------|-------------|---------|
|
||||
| `--validation_steps` | Run validation every N steps | 2000 |
|
||||
| `--validation_epochs` | Run validation every N epochs | 5 |
|
||||
| `--validation_prompts` | Prompts used for validation video generation | English prompt |
|
||||
| `--video_sample_size` | Validation sample size | 640 |
|
||||
| `--video_sample_n_frames` | Number of frames for validation videos | 81 |
|
||||
| `--fix_sample_size` | Fixed `[height, width]` used during validation | `480 832` |
|
||||
|
||||
### 4.3 Training with DeepSpeed-Zero-2 / FSDP
|
||||
|
||||
For multi-GPU training, the same memory-saving backends as the distillation stage are supported.
|
||||
@@ -313,7 +324,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
@@ -352,7 +363,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
@@ -400,7 +411,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
|
||||
@@ -225,7 +225,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_ode
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
@@ -301,47 +301,84 @@ bash scripts/wan2.1_self_forcing/train_ode.sh
|
||||
|
||||
### 4.3 使用 DeepSpeed-Zero-2 / FSDP 训练
|
||||
|
||||
多卡训练支持与蒸馏阶段相同的显存节约后端。将 4.1 中 `accelerate launch` 前缀替换为以下任意一种即可：
|
||||
多卡训练支持与蒸馏阶段相同的显存节约后端。
|
||||
|
||||
**DeepSpeed-Zero-2**（推荐默认）：
|
||||
|
||||
```bash
|
||||
accelerate launch \
|
||||
--use_deepspeed --deepspeed_config_file config/zero_stage2_config.json \
|
||||
--deepspeed_multinode_launcher standard \
|
||||
scripts/wan2.1_self_forcing/train_ode.py \
|
||||
... # 训练参数与 4.1 相同
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME=""
|
||||
export ODE_DATA_META="datasets/ode_pairs_output/outputs.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/wan2.1_self_forcing/train_ode.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$ODE_DATA_META \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.1_self_forcing_ode_regression" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--max_grad_norm=0.05 \
|
||||
--num_frame_per_block=3 \
|
||||
--train_sampling_steps=1000 \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--shift=8.0 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--trainable_modules "."
|
||||
```
|
||||
|
||||
**FSDP**（DeepSpeed-Zero-2 显存不足时使用）：
|
||||
|
||||
```bash
|
||||
accelerate launch --mixed_precision="bf16" \
|
||||
--use_fsdp \
|
||||
--fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
|
||||
--fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
|
||||
--fsdp_sharding_strategy "FULL_SHARD" \
|
||||
--fsdp_state_dict_type=SHARDED_STATE_DICT \
|
||||
--fsdp_backward_prefetch "BACKWARD_PRE" \
|
||||
--fsdp_cpu_ram_efficient_loading False \
|
||||
scripts/wan2.1_self_forcing/train_ode.py \
|
||||
... # 训练参数与 4.1 相同
|
||||
```
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME=""
|
||||
export ODE_DATA_META="datasets/ode_pairs_output/outputs.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
|
||||
|
||||
**DeepSpeed-Zero-3**（适用于超大模型，1.3B 通常不需要）：
|
||||
|
||||
```bash
|
||||
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true \
|
||||
--use_deepspeed --deepspeed_config_file config/zero_stage3_config.json \
|
||||
--deepspeed_multinode_launcher standard \
|
||||
scripts/wan2.1_self_forcing/train_ode.py \
|
||||
... # 训练参数与 4.1 相同
|
||||
|
||||
# 训练完成后将分片 checkpoint 转为单文件 bf16：
|
||||
python scripts/zero_to_bf16.py \
|
||||
output_dir_wan2.1_self_forcing_ode_regression/checkpoint-{N} \
|
||||
output_dir_wan2.1_self_forcing_ode_regression/checkpoint-{N}-outputs \
|
||||
--max_shard_size 80GB --safe_serialization
|
||||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_ode.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$ODE_DATA_META \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.1_self_forcing_ode_regression" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--max_grad_norm=0.05 \
|
||||
--num_frame_per_block=3 \
|
||||
--train_sampling_steps=1000 \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--shift=8.0 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--trainable_modules "."
|
||||
```
|
||||
|
||||
### 4.4 多机分布式训练
|
||||
@@ -351,25 +388,65 @@ python scripts/zero_to_bf16.py \
|
||||
**机器 0（Master）**：
|
||||
|
||||
```bash
|
||||
export MASTER_ADDR="192.168.1.100"
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME=""
|
||||
export ODE_DATA_META="datasets/ode_pairs_output/outputs.json"
|
||||
export MASTER_ADDR="192.168.1.100" # 主节点 IP
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2 # 机器总数
|
||||
export NUM_PROCESS=16 # 总进程数 = 机器数 × 8
|
||||
export RANK=0 # 本机 rank（0 或 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/wan2.1_self_forcing/train_ode.py \
|
||||
--config_path="config/wan2.1/wan_civitai.yaml" \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--train_data_dir=$DATASET_NAME \
|
||||
--train_data_meta=$ODE_DATA_META \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
--output_dir="output_dir_wan2.1_self_forcing_ode_regression" \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="bf16" \
|
||||
--adam_weight_decay=3e-2 \
|
||||
--adam_epsilon=1e-10 \
|
||||
--max_grad_norm=0.05 \
|
||||
--num_frame_per_block=3 \
|
||||
--train_sampling_steps=1000 \
|
||||
--denoising_step_indices_list 1000 750 500 250 \
|
||||
--shift=8.0 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--trainable_modules "."
|
||||
```
|
||||
|
||||
**机器 1（Worker）**：
|
||||
```bash
|
||||
export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
|
||||
export DATASET_NAME=""
|
||||
export ODE_DATA_META="datasets/ode_pairs_output/outputs.json"
|
||||
export MASTER_ADDR="192.168.1.100" # 与 Master 相同
|
||||
export MASTER_PORT=10086
|
||||
export WORLD_SIZE=2
|
||||
export NUM_PROCESS=16
|
||||
export RANK=0
|
||||
export RANK=1 # 注意此处为 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/wan2.1_self_forcing/train_ode.py \
|
||||
... # 训练参数与 4.1 相同
|
||||
# 与机器 0 使用完全相同的 accelerate launch 命令
|
||||
```
|
||||
|
||||
**机器 1（Worker）**：与 Master 完全相同，仅将 `export RANK=1`。
|
||||
|
||||
**注意事项**：
|
||||
- 优先使用 RDMA / InfiniBand。无 RDMA 时需设置 `NCCL_IB_DISABLE=1` 与 `NCCL_P2P_DISABLE=1`。
|
||||
- 所有机器必须共享同一份 `outputs.json` 与对应的 `.safetensors` 文件（NFS / 共享存储）。
|
||||
|
||||
@@ -126,8 +126,10 @@ def log_validation(transformer3d, args, config, accelerator, weight_dtype, globa
|
||||
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
|
||||
text_encoder = shard_fn(text_encoder)
|
||||
|
||||
scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs'])
|
||||
scheduler_kwargs['shift'] = args.shift
|
||||
scheduler = FlowMatchEulerDiscreteScheduler(
|
||||
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
|
||||
**filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs)
|
||||
)
|
||||
pipeline = WanSelfForcingPipeline(
|
||||
vae=vae,
|
||||
@@ -158,9 +160,11 @@ def log_validation(transformer3d, args, config, accelerator, weight_dtype, globa
|
||||
generator=generator,
|
||||
guidance_scale=1.0,
|
||||
num_inference_steps=len(args.denoising_step_indices_list),
|
||||
shift=args.shift,
|
||||
num_frame_per_block=args.num_frame_per_block,
|
||||
independent_first_frame=args.independent_first_frame,
|
||||
context_noise=args.context_noise,
|
||||
stochastic_sampling=True,
|
||||
).videos
|
||||
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
|
||||
save_videos_grid(
|
||||
@@ -188,7 +192,8 @@ def log_validation(transformer3d, args, config, accelerator, weight_dtype, globa
|
||||
|
||||
def get_timestep_for_ode(
|
||||
min_timestep, max_timestep, batch_size, num_frames,
|
||||
num_frame_per_block, independent_first_frame, device
|
||||
num_frame_per_block, independent_first_frame, device,
|
||||
generator=None,
|
||||
):
|
||||
"""
|
||||
Generate random timestep indices per frame/block.
|
||||
@@ -198,7 +203,8 @@ def get_timestep_for_ode(
|
||||
timestep = torch.randint(
|
||||
min_timestep, max_timestep,
|
||||
[batch_size, num_frames],
|
||||
device=device, dtype=torch.long
|
||||
device=device, dtype=torch.long,
|
||||
generator=generator,
|
||||
)
|
||||
if independent_first_frame:
|
||||
timestep_from_second = timestep[:, 1:]
|
||||
@@ -215,6 +221,38 @@ def get_timestep_for_ode(
|
||||
return timestep
|
||||
|
||||
|
||||
def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length,
|
||||
num_layers, num_heads, head_dim, dtype, device):
|
||||
"""Initialize KV cache for block-by-block training (mirrors train_distill)."""
|
||||
kv_cache_size = num_frames * frame_seq_length
|
||||
kv_cache = []
|
||||
for _ in range(num_layers):
|
||||
kv_cache.append({
|
||||
"k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim],
|
||||
dtype=dtype, device=device),
|
||||
"v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim],
|
||||
dtype=dtype, device=device),
|
||||
"global_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
"local_end_index": torch.tensor([0], dtype=torch.long, device=device),
|
||||
})
|
||||
return kv_cache
|
||||
|
||||
|
||||
def initialize_crossattn_cache_for_training(batch_size, text_len, num_layers,
|
||||
num_heads, head_dim, dtype, device):
|
||||
"""Initialize cross-attention cache for block-by-block training."""
|
||||
crossattn_cache = []
|
||||
for _ in range(num_layers):
|
||||
crossattn_cache.append({
|
||||
"k": torch.zeros([batch_size, text_len, num_heads, head_dim],
|
||||
dtype=dtype, device=device),
|
||||
"v": torch.zeros([batch_size, text_len, num_heads, head_dim],
|
||||
dtype=dtype, device=device),
|
||||
"is_init": False,
|
||||
})
|
||||
return crossattn_cache
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Args
|
||||
# ============================================================================
|
||||
@@ -477,6 +515,26 @@ def parse_args():
|
||||
default=8.0,
|
||||
help="Shift value for FlowMatchEulerDiscreteScheduler. Default: 8.0 (matches ODE data generation).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_kv_cache_training",
|
||||
action="store_true",
|
||||
help=(
|
||||
"If set, run block-by-block KV cache training that fully matches the "
|
||||
"pipeline_wan_self_forcing inference behavior. Otherwise fall back to "
|
||||
"the default one-shot causal-mask ODE regression (kept as baseline)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prob_full_zero_start",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help=(
|
||||
"Probability (per-sample) of forcing ALL frames in ALL blocks to use "
|
||||
"timestep index=0 (pure-noise start). Bridges the train-inference gap "
|
||||
"so the model also sees the real autoregressive rollout where every "
|
||||
"block starts from fresh noise. 0.0 disables (default)."
|
||||
),
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
||||
@@ -771,11 +829,63 @@ def main():
|
||||
RandomSampler(train_dataset, generator=batch_sampler_generator),
|
||||
batch_size=args.train_batch_size, drop_last=True
|
||||
)
|
||||
|
||||
def ode_safetensors_collate_fn(examples):
|
||||
"""Collate safetensors-loaded ODE samples into a batch.
|
||||
|
||||
Each sample is a dict with keys:
|
||||
- 'latents': [S, C, F, H, W]
|
||||
- 'prompt_embeds': [L, D]
|
||||
- 'prompt_attention_mask': [L]
|
||||
|
||||
The default torch collate fails when, across samples, the same key has
|
||||
slightly different dtypes/lengths (e.g. attention_mask saved as bool/int
|
||||
vs long, or prompt_embeds with different seq lengths). This custom
|
||||
collate normalizes dtypes and pads variable-length text fields so
|
||||
`torch.stack` always succeeds.
|
||||
"""
|
||||
out = {}
|
||||
|
||||
# ---- latents: assume identical shape across samples (fixed by pipeline) ----
|
||||
latents = [ex["latents"] for ex in examples]
|
||||
target_latent_dtype = latents[0].dtype
|
||||
latents = [t.to(target_latent_dtype) for t in latents]
|
||||
out["latents"] = torch.stack(latents, dim=0)
|
||||
|
||||
# ---- prompt_embeds: pad along seq dim, unify dtype ----
|
||||
embeds = [ex["prompt_embeds"] for ex in examples]
|
||||
embed_dtype = embeds[0].dtype
|
||||
max_len = max(e.shape[0] for e in embeds)
|
||||
padded_embeds = []
|
||||
for e in embeds:
|
||||
e = e.to(embed_dtype)
|
||||
if e.shape[0] < max_len:
|
||||
pad = torch.zeros(
|
||||
max_len - e.shape[0], *e.shape[1:], dtype=embed_dtype
|
||||
)
|
||||
e = torch.cat([e, pad], dim=0)
|
||||
padded_embeds.append(e)
|
||||
out["prompt_embeds"] = torch.stack(padded_embeds, dim=0)
|
||||
|
||||
# ---- prompt_attention_mask: pad along seq dim, force long dtype ----
|
||||
masks = [ex["prompt_attention_mask"].long() for ex in examples]
|
||||
max_len = max(m.shape[0] for m in masks)
|
||||
padded_masks = []
|
||||
for m in masks:
|
||||
if m.shape[0] < max_len:
|
||||
pad = torch.zeros(max_len - m.shape[0], dtype=torch.long)
|
||||
m = torch.cat([m, pad], dim=0)
|
||||
padded_masks.append(m)
|
||||
out["prompt_attention_mask"] = torch.stack(padded_masks, dim=0)
|
||||
|
||||
return out
|
||||
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset,
|
||||
batch_sampler=batch_sampler,
|
||||
persistent_workers=True if args.dataloader_num_workers != 0 else False,
|
||||
num_workers=args.dataloader_num_workers,
|
||||
collate_fn=ode_safetensors_collate_fn,
|
||||
)
|
||||
|
||||
# Scheduler and math around the number of training steps.
|
||||
@@ -802,9 +912,19 @@ def main():
|
||||
denoising_step_list = noise_scheduler.timesteps[
|
||||
args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)
|
||||
]
|
||||
num_denoising_steps = len(denoising_step_list)
|
||||
# Training denoising step list: append 0 (clean) for train-inference context alignment.
|
||||
# index=4 frames use clean latent as input but are excluded from loss via mask=(timestep!=0).
|
||||
# They serve as clean context for later blocks via causal attention.
|
||||
train_denoising_step_list = denoising_step_list
|
||||
if 0 not in denoising_step_list.tolist():
|
||||
train_denoising_step_list = torch.cat([
|
||||
denoising_step_list, torch.tensor([0], device=denoising_step_list.device)
|
||||
])
|
||||
num_denoising_steps = len(train_denoising_step_list)
|
||||
if accelerator.is_main_process:
|
||||
print(f"Denoising step list: {denoising_step_list.tolist()}")
|
||||
print(f"Denoising step list (inference): {denoising_step_list.tolist()}")
|
||||
print(f"Denoising step list (training): {train_denoising_step_list.tolist()}")
|
||||
print(f"num_denoising_steps (includes clean): {num_denoising_steps}")
|
||||
print(f"Dataset size: {len(train_dataset)}")
|
||||
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
@@ -910,77 +1030,226 @@ def main():
|
||||
# Target: clean endpoint (last timestep)
|
||||
target_latent = ode_latent[:, -1] # [B, C, F, H, W]
|
||||
num_frames = target_latent.shape[2]
|
||||
C_dim, F_dim, H_dim, W_dim = (
|
||||
ode_latent.shape[2], ode_latent.shape[3],
|
||||
ode_latent.shape[4], ode_latent.shape[5],
|
||||
)
|
||||
|
||||
# Random timestep index per frame/block
|
||||
index = get_timestep_for_ode(
|
||||
0, num_denoising_steps, bsz, num_frames,
|
||||
args.num_frame_per_block, args.independent_first_frame,
|
||||
accelerator.device
|
||||
) # [B, F]
|
||||
|
||||
# Gather noisy input from ODE trajectory
|
||||
# ode_latent: [B, S, C, F, H, W], index: [B, F] -> expand to gather
|
||||
C_dim, F_dim, H_dim, W_dim = ode_latent.shape[2], ode_latent.shape[3], ode_latent.shape[4], ode_latent.shape[5]
|
||||
gather_index = index.reshape(bsz, 1, 1, num_frames, 1, 1).expand(-1, -1, C_dim, -1, H_dim, W_dim)
|
||||
# Transpose ode_latent to [B, S, C, F, H, W] for gathering along dim=1
|
||||
noisy_input = torch.gather(ode_latent, dim=1, index=gather_index).squeeze(1) # [B, C, F, H, W]
|
||||
|
||||
# Compute actual timestep values: [B, F]
|
||||
timestep = denoising_step_list[index] # [B, F]
|
||||
|
||||
# --- Forward through causal generator ---
|
||||
# Create block mask for causal training
|
||||
patch_h, patch_w = accelerator.unwrap_model(transformer3d).config.patch_size[1:]
|
||||
frame_seqlen = (H_dim * W_dim) // (patch_h * patch_w)
|
||||
seq_len = frame_seqlen * num_frames
|
||||
|
||||
accelerator.unwrap_model(transformer3d).create_block_mask_for_training(
|
||||
num_frames=num_frames,
|
||||
frame_seqlen=frame_seqlen,
|
||||
num_frame_per_block=args.num_frame_per_block,
|
||||
independent_first_frame=args.independent_first_frame,
|
||||
device=accelerator.device
|
||||
)
|
||||
|
||||
# Convert to list format for transformer
|
||||
noisy_input_list = [noisy_input[i] for i in range(bsz)]
|
||||
|
||||
with accelerator.accumulate(transformer3d):
|
||||
with torch.cuda.amp.autocast(dtype=weight_dtype):
|
||||
# Pass per-frame timestep [B, F] so each frame gets its own time embedding.
|
||||
# This matches the original Self-Forcing: different frames are at different
|
||||
# noise levels and require independent time modulation.
|
||||
flow_pred = transformer3d(
|
||||
x=noisy_input_list,
|
||||
context=prompt_embeds,
|
||||
t=timestep,
|
||||
seq_len=seq_len,
|
||||
)
|
||||
if args.use_kv_cache_training:
|
||||
# ============================================================
|
||||
# Block-by-block KV cache training (autoregressive, single-step x0)
|
||||
# Starting timestep is randomly sampled per block — same as the
|
||||
# non-KV-cache (baseline) branch. Each block performs ONE forward
|
||||
# to predict x0; KV cache is then refreshed with pred_block +
|
||||
# context_noise to keep the autoregressive rollout intact.
|
||||
# ============================================================
|
||||
# 1) Block split (mirrors pipeline_wan_self_forcing)
|
||||
if not args.independent_first_frame:
|
||||
assert num_frames % args.num_frame_per_block == 0
|
||||
num_blocks_split = num_frames // args.num_frame_per_block
|
||||
all_num_frames = [args.num_frame_per_block] * num_blocks_split
|
||||
else:
|
||||
assert (num_frames - 1) % args.num_frame_per_block == 0
|
||||
num_blocks_split = (num_frames - 1) // args.num_frame_per_block
|
||||
all_num_frames = [1] + [args.num_frame_per_block] * num_blocks_split
|
||||
|
||||
# Convert flow prediction to x0 prediction (per-frame).
|
||||
# flow_pred: [B, C, F, H, W], xt: [B, C, F, H, W]
|
||||
# x0 = xt - sigma_t * flow_pred
|
||||
# Each frame has its own sigma from its own timestep.
|
||||
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=torch.float64)
|
||||
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
|
||||
# timestep: [B, F] -> flatten to [B*F] for per-frame sigma lookup
|
||||
step_indices = torch.argmin(
|
||||
(schedule_timesteps.unsqueeze(0) - timestep.reshape(-1).unsqueeze(1)).abs(), dim=1
|
||||
) # [B*F]
|
||||
sigma = sigmas[step_indices].to(weight_dtype)
|
||||
sigma = sigma.reshape(bsz, 1, num_frames, 1, 1) # [B, 1, F, 1, 1]
|
||||
# 2) Random timestep index per frame/block (same as baseline branch)
|
||||
index = get_timestep_for_ode(
|
||||
0, num_denoising_steps, bsz, num_frames,
|
||||
args.num_frame_per_block, args.independent_first_frame,
|
||||
accelerator.device,
|
||||
generator=torch_rng,
|
||||
) # [B, F]
|
||||
# Optional: force per-sample full-zero start to cover the
|
||||
# real inference rollout (all blocks starting from pure noise).
|
||||
if args.prob_full_zero_start > 0.0:
|
||||
zero_mask = (
|
||||
torch.rand(bsz, device=accelerator.device, generator=torch_rng)
|
||||
< args.prob_full_zero_start
|
||||
)
|
||||
if zero_mask.any():
|
||||
index[zero_mask] = 0
|
||||
gather_index = index.reshape(bsz, 1, 1, num_frames, 1, 1).expand(
|
||||
-1, -1, C_dim, -1, H_dim, W_dim
|
||||
)
|
||||
noisy_input_full = torch.gather(ode_latent, dim=1, index=gather_index).squeeze(1)
|
||||
timestep_full = train_denoising_step_list[index] # [B, F]
|
||||
|
||||
pred_x0 = noisy_input - sigma * flow_pred
|
||||
# 3) Initialize KV / cross-attention cache
|
||||
cfg = accelerator.unwrap_model(transformer3d).config
|
||||
num_layers_t = cfg.num_layers
|
||||
num_heads_t = cfg.num_heads
|
||||
head_dim_t = cfg.dim // num_heads_t
|
||||
text_len = 512 # T5 sequence length
|
||||
kv_cache = initialize_kv_cache_for_training(
|
||||
batch_size=bsz,
|
||||
num_frames=num_frames,
|
||||
frame_seq_length=frame_seqlen,
|
||||
num_layers=num_layers_t,
|
||||
num_heads=num_heads_t,
|
||||
head_dim=head_dim_t,
|
||||
dtype=weight_dtype,
|
||||
device=accelerator.device,
|
||||
)
|
||||
crossattn_cache = initialize_crossattn_cache_for_training(
|
||||
batch_size=bsz,
|
||||
text_len=text_len,
|
||||
num_layers=num_layers_t,
|
||||
num_heads=num_heads_t,
|
||||
head_dim=head_dim_t,
|
||||
dtype=weight_dtype,
|
||||
device=accelerator.device,
|
||||
)
|
||||
|
||||
# MSE loss (mask t=0 frames)
|
||||
# timestep: [B, F], mask frames where timestep != 0
|
||||
mask = (timestep != 0).unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # [B, 1, F, 1, 1]
|
||||
mask = mask.expand_as(target_latent).float()
|
||||
# 4) Sigma / timestep lookup tables (per-frame sigma)
|
||||
sigmas_full = noise_scheduler.sigmas.to(
|
||||
device=accelerator.device, dtype=torch.float64
|
||||
)
|
||||
schedule_timesteps_full = noise_scheduler.timesteps.to(accelerator.device)
|
||||
|
||||
if mask.sum() > 0:
|
||||
loss = F.mse_loss(pred_x0 * mask, target_latent * mask, reduction="sum") / mask.sum()
|
||||
current_start_frame = 0
|
||||
total_pred = torch.zeros_like(target_latent)
|
||||
full_seq_len = frame_seqlen * num_frames
|
||||
|
||||
# 5) Block-by-block rollout — single-step x0 prediction per block
|
||||
for block_idx, current_num_frames in enumerate(all_num_frames):
|
||||
start_idx = current_start_frame
|
||||
end_idx = current_start_frame + current_num_frames
|
||||
|
||||
noisy_input = noisy_input_full[:, :, start_idx:end_idx]
|
||||
timestep_block = timestep_full[:, start_idx:end_idx].to(torch.int64)
|
||||
|
||||
flow_pred = transformer3d(
|
||||
x=[noisy_input[i] for i in range(bsz)],
|
||||
context=prompt_embeds,
|
||||
t=timestep_block,
|
||||
seq_len=full_seq_len,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=current_start_frame * frame_seqlen,
|
||||
cache_start=None,
|
||||
)
|
||||
if isinstance(flow_pred, list):
|
||||
flow_pred = torch.stack(flow_pred, dim=0)
|
||||
|
||||
# Per-frame sigma -> x0 = xt - sigma * flow_pred
|
||||
step_indices_block = torch.argmin(
|
||||
(schedule_timesteps_full.unsqueeze(0)
|
||||
- timestep_block.reshape(-1).unsqueeze(1)).abs(), dim=1
|
||||
)
|
||||
sigma_block = sigmas_full[step_indices_block].to(weight_dtype)
|
||||
# timestep=0 (clean context) must use sigma=0 exactly.
|
||||
sigma_block[timestep_block.reshape(-1) == 0] = 0.0
|
||||
sigma_block = sigma_block.reshape(bsz, 1, current_num_frames, 1, 1)
|
||||
pred_block = noisy_input - sigma_block * flow_pred
|
||||
|
||||
total_pred[:, :, start_idx:end_idx] = pred_block
|
||||
|
||||
# 6) Update KV cache with student's pred_block + context_noise
|
||||
# (matches pipeline_wan_self_forcing L802-L839)
|
||||
if block_idx < len(all_num_frames) - 1:
|
||||
ctx_t = torch.full(
|
||||
[bsz, current_num_frames], args.context_noise,
|
||||
device=accelerator.device, dtype=torch.int64,
|
||||
)
|
||||
with torch.no_grad():
|
||||
transformer3d(
|
||||
x=[pred_block[i] for i in range(bsz)],
|
||||
context=prompt_embeds,
|
||||
t=ctx_t,
|
||||
seq_len=full_seq_len,
|
||||
kv_cache=kv_cache,
|
||||
crossattn_cache=crossattn_cache,
|
||||
current_start=current_start_frame * frame_seqlen,
|
||||
cache_start=None,
|
||||
)
|
||||
|
||||
current_start_frame += current_num_frames
|
||||
|
||||
# 7) ODE-endpoint MSE loss (mask out clean timestep=0 frames)
|
||||
mask = (timestep_full != 0).unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
|
||||
mask = mask.expand_as(target_latent).float()
|
||||
if mask.sum() > 0:
|
||||
loss = F.mse_loss(total_pred * mask, target_latent * mask, reduction="sum") / mask.sum()
|
||||
else:
|
||||
loss = F.mse_loss(total_pred, target_latent)
|
||||
else:
|
||||
loss = F.mse_loss(pred_x0, target_latent)
|
||||
# --- Baseline (one-shot causal-mask) preparation ---
|
||||
# Random timestep index per frame/block
|
||||
index = get_timestep_for_ode(
|
||||
0, num_denoising_steps, bsz, num_frames,
|
||||
args.num_frame_per_block, args.independent_first_frame,
|
||||
accelerator.device,
|
||||
generator=torch_rng,
|
||||
) # [B, F]
|
||||
# Optional: force per-sample full-zero start to cover the
|
||||
# real inference rollout (all blocks starting from pure noise).
|
||||
if args.prob_full_zero_start > 0.0:
|
||||
zero_mask = (
|
||||
torch.rand(bsz, device=accelerator.device, generator=torch_rng)
|
||||
< args.prob_full_zero_start
|
||||
)
|
||||
if zero_mask.any():
|
||||
index[zero_mask] = 0
|
||||
|
||||
# Gather noisy input from ODE trajectory
|
||||
gather_index = index.reshape(bsz, 1, 1, num_frames, 1, 1).expand(
|
||||
-1, -1, C_dim, -1, H_dim, W_dim
|
||||
)
|
||||
noisy_input = torch.gather(ode_latent, dim=1, index=gather_index).squeeze(1)
|
||||
|
||||
# Compute actual timestep values: [B, F]
|
||||
timestep = train_denoising_step_list[index] # [B, F]
|
||||
|
||||
# Build causal block mask
|
||||
accelerator.unwrap_model(transformer3d).create_block_mask_for_training(
|
||||
num_frames=num_frames,
|
||||
frame_seqlen=frame_seqlen,
|
||||
num_frame_per_block=args.num_frame_per_block,
|
||||
independent_first_frame=args.independent_first_frame,
|
||||
device=accelerator.device
|
||||
)
|
||||
|
||||
# Convert to list format for transformer
|
||||
noisy_input_list = [noisy_input[i] for i in range(bsz)]
|
||||
|
||||
# ============================================================
|
||||
# Baseline: one-shot causal-mask ODE regression
|
||||
# ============================================================
|
||||
flow_pred = transformer3d(
|
||||
x=noisy_input_list,
|
||||
context=prompt_embeds,
|
||||
t=timestep,
|
||||
seq_len=seq_len,
|
||||
)
|
||||
|
||||
# Convert flow prediction to x0 prediction (per-frame).
|
||||
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=torch.float64)
|
||||
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
|
||||
step_indices = torch.argmin(
|
||||
(schedule_timesteps.unsqueeze(0) - timestep.reshape(-1).unsqueeze(1)).abs(), dim=1
|
||||
)
|
||||
sigma = sigmas[step_indices].to(weight_dtype)
|
||||
# Fix: timestep=0 (clean context frames) should have sigma=0 exactly.
|
||||
sigma[timestep.reshape(-1) == 0] = 0.0
|
||||
sigma = sigma.reshape(bsz, 1, num_frames, 1, 1)
|
||||
|
||||
pred_x0 = noisy_input - sigma * flow_pred
|
||||
|
||||
# MSE loss (mask t=0 frames)
|
||||
mask = (timestep != 0).unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
|
||||
mask = mask.expand_as(target_latent).float()
|
||||
|
||||
if mask.sum() > 0:
|
||||
loss = F.mse_loss(pred_x0 * mask, target_latent * mask, reduction="sum") / mask.sum()
|
||||
else:
|
||||
loss = F.mse_loss(pred_x0, target_latent)
|
||||
|
||||
avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
|
||||
train_loss += avg_loss.item() / args.gradient_accumulation_steps
|
||||
|
||||
@@ -16,7 +16,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_ode
|
||||
--dataloader_num_workers=8 \
|
||||
--num_train_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--learning_rate=2e-05 \
|
||||
--learning_rate=2.0e-06 \
|
||||
--lr_scheduler="constant_with_warmup" \
|
||||
--lr_warmup_steps=100 \
|
||||
--seed=42 \
|
||||
|
||||
Reference in New Issue
Block a user