From 66189c804bbb6daae297721006bd5bc89ee8a346 Mon Sep 17 00:00:00 2001 From: Aryan Date: Fri, 20 Dec 2024 10:03:39 +0100 Subject: [PATCH] update --- README.md | 51 ++++ finetrainers/args.py | 43 +++ finetrainers/constants.py | 4 + finetrainers/dataset.py | 30 ++ finetrainers/ltx_video/ltx_video_lora.py | 87 ++++-- finetrainers/trainer.py | 372 ++++++++++++++++++----- finetrainers/utils/data_utils.py | 35 +++ finetrainers/utils/torch_utils.py | 21 ++ train.py | 1 + 9 files changed, 550 insertions(+), 94 deletions(-) create mode 100644 finetrainers/utils/data_utils.py diff --git a/README.md b/README.md index f1eee8c..2fe6216 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,57 @@ video = pipe("").frames[0] export_to_video(video, "output.mp4", fps=8) ``` +### Memory Usage + +LoRA with rank 128, batch size 1, gradient checkpointing, optimizer adamw, `49x512x768` resolution, **without precomputation**: + +``` +Memory before training start: { + "memory_allocated": 13.486, + "memory_reserved": 13.879, + "max_memory_allocated": 13.486, + "max_memory_reserved": 13.879 +} +Training configuration: { + "trainable parameters": 117440512, + "total samples": 69, + "train epochs": 1, + "train steps": 10, + "batches per device": 1, + "total batches observed per epoch": 69, + "train batch size": 1, + "gradient accumulation steps": 1 +} +Memory before validation start: { + "memory_allocated": 14.146, + "memory_reserved": 16.809, + "max_memory_allocated": 15.527, + "max_memory_reserved": 17.623 +} +Memory after validation end: { + "memory_allocated": 14.146, + "memory_reserved": 14.627, + "max_memory_allocated": 15.527, + "max_memory_reserved": 17.623 +} +Memory after epoch 1: { + "memory_allocated": 14.146, + "memory_reserved": 14.627, + "max_memory_allocated": 15.527, + "max_memory_reserved": 17.623 +} +Memory after training end: { + "memory_allocated": 4.461, + "memory_reserved": 5.014, + "max_memory_allocated": 15.527, + "max_memory_reserved": 17.623 +} +``` + +LoRA with rank 128, batch size 1, gradient checkpointing, optimizer adamw, `49x512x768` resolution, **with precomputation**: + +TODO +
diff --git a/finetrainers/args.py b/finetrainers/args.py index c172028..193f2a0 100644 --- a/finetrainers/args.py +++ b/finetrainers/args.py @@ -1,6 +1,8 @@ import argparse from typing import Any, Dict, List, Optional, Tuple +import torch + from .constants import DEFAULT_IMAGE_RESOLUTION_BUCKETS, DEFAULT_VIDEO_RESOLUTION_BUCKETS @@ -20,6 +22,12 @@ class Args: revision: Optional[str] = None variant: Optional[str] = None cache_dir: Optional[str] = None + text_encoder_dtype: torch.dtype = torch.bfloat16 + text_encoder_2_dtype: torch.dtype = torch.bfloat16 + text_encoder_3_dtype: torch.dtype = torch.bfloat16 + transformer_dtype: torch.dtype = torch.bfloat16 + unet_dtype: torch.dtype = torch.bfloat16 + vae_dtype: torch.dtype = torch.bfloat16 # Dataset arguments data_root: str = None @@ -32,6 +40,7 @@ class Args: video_reshape_mode: Optional[str] = None caption_dropout_p: float = 0.00 caption_dropout_technique: str = "empty" + precompute_conditions: bool = False # Dataloader arguments dataloader_num_workers: int = 0 @@ -113,6 +122,12 @@ class Args: "revision": self.revision, "variant": self.variant, "cache_dir": self.cache_dir, + "text_encoder_dtype": self.text_encoder_dtype, + "text_encoder_2_dtype": self.text_encoder_2_dtype, + "text_encoder_3_dtype": self.text_encoder_3_dtype, + "transformer_dtype": self.transformer_dtype, + "unet_dtype": self.unet_dtype, + "vae_dtype": self.vae_dtype, }, "dataset_arguments": { "data_root": self.data_root, @@ -124,6 +139,8 @@ class Args: "video_resolution_buckets": self.video_resolution_buckets, "video_reshape_mode": self.video_reshape_mode, "caption_dropout_p": self.caption_dropout_p, + "caption_dropout_technique": self.caption_dropout_technique, + "precompute_conditions": self.precompute_conditions, }, "dataloader_arguments": { "dataloader_num_workers": self.dataloader_num_workers, @@ -234,6 +251,12 @@ def _add_model_arguments(parser: argparse.ArgumentParser) -> None: default=None, help="The directory where the downloaded models and datasets will be stored.", ) + parser.add_argument("--text_encoder_dtype", type=str, default="bf16", help="Data type for the text encoder.") + parser.add_argument("--text_encoder_2_dtype", type=str, default="bf16", help="Data type for the text encoder 2.") + parser.add_argument("--text_encoder_3_dtype", type=str, default="bf16", help="Data type for the text encoder 3.") + parser.add_argument("--transformer_dtype", type=str, default="bf16", help="Data type for the transformer model.") + parser.add_argument("--unet_dtype", type=str, default="bf16", help="Data type for the U-Net model.") + parser.add_argument("--vae_dtype", type=str, default="bf16", help="Data type for the VAE model.") def _add_dataset_arguments(parser: argparse.ArgumentParser) -> None: @@ -317,6 +340,11 @@ def _add_dataset_arguments(parser: argparse.ArgumentParser) -> None: choices=["empty", "zero"], help="Technique to use for caption dropout.", ) + parser.add_argument( + "--precompute_conditions", + action="store_true", + help="Whether or not to precompute the conditionings for the model.", + ) def _add_dataloader_arguments(parser: argparse.ArgumentParser) -> None: @@ -645,6 +673,13 @@ def _add_miscellaneous_arguments(parser: argparse.ArgumentParser) -> None: ) +_DTYPE_MAP = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + def _map_to_args_type(args: Dict[str, Any]) -> Args: result_args = Args() @@ -654,6 +689,12 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args: result_args.revision = args.revision result_args.variant = args.variant result_args.cache_dir = args.cache_dir + result_args.text_encoder_dtype = _DTYPE_MAP[args.text_encoder_dtype] + result_args.text_encoder_2_dtype = _DTYPE_MAP[args.text_encoder_2_dtype] + result_args.text_encoder_3_dtype = _DTYPE_MAP[args.text_encoder_3_dtype] + result_args.transformer_dtype = _DTYPE_MAP[args.transformer_dtype] + result_args.unet_dtype = _DTYPE_MAP[args.unet_dtype] + result_args.vae_dtype = _DTYPE_MAP[args.vae_dtype] # Dataset arguments if args.data_root is None and args.dataset_file is None: @@ -668,6 +709,8 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args: result_args.video_resolution_buckets = args.video_resolution_buckets or DEFAULT_VIDEO_RESOLUTION_BUCKETS result_args.video_reshape_mode = args.video_reshape_mode result_args.caption_dropout_p = args.caption_dropout_p + result_args.caption_dropout_technique = args.caption_dropout_technique + result_args.precompute_conditions = args.precompute_conditions # Dataloader arguments result_args.dataloader_num_workers = args.dataloader_num_workers diff --git a/finetrainers/constants.py b/finetrainers/constants.py index bc050f6..26d64ff 100644 --- a/finetrainers/constants.py +++ b/finetrainers/constants.py @@ -19,6 +19,10 @@ for frames in DEFAULT_FRAME_BUCKETS: FINETRAINERS_LOG_LEVEL = os.environ.get("FINETRAINERS_LOG_LEVEL", "INFO") +PRECOMPUTED_DIR_NAME = "precomputed" +PRECOMPUTED_CONDITIONS_DIR_NAME = "conditions" +PRECOMPUTED_LATENTS_DIR_NAME = "latents" + MODEL_DESCRIPTION = r""" \# {model_id} {training_type} finetune diff --git a/finetrainers/dataset.py b/finetrainers/dataset.py index 8ca0bb6..7b016e0 100644 --- a/finetrainers/dataset.py +++ b/finetrainers/dataset.py @@ -1,3 +1,4 @@ +import os import random from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -19,6 +20,9 @@ import decord # isort:skip decord.bridge.set_bridge("torch") +from .constants import PRECOMPUTED_DIR_NAME, PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME + + logger = get_logger(__name__) @@ -257,6 +261,32 @@ class VideoDatasetWithResizeAndRectangleCrop(VideoDataset): return nearest_res[1], nearest_res[2] +class PrecomputedDataset(Dataset): + def __init__(self, data_root: str) -> None: + super().__init__() + + self.data_root = Path(data_root) + + self.latents_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME + self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME + + self.latent_conditions = sorted(os.listdir(self.latents_path)) + self.other_conditions = sorted(os.listdir(self.conditions_path)) + + assert len(self.latent_conditions) == len(self.other_conditions), "Number of captions and videos do not match" + + def __len__(self) -> int: + return len(self.latent_conditions) + + def __getitem__(self, index: int) -> Dict[str, Any]: + conditions = {} + latent_path = self.latents_path / self.latent_conditions[index] + condition_path = self.conditions_path / self.other_conditions[index] + conditions["latent_conditions"] = torch.load(latent_path, map_location="cpu", weights_only=True) + conditions["other_conditions"] = torch.load(condition_path, map_location="cpu", weights_only=True) + return conditions + + class BucketSampler(Sampler): r""" PyTorch Sampler that groups 3D data by height, width and frames. diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index f8109f9..e767f21 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -12,32 +12,45 @@ from PIL import Image logger = get_logger("finetrainers") # pylint: disable=invalid-name -def load_components( +def load_condition_models( model_id: str = "Lightricks/LTX-Video", text_encoder_dtype: torch.dtype = torch.bfloat16, - transformer_dtype: torch.dtype = torch.bfloat16, - vae_dtype: torch.dtype = torch.bfloat16, revision: Optional[str] = None, cache_dir: Optional[str] = None, + **kwargs, ) -> Dict[str, nn.Module]: tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir) text_encoder = T5EncoderModel.from_pretrained( model_id, subfolder="text_encoder", torch_dtype=text_encoder_dtype, revision=revision, cache_dir=cache_dir ) - transformer = LTXVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir - ) + return {"tokenizer": tokenizer, "text_encoder": text_encoder} + + +def load_latent_models( + model_id: str = "Lightricks/LTX-Video", + vae_dtype: torch.dtype = torch.bfloat16, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + **kwargs, +) -> Dict[str, nn.Module]: vae = AutoencoderKLLTXVideo.from_pretrained( model_id, subfolder="vae", torch_dtype=vae_dtype, revision=revision, cache_dir=cache_dir ) + return {"vae": vae} + + +def load_diffusion_models( + model_id: str = "Lightricks/LTX-Video", + transformer_dtype: torch.dtype = torch.bfloat16, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + **kwargs, +) -> Dict[str, nn.Module]: + transformer = LTXVideoTransformer3DModel.from_pretrained( + model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + ) scheduler = FlowMatchEulerDiscreteScheduler() - return { - "tokenizer": tokenizer, - "text_encoder": text_encoder, - "transformer": transformer, - "vae": vae, - "scheduler": scheduler, - } + return {"transformer": transformer, "scheduler": scheduler} def initialize_pipeline( @@ -114,19 +127,52 @@ def prepare_latents( device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, generator: Optional[torch.Generator] = None, + precompute: bool = False, ) -> torch.Tensor: device = device or vae.device - dtype = dtype or vae.dtype if image_or_video.ndim == 4: image_or_video = image_or_video.unsqueeze(2) assert image_or_video.ndim == 5, f"Expected 5D tensor, got {image_or_video.ndim}D tensor" - image_or_video = image_or_video.to(device=device, dtype=dtype) + image_or_video = image_or_video.to(device=device, dtype=vae.dtype) image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous() # [B, C, F, H, W] -> [B, F, C, H, W] - latents = vae.encode(image_or_video).latent_dist.sample(generator=generator) - _, _, num_frames, height, width = latents.shape - latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std) + if not precompute: + latents = vae.encode(image_or_video).latent_dist.sample(generator=generator) + latents = latents.to(dtype=dtype) + _, _, num_frames, height, width = latents.shape + latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std) + latents = _pack_latents(latents, patch_size, patch_size_t) + return {"latents": latents, "num_frames": num_frames, "height": height, "width": width} + else: + if vae.use_slicing and image_or_video.shape[0] > 1: + encoded_slices = [vae._encode(x_slice) for x_slice in image_or_video.split(1)] + h = torch.cat(encoded_slices) + else: + h = vae._encode(image_or_video) + _, _, num_frames, height, width = h.shape + # TODO(aryan): this is very very very stupid, but anything to make it work for now. refactor and design better later + return { + "latents": h, + "num_frames": num_frames, + "height": height, + "width": width, + "latents_mean": vae.latents_mean, + "latents_std": vae.latents_std, + } + + +def post_latent_preparation( + latents: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + latents = _normalize_latents(latents, latents_mean, latents_std) latents = _pack_latents(latents, patch_size, patch_size_t) return {"latents": latents, "num_frames": num_frames, "height": height, "width": width} @@ -260,10 +306,13 @@ def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int LTX_VIDEO_T2V_LORA_CONFIG = { "pipeline_cls": LTXPipeline, - "load_components": load_components, + "load_condition_models": load_condition_models, + "load_latent_models": load_latent_models, + "load_diffusion_models": load_diffusion_models, "initialize_pipeline": initialize_pipeline, "prepare_conditions": prepare_conditions, "prepare_latents": prepare_latents, + "post_latent_preparation": post_latent_preparation, "collate_fn": collate_fn_t2v, "forward_pass": forward_pass, "validation": validation, diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 71e4d60..85c740c 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -29,20 +29,27 @@ from diffusers.training_utils import ( compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3, ) +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.utils import export_to_video, load_image, load_video from huggingface_hub import create_repo, upload_folder from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict from tqdm import tqdm from .args import Args, validate_args -from .constants import FINETRAINERS_LOG_LEVEL -from .dataset import BucketSampler, VideoDatasetWithResizing +from .constants import ( + FINETRAINERS_LOG_LEVEL, + PRECOMPUTED_DIR_NAME, + PRECOMPUTED_CONDITIONS_DIR_NAME, + PRECOMPUTED_LATENTS_DIR_NAME, +) +from .dataset import BucketSampler, PrecomputedDataset, VideoDatasetWithResizing from .models import get_config_from_model_name from .state import State +from .utils.data_utils import should_perform_precomputation from .utils.file_utils import find_files, delete_files, string_to_filename from .utils.optimizer_utils import get_optimizer, gradient_norm from .utils.memory_utils import get_memory_statistics, free_memory, make_contiguous -from .utils.torch_utils import unwrap_model +from .utils.torch_utils import unwrap_model, align_device_and_dtype logger = get_logger("finetrainers") @@ -73,6 +80,9 @@ class Trainer: # Autoencoders self.vae = None + # Scheduler + self.scheduler = None + self._init_distributed() self._init_logging() self._init_directories_and_repositories() @@ -80,38 +90,8 @@ class Trainer: self.state.model_name = self.args.model_name self.model_config = get_config_from_model_name(self.args.model_name, self.args.training_type) - def prepare_models(self) -> None: - logger.info("Initializing models") - - # TODO(aryan): refactor in future - load_components_kwargs = { - "text_encoder_dtype": torch.bfloat16, - "transformer_dtype": torch.bfloat16, - "vae_dtype": torch.bfloat16, - "revision": self.args.revision, - "cache_dir": self.args.cache_dir, - } - if self.args.pretrained_model_name_or_path is not None: - load_components_kwargs["model_id"] = self.args.pretrained_model_name_or_path - components = self._model_config_call(self.model_config["load_components"], load_components_kwargs) - - self.tokenizer = components.get("tokenizer", None) - self.text_encoder = components.get("text_encoder", None) - self.tokenizer_2 = components.get("tokenizer_2", None) - self.text_encoder_2 = components.get("text_encoder_2", None) - self.transformer = components.get("transformer", None) - self.vae = components.get("vae", None) - self.scheduler = components.get("scheduler", None) - - if self.vae is not None: - if self.args.enable_slicing: - self.vae.enable_slicing() - if self.args.enable_tiling: - self.vae.enable_tiling() - - self.transformer_config = self.transformer.config if self.transformer is not None else None - def prepare_dataset(self) -> None: + # TODO(aryan): Make a background process for fetching logger.info("Initializing dataset and dataloader") self.dataset = VideoDatasetWithResizing( @@ -131,16 +111,238 @@ class Trainer: pin_memory=self.args.pin_memory, ) + def _get_load_components_kwargs(self) -> Dict[str, Any]: + load_component_kwargs = { + "text_encoder_dtype": self.args.text_encoder_dtype, + "text_encoder_2_dtype": self.args.text_encoder_2_dtype, + "text_encoder_3_dtype": self.args.text_encoder_3_dtype, + "transformer_dtype": self.args.transformer_dtype, + "unet_dtype": self.args.unet_dtype, + "vae_dtype": self.args.vae_dtype, + "revision": self.args.revision, + "cache_dir": self.args.cache_dir, + } + if self.args.pretrained_model_name_or_path is not None: + load_component_kwargs["model_id"] = self.args.pretrained_model_name_or_path + return load_component_kwargs + + def _set_components(self, components: Dict[str, Any]) -> None: + self.tokenizer = components.get("tokenizer", self.tokenizer) + self.tokenizer_2 = components.get("tokenizer_2", self.tokenizer_2) + self.tokenizer_3 = components.get("tokenizer_3", self.tokenizer_3) + self.text_encoder = components.get("text_encoder", self.text_encoder) + self.text_encoder_2 = components.get("text_encoder_2", self.text_encoder_2) + self.text_encoder_3 = components.get("text_encoder_3", self.text_encoder_3) + self.transformer = components.get("transformer", self.transformer) + self.unet = components.get("unet", self.unet) + self.vae = components.get("vae", self.vae) + self.scheduler = components.get("scheduler", self.scheduler) + + def _nuke_components(self) -> None: + self.tokenizer = None + self.tokenizer_2 = None + self.tokenizer_3 = None + self.text_encoder = None + self.text_encoder_2 = None + self.text_encoder_3 = None + self.transformer = None + self.unet = None + self.vae = None + self.scheduler = None + free_memory() + torch.cuda.synchronize(self.state.accelerator.device) + + def prepare_models(self) -> None: + logger.info("Initializing models") + + load_components_kwargs = self._get_load_components_kwargs() + condition_components, latent_components, diffusion_components = {}, {}, {} + if not self.args.precompute_conditions: + condition_components = self.model_config["load_condition_models"](**load_components_kwargs) + latent_components = self.model_config["load_latent_models"](**load_components_kwargs) + diffusion_components = self.model_config["load_diffusion_models"](**load_components_kwargs) + + components = {} + components.update(condition_components) + components.update(latent_components) + components.update(diffusion_components) + self._set_components(components) + + if self.vae is not None: + if self.args.enable_slicing: + self.vae.enable_slicing() + if self.args.enable_tiling: + self.vae.enable_tiling() + + self.transformer_config = self.transformer.config if self.transformer is not None else None + + def prepare_precomputations(self) -> None: + if not self.args.precompute_conditions: + return + + logger.info("Initializing precomputations") + + if self.args.batch_size != 1: + raise ValueError("Precomputation is only supported with batch size 1. This will be supported in future.") + + def collate_fn(batch): + latent_conditions = [x["latent_conditions"] for x in batch] + other_conditions = [x["other_conditions"] for x in batch] + batched_latent_conditions = {} + batched_other_conditions = {} + for key in list(latent_conditions[0].keys()): + if torch.is_tensor(latent_conditions[0][key]): + batched_latent_conditions[key] = torch.cat([x[key] for x in latent_conditions], dim=0) + else: + # TODO(aryan): implement batch sampler for precomputed latents + batched_latent_conditions[key] = [x[key] for x in latent_conditions][0] + for key in list(other_conditions[0].keys()): + if torch.is_tensor(other_conditions[0][key]): + batched_other_conditions[key] = torch.cat([x[key] for x in other_conditions], dim=0) + else: + # TODO(aryan): implement batch sampler for precomputed latents + batched_other_conditions[key] = [x[key] for x in other_conditions][0] + return {"latent_conditions": batched_latent_conditions, "other_conditions": batched_other_conditions} + + should_recompute = should_perform_precomputation(self.args.data_root) + if not should_recompute: + logger.info("Precomputed conditions and latents found. Loading precomputed data.") + self.dataloader = torch.utils.data.DataLoader( + PrecomputedDataset(self.args.data_root), + batch_size=self.args.batch_size, + shuffle=True, + collate_fn=collate_fn, + num_workers=self.args.dataloader_num_workers, + pin_memory=self.args.pin_memory, + ) + return + + logger.info("Precomputed conditions and latents not found. Running precomputation.") + + # At this point, no models are loaded, so we need to load and precompute conditions and latents + condition_components = self.model_config["load_condition_models"](**self._get_load_components_kwargs()) + self._set_components(condition_components) + self._move_components_to_device() + + if self.args.caption_dropout_p > 0 and self.args.caption_dropout_technique == "empty": + logger.warning( + "Caption dropout is not supported with precomputation yet. This will be supported in the future." + ) + + conditions_dir = Path(self.args.data_root) / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME + latents_dir = Path(self.args.data_root) / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME + conditions_dir.mkdir(parents=True, exist_ok=True) + latents_dir.mkdir(parents=True, exist_ok=True) + + # Precompute conditions + progress_bar = tqdm( + range(0, len(self.dataset)), + desc="Precomputing conditions", + disable=not self.state.accelerator.is_local_main_process, + ) + index = 0 + for i, data in enumerate(self.dataset): + if i % self.state.accelerator.num_processes != self.state.accelerator.process_index: + continue + + logger.debug( + f"Precomputing conditions and latents for batch {i + 1}/{len(self.dataset)} on process {self.state.accelerator.process_index}" + ) + + other_conditions = self.model_config["prepare_conditions"]( + tokenizer=self.tokenizer, + tokenizer_2=self.tokenizer_2, + tokenizer_3=self.tokenizer_3, + text_encoder=self.text_encoder, + text_encoder_2=self.text_encoder_2, + text_encoder_3=self.text_encoder_3, + prompt=data["prompt"], + device=self.state.accelerator.device, + dtype=self.state.weight_dtype, + ) + filename = conditions_dir / f"conditions-{i}-{index}.pt" + torch.save(other_conditions, filename.as_posix()) + index += 1 + progress_bar.update(1) + self._nuke_components() + + memory_statistics = get_memory_statistics() + logger.info(f"Memory after precomputing conditions: {json.dumps(memory_statistics, indent=4)}") + torch.cuda.reset_peak_memory_stats(self.state.accelerator.device) + + # Precompute latents + latent_components = self.model_config["load_latent_models"](**self._get_load_components_kwargs()) + self._set_components(latent_components) + self._move_components_to_device() + + if self.vae is not None: + if self.args.enable_slicing: + self.vae.enable_slicing() + if self.args.enable_tiling: + self.vae.enable_tiling() + + progress_bar = tqdm( + range(0, len(self.dataset)), + desc="Precomputing latents", + disable=not self.state.accelerator.is_local_main_process, + ) + index = 0 + for i, data in enumerate(self.dataset): + if i % self.state.accelerator.num_processes != self.state.accelerator.process_index: + continue + + logger.debug( + f"Precomputing latents for batch {i + 1}/{len(self.dataset)} on process {self.state.accelerator.process_index}" + ) + + latent_conditions = self.model_config["prepare_latents"]( + vae=self.vae, + image_or_video=data["video"].unsqueeze(0), + device=self.state.accelerator.device, + dtype=self.state.weight_dtype, + generator=self.state.generator, + precompute=True, + ) + filename = latents_dir / f"latents-{i}-{index}.pt" + torch.save(latent_conditions, filename.as_posix()) + index += 1 + progress_bar.update(1) + self._nuke_components() + + self.state.accelerator.wait_for_everyone() + logger.info("Precomputation complete") + + memory_statistics = get_memory_statistics() + logger.info(f"Memory after precomputing latents: {json.dumps(memory_statistics, indent=4)}") + torch.cuda.reset_peak_memory_stats(self.state.accelerator.device) + + # Update dataloader to use precomputed conditions and latents + self.dataloader = torch.utils.data.DataLoader( + PrecomputedDataset(self.args.data_root), + batch_size=self.args.batch_size, + shuffle=True, + collate_fn=collate_fn, + num_workers=self.args.dataloader_num_workers, + pin_memory=self.args.pin_memory, + ) + def prepare_trainable_parameters(self) -> None: logger.info("Initializing trainable parameters") - # TODO(aryan): refactor later. for now only lora is supported - self.text_encoder.requires_grad_(False) - self.transformer.requires_grad_(False) - self.vae.requires_grad_(False) + diffusion_components = self.model_config["load_diffusion_models"](**self._get_load_components_kwargs()) + self._set_components(diffusion_components) - if self.text_encoder_2 is not None: - self.text_encoder_2.requires_grad_(False) + # TODO(aryan): refactor later. for now only lora is supported + components_to_disable_grads = [ + self.text_encoder, + self.text_encoder_2, + self.text_encoder_3, + self.transformer, + self.vae, + ] + for component in components_to_disable_grads: + if component is not None: + component.requires_grad_(False) # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision # as these weights are only used for inference, keeping weights in full precision is not required. @@ -158,12 +360,8 @@ class Trainer: # TODO(aryan): handle torch dtype from accelerator vs model dtype; refactor self.state.weight_dtype = weight_dtype - self.text_encoder.to(self.state.accelerator.device, dtype=weight_dtype) - self.transformer.to(self.state.accelerator.device, dtype=weight_dtype) - self.vae.to(self.state.accelerator.device, dtype=weight_dtype) - - if self.text_encoder_2 is not None: - self.text_encoder_2.to(self.state.accelerator.device, dtype=weight_dtype) + self.transformer.to(dtype=weight_dtype) + self._move_components_to_device() if self.args.gradient_checkpointing: self.transformer.enable_gradient_checkpointing() @@ -364,34 +562,46 @@ class Trainer: logs = {} with accelerator.accumulate(models_to_accumulate): - videos = batch["videos"] - prompts = batch["prompts"] - batch_size = len(prompts) + if not self.args.precompute_conditions: + videos = batch["videos"] + prompts = batch["prompts"] + batch_size = len(prompts) - if self.args.caption_dropout_technique == "empty": - if random.random() < self.args.caption_dropout_p: - prompts = [""] * batch_size + if self.args.caption_dropout_technique == "empty": + if random.random() < self.args.caption_dropout_p: + prompts = [""] * batch_size + + latent_conditions = self.model_config["prepare_latents"]( + vae=self.vae, + image_or_video=videos, + patch_size=self.transformer_config.patch_size, + patch_size_t=self.transformer_config.patch_size_t, + device=accelerator.device, + dtype=weight_dtype, + generator=generator, + ) + other_conditions = self.model_config["prepare_conditions"]( + tokenizer=self.tokenizer, + text_encoder=self.text_encoder, + tokenizer_2=self.tokenizer_2, + text_encoder_2=self.text_encoder_2, + prompt=prompts, + device=accelerator.device, + dtype=weight_dtype, + ) + else: + latent_conditions = batch["latent_conditions"] + other_conditions = batch["other_conditions"] + latent_conditions["latents"] = DiagonalGaussianDistribution( + latent_conditions["latents"] + ).sample(generator) + if "post_latent_preparation" in self.model_config.keys(): + latent_conditions = self.model_config["post_latent_preparation"](**latent_conditions) + align_device_and_dtype(latent_conditions, accelerator.device, weight_dtype) + align_device_and_dtype(other_conditions, accelerator.device, weight_dtype) + batch_size = latent_conditions["latents"].shape[0] - latent_conditions = self.model_config["prepare_latents"]( - vae=self.vae, - image_or_video=videos, - patch_size=self.transformer_config.patch_size, - patch_size_t=self.transformer_config.patch_size_t, - device=accelerator.device, - dtype=weight_dtype, - generator=generator, - ) latent_conditions = make_contiguous(latent_conditions) - - other_conditions = self.model_config["prepare_conditions"]( - tokenizer=self.tokenizer, - text_encoder=self.text_encoder, - tokenizer_2=self.tokenizer_2, - text_encoder_2=self.text_encoder_2, - prompt=prompts, - device=accelerator.device, - dtype=weight_dtype, - ) other_conditions = make_contiguous(other_conditions) if self.args.caption_dropout_technique == "zero": @@ -629,9 +839,12 @@ class Trainer: tracker.log({"validation": all_artifacts}, step=step) accelerator.wait_for_everyone() + free_memory() memory_statistics = get_memory_statistics() logger.info(f"Memory after validation end: {json.dumps(memory_statistics, indent=4)}") + torch.cuda.reset_peak_memory_stats(accelerator.device) + self.transformer.train() def evaluate(self) -> None: @@ -692,7 +905,16 @@ class Trainer: repo_id = self.args.hub_model_id or Path(self.args.output_dir).name self.state.repo_id = create_repo(token=self.args.hub_token, name=repo_id).repo_id - def _model_config_call(self, fn, kwargs): - accepted_kwargs = inspect.signature(fn).parameters.keys() - kwargs = {k: v for k, v in kwargs.items() if k in accepted_kwargs} - return fn(**kwargs) + def _move_components_to_device(self): + if self.text_encoder is not None: + self.text_encoder = self.text_encoder.to(self.state.accelerator.device) + if self.text_encoder_2 is not None: + self.text_encoder_2 = self.text_encoder_2.to(self.state.accelerator.device) + if self.text_encoder_3 is not None: + self.text_encoder_3 = self.text_encoder_3.to(self.state.accelerator.device) + if self.transformer is not None: + self.transformer = self.transformer.to(self.state.accelerator.device) + if self.unet is not None: + self.unet = self.unet.to(self.state.accelerator.device) + if self.vae is not None: + self.vae = self.vae.to(self.state.accelerator.device) diff --git a/finetrainers/utils/data_utils.py b/finetrainers/utils/data_utils.py new file mode 100644 index 0000000..b05c49e --- /dev/null +++ b/finetrainers/utils/data_utils.py @@ -0,0 +1,35 @@ +from pathlib import Path +from typing import Union + +from accelerate.logging import get_logger + +from ..constants import PRECOMPUTED_DIR_NAME, PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME + + +logger = get_logger("finetrainers") + + +def should_perform_precomputation(data_root: Union[str, Path]) -> bool: + if isinstance(data_root, str): + data_root = Path(data_root) + conditions_dir = data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME + latents_dir = data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME + if conditions_dir.exists() and latents_dir.exists(): + num_files_conditions = len(list(conditions_dir.glob("*.pt"))) + num_files_latents = len(list(latents_dir.glob("*.pt"))) + if num_files_conditions != num_files_latents: + logger.warning( + f"Number of precomputed conditions ({num_files_conditions}) does not match number of precomputed latents ({num_files_latents})." + f"Cleaning up precomputed directories and re-running precomputation." + ) + # clean up precomputed directories + for file in conditions_dir.glob("*.pt"): + file.unlink() + for file in latents_dir.glob("*.pt"): + file.unlink() + return True + if num_files_conditions > 0: + logger.info(f"Found {num_files_conditions} precomputed conditions and latents.") + return False + logger.info("Precomputed data not found. Running precomputation.") + return True diff --git a/finetrainers/utils/torch_utils.py b/finetrainers/utils/torch_utils.py index 32190bb..2aad077 100644 --- a/finetrainers/utils/torch_utils.py +++ b/finetrainers/utils/torch_utils.py @@ -1,3 +1,6 @@ +from typing import Dict, Optional, Union + +import torch from accelerate import Accelerator from diffusers.utils.torch_utils import is_compiled_module @@ -6,3 +9,21 @@ def unwrap_model(accelerator: Accelerator, model): model = accelerator.unwrap_model(model) model = model._orig_mod if is_compiled_module(model) else model return model + + +def align_device_and_dtype( + x: Union[torch.Tensor, Dict[str, torch.Tensor]], + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, +): + if isinstance(x, torch.Tensor): + if device is not None: + x = x.to(device) + if dtype is not None: + x = x.to(dtype) + elif isinstance(x, dict): + if device is not None: + x = {k: align_device_and_dtype(v, device, dtype) for k, v in x.items()} + if dtype is not None: + x = {k: align_device_and_dtype(v, device, dtype) for k, v in x.items()} + return x diff --git a/train.py b/train.py index 0a45987..e12d47c 100644 --- a/train.py +++ b/train.py @@ -27,6 +27,7 @@ def main(): trainer.prepare_dataset() trainer.prepare_models() + trainer.prepare_precomputations() trainer.prepare_trainable_parameters() trainer.prepare_optimizer() trainer.prepare_for_training()