diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index df50350fd..c03a9f6d7 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -39,6 +39,11 @@ on: required: false default: false type: boolean + run_training_test: + description: "Run training-test" + required: false + default: false + type: boolean env: PYTHONUNBUFFERED: "1" @@ -160,6 +165,25 @@ jobs: secrets: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + + training-test: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true') + uses: ./.github/workflows/runpod-test.yml + with: + job_id: "training-test" + gpu_type: "NVIDIA A40" + gpu_count: 4 + volume_size: 100 + disk_size: 100 + image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}" + test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/training -srP" + timeout_minutes: 30 + secrets: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} runpod-cleanup: needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index c763d0ddb..942095fc2 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -529,6 +529,7 @@ class TrainingArgs(FastVideoArgs): validation_steps: float = 0.0 log_validation: bool = False tracker_project_name: str = "" + wandb_run_name: str = "" seed: Optional[int] = None # output @@ -689,6 +690,9 @@ class TrainingArgs(FastVideoArgs): parser.add_argument("--tracker-project-name", type=str, help="Project name for tracking") + parser.add_argument("--wandb-run-name", + type=str, + help="Run name for wandb") parser.add_argument("--seed", type=int, default=42, diff --git a/fastvideo/v1/tests/conftest.py b/fastvideo/v1/tests/conftest.py index 9aa6bfd4f..c7547ad6b 100644 --- a/fastvideo/v1/tests/conftest.py +++ b/fastvideo/v1/tests/conftest.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 import pytest -import torch.distributed as dist -import pytest import torch import numpy as np diff --git a/fastvideo/v1/tests/training/reference_wandb_summary.json b/fastvideo/v1/tests/training/reference_wandb_summary.json new file mode 100644 index 000000000..dd7e64c0f --- /dev/null +++ b/fastvideo/v1/tests/training/reference_wandb_summary.json @@ -0,0 +1 @@ +{"_timestamp":1.7496170016478686e+09,"validation_videos_8_steps":{"_type":"videos","count":5,"videos":[{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"},{"sha256":"cd72a3d513eca6b41b03b80e6fa044ce7219c35e969d2ca20b9cb48c91e585c6","size":477837,"path":"media/videos/validation_videos_8_steps_0_cd72a3d513eca6b41b03.mp4","_type":"video-file"},{"_type":"video-file","sha256":"43d47c211a69bf0be3544738e76e7d8fa58bb108c00cb21dc34eaad3c6ce7cc3","size":409419,"path":"media/videos/validation_videos_8_steps_0_43d47c211a69bf0be354.mp4"},{"_type":"video-file","sha256":"ea674ec9e200bc97563c9d87d9dc07110c3f42c0ab7277dd237ae96dd8f90a10","size":333966,"path":"media/videos/validation_videos_8_steps_0_ea674ec9e200bc97563c.mp4"},{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"}],"captions":false},"step_time":2.5065076276659966,"_wandb":{"runtime":53},"learning_rate":1e-06,"_step":5,"_runtime":53.172758961,"grad_norm":5.65625,"train_loss":0.3915919363498688,"avg_step_time":2.8116052336990833} \ No newline at end of file diff --git a/fastvideo/v1/tests/training/test_training_loss.py b/fastvideo/v1/tests/training/test_training_loss.py new file mode 100644 index 000000000..b38dda5da --- /dev/null +++ b/fastvideo/v1/tests/training/test_training_loss.py @@ -0,0 +1,142 @@ +import os +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 + +# 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 + +wandb_name = "test_training_loss" +reference_wandb_summary_file = "fastvideo/v1/tests/training/reference_wandb_summary.json" + +NUM_NODES = "1" +NUM_GPUS_PER_NODE = "4" + + +def run_worker(): + """Worker function that will be run on each GPU""" + # Create and populate args + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + + # Set the arguments as they are in finetune_v1_test.sh + args = parser.parse_args([ + "--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", + "--cache_dir", "/home/.cache", + "--data_path", "data/crush-smol_parq/combined_parquet_dataset", + "--validation_prompt_dir", "data/crush-smol_parq/validation_parquet_dataset", + "--train_batch_size", "2", + "--num_latent_t", "4", + "--num_gpus", "4", + "--sp_size", "4", + "--tp_size", "4", + "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", "4", + "--train_sp_batch_size", "1", + "--dataloader_num_workers", "1", + "--gradient_accumulation_steps", "2", + "--max_train_steps", "5", + "--learning_rate", "1e-6", + "--mixed_precision", "bf16", + "--checkpointing_steps", "30", + "--validation_steps", "10", + "--validation_sampling_steps", "8", + "--log_validation", + "--checkpoints_total_limit", "3", + "--allow_tf32", + "--ema_start_step", "0", + "--cfg", "0.0", + "--output_dir", "data/wan_finetune_test", + "--tracker_project_name", "wan_finetune_ci", + "--wandb_run_name", wandb_name, + "--num_height", "480", + "--num_width", "832", + "--num_frames", "81", + "--shift", "3", + "--validation_guidance_scale", "1.0", + "--num_euler_timesteps", "50", + "--multi_phased_distill_schedule", "4000-1", + "--weight_decay", "0.01", + "--not_apply_cfg_solver", + "--master_weight_type", "fp32", + "--max_grad_norm", "1.0" + ]) + + # Call the main training function + main(args) + +def test_distributed_training(): + """Test the distributed training setup""" + os.environ["WANDB_API_KEY"] = "8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" + os.environ["WANDB_MODE"] = "online" + + data_dir = Path("data/crush-smol_parq") + + if not data_dir.exists(): + print(f"Downloading test dataset to {data_dir}...") + snapshot_download( + repo_id="PY007/crush-smol", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False + ) + + # Get the current file path + current_file = Path(__file__).resolve() + + # Run torchrun command + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE, + str(current_file) + ] + + process = subprocess.run(cmd, check=True) + + wandb_dir = f"wandb/latest-run" + summary_file = f"{wandb_dir}/files/wandb-summary.json" + + reference_wandb_summary = json.load(open(reference_wandb_summary_file)) + wandb_summary = json.load(open(summary_file)) + + fields_and_thresholds = { + 'avg_step_time': 1.0, + 'grad_norm': 0.1, + 'step_time': 0.5, + 'train_loss': 0.001 + } + + failures = [] + for field, threshold in fields_and_thresholds.items(): + 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})") + + if failures: + raise AssertionError("\n".join(failures)) + +if __name__ == "__main__": + if os.environ.get("LOCAL_RANK") is not None: + # We're being run by torchrun + run_worker() + else: + # We're being run directly + test_distributed_training() diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index eca1ebc8a..f8e66e837 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -121,7 +121,9 @@ class TrainingPipeline(ComposedPipelineBase, ABC): if self.global_rank == 0: project = training_args.tracker_project_name or "fastvideo" - wandb.init(project=project, config=training_args) + wandb.init(project=project, + config=training_args, + name=training_args.wandb_run_name) @abstractmethod def initialize_validation_pipeline(self, training_args: TrainingArgs): @@ -177,115 +179,117 @@ class TrainingPipeline(ComposedPipelineBase, ABC): transformer.eval() - # Process each validation prompt - videos: List[np.ndarray] = [] - captions: List[str | None] = [] - for _, embeddings, masks, infos in validation_dataloader: - captions.extend([None]) # TODO(peiyuan): add caption - prompt_embeds = embeddings.to(get_torch_device()) - prompt_attention_mask = masks.to(get_torch_device()) + validation_steps = training_args.validation_sampling_steps.split(",") + validation_steps = [int(step) for step in validation_steps] + validation_steps = [step for step in validation_steps if step > 0] - # Calculate sizes - latents_size = [(sampling_param.num_frames - 1) // 4 + 1, - sampling_param.height // 8, - sampling_param.width // 8] - n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + # Process each validation prompt for each validation step + for num_inference_steps in validation_steps: + step_videos: List[np.ndarray] = [] + step_captions: List[str | None] = [] - temporal_compression_factor = training_args.vae_config.arch_config.temporal_compression_ratio - num_frames = (training_args.num_latent_t - - 1) * temporal_compression_factor + 1 + for _, embeddings, masks, infos in validation_dataloader: + step_captions.extend([None]) # TODO(peiyuan): add caption + prompt_embeds = embeddings.to(get_torch_device()) + prompt_attention_mask = masks.to(get_torch_device()) - # Prepare batch for validation - batch = ForwardBatch( - data_type="video", - latents=None, - seed=validation_seed, # Use deterministic seed - generator=torch.Generator( - device="cpu").manual_seed(validation_seed), - prompt_embeds=[prompt_embeds], - prompt_attention_mask=[prompt_attention_mask], - negative_prompt_embeds=[negative_prompt_embeds], - negative_attention_mask=[negative_prompt_attention_mask], - # make sure we use the same height, width, and num_frames as the training pipeline - height=training_args.num_height, - width=training_args.num_width, - num_frames=num_frames, - # TODO(will): validation_sampling_steps and - # validation_guidance_scale are actually passed in as a list of - # values, like "10,20,30". The validation should be run for each - # combination of values. - # num_inference_steps=fastvideo_args.validation_sampling_steps, - num_inference_steps=sampling_param.num_inference_steps, - # guidance_scale=fastvideo_args.validation_guidance_scale, - guidance_scale=sampling_param.guidance_scale, - n_tokens=n_tokens, - eta=0.0, - ) + # Calculate sizes + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, + sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - # Run validation inference - with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): - output_batch = self.validation_pipeline.forward( - batch, training_args) - samples = output_batch.output + temporal_compression_factor = training_args.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 - # Re-enable gradients for training - transformer.requires_grad_(True) - transformer.train() + # Prepare batch for validation + batch = ForwardBatch( + data_type="video", + latents=None, + seed=validation_seed, # Use deterministic seed + generator=torch.Generator( + device="cpu").manual_seed(validation_seed), + prompt_embeds=[prompt_embeds], + prompt_attention_mask=[prompt_attention_mask], + negative_prompt_embeds=[negative_prompt_embeds], + negative_attention_mask=[negative_prompt_attention_mask], + height=training_args.num_height, + width=training_args.num_width, + num_frames=num_frames, + num_inference_steps= + num_inference_steps, # Use the current validation step + guidance_scale=sampling_param.guidance_scale, + n_tokens=n_tokens, + eta=0.0, + ) - if self.rank_in_sp_group != 0: - continue + # Run validation inference + with torch.no_grad(), torch.autocast("cuda", + dtype=torch.bfloat16): + output_batch = self.validation_pipeline.forward( + batch, training_args) + samples = output_batch.output - # Process outputs - video = rearrange(samples, "b c t h w -> t b c h w") - frames = [] - for x in video: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - videos.append(frames) + if self.rank_in_sp_group != 0: + continue - # Log validation results - world_group = get_world_group() - num_sp_groups = world_group.world_size // self.sp_group.world_size + # Process outputs + video = rearrange(samples, "b c t h w -> t b c h w") + frames = [] + for x in video: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + step_videos.append(frames) - # Only sp_group leaders (rank_in_sp_group == 0) need to send their - # results to global rank 0 - if self.rank_in_sp_group == 0: - if self.global_rank == 0: - # Global rank 0 collects results from all sp_group leaders - all_videos = videos # Start with own results - all_captions = captions + # Log validation results for this step + world_group = get_world_group() + num_sp_groups = world_group.world_size // self.sp_group.world_size - # Receive from other sp_group leaders - for sp_group_idx in range(1, num_sp_groups): - src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders - recv_videos = world_group.recv_object(src=src_rank) - recv_captions = world_group.recv_object(src=src_rank) - all_videos.extend(recv_videos) - all_captions.extend(recv_captions) + # Only sp_group leaders (rank_in_sp_group == 0) need to send their + # results to global rank 0 + if self.rank_in_sp_group == 0: + if self.global_rank == 0: + # Global rank 0 collects results from all sp_group leaders + all_videos = step_videos # Start with own results + all_captions = step_captions - video_filenames = [] - for i, (video, - caption) in enumerate(zip(all_videos, all_captions)): - os.makedirs(training_args.output_dir, exist_ok=True) - filename = os.path.join( - training_args.output_dir, - f"validation_step_{global_step}_video_{i}.mp4") - imageio.mimsave(filename, video, fps=sampling_param.fps) - video_filenames.append(filename) + # Receive from other sp_group leaders + for sp_group_idx in range(1, num_sp_groups): + src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders + recv_videos = world_group.recv_object(src=src_rank) + recv_captions = world_group.recv_object(src=src_rank) + all_videos.extend(recv_videos) + all_captions.extend(recv_captions) - logs = { - "validation_videos": [ - wandb.Video(filename, caption=caption) for filename, - caption in zip(video_filenames, all_captions) - ] - } - wandb.log(logs, step=global_step) - else: - # Other sp_group leaders send their results to global rank 0 - world_group.send_object(videos, dst=0) - world_group.send_object(captions, dst=0) + video_filenames = [] + for i, (video, + caption) in enumerate(zip(all_videos, + all_captions)): + os.makedirs(training_args.output_dir, exist_ok=True) + filename = os.path.join( + training_args.output_dir, + f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" + ) + imageio.mimsave(filename, video, fps=sampling_param.fps) + video_filenames.append(filename) + logs = { + f"validation_videos_{num_inference_steps}_steps": [ + wandb.Video(filename, caption=caption) + for filename, caption in zip( + video_filenames, all_captions) + ] + } + wandb.log(logs, step=global_step) + else: + # Other sp_group leaders send their results to global rank 0 + world_group.send_object(step_videos, dst=0) + world_group.send_object(step_captions, dst=0) + + # Re-enable gradients for training + transformer.train() gc.collect() torch.cuda.empty_cache() diff --git a/fastvideo/v1/training/wan_training_pipeline.py b/fastvideo/v1/training/wan_training_pipeline.py index a879dd2d6..d2e9637d9 100644 --- a/fastvideo/v1/training/wan_training_pipeline.py +++ b/fastvideo/v1/training/wan_training_pipeline.py @@ -338,6 +338,7 @@ class WanTrainingPipeline(TrainingPipeline): logger.info("GPU memory usage after validation: %s MB", gpu_memory_usage) + wandb.finish() save_checkpoint(self.transformer, self.global_rank, self.training_args.output_dir, self.training_args.max_train_steps, self.optimizer, diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index b5fb9ec58..da0d38923 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -22,7 +22,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --sp_size 4 \ --tp_size 4 \ --hsdp_replicate_dim 1 \ - --hsdp_shards 4 \ + --hsdp_shard_dim 4 \ --num_gpus $NUM_GPUS \ --train_sp_batch_size 1\ --dataloader_num_workers 10\