CogVideoX LoRA and full finetuning (#1)

* update * update * update * update * update * update * update * update * update * update * update * update * update * update * update
2026-10-17 00:23:35 +00:00 · 2024-10-01 08:05:10 +05:30
parent 95774133eb
commit 512a8a6854
26 changed files with 3823 additions and 1 deletions
@@ -1 +1,101 @@
-# cogvideox-distillation
+# Finetuning CogVideoX
+
+
+## Dataset Preparation
+
+Create two files where one file contains line-separated prompts and another file contains line-separated paths to video data (the path to video files must be relative to the path you pass when specifying `--data_root`). Let's take a look at an example to understand this better!
+
+Assume you've specified `--data_root` as `/dataset`, and that this directory contains the files: `prompts.txt` and `videos.txt`.
+
+The `prompts.txt` file should contain line-separated prompts:
+
+```
+A black and white animated sequence featuring a rabbit, named Rabbity Ribfried, and an anthropomorphic goat in a musical, playful environment, showcasing their evolving interaction.
+A black and white animated sequence on a ship's deck features a bulldog character, named Bully Bulldoger, showcasing exaggerated facial expressions and body language. The character progresses from confident to focused, then to strained and distressed, displaying a range of emotions as it navigates challenges. The ship's interior remains static in the background, with minimalistic details such as a bell and open door. The character's dynamic movements and changing expressions drive the narrative, with no camera movement to distract from its evolving reactions and physical gestures.
+...
+```
+
+The `videos.txt` file should contain line-separate paths to video files. Note that the path should be _relative_ to the `--data_root` directory.
+
+```bash
+videos/00000.mp4
+videos/00001.mp4
+...
+```
+
+Overall, this is how your dataset would look like if you ran the `tree` command on the dataset root directory:
+
+```bash
+/dataset
+├── prompts.txt
+├── videos.txt
+├── videos
+    ├── videos/00000.mp4
+    ├── videos/00001.mp4
+    ├── ...
+```
+
+When using this format, the `--caption_column` must be `prompts.txt` and `--video_column` must be `videos.txt`. If you, instead, have your data stored in a CSV file, you can also specify `--dataset_file` as the path to CSV, the `--caption_column` and `--video_column` as the actual column names in the CSV file.
+
+As an example, let's use [this](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) Disney dataset for finetuning. To download, one can use the 🤗 Hugging Face CLI.
+
+```bash
+huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir video-dataset-disney
+```
+
+#### Rough notes and TODOs:
+
+- Uncompiled SFT works end-to-end on dummy example. Need to test on larger dataset (not priority at the moment)
+- Compiled SFT fails with `THUDM/CogVideoX-2b` throwing the following error (by error, it's more of a graph break situation due to mixin numpy/cpu device when getting sincos positional embeddings).
+
+## Training
+
+TODO
+
+Take a look at `training/*.sh`
+
+Note: Untested on MPS
+
+## Memory requirements
+
+|       model        | lora rank | optimizer | gradient_checkpointing | memory_before_training | memory_after_validation | memory_after_testing |
+|:------------------:|:---------:|:---------:|:----------------------:|:----------------------:|:-----------------------:|:--------------------:|
+| THUDM/CogVideoX-2b |    16     |  adamw    |          False         |         12.945         |         39.553          |       23.148         |
+| THUDM/CogVideoX-2b |    16     |  adamw    |          True          |         12.946         |         18.436          |       23.160         |
+| THUDM/CogVideoX-2b |    64     |  adamw    |          False         |         13.035         |         40.051          |       23.430         |
+| THUDM/CogVideoX-2b |    64     |  adamw    |          True          |         13.035         |         18.883          |       23.414         |
+| THUDM/CogVideoX-2b |    256    |  adamw    |          False         |         13.095         |         42.004          |       24.385         |
+| THUDM/CogVideoX-2b |    256    |  adamw    |          True          |         13.095         |         19.307          |       24.381         |
+
+**Note:** `memory_after_validation` is indicative of the peak memory required for training.
+
+<details>
+<summary> stack trace </summary>
+
+```
+skipping cudagraphs due to skipping cudagraphs due to cpu device (cat_3). Found from : 
+   File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 820, in forward
+    return model_forward(*args, **kwargs)
+  File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/accelerate/utils/operations.py", line 808, in __call__
+    return convert_to_fp32(self.model_forward(*args, **kwargs))
+  File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
+    return func(*args, **kwargs)
+  File "/home/aryan/work/diffusers/src/diffusers/models/transformers/cogvideox_transformer_3d.py", line 446, in forward
+    hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+  File "/home/aryan/work/diffusers/src/diffusers/models/embeddings.py", line 435, in forward
+    pos_embedding = self._get_positional_embeddings(height, width, pre_time_compression_frames)
+  File "/home/aryan/work/diffusers/src/diffusers/models/embeddings.py", line 385, in _get_positional_embeddings
+    pos_embedding = get_3d_sincos_pos_embed(
+  File "/home/aryan/work/diffusers/src/diffusers/models/embeddings.py", line 108, in get_3d_sincos_pos_embed
+    grid = np.stack(grid, axis=0)
+```
+</details>
+
+- Make T2V LoRA script up-to-date
+- Make I2V LoRA script up-to-date
+- Make scripts compatible with DDP
+- Make scripts compatible with FSDP
+- Make scripts compatible with DeepSpeed
+- Test scripts with memory-efficient optimizer
+- Test scripts with quantization using torchao, CPUOffloadOptimizer, etc.
+- Make 5B lora finetuning work in under 24GB