[WIP][tests] add precomputation tests (#234)

sayakpaul · a-r-r-o-w · web-flow · commit 61d14a7ceb73 · 2025-02-24T09:34:39.000+05:30
* add precomputation tests

* update

* updates

* changes

* quality

* remove do_not_run_validation.

* make get_memory_stat method leaner.

* reset memory utils.

* sync util.

* updates

* typo.

Co-authored-by: a-r-r-o-w &lt;aryan@huggingface.co&gt;

* resolve imports.

* updates

* update ltx

* updates

* updates

* updates

* updates

* fixes

* fixes

---------

Co-authored-by: a-r-r-o-w &lt;aryan@huggingface.co&gt;
diff --git a/Makefile b/Makefile
@@ -3,9 +3,9 @@
 check_dirs := finetrainers tests examples
 
 quality:
-	ruff check $(check_dirs)
-	ruff format --check $(check_dirs)
+	ruff check $(check_dirs) --exclude examples/_legacy
+	ruff format --check $(check_dirs) --exclude examples/_legacy
 
 style:
-	ruff check $(check_dirs) --fix
-	ruff format $(check_dirs)
+	ruff check $(check_dirs) --fix --exclude examples/_legacy
+	ruff format $(check_dirs) --exclude examples/_legacy
diff --git a/examples/_legacy/cogvideox/cogvideox_image_to_video_sft.py b/examples/_legacy/cogvideox/cogvideox_image_to_video_sft.py
@@ -44,7 +44,7 @@
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import cast_training_params
-from diffusers.utils import convert_unet_state_dict_to_peft, export_to_video, load_image
+from diffusers.utils import export_to_video, load_image
 from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from huggingface_hub import create_repo, upload_folder
 from torch.utils.data import DataLoader
diff --git a/examples/_legacy/cogvideox/cogvideox_text_to_video_lora.py b/examples/_legacy/cogvideox/cogvideox_text_to_video_lora.py
@@ -62,7 +62,7 @@
     print_memory,
     reset_memory,
     unwrap_model,
-)  # isort:skip
+)
 
 
 logger = get_logger(__name__)
diff --git a/examples/_legacy/cogvideox/cogvideox_text_to_video_sft.py b/examples/_legacy/cogvideox/cogvideox_text_to_video_sft.py
@@ -61,7 +61,7 @@
     print_memory,
     reset_memory,
     unwrap_model,
-)  # isort:skip
+)
 
 
 logger = get_logger(__name__)
diff --git a/examples/_legacy/cogvideox/dataset.py b/examples/_legacy/cogvideox/dataset.py
@@ -425,4 +425,4 @@ def __iter__(self):
                 random.shuffle(bucket)
                 yield bucket
                 del self.buckets[fhw]
-                self.buckets[fhw] = []
+                self.buckets[fhw] = []
diff --git a/examples/_legacy/mochi-1/args.py b/examples/_legacy/mochi-1/args.py
@@ -1,5 +1,5 @@
 """
-Default values taken from 
+Default values taken from
 https://github.com/genmoai/mochi/blob/aba74c1b5e0755b1fa3343d9e4bd22e89de77ab1/demos/fine_tuner/configs/lora.yaml
 when applicable.
 """
diff --git a/examples/_legacy/mochi-1/dataset_simple.py b/examples/_legacy/mochi-1/dataset_simple.py
@@ -37,7 +37,7 @@ def __getitem__(self, idx):
 def process_videos(directory):
     dir_path = Path(directory)
     mp4_files = [str(f) for f in dir_path.glob("**/*.mp4") if not f.name.endswith(".recon.mp4")]
-    assert mp4_files, f"No mp4 files found"
+    assert mp4_files, "No mp4 files found"
 
     dataset = LatentEmbedDataset(mp4_files)
     dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
