From cc1d2e759f32058d329937022b5701a450e84a23 Mon Sep 17 00:00:00 2001 From: "Yuxuan.Zhang" <2448370773@qq.com> Date: Wed, 9 Oct 2024 06:49:59 +0800 Subject: [PATCH] Multi-GPU parallel encoding support for training videos. (#6) * Multi-GPU parallel encoding support for training videos. * revert * make style * update --------- Co-authored-by: Aryan --- .gitignore | 3 + README.md | 2 +- README_zh.md | 95 ++++++++++++++++++++ prepare_dataset.sh | 33 +++---- training/prepare_dataset.py | 168 +++++++++++++++++++++++++----------- 5 files changed, 236 insertions(+), 65 deletions(-) create mode 100644 README_zh.md diff --git a/.gitignore b/.gitignore index e652a10..bcc6640 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,9 @@ __pycache__/ *.py[cod] *$py.class +# JetBrains +.idea + # C extensions *.so diff --git a/README.md b/README.md index 81916ad..5da5fda 100644 --- a/README.md +++ b/README.md @@ -402,4 +402,4 @@ With `train_batch_size = 4`: - [x] Make 5B lora finetuning work in under 24GB > [!IMPORTANT] -> Since our goal is to make the scripts as memory-friendly as possible we don't guarantee multi-GPU training. \ No newline at end of file +> Since our goal is to make the scripts as memory-friendly as possible we don't guarantee multi-GPU training. diff --git a/README_zh.md b/README_zh.md new file mode 100644 index 0000000..3758d84 --- /dev/null +++ b/README_zh.md @@ -0,0 +1,95 @@ +# CogVideoX Factory + +## 简介 + +这是用于 CogVideoX 微调的仓库。 + +## 数据集准备 + +创建两个文件，一个文件包含以换行符分隔的提示词，另一个文件包含以换行符分隔的视频数据路径（视频文件的路径必须相对于您在指定 `--data_root` 时传递的路径）。让我们通过一个例子来更好地理解这一点！ + +假设您将 `--data_root` 指定为 `/dataset`，并且该目录包含文件：`prompts.txt` 和 `videos.txt`。 + +`prompts.txt` 文件应包含以换行符分隔的提示词： + +``` +一段黑白动画序列，主角是一只名为 Rabbity Ribfried 的兔子和一只拟人化的山羊，在一个充满音乐和趣味的环境中，展示他们不断发展的互动。 +一段黑白动画序列，场景在船甲板上，主角是一只名为 Bully Bulldoger 的斗牛犬角色，展示了夸张的面部表情和肢体语言。角色从自信到专注，再到紧张和痛苦，展示了一系列情绪，随着它克服挑战。船的内部在背景中保持静止，只有简单的细节，如钟声和开着的门。角色的动态动作和变化的表情推动了故事的发展，没有镜头移动，确保观众专注于其不断变化的反应和肢体动作。 +... +``` + +`videos.txt` 文件应包含以换行符分隔的视频文件路径。请注意，路径应相对于 `--data_root` 目录。 + +```bash +videos/00000.mp4 +videos/00001.mp4 +... +``` + +总体而言，如果您在数据集根目录运行 `tree` 命令，您的数据集应如下所示： + +```bash +/dataset +├── prompts.txt +├── videos.txt +├── videos + ├── videos/00000.mp4 + ├── videos/00001.mp4 + ├── ... +``` + +使用此格式时，`--caption_column` 必须是 `prompts.txt`，`--video_column` 必须是 `videos.txt`。如果您的数据存储在 CSV 文件中，您也可以指定 `--dataset_file` 为 CSV 的路径，`--caption_column` 和 `--video_column` 为 CSV 文件中的实际列名。 + +例如，让我们使用这个 [Disney 数据集](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) 进行微调。要下载，可以使用 🤗 Hugging Face CLI。 + +```bash +huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir video-dataset-disney +``` + +## 训练 + +TODO + +请查看 `training/*.sh` + +注意：未在 MPS 上测试 + +## 内存需求 + +训练支持并验证的内存优化包括： + +- 来自 [TorchAO](https://github.com/pytorch/ao) 的 `CPUOffloadOptimizer`。 +- 来自 [bitsandbytes](https://huggingface.co/docs/bitsandbytes/optimizers) 的低位优化器。 + +### LoRA 微调 + +
+ AdamW + +With `train_batch_size = 1`: + +| model | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing | +|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:| +| THUDM/CogVideoX-2b | 16 | False | 12.945 | 43.764 | 46.918 | 24.234 | +| THUDM/CogVideoX-2b | 16 | True | 12.945 | 12.945 | 21.121 | 24.234 | +| THUDM/CogVideoX-2b | 64 | False | 13.035 | 44.314 | 47.469 | 24.469 | +| THUDM/CogVideoX-2b | 64 | True | 13.036 | 13.035 | 21.564 | 24.500 | +| THUDM/CogVideoX-2b | 256 | False | 13.095 | 45.826 | 48.990 | 25.543 | +| THUDM/CogVideoX-2b | 256 | True | 13.094 | 13.095 | 22.344 | 25.537 | +| THUDM/CogVideoX-5b | 16 | True | 19.742 | 19.742 | 28.746 | 38.123 | +| THUDM/CogVideoX-5b | 64 | True | 20.006 | 20.818 | 30.338 | 38.738 | +| THUDM/CogVideoX-5b | 256 | True | 20.771 | 22.119 | 31.939 | 41.537 | + +With `train_batch_size = 4`: + +| model | lora rank | gradient_checkpointing | memory_before_training | memory_before_validation | memory_after_validation | memory_after_testing | +|:------------------:|:---------:|:----------------------:|:----------------------:|:------------------------:|:-----------------------:|:--------------------:| +| THUDM/CogVideoX-2b | 16 | True | 12.945 | 21.803 | 21.814 | 24.322 | +| THUDM/CogVideoX-2b | 64 | True | 13.035 | 22.254 | 22.254 | 24.572 | +| THUDM/CogVideoX-2b | 256 | True | 13.094 | 22.020 | 22.033 | 25.574 | +| THUDM/CogVideoX-5b | 16 | True | 19.742 | 46.492 | 46.492 | 38.197 | +| THUDM/CogVideoX-5b | 64 | True | 20.006 | 47.805 | 47.805 | 39.365 | +| THUDM/CogVideoX-5b | 256 | True | 20.771 | 47.268 | 47.332 | 41.008 | + +> [!NOTE] +> \ No newline at end of file diff --git a/prepare_dataset.sh b/prepare_dataset.sh index b0414e9..44cba32 100755 --- a/prepare_dataset.sh +++ b/prepare_dataset.sh @@ -2,11 +2,13 @@ MODEL_ID="THUDM/CogVideoX-2b" +NUM_GPUS=8 + # For more details on the expected data format, please refer to the README. -DATA_ROOT="/raid/aryan/video-dataset-tom-and-jerry" # This needs to be the path to the base directory where your videos are located. +DATA_ROOT="/path/to/my/datasets/video-dataset" # This needs to be the path to the base directory where your videos are located. CAPTION_COLUMN="prompts.txt" VIDEO_COLUMN="videos.txt" -OUTPUT_DIR="/raid/aryan/video-dataset-tom-and-jerry-encoded" +OUTPUT_DIR="/path/to/my/datasets/preprocessed-dataset" HEIGHT=480 WIDTH=720 MAX_NUM_FRAMES=49 @@ -17,19 +19,20 @@ DTYPE=fp32 # To create a folder-style dataset structure without pre-encoding videos and captions' CMD_WITHOUT_PRE_ENCODING="\ - python3 training/prepare_dataset.py \ - --model_id $MODEL_ID \ - --data_root $DATA_ROOT \ - --caption_column $CAPTION_COLUMN \ - --video_column $VIDEO_COLUMN \ - --output_dir $OUTPUT_DIR \ - --height $HEIGHT \ - --width $WIDTH \ - --max_num_frames $MAX_NUM_FRAMES \ - --max_sequence_length $MAX_SEQUENCE_LENGTH \ - --target_fps $TARGET_FPS \ - --batch_size $BATCH_SIZE \ - --dtype $DTYPE + torchrun --nproc_per_node=$NUM_GPUS \ + training/prepare_dataset.py \ + --model_id $MODEL_ID \ + --data_root $DATA_ROOT \ + --caption_column $CAPTION_COLUMN \ + --video_column $VIDEO_COLUMN \ + --output_dir $OUTPUT_DIR \ + --height $HEIGHT \ + --width $WIDTH \ + --max_num_frames $MAX_NUM_FRAMES \ + --max_sequence_length $MAX_SEQUENCE_LENGTH \ + --target_fps $TARGET_FPS \ + --batch_size $BATCH_SIZE \ + --dtype $DTYPE " CMD_WITH_PRE_ENCODING="$CMD_WITHOUT_PRE_ENCODING --save_tensors" diff --git a/training/prepare_dataset.py b/training/prepare_dataset.py index 9d70884..42f0b3a 100644 --- a/training/prepare_dataset.py +++ b/training/prepare_dataset.py @@ -1,23 +1,22 @@ #!/usr/bin/env python3 -# For folder structure dataset: python3 prepare_dataset.py --model_id THUDM/CogVideoX-2b --data_root /raid/aryan/video-dataset-disney/ --caption_column prompts.txt --video_column videos.txt --output_dir dump --height 480 --width 720 --max_num_frames 49 --max_sequence_length 226 --target_fps 8 --batch_size 1 --dtype fp32 -# For latent/embed structure dataset: python3 prepare_dataset.py --model_id THUDM/CogVideoX-2b --data_root /raid/aryan/video-dataset-disney/ --caption_column prompts.txt --video_column videos.txt --output_dir dump --height 480 --width 720 --max_num_frames 49 --max_sequence_length 226 --target_fps 8 --batch_size 1 --dtype fp32 --save_tensors - import argparse import gc +import os import pathlib import traceback from typing import Any, Dict, List, Optional, Tuple, Union import pandas as pd import torch +import torch.distributed as dist from diffusers import AutoencoderKLCogVideoX from diffusers.utils import export_to_video, get_logger from torchvision import transforms +from tqdm import tqdm from transformers import T5EncoderModel, T5Tokenizer -# Must import after importing torch, otherwise there's a nasty segfault when loading text_encoder/vae import decord # isort:skip decord.bridge.set_bridge("torch") @@ -104,7 +103,7 @@ def get_args() -> Dict[str, Any]: def load_dataset_from_local_path( data_root: pathlib.Path, caption_column: str, video_column: str -) -> Tuple[List[str], List[str]]: +) -> Tuple[List[str], List[pathlib.Path]]: if not data_root.exists(): raise ValueError("Root folder for videos does not exist") @@ -127,7 +126,7 @@ def load_dataset_from_local_path( if any(not path.is_file() for path in video_paths): raise ValueError( - f"Expected `{video_column=}` to be a path to a file in `{data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file." + f"Expected `{video_column}` to be a path to a file in `{data_root}` containing line-separated paths to video data but found at least one path that is not a valid file." ) return prompts, video_paths @@ -135,7 +134,7 @@ def load_dataset_from_local_path( def load_dataset_from_csv( data_root: pathlib.Path, dataset_file: pathlib.Path, caption_column: str, video_column: str -) -> Tuple[List[str], List[str]]: +) -> Tuple[List[str], List[pathlib.Path]]: df = pd.read_csv(dataset_file) prompts = df[caption_column].tolist() video_paths = df[video_column].tolist() @@ -143,7 +142,7 @@ def load_dataset_from_csv( if any(not path.is_file() for path in video_paths): raise ValueError( - f"Expected `{video_column=}` to be a path to a file in `{data_root=}` containing line-separated paths to video data but found atleast one path that is not a valid file." + f"Expected `{video_column}` to be a path to a file in `{data_root}` containing line-separated paths to video data but found at least one path that is not a valid file." ) return prompts, video_paths @@ -151,7 +150,7 @@ def load_dataset_from_csv( def load_and_preprocess_video( path: pathlib.Path, height: int, width: int, max_num_frames: int, video_transforms, num_threads: int = 0 -) -> torch.Tensor: +) -> Optional[torch.Tensor]: frames = None try: @@ -160,11 +159,11 @@ def load_and_preprocess_video( if video_num_frames < max_num_frames: logger.warning( - f"Video at `{path.as_posix()}` should have atleast `{max_num_frames=}`, but got only `{video_num_frames=}`. Skipping it." + f"Video at `{path.as_posix()}` should have at least `{max_num_frames=}`, but got only `{video_num_frames=}`. Skipping it." ) - return + return None - indices = list(range(0, video_num_frames, video_num_frames // max_num_frames)) + indices = list(range(0, video_num_frames, max(video_num_frames // max_num_frames, 1))) frames: torch.Tensor = video_reader.get_batch(indices) frames = frames[:max_num_frames].float() frames = frames.permute(0, 3, 1, 2).contiguous() @@ -241,7 +240,7 @@ def encode_prompt( def compute_prompt_embeddings( tokenizer: T5Tokenizer, text_encoder: T5EncoderModel, - prompt: str, + prompts: List[str], max_sequence_length: int, device: torch.device, dtype: torch.dtype, @@ -251,7 +250,7 @@ def compute_prompt_embeddings( prompt_embeds = encode_prompt( tokenizer, text_encoder, - prompt, + prompts, num_videos_per_prompt=1, max_sequence_length=max_sequence_length, device=device, @@ -262,7 +261,7 @@ def compute_prompt_embeddings( prompt_embeds = encode_prompt( tokenizer, text_encoder, - prompt, + prompts, num_videos_per_prompt=1, max_sequence_length=max_sequence_length, device=device, @@ -272,7 +271,11 @@ def compute_prompt_embeddings( def save_videos( - videos: torch.Tensor, video_paths: List[str], prompts: List[str], output_dir: pathlib.Path, target_fps: int = 8 + videos: torch.Tensor, + video_paths: List[pathlib.Path], + prompts: List[str], + output_dir: pathlib.Path, + target_fps: int = 8, ) -> None: assert videos.size(0) == len(video_paths) @@ -289,13 +292,13 @@ def save_videos( videos_pil = [[to_pil_image(frame) for frame in video] for video in videos] for video, video_path in zip(videos_pil, video_paths): - filename = video_dir.joinpath(pathlib.Path(video_path).name) + filename = video_dir.joinpath(video_path.name) logger.debug(f"Saving video to `{filename}`") export_to_video(video, filename.as_posix(), fps=target_fps) with open(output_dir.joinpath("videos.txt").as_posix(), "w", encoding="utf-8") as file: for video_path in video_paths: - file.write(f"videos/{pathlib.Path(video_path).name}\n") + file.write(f"videos/{video_path.name}\n") with open(output_dir.joinpath("prompts.txt").as_posix(), "w", encoding="utf-8") as file: for prompt in prompts: @@ -305,7 +308,7 @@ def save_videos( def save_latents_and_embeddings( latents: torch.Tensor, prompt_embeds: torch.Tensor, - video_paths: List[str], + video_paths: List[pathlib.Path], prompts: List[str], output_dir: pathlib.Path, ) -> None: @@ -321,27 +324,20 @@ def save_latents_and_embeddings( embeds_dir.mkdir(parents=True, exist_ok=True) for latent, embed, video_path in zip(latents, prompt_embeds, video_paths): - # Need to perform the clone, otherwise the entire `latents` or `prompt_embeds` tensor is - # saved for every single video/prompt embedding. This is due to us viewing a slice of a - # large tensor when iteratively saving stuff here. latent = latent.clone() embed = embed.clone() - video_path = pathlib.Path(video_path) - filename_without_ext = video_path.name.split(".")[0] + filename_without_ext = video_path.stem - latent_filename = latents_dir.joinpath(filename_without_ext) - embed_filename = embeds_dir.joinpath(filename_without_ext) - - latent_filename = f"{latent_filename}.pt" - embed_filename = f"{embed_filename}.pt" + latent_filename = latents_dir.joinpath(f"{filename_without_ext}.pt") + embed_filename = embeds_dir.joinpath(f"{filename_without_ext}.pt") torch.save(latent, latent_filename) torch.save(embed, embed_filename) with open(output_dir.joinpath("videos.txt").as_posix(), "w", encoding="utf-8") as file: for video_path in video_paths: - file.write(f"videos/{pathlib.Path(video_path).name}\n") + file.write(f"videos/{video_path.name}\n") with open(output_dir.joinpath("prompts.txt").as_posix(), "w", encoding="utf-8") as file: for prompt in prompts: @@ -349,7 +345,23 @@ def save_latents_and_embeddings( @torch.no_grad() -def main(args: Dict[str, Any]) -> None: +def main(): + args = get_args() + + # Initialize distributed processing + if "LOCAL_RANK" in os.environ: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + world_size = dist.get_world_size() + rank = dist.get_rank() + else: + # Single GPU + local_rank = 0 + world_size = 1 + rank = 0 + torch.cuda.set_device(local_rank) + data_root = pathlib.Path(args.data_root) dataset_file = None if args.dataset_file: @@ -367,10 +379,18 @@ def main(args: Dict[str, Any]) -> None: ] ) + # Preprocess videos with progress bar prompts_usable = [] video_paths_usable = [] videos = [] - for prompt, path in zip(prompts, video_paths): + + # Only show progress bar on the main process + if rank == 0: + iterator = tqdm(zip(prompts, video_paths), total=len(prompts), desc="Load and Preprocess videos") + else: + iterator = zip(prompts, video_paths) + + for prompt, path in iterator: video = load_and_preprocess_video( path, args.height, args.width, args.max_num_frames, video_transforms, args.num_decode_threads ) @@ -378,18 +398,47 @@ def main(args: Dict[str, Any]) -> None: prompts_usable.append(prompt) video_paths_usable.append(path) videos.append(video) + + if len(videos) == 0: + logger.error("No usable videos found after preprocessing.") + return + videos = torch.stack(videos) + # Split data among GPUs + if world_size > 1: + total_samples = len(prompts_usable) + samples_per_gpu = total_samples // world_size + start_index = rank * samples_per_gpu + end_index = start_index + samples_per_gpu + if rank == world_size - 1: + end_index = total_samples # Make sure the last GPU gets the remaining data + + # Slice the data + prompts_usable = prompts_usable[start_index:end_index] + video_paths_usable = video_paths_usable[start_index:end_index] + videos = videos[start_index:end_index] + else: + pass + + device = torch.device(f"cuda:{local_rank}") + if not args.save_tensors: save_videos(videos, video_paths_usable, prompts_usable, pathlib.Path(args.output_dir), args.target_fps) else: dtype = DTYPE_MAPPING[args.dtype] tokenizer = T5Tokenizer.from_pretrained(args.model_id, subfolder="tokenizer") text_encoder = T5EncoderModel.from_pretrained(args.model_id, subfolder="text_encoder", torch_dtype=dtype) - text_encoder = text_encoder.to("cuda") + text_encoder = text_encoder.to(device) prompt_embeds_list = [] - for start_index in range(0, len(prompts_usable), args.batch_size): + + if rank == 0: + iterator = tqdm(range(0, len(prompts_usable), args.batch_size), desc="Encoding prompts") + else: + iterator = range(0, len(prompts_usable), args.batch_size) + + for start_index in iterator: end_index = min(len(prompts_usable), start_index + args.batch_size) batch_prompts = prompts_usable[start_index:end_index] @@ -398,32 +447,40 @@ def main(args: Dict[str, Any]) -> None: text_encoder, batch_prompts, max_sequence_length=args.max_sequence_length, - device="cuda", + device=device, dtype=dtype, ) - prompt_embeds_list.append(prompt_embeds) + prompt_embeds_list.append(prompt_embeds.to("cpu")) - prompt_embeds = torch.cat(prompt_embeds_list).to("cpu") + prompt_embeds = None + if len(prompt_embeds_list) > 0: + prompt_embeds = torch.cat(prompt_embeds_list) del tokenizer, text_encoder gc.collect() torch.cuda.empty_cache() - torch.cuda.synchronize("cuda") + torch.cuda.synchronize(device) vae = AutoencoderKLCogVideoX.from_pretrained(args.model_id, subfolder="vae", torch_dtype=dtype) - vae = vae.to("cuda") + vae = vae.to(device) if args.use_slicing: vae.enable_slicing() if args.use_tiling: vae.enable_tiling() - encoded_videos = [] - for start_index in range(0, len(video_paths_usable), args.batch_size): + encoded_videos_list = [] + + if rank == 0: + iterator = tqdm(range(0, len(video_paths_usable), args.batch_size), desc="Encoding videos") + else: + iterator = range(0, len(video_paths_usable), args.batch_size) + + for start_index in iterator: end_index = min(len(video_paths_usable), start_index + args.batch_size) batch_videos = videos[start_index:end_index] - batch_videos = batch_videos.to("cuda") + batch_videos = batch_videos.to(device) batch_videos = batch_videos.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] if args.use_slicing: @@ -432,18 +489,31 @@ def main(args: Dict[str, Any]) -> None: else: encoded_video = vae._encode(batch_videos) - encoded_videos.append(encoded_video) + encoded_videos_list.append(encoded_video.to("cpu")) - encoded_videos = torch.cat(encoded_videos).to("cpu") + encoded_videos = None + if len(encoded_videos_list) > 0: + encoded_videos = torch.cat(encoded_videos_list) del vae gc.collect() torch.cuda.empty_cache() - torch.cuda.synchronize("cuda") + torch.cuda.synchronize(device) - save_latents_and_embeddings( - encoded_videos, prompt_embeds, video_paths_usable, prompts_usable, pathlib.Path(args.output_dir) - ) + # Ensure that only one process creates the output directories + if world_size > 1: + dist.barrier() + + if prompt_embeds is not None: + assert encoded_videos is not None + save_latents_and_embeddings( + encoded_videos, prompt_embeds, video_paths_usable, prompts_usable, pathlib.Path(args.output_dir) + ) + + # Finalize distributed processing + if world_size > 1: + dist.barrier() + dist.destroy_process_group() if __name__ == "__main__": @@ -455,4 +525,4 @@ if __name__ == "__main__": args.max_num_frames % 4 == 0 or args.max_num_frames % 4 == 1 ), "`--max_num_frames` must be of form 4 * k or 4 * k + 1 to be compatible with VAE." - main(args) + main()