nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[LoRA] Support v1 LoRA training (#576)

Browse Source
This commit is contained in:
Wenxuan Tan
2025-07-17 15:52:09 -05:00
committed by GitHub
parent 576c214f28
commit bb8c769c8e
38 changed files with 757 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/pipeline.yml
+11
View File
@@ -91,6 +91,17 @@ steps:
- TEST_TYPE=training
agents:
queue: "default"
- path:
- "fastvideo/v1/**"
- "pyproject.toml"
- "docker/Dockerfile.python3.12"
config:
command: "timeout 15m .buildkite/scripts/pr_test.sh"
label: "LoRA Training Tests"
env:
- TEST_TYPE=training_lora
agents:
queue: "default"
- path:
- "fastvideo/v1/**"
- "csrc/attn/vsa/**"
.buildkite/scripts/pr_test.sh
+4
View File
@@ -81,6 +81,10 @@ case "$TEST_TYPE" in
log "Running training tests..."
MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
;;
"training_lora")
log "Running LoRA training tests..."
MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_lora_tests"
;;
"training_vsa")
log "Running training VSA tests..."
MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
docs/source/training/data_preprocess.md
+1 -1
View File
@@ -7,7 +7,7 @@ To save GPU memory, we precompute text embeddings and VAE latents to eliminate t
We provide a sample dataset to help you get started. Download the source media using the following command:
```bash
python scripts/huggingface/download_hf.py --repo_id=FastVideo/mini_i2v_dataset --local_dir=FastVideo/mini_i2v_dataset --repo_type=dataset
python scripts/huggingface/download_hf.py --repo_id=FastVideo/mini_i2v_dataset --local_dir=data/mini_i2v_dataset --repo_type=dataset
```
The folder `crush-smol_raw/` contains raw videos and captions for testing preprocessing, while `crush-smol_preprocessed/` contains latents prepared for testing training.
examples/inference/lora/wan_lora_inference.py
+2 -2
View File
@@ -1,7 +1,7 @@
from fastvideo import VideoGenerator
from fastvideo.v1.configs.sample import SamplingParam
OUTPUT_PATH = "./lora"
OUTPUT_PATH = "./lora_out"
def main():
# Initialize VideoGenerator with the Wan model
generator = VideoGenerator.from_pretrained(
@@ -32,7 +32,7 @@ def main():
)
del generator
# Until FSDP resharding bug is fixed, multi-lora requires reloading the model
# Until FSDP resharding bug is fixed, multi-lora requires reloading the model or disabling FSDP
# see https://github.com/pytorch/pytorch/issues/157209
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh
+93
View File
@@ -0,0 +1,93 @@
#!/bin/bash
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
--tracker_project_name "wan_i2v_finetune"
--output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
--max_train_steps 2000
--train_batch_size 1
--train_sp_batch_size 1
--gradient_accumulation_steps 1
--num_latent_t 8
--num_height 480
--num_width 832
--num_frames 77
--lora_rank 32
--lora_training True
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size $NUM_GPUS
--tp_size $NUM_GPUS
--hsdp_replicate_dim 1
--hsdp_shard_dim $NUM_GPUS
)
# Model arguments
model_args=(
--model_path $MODEL_PATH
--pretrained_model_name_or_path $MODEL_PATH
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 1
)
# Validation arguments
validation_args=(
--log_validation
--validation_preprocessed_path "$VALIDATION_DIR"
--validation_steps 100
--validation_sampling_steps "40"
--validation_guidance_scale "1.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 1e-5
--mixed_precision "bf16"
--checkpointing_steps 1000
--weight_decay 1e-4
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--allow_tf32
--checkpoints_total_limit 3
--training_cfg_rate 0.1
--multi_phased_distill_schedule "4000-1"
--not_apply_cfg_solver
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
fastvideo/v1/training/wan_i2v_training_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}"
examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh
+92
View File
@@ -0,0 +1,92 @@
#!/bin/bash
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
NUM_GPUS=2
# export CUDA_VISIBLE_DEVICES=4,5
# Training arguments
training_args=(
--tracker_project_name "wan_t2v_finetune"
--output_dir "outputs/wan_t2v_finetune"
--max_train_steps 5
--train_batch_size 1
--train_sp_batch_size 1
--gradient_accumulation_steps 8
--num_latent_t 8
--num_height 480
--num_width 832
--num_frames 77
--lora_rank 32
--lora_training True
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size $NUM_GPUS
--tp_size $NUM_GPUS
--hsdp_replicate_dim 1
--hsdp_shard_dim $NUM_GPUS
)
# Model arguments
model_args=(
--model_path $MODEL_PATH
--pretrained_model_name_or_path $MODEL_PATH
)
# Dataset arguments
dataset_args=(
--data_path $DATA_DIR
--dataloader_num_workers 1
)
# Validation arguments
validation_args=(
--log_validation
--validation_dataset_file $VALIDATION_DATASET_FILE
--validation_steps 50
--validation_sampling_steps "50"
--validation_guidance_scale "1.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 5e-5
--mixed_precision "bf16"
--checkpointing_steps 6000
--weight_decay 1e-4
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--allow_tf32
--checkpoints_total_limit 3
--training_cfg_rate 0.1
--multi_phased_distill_schedule "4000-1"
--not_apply_cfg_solver
--dit_precision "fp32"
--num_euler_timesteps 50
--ema_start_step 0
)
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
fastvideo/v1/training/wan_training_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${validation_args[@]}" \
"${miscellaneous_args[@]}"
fastvideo/v1/configs/pipelines/base.py
-6
View File
@@ -73,12 +73,6 @@ class PipelineConfig:
...] = field(default_factory=lambda:
(postprocess_text, ))
# LoRA parameters
lora_path: str | None = None
lora_nickname: str | None = "default" # for swapping adapters in the pipeline
lora_target_names: list[
str] | None = None # can restrict list of layers to adapt, e.g. ["q_proj"]
# StepVideo specific parameters
pos_magic: str | None = None
neg_magic: str | None = None
fastvideo/v1/dataset/__init__.py
+2 -1
View File
@@ -36,7 +36,8 @@ def getdataset(args) -> VideoCaptionMergedDataset:
args=args,
transform=transform,
temporal_sample=temporal_sample,
transform_topcrop=transform_topcrop)
transform_topcrop=transform_topcrop,
seed=args.seed)
__all__ = [
fastvideo/v1/dataset/latent_datasets.py
+5 -1
View File
@@ -14,10 +14,14 @@ class LatentDataset(Dataset):
json_path,
num_latent_t,
cfg_rate,
seed: int = 42,
) -> None:
# data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
self.json_path = json_path
self.cfg_rate = cfg_rate
self.seed = seed
# Create a seeded random generator for deterministic CFG
self.rng = random.Random(seed)
self.datase_dir_path = os.path.dirname(json_path)
self.video_dir = os.path.join(self.datase_dir_path, "video")
self.latent_dir = os.path.join(self.datase_dir_path, "latent")
@@ -50,7 +54,7 @@ class LatentDataset(Dataset):
weights_only=True,
)
latent = latent.squeeze(0)[:, -self.num_latent_t:]
if random.random() < self.cfg_rate:
if self.rng.random() < self.cfg_rate:
prompt_embed = self.uncond_prompt_embed
prompt_attention_mask = self.uncond_prompt_mask
else:
fastvideo/v1/dataset/parquet_dataset_iterable_style.py
+2 -2
View File
@@ -236,8 +236,8 @@ def shard_parquet_files_across_sp_groups_and_workers(
shard_parquet_lengths[target_shard][file] = length
#randomize each shard
for shard in shard_parquet_files:
random.seed(seed)
random.shuffle(shard)
rng = random.Random(seed)
rng.shuffle(shard)
# Save the sharding plan
os.makedirs(sharding_info_dir, exist_ok=True)
fastvideo/v1/dataset/parquet_dataset_map_style.py
+8 -2
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
import os
import pickle
import random
from typing import Any
import pyarrow as pa
@@ -235,6 +236,9 @@ class LatentsParquetMapStyleDataset(Dataset):
self.path = path
self.cfg_rate = cfg_rate
self.parquet_schema = parquet_schema
self.seed = seed
# Create a seeded random generator for deterministic CFG
self.rng = random.Random(seed)
logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
path)
self.parquet_files, self.lengths = get_parquet_files_and_length(path)
@@ -271,7 +275,8 @@ class LatentsParquetMapStyleDataset(Dataset):
batch = collate_rows_from_parquet_schema([row_dict],
self.parquet_schema,
self.text_padding_length,
cfg_rate=0.0)
cfg_rate=0.0,
rng=self.rng)
negative_prompt = batch['info_list'][0]['prompt']
negative_prompt_embedding = batch['text_embedding']
negative_prompt_attention_mask = batch['text_attention_mask']
@@ -296,7 +301,8 @@ class LatentsParquetMapStyleDataset(Dataset):
batch = collate_rows_from_parquet_schema(rows,
self.parquet_schema,
self.text_padding_length,
cfg_rate=self.cfg_rate)
cfg_rate=self.cfg_rate,
rng=self.rng)
return batch
def __len__(self):
fastvideo/v1/dataset/preprocessing_datasets.py
+21 -8
View File
@@ -211,12 +211,15 @@ class FrameSamplingStage(DatasetFilterStage):
train_fps: int,
speed_factor: int = 1,
video_length_tolerance_range: float = 5.0,
drop_short_ratio: float = 0.0):
drop_short_ratio: float = 0.0,
seed: int = 42):
self.num_frames = num_frames
self.train_fps = train_fps
self.speed_factor = speed_factor
self.video_length_tolerance_range = video_length_tolerance_range
self.drop_short_ratio = drop_short_ratio
# Create a seeded random generator for deterministic sampling
self.rng = random.Random(seed)
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
"""
@@ -249,7 +252,7 @@ class FrameSamplingStage(DatasetFilterStage):
# Filter short videos
return not (len(frame_indices) < self.num_frames
and random.random() < self.drop_short_ratio)
and self.rng.random() < self.drop_short_ratio)
def process(self,
batch: PreprocessBatch,
@@ -372,10 +375,16 @@ class ImageTransformStage(DatasetStage):
class TextEncodingStage(DatasetStage):
"""Stage for text tokenization and encoding."""
def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0):
def __init__(self,
tokenizer,
text_max_length: int,
cfg_rate: float = 0.0,
seed: int = 42):
self.tokenizer = tokenizer
self.text_max_length = text_max_length
self.cfg_rate = cfg_rate
# Create a seeded random generator for deterministic CFG
self.rng = random.Random(seed)
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
"""
@@ -390,9 +399,9 @@ class TextEncodingStage(DatasetStage):
text = batch.cap
if not isinstance(text, list):
text = [text]
text = [random.choice(text)]
text = [self.rng.choice(text)]
text = text[0] if random.random() > self.cfg_rate else ""
text = text[0] if self.rng.random() > self.cfg_rate else ""
text_tokens_and_mask = self.tokenizer(
text,
max_length=self.text_max_length,
@@ -451,11 +460,13 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
transform,
temporal_sample,
transform_topcrop,
start_idx: int = 0):
start_idx: int = 0,
seed: int = 42):
self.data_merge_path = data_merge_path
self.start_idx = start_idx
self.args = args
self.temporal_sample = temporal_sample
self.seed = seed
# Initialize tokenizer
tokenizer_path = os.path.join(args.model_path, "tokenizer")
@@ -479,14 +490,16 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
train_fps=args.train_fps,
speed_factor=args.speed_factor,
video_length_tolerance_range=args.video_length_tolerance_range,
drop_short_ratio=args.drop_short_ratio)
drop_short_ratio=args.drop_short_ratio,
seed=self.seed)
self.video_transform_stage = VideoTransformStage(transform)
self.image_transform_stage = ImageTransformStage(
transform, transform_topcrop)
self.text_encoding_stage = TextEncodingStage(
tokenizer=tokenizer,
text_max_length=args.text_max_length,
cfg_rate=args.training_cfg_rate)
cfg_rate=args.training_cfg_rate,
seed=self.seed)
def _load_raw_data(self) -> list[dict]:
"""Load raw data from JSON files."""
fastvideo/v1/dataset/utils.py
+13 -8
View File
@@ -21,7 +21,10 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
return t[:padding_length], torch.ones(padding_length)
def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> dict[str, Any]:
def get_torch_tensors_from_row_dict(row_dict,
keys,
cfg_rate,
rng=None) -> dict[str, Any]:
"""
Get the latents and prompts from a row dictionary.
"""
@@ -43,7 +46,8 @@ def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> dict[str, Any]:
bytes = row_dict[f"{key}_bytes"]
# TODO (peiyuan): read precision
if key == 'text_embedding' and random.random() < cfg_rate:
if key == 'text_embedding' and (rng.random()
if rng else random.random()) < cfg_rate:
data = np.zeros((512, 4096), dtype=np.float32)
else:
data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
@@ -60,8 +64,8 @@ def collate_latents_embs_masks(
batch_to_process,
text_padding_length,
keys,
cfg_rate=0.0
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
cfg_rate=0.0,
rng=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
# Initialize tensors to hold padded embeddings and masks
all_latents = []
all_embs = []
@@ -70,7 +74,7 @@ def collate_latents_embs_masks(
# Process each row individually
for i, row in enumerate(batch_to_process):
# Get tensors from row
data = get_torch_tensors_from_row_dict(row, keys, cfg_rate)
data = get_torch_tensors_from_row_dict(row, keys, cfg_rate, rng)
latents, emb = data["vae_latent"], data["text_embedding"]
padded_emb, mask = pad(emb, text_padding_length)
@@ -95,7 +99,8 @@ def collate_latents_embs_masks(
def collate_rows_from_parquet_schema(rows,
parquet_schema,
text_padding_length,
cfg_rate=0.0) -> dict[str, Any]:
cfg_rate=0.0,
rng=None) -> dict[str, Any]:
"""
Collate rows from parquet files based on the provided schema.
Dynamically processes tensor fields based on schema and returns batched data.
@@ -145,8 +150,8 @@ def collate_rows_from_parquet_schema(rows,
tensor = torch.zeros(0, dtype=torch.bfloat16)
else:
# Convert bytes to tensor using float32 as default
if tensor_name == 'text_embedding' and random.random(
) < cfg_rate:
if tensor_name == 'text_embedding' and (rng.random(
) if rng else random.random()) < cfg_rate:
data = np.zeros((512, 4096), dtype=np.float32)
else:
data = np.frombuffer(
fastvideo/v1/distributed/parallel_state.py
+10 -5
View File
@@ -177,9 +177,12 @@ class GroupCoordinator:
self.rank_in_group = ranks.index(self.rank)
self.device_group = device_group
self.cpu_group = cpu_group
assert self.cpu_group is not None
assert self.device_group is not None
try:
assert self.cpu_group is not None
assert self.device_group is not None
except Exception as e:
print(f"rank: {self.rank} group not found")
raise e
from fastvideo.v1.platforms import current_platform
@@ -858,7 +861,7 @@ def initialize_model_parallel(
world_size: int = get_world_size()
backend = backend or torch.distributed.get_backend(
get_world_group().device_group)
assert world_size >= tensor_model_parallel_size, f"world_size({world_size}) must be greater than or equal to tensor_model_parallel_size({tensor_model_parallel_size})"
num_tensor_model_parallel_groups: int = (world_size //
tensor_model_parallel_size)
global _TP
@@ -962,8 +965,10 @@ def maybe_init_distributed_environment_and_model_parallel(
local_rank = int(os.environ.get("LOCAL_RANK", 0))
world_size = int(os.environ.get("WORLD_SIZE", 1))
rank = int(os.environ.get("RANK", 0))
device = get_local_torch_device()
logger.info(
"Initializing distributed environment with world_size=%d, device=%s",
world_size, device)
init_distributed_environment(
world_size=world_size,
fastvideo/v1/fastvideo_args.py
+19
View File
@@ -57,6 +57,15 @@ class FastVideoArgs:
pipeline_config: PipelineConfig = field(default_factory=PipelineConfig)
# LoRA parameters
# (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
lora_path: str | None = None
lora_nickname: str = "default" # for swapping adapters in the pipeline
# can restrict layers to adapt, e.g. ["q_proj"]
# For inference, will be consistent with loaded adapter by default
# For training, will adapt only q, k, v, o by default.
lora_target_modules: list[str] | None = None
output_type: str = "pil"
use_cpu_offload: bool = True # For DiT
@@ -475,6 +484,11 @@ class TrainingArgs(FastVideoArgs):
VSA_decay_rate: float = 0.01 # decay rate -> 0.02
VSA_decay_interval_steps: int = 1 # decay interval steps -> 50
# LoRA training parameters
lora_rank: int | None = None
lora_alpha: int | None = None
lora_training: bool = False
@classmethod
def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
provided_args = clean_cli_args(args)
@@ -755,5 +769,10 @@ class TrainingArgs(FastVideoArgs):
type=int,
default=TrainingArgs.VSA_decay_interval_steps,
help="VSA decay interval steps")
parser.add_argument("--lora-training",
action=StoreBoolean,
help="Whether to use LoRA training")
parser.add_argument("--lora-rank", type=int, help="LoRA rank")
parser.add_argument("--lora-alpha", type=int, help="LoRA alpha")
return parser
fastvideo/v1/forward_context.py
+1 -3
View File
@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Optional
import torch
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.logger import init_logger
if TYPE_CHECKING:
@@ -55,8 +54,7 @@ def get_forward_context() -> "ForwardContext":
@contextmanager
def set_forward_context(current_timestep,
attn_metadata,
forward_batch: Optional["ForwardBatch"] = None,
fastvideo_args: FastVideoArgs | None = None):
forward_batch: Optional["ForwardBatch"] = None):
"""A context manager that stores the current forward context,
can be attention metadata, etc.
Here we can inject common logic for every model forward pass.
fastvideo/v1/layers/lora/linear.py
+94 -18
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
import math
import torch
from torch import nn
from torch.distributed._composable.fsdp import (CPUOffloadPolicy, OffloadPolicy,
@@ -18,26 +20,76 @@ from fastvideo.v1.layers.linear import (ColumnParallelLinear, LinearBase,
from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
from fastvideo.v1.utils import get_mixed_precision_state
torch._dynamo.config.recompile_limit = 16
class BaseLayerWithLoRA(nn.Module):
def __init__(
self,
base_layer: nn.Module,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False,
):
super().__init__()
self.base_layer: nn.Module = base_layer
self.lora_A: torch.Tensor = None
self.lora_B: torch.Tensor = None
self.merged: bool = False
self.cpu_weight = base_layer.weight.to("cpu")
# indicates adapter weights don't contain this layer
# (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
self.disable_lora: bool = False
self.lora_rank = lora_rank
self.lora_alpha = lora_alpha
self.training_mode = training_mode
self.lora_path: str | None = None
if training_mode:
assert self.lora_rank is not None, "LoRA rank must be set for training mode"
if self.lora_rank is None or self.lora_alpha is None:
self.lora_alpha = lora_rank
self.base_layer.requires_grad_(False)
in_dim = self.base_layer.weight.shape[1]
out_dim = self.base_layer.weight.shape[0]
self.lora_A = nn.Parameter(
torch.zeros(self.lora_rank,
in_dim,
device=self.base_layer.weight.device,
dtype=self.base_layer.weight.dtype))
self.lora_B = nn.Parameter(
torch.zeros(out_dim,
self.lora_rank,
device=self.base_layer.weight.device,
dtype=self.base_layer.weight.dtype))
torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
torch.nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5))
else:
self.lora_A = None
self.lora_B = None
@torch.compile()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.base_layer.forward(x)
lora_A = self.lora_A
lora_B = self.lora_B
if isinstance(self.lora_B, DTensor):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()
if self.training_mode or not self.merged:
delta = x @ (
self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
@ self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)))
if self.lora_alpha != self.lora_rank:
delta = delta * (
self.lora_alpha / self.lora_rank # type: ignore
) # type: ignore
out, output_bias = self.base_layer(x)
return out + delta, output_bias
else:
# already merged
out, output_bias = self.base_layer(x)
return out.to(x), output_bias
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
return A
@@ -58,6 +110,7 @@ class BaseLayerWithLoRA(nn.Module):
self.lora_path = lora_path
@torch.no_grad()
# @torch.compile()
def merge_lora_weights(self) -> None:
if self.disable_lora:
return
@@ -71,16 +124,17 @@ class BaseLayerWithLoRA(nn.Module):
current_device = self.base_layer.weight.data.device
data = self.base_layer.weight.data.to(
get_local_torch_device()).full_tensor()
data += (self.slice_lora_b_weights(self.lora_B).to(data)
@ self.slice_lora_a_weights(self.lora_A).to(data))
data += (self.slice_lora_b_weights(self.lora_B).to(
data, non_blocking=True) @ self.slice_lora_a_weights(
self.lora_A).to(data, non_blocking=True))
# Must re-register updated weights for FSDP to recognize them
self.base_layer.weight = nn.Parameter(data.to(current_device))
if isinstance(getattr(self.base_layer, "bias", None), DTensor):
self.base_layer.bias = nn.Parameter(
self.base_layer.bias.to(
get_local_torch_device(),
non_blocking=True).full_tensor().to(current_device))
get_local_torch_device()).full_tensor().to(
current_device))
offload_policy = CPUOffloadPolicy() if "cpu" in str(
current_device) else OffloadPolicy()
@@ -96,14 +150,17 @@ class BaseLayerWithLoRA(nn.Module):
else:
current_device = self.base_layer.weight.data.device
data = self.base_layer.weight.data.to(get_local_torch_device())
data += \
(self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data)))
data += (self.slice_lora_b_weights(
self.lora_B.to(data, non_blocking=True))
@ self.slice_lora_a_weights(
self.lora_A.to(data, non_blocking=True)))
self.base_layer.weight.data = data.to(current_device,
non_blocking=True)
self.merged = True
@torch.no_grad()
# @torch.compile(dynamic=True)
def unmerge_lora_weights(self) -> None:
if self.disable_lora:
return
@@ -113,13 +170,14 @@ class BaseLayerWithLoRA(nn.Module):
"LoRA weights not merged. Please merge them first before unmerging."
)
# To avoid precision loss we do not subtract the LoRA weights here
# avoid precision loss
if isinstance(self.base_layer.weight, DTensor):
device = self.base_layer.weight.data.device
self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device))
self.base_layer.weight = nn.Parameter(
self.cpu_weight.to(device, non_blocking=True))
else:
self.base_layer.weight.data = self.cpu_weight.data.to(
self.base_layer.weight)
self.base_layer.weight, non_blocking=True)
self.merged = False
@@ -149,8 +207,11 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
def __init__(
self,
base_layer: ColumnParallelLinear,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False,
) -> None:
super().__init__(base_layer)
super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
def forward(self, input_: torch.Tensor) -> torch.Tensor:
# duplicate the logic in ColumnParallelLinear
@@ -181,8 +242,11 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
def __init__(
self,
base_layer: MergedColumnParallelLinear,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False,
) -> None:
super().__init__(base_layer)
super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
return A.to(self.base_layer.weight)
@@ -201,8 +265,11 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
def __init__(
self,
base_layer: QKVParallelLinear,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False,
) -> None:
super().__init__(base_layer)
super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
return A
@@ -232,8 +299,11 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
def __init__(
self,
base_layer: RowParallelLinear,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False,
) -> None:
super().__init__(base_layer)
super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
def forward(self, input_: torch.Tensor):
# duplicate the logic in RowParallelLinear
@@ -276,7 +346,10 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
return B
def get_lora_layer(layer: nn.Module) -> BaseLayerWithLoRA | None:
def get_lora_layer(layer: nn.Module,
lora_rank: int | None = None,
lora_alpha: int | None = None,
training_mode: bool = False) -> BaseLayerWithLoRA | None:
supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = {
# the order matters
# VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
@@ -288,7 +361,10 @@ def get_lora_layer(layer: nn.Module) -> BaseLayerWithLoRA | None:
}
for src_layer_type, lora_layer_type in supported_layer_types.items():
if isinstance(layer, src_layer_type): # pylint: disable=unidiomatic-typecheck
ret = lora_layer_type(layer)
ret = lora_layer_type(layer,
lora_rank=lora_rank,
lora_alpha=lora_alpha,
training_mode=training_mode)
return ret
return None
fastvideo/v1/models/loader/fsdp_load.py
+1 -1
View File
@@ -134,9 +134,9 @@ def maybe_load_fsdp_model(
if p.is_meta:
raise RuntimeError(
f"Unexpected param or buffer {n} on meta device.")
# Avoid unintended computation graph accumulation during inference
if isinstance(p, torch.nn.Parameter):
p.requires_grad = False
return model
fastvideo/v1/pipelines/composed_pipeline_base.py
+27 -10
View File
@@ -38,8 +38,9 @@ class ComposedPipelineBase(ABC):
is_video_pipeline: bool = False # To be overridden by video pipelines
_required_config_modules: list[str] = []
training_args: TrainingArgs | None = None
fastvideo_args: FastVideoArgs | None = None
fastvideo_args: FastVideoArgs | TrainingArgs | None = None
modules: dict[str, torch.nn.Module] = {}
post_init_called: bool = False
# TODO(will): args should support both inference args and training args
def __init__(self,
@@ -70,20 +71,33 @@ class ComposedPipelineBase(ABC):
# Load modules directly in initialization
logger.info("Loading pipeline modules...")
self.modules = self.load_modules(fastvideo_args, loaded_modules)
# Only train DiT
for name, module in self.modules.items():
if not isinstance(module, torch.nn.Module):
continue
if name == "transformer":
module.requires_grad_(True)
else:
module.requires_grad_(False)
if fastvideo_args.training_mode:
assert isinstance(fastvideo_args, TrainingArgs)
self.training_args = fastvideo_args
def post_init(self) -> None:
assert self.fastvideo_args is not None, "fastvideo_args must be set"
if self.post_init_called:
return
self.post_init_called = True
if getattr(self.fastvideo_args, "training_mode", False):
assert isinstance(self.fastvideo_args, TrainingArgs)
self.training_args = self.fastvideo_args
assert self.training_args is not None
if self.training_args.log_validation:
self.initialize_validation_pipeline(self.training_args)
self.initialize_training_pipeline(self.training_args)
self.initialize_pipeline(fastvideo_args)
self.initialize_pipeline(self.fastvideo_args)
if not fastvideo_args.training_mode:
if not getattr(self.fastvideo_args, "training_mode", False):
logger.info("Creating pipeline stages...")
self.create_pipeline_stages(fastvideo_args)
self.create_pipeline_stages(self.fastvideo_args)
def initialize_training_pipeline(self, training_args: TrainingArgs):
raise NotImplementedError(
@@ -123,8 +137,6 @@ class ComposedPipelineBase(ABC):
setattr(fastvideo_args, key, value)
fastvideo_args.use_cpu_offload = False
# make sure we are in training mode
fastvideo_args.inference_mode = False
# we hijack the precision to be the master weight type so that the
# model is loaded with the correct precision. Subsequently we will
# use FSDP2's MixedPrecisionPolicy to set the precision for the
@@ -135,10 +147,12 @@ class ComposedPipelineBase(ABC):
logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
return cls(model_path,
pipe = cls(model_path,
fastvideo_args,
required_config_modules=required_config_modules,
loaded_modules=loaded_modules)
pipe.post_init()
return pipe
def get_module(self, module_name: str, default_value: Any = None) -> Any:
if module_name not in self.modules:
@@ -286,6 +300,9 @@ class ComposedPipelineBase(ABC):
Returns:
ForwardBatch: The batch with the generated video or image.
"""
if not self.post_init_called:
self.post_init()
# Execute each stage
logger.info("Running pipeline stages: %s",
self._stage_name_mapping.keys())
fastvideo/v1/pipelines/lora_pipeline.py
+76 -24
View File
@@ -5,10 +5,13 @@ from typing import Any
import torch
import torch.distributed as dist
import torch.nn as nn
from safetensors.torch import load_file
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor
from fastvideo.v1.distributed.parallel_state import get_local_torch_device
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.distributed import get_local_torch_device
from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.v1.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer,
replace_submodule)
from fastvideo.v1.logger import init_logger
@@ -27,34 +30,67 @@ class LoRAPipeline(ComposedPipelineBase):
lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
dict) # state dicts of loaded lora adapters
cur_adapter_name: str = ""
cur_adapter_path: str = ""
lora_layers: dict[str, BaseLayerWithLoRA] = {}
fastvideo_args: FastVideoArgs
fastvideo_args: FastVideoArgs | TrainingArgs
exclude_lora_layers: list[str] = []
device: torch.device | None = None
device: torch.device = get_local_torch_device()
lora_target_modules: list[str] | None = None
lora_path: str | None = None
lora_nickname: str = "default"
lora_rank: int | None = None
lora_alpha: int | None = None
lora_initialized: bool = False
def __init__(self, *args, **kwargs):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.device = get_local_torch_device()
self.exclude_lora_layers = self.modules[
"transformer"].config.arch_config.exclude_lora_layers
self.device = get_local_torch_device()
self.convert_to_lora_layers()
if self.fastvideo_args.pipeline_config.lora_path is not None:
self.lora_target_modules = self.fastvideo_args.lora_target_modules
self.lora_path = self.fastvideo_args.lora_path
self.lora_nickname = self.fastvideo_args.lora_nickname
self.training_mode = self.fastvideo_args.training_mode
if self.training_mode and getattr(self.fastvideo_args, "lora_training",
False):
assert isinstance(self.fastvideo_args, TrainingArgs)
if self.fastvideo_args.lora_alpha is None:
self.fastvideo_args.lora_alpha = self.fastvideo_args.lora_rank
self.lora_rank = self.fastvideo_args.lora_rank # type: ignore
self.lora_alpha = self.fastvideo_args.lora_alpha # type: ignore
logger.info("Using LoRA training with rank %d and alpha %d",
self.lora_rank, self.lora_alpha)
if self.lora_target_modules is None:
# Possible names for q, k, v, o
self.lora_target_modules = [
"q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k",
"to_v", "to_out", "to_qkv"
]
self.convert_to_lora_layers()
# Inference
elif not self.training_mode and self.lora_path is not None:
self.convert_to_lora_layers()
self.set_lora_adapter(
self.fastvideo_args.pipeline_config.
lora_nickname, # type: ignore
self.fastvideo_args.pipeline_config.lora_path)
self.lora_nickname, # type: ignore
self.lora_path) # type: ignore
def is_target_layer(self, module_name: str) -> bool:
if self.fastvideo_args.pipeline_config.lora_target_names is None:
if self.lora_target_modules is None:
return True
return any(target_name in module_name for target_name in
self.fastvideo_args.pipeline_config.lora_target_names)
return any(target_name in module_name
for target_name in self.lora_target_modules)
def convert_to_lora_layers(self) -> None:
"""
Converts the transformer to a LoRA transformer.
Unified method to convert the transformer to a LoRA transformer.
"""
if self.lora_initialized:
return
self.lora_initialized = True
is_lora_training = self.training_mode and getattr(
self.fastvideo_args, "lora_training", False)
if is_lora_training:
self.modules["transformer"].requires_grad_(False)
for name, layer in self.modules["transformer"].named_modules():
if not self.is_target_layer(name):
continue
@@ -67,16 +103,31 @@ class LoRAPipeline(ComposedPipelineBase):
if excluded:
continue
layer = get_lora_layer(layer)
layer = get_lora_layer(layer,
lora_rank=self.lora_rank,
lora_alpha=self.lora_alpha,
training_mode=self.training_mode)
if layer is not None:
self.lora_layers[name] = layer
replace_submodule(self.modules["transformer"], name, layer)
if is_lora_training:
device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1),
mesh_dim_names=["fake", "replicate"])
for name, layer in self.lora_layers.items():
# Enable grads for lora weights only
layer.lora_A.requires_grad_(True)
layer.lora_B.requires_grad_(True)
# Must convert to DTensor for compatibility with other FSDP modules in grad calculation
layer.lora_A = nn.Parameter(
DTensor.from_local(layer.lora_A, device_mesh=device_mesh))
layer.lora_B = nn.Parameter(
DTensor.from_local(layer.lora_B, device_mesh=device_mesh))
def set_lora_adapter(self,
lora_nickname: str,
lora_path: str | None = None): # type: ignore
"""
Loads a LoRA adapter into the pipeline and applies it to the transformer.
Load a LoRA adapter into the pipeline and merge it into the transformer.
Args:
lora_nickname: The "nick name" of the adapter when referenced in the pipeline.
lora_path: The path to the adapter, either a local path or a Hugging Face repo id.
@@ -86,13 +137,14 @@ class LoRAPipeline(ComposedPipelineBase):
raise ValueError(
f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
)
if not self.lora_initialized:
self.convert_to_lora_layers()
adapter_updated = False
rank = dist.get_rank()
if lora_path is not None:
if lora_path is not None and lora_path != self.cur_adapter_path:
lora_local_path = maybe_download_lora(lora_path)
lora_state_dict = load_file(lora_local_path,
device=str(self.device))
lora_state_dict = load_file(lora_local_path)
# Map the hf layer names to our custom layer names
param_names_mapping_fn = get_param_names_mapping(
self.modules["transformer"].param_names_mapping)
@@ -130,10 +182,12 @@ class LoRAPipeline(ComposedPipelineBase):
self.lora_adapters[lora_nickname][target_name] = weight.to(
self.device)
adapter_updated = True
self.cur_adapter_path = lora_path
logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)
if not adapter_updated and lora_nickname == self.cur_adapter_name:
if not adapter_updated and self.cur_adapter_name == lora_nickname:
return
self.cur_adapter_name = lora_nickname
# Merge the new adapter
adapted_count = 0
@@ -156,5 +210,3 @@ class LoRAPipeline(ComposedPipelineBase):
layer.disable_lora = True
logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
lora_path, adapted_count)
self.cur_adapter_name = lora_nickname
fastvideo/v1/pipelines/stages/conditioning.py
+1
View File
@@ -23,6 +23,7 @@ class ConditioningStage(PipelineStage):
to the diffusion process.
"""
@torch.no_grad()
def forward(
self,
batch: ForwardBatch,
fastvideo/v1/pipelines/stages/decoding.py
+1
View File
@@ -51,6 +51,7 @@ class DecodingStage(PipelineStage):
result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
return result
@torch.no_grad()
def forward(
self,
batch: ForwardBatch,
fastvideo/v1/pipelines/stages/encoding.py
+1
View File
@@ -33,6 +33,7 @@ class EncodingStage(PipelineStage):
def __init__(self, vae: ParallelTiledVAE) -> None:
self.vae: ParallelTiledVAE = vae
@torch.no_grad()
def forward(
self,
batch: ForwardBatch,
fastvideo/v1/pipelines/stages/image_encoding.py
+3
View File
@@ -5,6 +5,8 @@ Image encoding stages for I2V diffusion pipelines.
This module contains implementations of image encoding stages for diffusion pipelines.
"""
import torch
from fastvideo.v1.distributed import get_local_torch_device
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.forward_context import set_forward_context
@@ -37,6 +39,7 @@ class ImageEncodingStage(PipelineStage):
self.image_processor = image_processor
self.image_encoder = image_encoder
@torch.no_grad()
def forward(
self,
batch: ForwardBatch,
fastvideo/v1/pipelines/stages/stepvideo_encoding.py
+1
View File
@@ -31,6 +31,7 @@ class StepvideoPromptEncodingStage(PipelineStage):
self.stepllm = stepllm
self.clip = clip
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args) -> ForwardBatch:
prompts = [batch.prompt + fastvideo_args.pipeline_config.pos_magic]
fastvideo/v1/pipelines/stages/text_encoding.py
+3
View File
@@ -5,6 +5,8 @@ Prompt encoding stages for diffusion pipelines.
This module contains implementations of prompt encoding stages for diffusion pipelines.
"""
import torch
from fastvideo.v1.distributed import get_local_torch_device
from fastvideo.v1.fastvideo_args import FastVideoArgs
from fastvideo.v1.forward_context import set_forward_context
@@ -37,6 +39,7 @@ class TextEncodingStage(PipelineStage):
self.tokenizers = tokenizers
self.text_encoders = text_encoders
@torch.no_grad()
def forward(
self,
batch: ForwardBatch,
fastvideo/v1/tests/modal/pr_test.py
+4
View File
@@ -82,6 +82,10 @@ def run_ssim_tests():
def run_training_tests():
run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/Vanilla -srP")
@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
def run_training_lora_tests():
run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/lora/test_lora_training.py -srP")
@app.function(gpu="H100:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
def run_training_tests_VSA():
run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/v1/tests/training/VSA -srP")
fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json
+1 -1
View File
@@ -1 +1 @@
{"step_time":5.501357046999999,"grad_norm":0.384765625,"train_loss":0.07890288904309273,"avg_step_time":5.831571423200001}
{"train_loss":0.10545400530099869,"_step":5,"_wandb":{"runtime":35},"step_time":2.2575189135968685,"grad_norm":0.53125,"_runtime":35.701958502,"avg_step_time":2.4754387199878694,"_timestamp":1.7525528420185745e+09,"learning_rate":1e-06,"vsa_sparsity":0,"validation_videos_8_steps":{"videos":[{"caption":"A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","_type":"video-file","sha256":"ee83e83df073a648f89dcd288cccaed9af765e11fc35c77a6e0cb2ebaa1be5b0","size":475248,"path":"media/videos/validation_videos_8_steps_0_ee83e83df073a648f89d.mp4"},{"size":341490,"path":"media/videos/validation_videos_8_steps_0_d0b2758549d5c82845ca.mp4","caption":"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","_type":"video-file","sha256":"d0b2758549d5c82845ca3c1ac0db6812877631495e77ac041acdc8ea31f0a5ee"},{"caption":"A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.","_type":"video-file","sha256":"08638381f1607d6ab10684772be38a58b18628fd402a208b25f6af96e765d454","size":436814,"path":"media/videos/validation_videos_8_steps_0_08638381f1607d6ab106.mp4"}],"captions":["A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."],"_type":"videos","count":3}}
fastvideo/v1/tests/training/Vanilla/test_training_loss.py
+27 -15
View File
@@ -1,28 +1,26 @@
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29512"
import sys
import subprocess
from pathlib import Path
import torch.distributed.elastic.multiprocessing.errors as errors
from torch.distributed.elastic.multiprocessing.errors import record
from torch.utils.data import DataLoader
import torch
import pytest
import wandb
import json
from huggingface_hub import snapshot_download
from fastvideo.v1.utils import logger
# Import the training pipeline
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
from fastvideo.v1.training.wan_training_pipeline import main
from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.v1.utils import FlexibleArgumentParser
from fastvideo.v1.training.wan_training_pipeline import WanTrainingPipeline
wandb_name = "test_training_loss"
a40_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/a40_reference_wandb_summary.json"
l40s_reference_wandb_summary_file = "fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json"
NUM_NODES = "1"
NUM_GPUS_PER_NODE = "4"
NUM_GPUS_PER_NODE = "2"
def run_worker():
@@ -41,11 +39,11 @@ def run_worker():
"--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
"--train_batch_size", "2",
"--num_latent_t", "4",
"--num_gpus", "4",
"--sp_size", "4",
"--tp_size", "1",
"--num_gpus", "2",
"--sp_size", "2",
"--tp_size", "2",
"--hsdp_replicate_dim", "1",
"--hsdp_shard_dim", "4",
"--hsdp_shard_dim", "2",
"--train_sp_batch_size", "1",
"--dataloader_num_workers", "1",
"--gradient_accumulation_steps", "2",
@@ -75,9 +73,12 @@ def run_worker():
"--dit_precision", "fp32",
"--max_grad_norm", "1.0"
])
# Call the main training function
main(args)
pipeline = WanTrainingPipeline.from_pretrained(
args.pretrained_model_name_or_path, args=args)
args = pipeline.training_args
pipeline.train()
logger.info("Training pipeline done")
def test_distributed_training():
"""Test the distributed training setup"""
@@ -93,7 +94,7 @@ def test_distributed_training():
repo_type="dataset",
local_dir_use_symlinks=False
)
# Get the current file path
current_file = Path(__file__).resolve()
@@ -102,10 +103,21 @@ def test_distributed_training():
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE,
"--master_port", os.environ["MASTER_PORT"],
str(current_file)
]
process = subprocess.run(cmd, capture_output=True, text=True)
process = subprocess.run(cmd, check=True)
# Print stdout and stderr for debugging
if process.stdout:
print("STDOUT:", process.stdout)
if process.stderr:
print("STDERR:", process.stderr)
# Check if the process failed
if process.returncode != 0:
print(f"Process failed with return code: {process.returncode}")
raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr)
summary_file = 'wandb/latest-run/files/wandb-summary.json'
fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json
+44
View File
@@ -0,0 +1,44 @@
{
"_runtime": 192.227662299,
"_step": 5,
"_timestamp": 1751927956.4331577,
"_wandb.runtime": 192,
"avg_step_time": 12.131943785200008,
"grad_norm": 0.48828125,
"learning_rate": 0.00005,
"step_time": 11.57419652599998,
"train_loss": 0.1554037211462855,
"validation_videos_50_steps": {
"_type": "videos",
"captions": [
"A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."
],
"count": 3,
"videos": [
{
"_type": "video-file",
"caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
"path": "media/videos/validation_videos_50_steps_0_54dafffa2c2dea9fcb21.mp4",
"sha256": "54dafffa2c2dea9fcb21b734a6b4abb072b99a32e57603899d12b8eaf4324434",
"size": 611595
},
{
"_type": "video-file",
"caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"path": "media/videos/validation_videos_50_steps_0_6d82abfb6ab979f9df94.mp4",
"sha256": "6d82abfb6ab979f9df94e96fcd1ac344f40f03cb3f106a6b1d9c510dc17e8f91",
"size": 460476
},
{
"_type": "video-file",
"caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
"path": "media/videos/validation_videos_50_steps_0_cedf3954f184052149c8.mp4",
"sha256": "cedf3954f184052149c83ac23776ad83f75cf911351fcb6f8e05854ba4d4c614",
"size": 738913
}
]
},
"vsa_sparsity": 0
}
fastvideo/v1/tests/training/lora/test_lora_training.py
+120
View File
@@ -0,0 +1,120 @@
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29514"
import sys
import subprocess
from pathlib import Path
import torch
import json
from huggingface_hub import snapshot_download
wandb_name = "test_lora_training"
l40s_reference_wandb_summary_file = "fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json"
NUM_NODES = "1"
NUM_GPUS_PER_NODE = "2"
def test_lora_training():
"""Test the LoRA training setup"""
os.environ["WANDB_MODE"] = "online"
data_dir = Path("data/crush-smol_processed_t2v")
if not data_dir.exists():
print(f"Downloading test dataset to {data_dir}...")
snapshot_download(
repo_id="wlsaidhi/crush-smol_processed_t2v",
local_dir=str(data_dir),
repo_type="dataset",
local_dir_use_symlinks=False
)
# Run torchrun command directly on the training pipeline like the shell script
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE,
"--master_port", os.environ["MASTER_PORT"],
"fastvideo/v1/training/wan_training_pipeline.py",
"--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--inference_mode", "False",
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset",
"--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
"--train_batch_size", "1",
"--num_latent_t", "8",
"--num_gpus", NUM_GPUS_PER_NODE,
"--sp_size", NUM_GPUS_PER_NODE,
"--tp_size", NUM_GPUS_PER_NODE,
"--hsdp_replicate_dim", "1",
"--hsdp_shard_dim", "2",
"--train_sp_batch_size", "1",
"--dataloader_num_workers", "1",
"--gradient_accumulation_steps", "8",
"--max_train_steps", "5",
"--learning_rate", "5e-5",
"--mixed_precision", "bf16",
"--checkpointing_steps", "6000",
"--validation_steps", "50",
"--validation_sampling_steps", "50",
"--log_validation",
"--checkpoints_total_limit", "3",
"--allow_tf32",
"--ema_start_step", "0",
"--training_cfg_rate", "0.1",
"--output_dir", "/workspace",
"--tracker_project_name", "wan_lora_finetune_ci",
"--wandb_run_name", wandb_name,
"--num_height", "480",
"--num_width", "832",
"--num_frames", "77",
"--validation_guidance_scale", "1.0",
"--num_euler_timesteps", "50",
"--multi_phased_distill_schedule", "4000-1",
"--weight_decay", "1e-4",
"--not_apply_cfg_solver",
"--dit_precision", "fp32",
"--max_grad_norm", "1.0",
"--lora_rank", "32",
"--lora_training", "True",
"--seed", "42"
]
process = subprocess.run(cmd, check=True)
summary_file = 'wandb/latest-run/files/wandb-summary.json'
device_name = torch.cuda.get_device_name()
assert "L40S" in device_name, "Test must be run on L40S"
reference_wandb_summary_file = l40s_reference_wandb_summary_file
reference_wandb_summary = json.load(open(reference_wandb_summary_file))
wandb_summary = json.load(open(summary_file))
# Define thresholds for LoRA training based on the provided console outputs
fields_and_thresholds = {
'avg_step_time': 1.0, # Allow 1 second variance
# 'grad_norm': 0.05, # too volatile for now. TODO: fix nondeterminism in training
'step_time': 1.0, # Allow 1 second variance
'train_loss': 0.01 # Allow 0.05 variance in training loss
}
failures = []
for field, threshold in fields_and_thresholds.items():
if field in reference_wandb_summary and field in wandb_summary:
ref_value = reference_wandb_summary[field]
current_value = wandb_summary[field]
diff = abs(ref_value - current_value)
print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}")
if diff > threshold:
failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})")
else:
print(f"WARNING: Field {field} not found in one or both summary files")
if failures:
raise AssertionError("\n".join(failures))
if __name__ == "__main__":
test_lora_training()
fastvideo/v1/training/checkpointing_utils.py
+5 -1
View File
@@ -18,7 +18,11 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
self.model = model
def state_dict(self) -> dict[str, Any]:
return get_model_state_dict(self.model) # type: ignore[no-any-return]
state_dict = get_model_state_dict(
self.model) # type: ignore[no-any-return]
# filter out non-trainable parameters
state_dict = {k: v for k, v in state_dict.items() if v.requires_grad}
return state_dict # type: ignore
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
set_model_state_dict(
fastvideo/v1/training/training_pipeline.py
+51 -18
View File
@@ -10,6 +10,7 @@ from typing import Any
import imageio
import numpy as np
import torch
import torch.distributed as dist
import torchvision
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.optimization import get_scheduler
@@ -33,7 +34,7 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
from fastvideo.v1.forward_context import set_forward_context
from fastvideo.v1.logger import init_logger
from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
TrainingBatch)
LoRAPipeline, TrainingBatch)
from fastvideo.v1.training.activation_checkpoint import (
apply_activation_checkpointing)
from fastvideo.v1.training.training_utils import (
@@ -49,7 +50,7 @@ vsa_available = is_vsa_available()
logger = init_logger(__name__)
class TrainingPipeline(ComposedPipelineBase, ABC):
class TrainingPipeline(LoRAPipeline, ABC):
"""
A pipeline for training a model. All training pipelines should inherit from this class.
All reusable components and code should be implemented in this class.
@@ -60,6 +61,20 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
train_loader_iter: Iterator[dict[str, Any]]
current_epoch: int = 0
def __init__(
self,
model_path: str,
fastvideo_args: TrainingArgs,
required_config_modules: list[str] | None = None,
loaded_modules: dict[str, torch.nn.Module] | None = None) -> None:
fastvideo_args.inference_mode = False
self.lora_training = fastvideo_args.lora_training
if self.lora_training and fastvideo_args.lora_rank is None:
raise ValueError("lora rank must be set when using lora training")
super().__init__(model_path, fastvideo_args, required_config_modules,
loaded_modules) # type: ignore
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
raise RuntimeError(
"create_pipeline_stages should not be called for training pipeline")
@@ -87,12 +102,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
# Set random seeds for deterministic training
set_random_seed(self.seed)
self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
self.seed)
self.validation_random_generator = torch.Generator(
device="cpu").manual_seed(self.seed)
logger.info("Initialized random seeds with seed: %s", self.seed)
self.transformer.requires_grad_(True)
self.transformer.train()
if training_args.enable_gradient_checkpointing_type is not None:
@@ -105,7 +114,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
params_to_optimize = self.transformer.parameters()
params_to_optimize = list(
filter(lambda p: p.requires_grad, params_to_optimize))
self.optimizer = torch.optim.AdamW(
params_to_optimize,
lr=training_args.learning_rate,
@@ -167,7 +175,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
"Training pipelines must implement this method")
def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
self.transformer.requires_grad_(True)
self.transformer.train()
self.optimizer.zero_grad()
training_batch.total_loss = 0.0
@@ -218,9 +225,12 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
assert training_batch.encoder_hidden_states is not None
assert training_batch.encoder_attention_mask is not None
assert self.noise_random_generator is not None
batch_size = training_batch.latents.shape[0]
noise = torch.randn_like(training_batch.latents)
latents = training_batch.latents
batch_size = latents.shape[0]
noise = torch.randn(latents.shape,
generator=self.noise_gen_cuda,
device=latents.device,
dtype=latents.dtype)
u = compute_density_for_timestep_sampling(
weighting_scheme=self.training_args.weighting_scheme,
batch_size=batch_size,
@@ -231,17 +241,17 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
)
indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
timesteps = self.noise_scheduler.timesteps[indices].to(
device=training_batch.latents.device)
device=latents.device)
if self.training_args.sp_size > 1:
# Make sure that the timesteps are the same across all sp processes.
sp_group = get_sp_group()
sp_group.broadcast(timesteps, src=0)
sigmas = get_sigmas(
self.noise_scheduler,
training_batch.latents.device,
latents.device,
timesteps,
n_dim=training_batch.latents.ndim,
dtype=training_batch.latents.dtype,
n_dim=latents.ndim,
dtype=latents.dtype,
)
noisy_model_input = (1.0 -
sigmas) * training_batch.latents + sigmas * noise
@@ -342,7 +352,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
# logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
# local_main_process_only=False)
world_group = get_world_group()
world_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
world_group.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
training_batch.total_loss += avg_loss.item()
return training_batch
@@ -426,10 +436,30 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
self.init_steps = 0
def train(self) -> None:
set_random_seed(self.seed)
logger.info('rank: %s: start training',
self.global_rank,
local_main_process_only=False)
assert self.training_args is not None
if not self.post_init_called:
self.post_init()
num_trainable_params = 0
for name, param in self.transformer.named_parameters():
if param.requires_grad:
num_trainable_params += param.numel()
logger.info("Starting training with %s B trainable parameters",
round(num_trainable_params / 1e9, 3))
# Set random seeds for deterministic training
self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
self.seed)
self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(
self.seed)
self.validation_random_generator = torch.Generator(
device="cpu").manual_seed(self.seed)
logger.info("Initialized random seeds with seed: %s", self.seed)
self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
@@ -586,6 +616,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
@torch.no_grad()
def _log_validation(self, transformer, training_args, global_step) -> None:
"""
Generate a validation video and log it to wandb to check the quality during training.
"""
assert training_args is not None
training_args.inference_mode = True
training_args.use_cpu_offload = True
fastvideo/v1/training/wan_i2v_training_pipeline.py
+5 -4
View File
@@ -52,18 +52,19 @@ class WanI2VTrainingPipeline(TrainingPipeline):
args_copy.use_cpu_offload = True
# args_copy.pipeline_config.vae_config.load_encoder = False
# validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
validation_pipeline = WanImageToVideoPipeline.from_pretrained(
self.validation_pipeline = WanImageToVideoPipeline.from_pretrained(
training_args.model_path,
args=None,
inference_mode=True,
loaded_modules={"transformer": self.get_module("transformer")},
loaded_modules={
"transformer": self.get_module("transformer"),
"vae": self.get_module("vae")
},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus,
use_cpu_offload=True)
self.validation_pipeline = validation_pipeline
def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
assert self.training_args is not None
assert self.train_dataloader is not None
fastvideo/v1/training/wan_training_pipeline.py
+4 -1
View File
@@ -42,7 +42,10 @@ class WanTrainingPipeline(TrainingPipeline):
training_args.model_path,
args=None,
inference_mode=True,
loaded_modules={"transformer": self.get_module("transformer")},
loaded_modules={
"transformer": self.get_module("transformer"),
# "vae": self.get_module("vae")
},
tp_size=training_args.tp_size,
sp_size=training_args.sp_size,
num_gpus=training_args.num_gpus)
fastvideo/v1/worker/gpu_worker.py
+1 -4
View File
@@ -63,7 +63,6 @@ class Worker:
# Related issue:
# https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
@@ -95,9 +94,7 @@ class Worker:
output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args)
return cast(ForwardBatch, output_batch)
def set_lora_adapter(self,
lora_nickname: str,
lora_path: str | None = None) -> None:
def set_lora_adapter(self, lora_nickname: str, lora_path: str) -> None:
self.pipeline.set_lora_adapter(lora_nickname, lora_path)
def shutdown(self) -> dict[str, Any]:
fastvideo/v1/worker/multiproc_executor.py
-1
View File
@@ -45,7 +45,6 @@ class MultiprocExecutor(Executor):
for rank in range(self.world_size):
executor_pipe, worker_pipe = mp.Pipe(duplex=True)
self.worker_pipes.append(executor_pipe)
worker = mp.Process(target=run_worker_process,
name=f"FVWorkerProc-{rank}",
kwargs=dict(fastvideo_args=self.fastvideo_args,
scripts/preprocess/v1_preprocess_wan_data_t2v.sh
+3 -3
View File
@@ -2,9 +2,9 @@
GPU_NUM=1 # 2,4,8
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
MODEL_TYPE="wan"
DATA_MERGE_PATH="data/crush-smol/merge.txt"
OUTPUT_DIR="data/crush-smol/latents"
VALIDATION_PATH="assets/prompt.txt"
DATA_MERGE_PATH="data/mini_i2v_dataset/crush-smol_raw/merge.txt"
OUTPUT_DIR="data/mini_i2v_dataset/crush-smol_preprocessed"
VALIDATION_PATH="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
torchrun --nproc_per_node=$GPU_NUM \
fastvideo/v1/pipelines/preprocess/v1_preprocess.py \