diff --git a/examples/_legacy/mochi-1/embed.py b/examples/_legacy/mochi-1/embed.py
@@ -4,13 +4,14 @@
 https://github.com/genmoai/mochi/blob/main/demos/fine_tuner/embed_captions.py
 """
 
+from pathlib import Path
+
 import click
 import torch
 import torchvision
-from pathlib import Path
 from diffusers import AutoencoderKLMochi, MochiPipeline
-from transformers import T5EncoderModel, T5Tokenizer
 from tqdm.auto import tqdm
+from transformers import T5EncoderModel, T5Tokenizer
 
 
 def encode_videos(model: torch.nn.Module, vid_path: Path, shape: str):
diff --git a/examples/_legacy/mochi-1/text_to_video_lora.py b/examples/_legacy/mochi-1/text_to_video_lora.py
@@ -14,16 +14,16 @@
 # limitations under the License.
 
 import gc
-import random
-from glob import glob
 import math
 import os
-import torch.nn.functional as F
-import numpy as np
+import random
+from glob import glob
 from pathlib import Path
-from typing import Any, Dict, Tuple, List
+from typing import Any, Dict, List, Tuple
 
+import numpy as np
 import torch
+import torch.nn.functional as F
 import wandb
 from diffusers import FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
@@ -37,9 +37,10 @@
 
 
 from args import get_args  # isort:skip
+
 from dataset_simple import LatentEmbedDataset
 
-import sys
+
 from utils import print_memory, reset_memory  # isort:skip
 
 
@@ -100,7 +101,7 @@ def save_model_card(
 ```py
 from diffusers import MochiPipeline
 from diffusers.utils import export_to_video
-import torch 
+import torch
 
 pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview")
 pipe.load_lora_weights("CHANGE_ME")
@@ -315,7 +316,7 @@ def main(args):
     optimizer = torch.optim.AdamW(transformer_lora_parameters, lr=args.learning_rate, weight_decay=args.weight_decay)
 
     # Dataset and DataLoader
-    train_vids = list(sorted(glob(f"{args.data_root}/*.mp4")))
+    train_vids = sorted(glob(f"{args.data_root}/*.mp4"))
     train_vids = [v for v in train_vids if not v.endswith(".recon.mp4")]
     print(f"Found {len(train_vids)} training videos in {args.data_root}")
     assert len(train_vids) > 0, f"No training data found in {args.data_root}"
diff --git a/examples/_legacy/mochi-1/trim_and_crop_videos.py b/examples/_legacy/mochi-1/trim_and_crop_videos.py
@@ -3,8 +3,8 @@
 https://github.com/genmoai/mochi/blob/main/demos/fine_tuner/trim_and_crop_videos.py
 """
 
-from pathlib import Path
 import shutil
+from pathlib import Path
 
 import click
 from moviepy.editor import VideoFileClip
diff --git a/examples/_legacy/mochi-1/utils.py b/examples/_legacy/mochi-1/utils.py
@@ -1,9 +1,9 @@
 import gc
-import inspect
-from typing import Optional, Tuple, Union
+from typing import Union
 
 import torch
 
+
 logger = get_logger(__name__)
 
 def reset_memory(device: Union[str, torch.device]) -> None:
diff --git a/finetrainers/models/cogvideox/lora.py b/finetrainers/models/cogvideox/lora.py
@@ -3,7 +3,7 @@
 import torch
 from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel
 from PIL import Image
-from transformers import T5EncoderModel, T5Tokenizer
+from transformers import AutoTokenizer, T5EncoderModel, T5Tokenizer
 
 from .utils import prepare_rotary_positional_embeddings
 
@@ -15,7 +15,14 @@ def load_condition_models(
     cache_dir: Optional[str] = None,
     **kwargs,
 ):
-    tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir)
+    try:
+        tokenizer = T5Tokenizer.from_pretrained(
+            model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir
+        )
+    except:  # noqa
+        tokenizer = AutoTokenizer.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
     )
diff --git a/finetrainers/models/ltx_video/lora.py b/finetrainers/models/ltx_video/lora.py
@@ -5,7 +5,7 @@
 from accelerate.logging import get_logger
 from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
 from PIL import Image
-from transformers import T5EncoderModel, T5Tokenizer
+from transformers import AutoTokenizer, T5EncoderModel, T5Tokenizer
 
 
 logger = get_logger("finetrainers")  # pylint: disable=invalid-name
