diff --git a/README.md b/README.md index 2fe6216..782d625 100644 --- a/README.md +++ b/README.md @@ -192,7 +192,60 @@ Memory after training end: { LoRA with rank 128, batch size 1, gradient checkpointing, optimizer adamw, `49x512x768` resolution, **with precomputation**: -TODO +``` +Memory after precomputing conditions: { + "memory_allocated": 8.88, + "memory_reserved": 8.895, + "max_memory_allocated": 8.897, + "max_memory_reserved": 8.92 +} +Memory after precomputing latents: { + "memory_allocated": 9.684, + "memory_reserved": 9.807, + "max_memory_allocated": 11.155, + "max_memory_reserved": 11.613 +} +Memory before training start: { + "memory_allocated": 3.809, + "memory_reserved": 10.01, + "max_memory_allocated": 9.684, + "max_memory_reserved": 10.01 +} +Training configuration: { + "trainable parameters": 117440512, + "total samples": 1, + "train epochs": 10, + "train steps": 10, + "batches per device": 1, + "total batches observed per epoch": 1, + "train batch size": 1, + "gradient accumulation steps": 1 +} +Memory after epoch 1: { + "memory_allocated": 4.26, + "memory_reserved": 10.916, + "max_memory_allocated": 9.684, + "max_memory_reserved": 10.916 +} +Memory before validation start: { + "memory_allocated": 4.26, + "memory_reserved": 10.916, + "max_memory_allocated": 9.684, + "max_memory_reserved": 10.916 +} +Memory after validation end: { + "memory_allocated": 13.924, + "memory_reserved": 14.209, + "max_memory_allocated": 15.083, + "max_memory_reserved": 17.262 +} +Memory after training end: { + "memory_allocated": 4.26, + "memory_reserved": 4.602, + "max_memory_allocated": 13.923, + "max_memory_reserved": 14.314 +} +``` @@ -327,6 +380,11 @@ export_to_video(output, "output.mp4", fps=15) If you would like to use a custom dataset, refer to the dataset preparation guide [here](./assets/dataset.md). +> [!NOTE] +> To lower memory requirements: +> - Pass `--precompute_conditions` when launching training. +> - Do not perform validation/testing. This saves a significant amount of memory, which can be used to focus solely on training if you're on smaller VRAM GPUs. + ## Memory requirements diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 85c740c..87c8b4c 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -224,6 +224,16 @@ class Trainer: self._set_components(condition_components) self._move_components_to_device() + # 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, + ] + for component in components_to_disable_grads: + if component is not None: + component.requires_grad_(False) + 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." @@ -275,6 +285,12 @@ class Trainer: self._set_components(latent_components) self._move_components_to_device() + # TODO(aryan): refactor later + components_to_disable_grads = [self.vae] + for component in components_to_disable_grads: + if component is not None: + component.requires_grad_(False) + if self.vae is not None: if self.args.enable_slicing: self.vae.enable_slicing()