From bb8c769c8e11bf64d2889427389b062ef95dc326 Mon Sep 17 00:00:00 2001 From: Wenxuan Tan Date: Thu, 17 Jul 2025 15:52:09 -0500 Subject: [PATCH] [LoRA] Support v1 LoRA training (#576) --- .buildkite/pipeline.yml | 11 ++ .buildkite/scripts/pr_test.sh | 4 + docs/source/training/data_preprocess.md | 2 +- examples/inference/lora/wan_lora_inference.py | 4 +- .../crush_smol/finetune_i2v_lora.sh | 93 ++++++++++++++ .../crush_smol/finetune_t2v_lora.sh | 92 ++++++++++++++ fastvideo/v1/configs/pipelines/base.py | 6 - fastvideo/v1/dataset/__init__.py | 3 +- fastvideo/v1/dataset/latent_datasets.py | 6 +- .../dataset/parquet_dataset_iterable_style.py | 4 +- .../v1/dataset/parquet_dataset_map_style.py | 10 +- .../v1/dataset/preprocessing_datasets.py | 29 +++-- fastvideo/v1/dataset/utils.py | 21 +-- fastvideo/v1/distributed/parallel_state.py | 15 ++- fastvideo/v1/fastvideo_args.py | 19 +++ fastvideo/v1/forward_context.py | 4 +- fastvideo/v1/layers/lora/linear.py | 112 +++++++++++++--- fastvideo/v1/models/loader/fsdp_load.py | 2 +- .../v1/pipelines/composed_pipeline_base.py | 37 ++++-- fastvideo/v1/pipelines/lora_pipeline.py | 100 +++++++++++---- fastvideo/v1/pipelines/stages/conditioning.py | 1 + fastvideo/v1/pipelines/stages/decoding.py | 1 + fastvideo/v1/pipelines/stages/encoding.py | 1 + .../v1/pipelines/stages/image_encoding.py | 3 + .../v1/pipelines/stages/stepvideo_encoding.py | 1 + .../v1/pipelines/stages/text_encoding.py | 3 + fastvideo/v1/tests/modal/pr_test.py | 4 + .../Vanilla/l40s_reference_wandb_summary.json | 2 +- .../training/Vanilla/test_training_loss.py | 42 +++--- .../l40s_reference_lora_wandb_summary.json | 44 +++++++ .../tests/training/lora/test_lora_training.py | 120 ++++++++++++++++++ fastvideo/v1/training/checkpointing_utils.py | 6 +- fastvideo/v1/training/training_pipeline.py | 69 +++++++--- .../v1/training/wan_i2v_training_pipeline.py | 9 +- .../v1/training/wan_training_pipeline.py | 5 +- fastvideo/v1/worker/gpu_worker.py | 5 +- fastvideo/v1/worker/multiproc_executor.py | 1 - .../preprocess/v1_preprocess_wan_data_t2v.sh | 6 +- 38 files changed, 757 insertions(+), 140 deletions(-) create mode 100644 examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh create mode 100644 examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh create mode 100644 fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json create mode 100644 fastvideo/v1/tests/training/lora/test_lora_training.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index bbb11c2d2..155d67704 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -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/**" diff --git a/.buildkite/scripts/pr_test.sh b/.buildkite/scripts/pr_test.sh index f43f08711..2d611fa35 100755 --- a/.buildkite/scripts/pr_test.sh +++ b/.buildkite/scripts/pr_test.sh @@ -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" diff --git a/docs/source/training/data_preprocess.md b/docs/source/training/data_preprocess.md index e302191a3..ede86cbff 100644 --- a/docs/source/training/data_preprocess.md +++ b/docs/source/training/data_preprocess.md @@ -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. diff --git a/examples/inference/lora/wan_lora_inference.py b/examples/inference/lora/wan_lora_inference.py index b96599e30..e392fed5f 100644 --- a/examples/inference/lora/wan_lora_inference.py +++ b/examples/inference/lora/wan_lora_inference.py @@ -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", diff --git a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh new file mode 100644 index 000000000..e2889fb12 --- /dev/null +++ b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v_lora.sh @@ -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[@]}" diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh new file mode 100644 index 000000000..22911adef --- /dev/null +++ b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v_lora.sh @@ -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[@]}" diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py index 9fc5ad5e6..5875370e8 100644 --- a/fastvideo/v1/configs/pipelines/base.py +++ b/fastvideo/v1/configs/pipelines/base.py @@ -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 diff --git a/fastvideo/v1/dataset/__init__.py b/fastvideo/v1/dataset/__init__.py index e1e9bdb53..fde109c9b 100644 --- a/fastvideo/v1/dataset/__init__.py +++ b/fastvideo/v1/dataset/__init__.py @@ -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__ = [ diff --git a/fastvideo/v1/dataset/latent_datasets.py b/fastvideo/v1/dataset/latent_datasets.py index 6bd43e116..dc01171d6 100644 --- a/fastvideo/v1/dataset/latent_datasets.py +++ b/fastvideo/v1/dataset/latent_datasets.py @@ -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: diff --git a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py index f71608732..256224ef6 100644 --- a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py +++ b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py @@ -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) diff --git a/fastvideo/v1/dataset/parquet_dataset_map_style.py b/fastvideo/v1/dataset/parquet_dataset_map_style.py index a1ada0a98..4783a1dd0 100644 --- a/fastvideo/v1/dataset/parquet_dataset_map_style.py +++ b/fastvideo/v1/dataset/parquet_dataset_map_style.py @@ -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): diff --git a/fastvideo/v1/dataset/preprocessing_datasets.py b/fastvideo/v1/dataset/preprocessing_datasets.py index fdbca4ae7..24660f6eb 100644 --- a/fastvideo/v1/dataset/preprocessing_datasets.py +++ b/fastvideo/v1/dataset/preprocessing_datasets.py @@ -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.""" diff --git a/fastvideo/v1/dataset/utils.py b/fastvideo/v1/dataset/utils.py index 49f1152da..ab1bf128a 100644 --- a/fastvideo/v1/dataset/utils.py +++ b/fastvideo/v1/dataset/utils.py @@ -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( diff --git a/fastvideo/v1/distributed/parallel_state.py b/fastvideo/v1/distributed/parallel_state.py index 30ef4815c..a2bace07a 100644 --- a/fastvideo/v1/distributed/parallel_state.py +++ b/fastvideo/v1/distributed/parallel_state.py @@ -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, diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 36f681e7f..319a5a658 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -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 diff --git a/fastvideo/v1/forward_context.py b/fastvideo/v1/forward_context.py index 567bc54bd..11b5ead18 100644 --- a/fastvideo/v1/forward_context.py +++ b/fastvideo/v1/forward_context.py @@ -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. diff --git a/fastvideo/v1/layers/lora/linear.py b/fastvideo/v1/layers/lora/linear.py index 76fa21557..309030c10 100644 --- a/fastvideo/v1/layers/lora/linear.py +++ b/fastvideo/v1/layers/lora/linear.py @@ -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 diff --git a/fastvideo/v1/models/loader/fsdp_load.py b/fastvideo/v1/models/loader/fsdp_load.py index f35aa9485..9a8bce7ba 100644 --- a/fastvideo/v1/models/loader/fsdp_load.py +++ b/fastvideo/v1/models/loader/fsdp_load.py @@ -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 diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index 7444e2031..79b4a1687 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -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()) diff --git a/fastvideo/v1/pipelines/lora_pipeline.py b/fastvideo/v1/pipelines/lora_pipeline.py index c65b9806f..ca82d7bce 100644 --- a/fastvideo/v1/pipelines/lora_pipeline.py +++ b/fastvideo/v1/pipelines/lora_pipeline.py @@ -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 diff --git a/fastvideo/v1/pipelines/stages/conditioning.py b/fastvideo/v1/pipelines/stages/conditioning.py index 8355c47b6..cbed69d7a 100644 --- a/fastvideo/v1/pipelines/stages/conditioning.py +++ b/fastvideo/v1/pipelines/stages/conditioning.py @@ -23,6 +23,7 @@ class ConditioningStage(PipelineStage): to the diffusion process. """ + @torch.no_grad() def forward( self, batch: ForwardBatch, diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index 1f98e0b7d..7ccde3102 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -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, diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index a99bf9c2f..8c9ab6193 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -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, diff --git a/fastvideo/v1/pipelines/stages/image_encoding.py b/fastvideo/v1/pipelines/stages/image_encoding.py index 62370c3e0..c98bd9b2e 100644 --- a/fastvideo/v1/pipelines/stages/image_encoding.py +++ b/fastvideo/v1/pipelines/stages/image_encoding.py @@ -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, diff --git a/fastvideo/v1/pipelines/stages/stepvideo_encoding.py b/fastvideo/v1/pipelines/stages/stepvideo_encoding.py index df9e8cee3..93a2f07f5 100644 --- a/fastvideo/v1/pipelines/stages/stepvideo_encoding.py +++ b/fastvideo/v1/pipelines/stages/stepvideo_encoding.py @@ -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] diff --git a/fastvideo/v1/pipelines/stages/text_encoding.py b/fastvideo/v1/pipelines/stages/text_encoding.py index 9eff83d0f..03c217e2c 100644 --- a/fastvideo/v1/pipelines/stages/text_encoding.py +++ b/fastvideo/v1/pipelines/stages/text_encoding.py @@ -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, diff --git a/fastvideo/v1/tests/modal/pr_test.py b/fastvideo/v1/tests/modal/pr_test.py index 75422d157..2bf879353 100644 --- a/fastvideo/v1/tests/modal/pr_test.py +++ b/fastvideo/v1/tests/modal/pr_test.py @@ -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") diff --git a/fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json b/fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json index 952d0b9f6..fc3d85a22 100644 --- a/fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json +++ b/fastvideo/v1/tests/training/Vanilla/l40s_reference_wandb_summary.json @@ -1 +1 @@ -{"step_time":5.501357046999999,"grad_norm":0.384765625,"train_loss":0.07890288904309273,"avg_step_time":5.831571423200001} \ No newline at end of file +{"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}} \ No newline at end of file diff --git a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py index b07ade0ee..35652443a 100644 --- a/fastvideo/v1/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/v1/tests/training/Vanilla/test_training_loss.py @@ -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' diff --git a/fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json b/fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json new file mode 100644 index 000000000..b0859886b --- /dev/null +++ b/fastvideo/v1/tests/training/lora/l40s_reference_lora_wandb_summary.json @@ -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 + } \ No newline at end of file diff --git a/fastvideo/v1/tests/training/lora/test_lora_training.py b/fastvideo/v1/tests/training/lora/test_lora_training.py new file mode 100644 index 000000000..483d827f5 --- /dev/null +++ b/fastvideo/v1/tests/training/lora/test_lora_training.py @@ -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() \ No newline at end of file diff --git a/fastvideo/v1/training/checkpointing_utils.py b/fastvideo/v1/training/checkpointing_utils.py index e55ea8862..b3b0f2c27 100644 --- a/fastvideo/v1/training/checkpointing_utils.py +++ b/fastvideo/v1/training/checkpointing_utils.py @@ -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( diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 9b355e251..f51063e4e 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -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 diff --git a/fastvideo/v1/training/wan_i2v_training_pipeline.py b/fastvideo/v1/training/wan_i2v_training_pipeline.py index e5b37e6b4..8c1433039 100644 --- a/fastvideo/v1/training/wan_i2v_training_pipeline.py +++ b/fastvideo/v1/training/wan_i2v_training_pipeline.py @@ -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 diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index 971aa7845..99666e24d 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -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) diff --git a/fastvideo/v1/worker/gpu_worker.py b/fastvideo/v1/worker/gpu_worker.py index 546103851..9df0f14d7 100644 --- a/fastvideo/v1/worker/gpu_worker.py +++ b/fastvideo/v1/worker/gpu_worker.py @@ -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]: diff --git a/fastvideo/v1/worker/multiproc_executor.py b/fastvideo/v1/worker/multiproc_executor.py index 2a2d48856..928aeaea0 100644 --- a/fastvideo/v1/worker/multiproc_executor.py +++ b/fastvideo/v1/worker/multiproc_executor.py @@ -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, diff --git a/scripts/preprocess/v1_preprocess_wan_data_t2v.sh b/scripts/preprocess/v1_preprocess_wan_data_t2v.sh index 56e384ebc..f587f8a95 100644 --- a/scripts/preprocess/v1_preprocess_wan_data_t2v.sh +++ b/scripts/preprocess/v1_preprocess_wan_data_t2v.sh @@ -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 \