@@ -18,7 +18,14 @@ def load_condition_models(
     cache_dir: Optional[str] = None,
     **kwargs,
 ) -> Dict[str, nn.Module]:
-    tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir)
+    try:
+        tokenizer = T5Tokenizer.from_pretrained(
+            model_id, subfolder="tokenizer", revision=revision, cache_dir=cache_dir
+        )
+    except:  # noqa
+        tokenizer = AutoTokenizer.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
     )
diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py
@@ -54,7 +54,13 @@
 )
 from .utils.file_utils import string_to_filename
 from .utils.hub_utils import save_model_card
-from .utils.memory_utils import free_memory, get_memory_statistics, make_contiguous
+from .utils.memory_utils import (
+    free_memory,
+    get_memory_statistics,
+    make_contiguous,
+    reset_memory_stats,
+    synchronize_device,
+)
 from .utils.model_utils import resolve_vae_cls_from_ckpt_path
 from .utils.optimizer_utils import get_optimizer
 from .utils.torch_utils import align_device_and_dtype, expand_tensor_dims, unwrap_model
@@ -259,7 +265,7 @@ def collate_fn(batch):
 
         memory_statistics = get_memory_statistics()
         logger.info(f"Memory after precomputing conditions: {json.dumps(memory_statistics, indent=4)}")
-        torch.cuda.reset_peak_memory_stats(accelerator.device)
+        reset_memory_stats(accelerator.device)
 
         # Precompute latents
         with self.state.accelerator.main_process_first():
@@ -307,7 +313,7 @@ def collate_fn(batch):
 
         memory_statistics = get_memory_statistics()
         logger.info(f"Memory after precomputing latents: {json.dumps(memory_statistics, indent=4)}")
-        torch.cuda.reset_peak_memory_stats(accelerator.device)
+        reset_memory_stats(accelerator.device)
 
         # Update dataloader to use precomputed conditions and latents
         self.dataloader = torch.utils.data.DataLoader(
@@ -997,7 +1003,7 @@ def validate(self, step: int, final_validation: bool = False) -> None:
         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)
+        reset_memory_stats(accelerator.device)
 
         if not final_validation:
             self.transformer.train()
@@ -1120,7 +1126,7 @@ def _delete_components(self) -> None:
         self.vae = None
         self.scheduler = None
         free_memory()
-        torch.cuda.synchronize(self.state.accelerator.device)
+        synchronize_device(self.state.accelerator.device)
 
     def _get_and_prepare_pipeline_for_validation(self, final_validation: bool = False) -> DiffusionPipeline:
         accelerator = self.state.accelerator
diff --git a/finetrainers/utils/memory_utils.py b/finetrainers/utils/memory_utils.py
@@ -9,29 +9,33 @@
 
 
 def get_memory_statistics(precision: int = 3) -> Dict[str, Any]:
-    memory_allocated = None
-    memory_reserved = None
-    max_memory_allocated = None
-    max_memory_reserved = None
+    memory_stats = {
+        "memory_allocated": None,
+        "memory_reserved": None,
+        "max_memory_allocated": None,
+        "max_memory_reserved": None,
+    }
 
     if torch.cuda.is_available():
         device = torch.cuda.current_device()
-        memory_allocated = torch.cuda.memory_allocated(device)
-        memory_reserved = torch.cuda.memory_reserved(device)
-        max_memory_allocated = torch.cuda.max_memory_allocated(device)
-        max_memory_reserved = torch.cuda.max_memory_reserved(device)
+        memory_stats.update(
+            {
+                "memory_allocated": torch.cuda.memory_allocated(device),
+                "memory_reserved": torch.cuda.memory_reserved(device),
+                "max_memory_allocated": torch.cuda.max_memory_allocated(device),
+                "max_memory_reserved": torch.cuda.max_memory_reserved(device),
+            }
+        )
 
     elif torch.backends.mps.is_available():
-        memory_allocated = torch.mps.current_allocated_memory()
+        memory_stats["memory_allocated"] = torch.mps.current_allocated_memory()
 
     else:
         logger.warning("No CUDA, MPS, or ROCm device found. Memory statistics are not available.")
 
     return {
-        "memory_allocated": round(bytes_to_gigabytes(memory_allocated), ndigits=precision),
-        "memory_reserved": round(bytes_to_gigabytes(memory_reserved), ndigits=precision),
-        "max_memory_allocated": round(bytes_to_gigabytes(max_memory_allocated), ndigits=precision),
-        "max_memory_reserved": round(bytes_to_gigabytes(max_memory_reserved), ndigits=precision),
+        key: (round(bytes_to_gigabytes(value), ndigits=precision) if value else None)
+        for key, value in memory_stats.items()
     }
 
 
@@ -49,6 +53,21 @@ def free_memory() -> None:
     # TODO(aryan): handle non-cuda devices
 
 
+def reset_memory_stats(device: torch.device):
+    # TODO: handle for non-cuda devices
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+    else:
+        logger.warning("No CUDA, device found. Nothing to reset memory of.")
+
+
+def synchronize_device(device: torch.device):
+    if torch.cuda.is_available():
+        torch.cuda.synchronize(device)
+    else:
+        logger.warning("No CUDA, device found. Nothing to synchronize.")
+
+
 def make_contiguous(x: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
     if isinstance(x, torch.Tensor):
         return x.contiguous()
diff --git a/tests/trainers/__init__.py b/tests/trainers/__init__.py
diff --git a/tests/trainers/cogvideox/__init__.py b/tests/trainers/cogvideox/__init__.py
diff --git a/tests/trainers/cogvideox/test_cogvideox.py b/tests/trainers/cogvideox/test_cogvideox.py
@@ -0,0 +1,44 @@
+import sys
+import unittest
+from pathlib import Path
+
+
+current_file = Path(__file__).resolve()
+root_dir = current_file.parents[3]
+sys.path.append(str(root_dir))
+
+from finetrainers import Args  # noqa
+from ..test_trainers_common import TrainerTestMixin, parse_resolution_bucket  # noqa
+
+
+class CogVideoXTester(unittest.TestCase, TrainerTestMixin):
+    MODEL_NAME = "cogvideox"
+    EXPECTED_PRECOMPUTATION_LATENT_KEYS = {"latents"}
+    EXPECTED_PRECOMPUTATION_CONDITION_KEYS = {"prompt_embeds"}
+
+    def get_training_args(self):
+        args = Args()
+        args.model_name = self.MODEL_NAME
+        args.training_type = "lora"
+        args.pretrained_model_name_or_path = "finetrainers/dummy-cogvideox"
+        args.data_root = ""  # will be set from the tester method.
+        args.video_resolution_buckets = [parse_resolution_bucket("9x16x16")]
+        args.precompute_conditions = True
+        args.validation_prompts = []
+        args.validation_heights = []
+        args.validation_widths = []
+        return args
+
+    @property
+    def latent_output_shape(self):
+        return (8, 3, 2, 2)
+
+    @property
+    def condition_output_shape(self):
+        return (226, 32)
+
+    def populate_shapes(self):
+        for k in self.EXPECTED_PRECOMPUTATION_LATENT_KEYS:
+            self.EXPECTED_LATENT_SHAPES[k] = self.latent_output_shape
+        for k in self.EXPECTED_PRECOMPUTATION_CONDITION_KEYS:
+            self.EXPECTED_CONDITION_SHAPES[k] = self.condition_output_shape
diff --git a/tests/trainers/hunyaun_video/__init__.py b/tests/trainers/hunyaun_video/__init__.py
diff --git a/tests/trainers/hunyaun_video/test_hunyaun_video.py b/tests/trainers/hunyaun_video/test_hunyaun_video.py
diff --git a/tests/trainers/ltx_video/__init__.py b/tests/trainers/ltx_video/__init__.py
diff --git a/tests/trainers/ltx_video/test_ltx_video.py b/tests/trainers/ltx_video/test_ltx_video.py
diff --git a/tests/trainers/test_trainers_common.py b/tests/trainers/test_trainers_common.py
