storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-09 00:09:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Low-bit memory optimizers, CpuOffloadOptimizer, Memory Reports (#3)

Browse Source 
* update

* update

* update

* update

* update

* update

* update

* update

* update
This commit is contained in:
Aryan
2024-10-04 18:18:54 +05:30
committed by GitHub
parent 512a8a6854
commit d81209bf49
11 changed files with 434 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -165,3 +165,5 @@ cython_debug/
wandb/
*.txt
dump*
!requirements.txt
README.md
+195 -39
View File
@@ -43,11 +43,6 @@ As an example, let's use [this](https://huggingface.co/datasets/Wild-Heart/Disne
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
@@ -58,44 +53,205 @@ 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 |
<table align="center">
<tr>
<td align="center"><a href="https://www.youtube.com/watch?v=UvRl4ansfCg"> Slaying OOMs with PyTorch</a></td>
</tr>
<tr>
<td align="center"><img src="assets/slaying-ooms.png" style="width: 480px; height: 480px;"></td>
</tr>
</table>
**Note:** `memory_after_validation` is indicative of the peak memory required for training.
Supported and verified memory optimizations for training include:
- `CPUOffloadOptimizer` from [TorchAO](https://github.com/pytorch/ao). You can read about its capabilities and limitations [here](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload). In short, it allows you to use the CPU for storing trainable parameters and gradients. This results in the optimizer step happening on the CPU, which requires a fast CPU optimizer, such as `torch.AdamW(fused=True)` or applying `torch.compile` on the optimizer step. Additionally, it is recommended to not `torch.compile` your model for training. Gradient clipping and accumulation is not supported yet either.
- Low-bit optimizers from [bitsandbytes](https://huggingface.co/docs/bitsandbytes/optimizers). TODO: to test and make [TorchAO](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim) ones work
- TODO: DeepSpeed ZeRO
> [!IMPORTANT]
> The memory requirements are reported after running the `training/prepare_dataset.py`, which converts the videos and captions to latents and embeddings. During training, we directly load the latents and embeddings, and do not require the VAE or the T5 text encoder. However, if you perform validation/testing, these must be loaded and increase the amount of required memory. Not performing validation/testing saves a significant amount of memory, which can be used to focus solely on training if you're on smaller VRAM GPUs.
### LoRA finetuning
<details>
<summary> stack trace </summary>
<summary> AdamW </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)
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]
> Trying to run CogVideoX-5b without gradient checkpointing OOMs even on an A100 (80 GB), so the memory measurements have not been specified.
</details>
<details>
<summary> AdamW (8-bit bitsandbytes) </summary>
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.732 | 46.887 | 24.195 |
| THUDM/CogVideoX-2b | 16 | True | 12.945 | 12.945 | 21.430 | 24.195 |
| THUDM/CogVideoX-2b | 64 | False | 13.035 | 44.004 | 47.158 | 24.369 |
| THUDM/CogVideoX-2b | 64 | True | 13.035 | 13.035 | 21.297 | 24.357 |
| THUDM/CogVideoX-2b | 256 | False | 13.035 | 45.291 | 48.455 | 24.836 |
| THUDM/CogVideoX-2b | 256 | True | 13.035 | 13.035 | 21.625 | 24.869 |
| THUDM/CogVideoX-5b | 16 | True | 19.742 | 19.742 | 28.602 | 38.049 |
| THUDM/CogVideoX-5b | 64 | True | 20.006 | 20.818 | 29.359 | 38.520 |
| THUDM/CogVideoX-5b | 256 | True | 20.771 | 21.352 | 30.727 | 39.596 |
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.734 | 21.775 | 24.281 |
| THUDM/CogVideoX-2b | 64 | True | 13.036 | 21.941 | 21.941 | 24.445 |
| THUDM/CogVideoX-2b | 256 | True | 13.094 | 22.020 | 22.266 | 24.943 |
| THUDM/CogVideoX-5b | 16 | True | 19.742 | 46.320 | 46.326 | 38.104 |
| THUDM/CogVideoX-5b | 64 | True | 20.006 | 46.820 | 46.820 | 38.588 |
| THUDM/CogVideoX-5b | 256 | True | 20.771 | 47.920 | 47.980 | 40.002 |
</details>
<details>
<summary> AdamW (8-bit torchao) </summary>
Currently, errors out with following stack-trace:
```python
Traceback (most recent call last):
File "/raid/aryan/cogvideox-distillation/training/cogvideox_text_to_video_lora.py", line 915, in <module>
main(args)
File "/raid/aryan/cogvideox-distillation/training/cogvideox_text_to_video_lora.py", line 719, in main
optimizer.step()
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/accelerate/optimizer.py", line 159, in step
self.scaler.step(self.optimizer, closure)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 457, in step
retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/amp/grad_scaler.py", line 352, in _maybe_opt_step
retval = optimizer.step(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/accelerate/optimizer.py", line 214, in patched_step
return method(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 137, in wrapper
return func.__get__(opt, opt.__class__)(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 487, in wrapper
out = func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torchao/prototype/low_bit_optim/adam.py", line 87, in step
raise RuntimeError(
RuntimeError: lr was changed to a non-Tensor object. If you want to update lr, please use optim.param_groups[0]['lr'].fill_(new_lr)
```
</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
<details>
<summary> AdamW (4-bit torchao) </summary>
Same error as AdamW (8-bit torchao)
</details>
<details>
<summary> AdamW + CPUOffloadOptimizer (with gradient offloading) </summary>
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.705 | 46.859 | 24.180 |
| THUDM/CogVideoX-2b | 16 | True | 12.945 | 12.945 | 21.395 | 24.180 |
| THUDM/CogVideoX-2b | 64 | False | 13.035 | 43.916 | 47.070 | 24.234 |
| THUDM/CogVideoX-2b | 64 | True | 13.035 | 13.035 | 20.887 | 24.266 |
| THUDM/CogVideoX-2b | 256 | False | 13.095 | 44.947 | 48.111 | 24.607 |
| THUDM/CogVideoX-2b | 256 | True | 13.095 | 13.095 | 21.391 | 24.635 |
| THUDM/CogVideoX-5b | 16 | True | 19.742 | 19.742 | 28.533 | 38.002 |
| THUDM/CogVideoX-5b | 64 | True | 20.006 | 20.006 | 29.107 | 38.785 |
| THUDM/CogVideoX-5b | 256 | True | 20.771 | 20.771 | 30.078 | 39.559 |
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.709 | 21.762 | 24.254 |
| THUDM/CogVideoX-2b | 64 | True | 13.035 | 21.844 | 21.855 | 24.338 |
| THUDM/CogVideoX-2b | 256 | True | 13.094 | 22.020 | 22.031 | 24.709 |
| THUDM/CogVideoX-5b | 16 | True | 19.742 | 46.262 | 46.297 | 38.400 |
| THUDM/CogVideoX-5b | 64 | True | 20.006 | 46.561 | 46.574 | 38.840 |
| THUDM/CogVideoX-5b | 256 | True | 20.771 | 47.268 | 47.332 | 39.623 |
> [!NOTE]
> Trying to run CogVideoX-5b without gradient checkpointing OOMs even on an A100 (80 GB), so the memory measurements have not been specified.
</details>
<details>
<summary> AdamW (8-bit bitsandbytes) + CPUOffloadOptimizer (with gradient offloading) </summary>
Currently, errors out with the following stack-trace:
```python
File "/raid/aryan/cogvideox-distillation/training/cogvideox_text_to_video_lora.py", line 925, in <module>
main(args)
File "/raid/aryan/cogvideox-distillation/training/cogvideox_text_to_video_lora.py", line 727, in main
optimizer.step()
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torchao/prototype/low_bit_optim/cpu_offload.py", line 87, in step
self.optim_dict[p_cuda].step()
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 487, in wrapper
out = func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py", line 287, in step
self.update_step(group, p, gindex, pindex)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py", line 546, in update_step
F.optimizer_update_8bit_blockwise(
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/bitsandbytes/functional.py", line 1774, in optimizer_update_8bit_blockwise
prev_device = pre_call(g.device)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/bitsandbytes/functional.py", line 463, in pre_call
torch.cuda.set_device(device)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/cuda/__init__.py", line 476, in set_device
device = _get_device_index(device)
File "/raid/aryan/nightly-venv/lib/python3.10/site-packages/torch/cuda/_utils.py", line 34, in _get_device_index
raise ValueError(f"Expected a cuda device, but got: {device}")
ValueError: Expected a cuda device, but got: cpu
```
</details>
### Full finetuning
> [!NOTE]
> `memory_after_validation` is indicative of the peak memory required for training. This is because apart from the activations, parameters and gradients stored for training, you also need to load the vae and text encoder in memory and spend some memory to perform inference. In order to reduce total memory required to perform training, one can choose to not perform validation/testing as part of the training script.
- [ ] Make scripts compatible with DDP
- [ ] Make scripts compatible with FSDP
- [ ] Make scripts compatible with DeepSpeed
- [x] Test scripts with memory-efficient optimizer from bitsandbytes
- [x] Test scripts with CPUOffloadOptimizer, etc.
- [ ] Test scripts with torchao quantization, and low bit memory optimizers, etc.
- [x] Make 5B lora finetuning work in under 24GB
assets/slaying-ooms.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 164 KiB

metrics.sh
-155
View File
@@ -1,155 +0,0 @@
# export TORCH_LOGS="+dynamo,recompiles,graph_breaks"
# export TORCHDYNAMO_VERBOSE=1
# export WANDB_MODE="offline"
# export NCCL_P2P_DISABLE=1
# export TORCH_NCCL_ENABLE_MONITORING=0
# GPU_IDS="1"
# LEARNING_RATES=("1e-4")
# LR_SCHEDULES=("cosine_with_restarts")
# OPTIMIZERS=("adamw")
# MAX_TRAIN_STEPS=("2")
# RANK=("16" "64" "256")
# GRADIENT_CHECKPOINTING=("" "--gradient_checkpointing")
# DATA_ROOT="/raid/aryan/video-dataset-disney/"
# CAPTION_COLUMN="prompts.txt"
# VIDEO_COLUMN="videos.txt"
# for learning_rate in "${LEARNING_RATES[@]}"; do
# for lr_schedule in "${LR_SCHEDULES[@]}"; do
# for optimizer in "${OPTIMIZERS[@]}"; do
# for steps in "${MAX_TRAIN_STEPS[@]}"; do
# for rank in "${RANK[@]}"; do
# for gradient_checkpointing in "${GRADIENT_CHECKPOINTING[@]}"; do
# cache_dir="/raid/aryan/cogvideox-lora/"
# output_dir="/raid/aryan/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
# cmd="accelerate launch --config_file accelerate_configs/uncompiled_1.yaml --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
# --pretrained_model_name_or_path THUDM/CogVideoX-2b \
# --cache_dir $cache_dir \
# --data_root $DATA_ROOT \
# --caption_column $CAPTION_COLUMN \
# --video_column $VIDEO_COLUMN \
# --id_token BW_STYLE \
# --height_buckets 480 \
# --width_buckets 720 \
# --frame_buckets 49 \
# --validation_prompt \"BW_STYLE A black and white animated scene unfolds with an anthropomorphic goat surrounded by musical notes and symbols, suggesting a playful environment. Mickey Mouse appears, leaning forward in curiosity as the goat remains still. The goat then engages with Mickey, who bends down to converse or react. The dynamics shift as Mickey grabs the goat, potentially in surprise or playfulness, amidst a minimalistic background. The scene captures the evolving relationship between the two characters in a whimsical, animated setting, emphasizing their interactions and emotions:::BW_STYLE A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance\" \
# --validation_prompt_separator ::: \
# --num_validation_videos 1 \
# --validation_epochs 1 \
# --seed 42 \
# --rank $rank \
# --lora_alpha 64 \
# --mixed_precision fp16 \
# --output_dir $output_dir \
# --max_num_frames 49 \
# --train_batch_size 1 \
# --max_train_steps $steps \
# --checkpointing_steps 1000 \
# --gradient_accumulation_steps 1 \
# $gradient_checkpointing \
# --learning_rate $learning_rate \
# --lr_scheduler $lr_schedule \
# --lr_warmup_steps 200 \
# --lr_num_cycles 1 \
# --enable_slicing \
# --enable_tiling \
# --optimizer $optimizer \
# --beta1 0.9 \
# --beta2 0.95 \
# --weight_decay 0.001 \
# --max_grad_norm 1.0 \
# --allow_tf32 \
# --report_to wandb \
# --nccl_timeout 1800"
# echo "Running command: $cmd"
# eval $cmd
# echo -ne "-------------------- Finished executing script --------------------\n\n"
# done
# done
# done
# done
# done
# done
# For testing load from tensor data
export TORCH_LOGS="+dynamo,recompiles,graph_breaks"
export TORCHDYNAMO_VERBOSE=1
export WANDB_MODE="offline"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
GPU_IDS="1"
LEARNING_RATES=("1e-4")
LR_SCHEDULES=("cosine_with_restarts")
OPTIMIZERS=("adamw")
MAX_TRAIN_STEPS=("2")
RANK=("16" "64" "256")
GRADIENT_CHECKPOINTING=("" "--gradient_checkpointing")
DATA_ROOT="training/dump"
CAPTION_COLUMN="prompts.txt"
VIDEO_COLUMN="videos.txt"
for learning_rate in "${LEARNING_RATES[@]}"; do
for lr_schedule in "${LR_SCHEDULES[@]}"; do
for optimizer in "${OPTIMIZERS[@]}"; do
for steps in "${MAX_TRAIN_STEPS[@]}"; do
for rank in "${RANK[@]}"; do
for gradient_checkpointing in "${GRADIENT_CHECKPOINTING[@]}"; do
cache_dir="/raid/aryan/cogvideox-lora/"
output_dir="/raid/aryan/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
cmd="accelerate launch --config_file accelerate_configs/uncompiled_1.yaml --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
--pretrained_model_name_or_path THUDM/CogVideoX-2b \
--cache_dir $cache_dir \
--data_root $DATA_ROOT \
--caption_column $CAPTION_COLUMN \
--video_column $VIDEO_COLUMN \
--id_token BW_STYLE \
--height_buckets 480 \
--width_buckets 720 \
--frame_buckets 49 \
--load_tensors \
--validation_prompt \"BW_STYLE A black and white animated scene unfolds with an anthropomorphic goat surrounded by musical notes and symbols, suggesting a playful environment. Mickey Mouse appears, leaning forward in curiosity as the goat remains still. The goat then engages with Mickey, who bends down to converse or react. The dynamics shift as Mickey grabs the goat, potentially in surprise or playfulness, amidst a minimalistic background. The scene captures the evolving relationship between the two characters in a whimsical, animated setting, emphasizing their interactions and emotions\" \
--validation_prompt_separator ::: \
--num_validation_videos 1 \
--validation_epochs 2 \
--seed 42 \
--rank $rank \
--lora_alpha 64 \
--mixed_precision fp16 \
--output_dir $output_dir \
--max_num_frames 49 \
--train_batch_size 1 \
--max_train_steps $steps \
--checkpointing_steps 1000 \
--gradient_accumulation_steps 1 \
$gradient_checkpointing \
--learning_rate $learning_rate \
--lr_scheduler $lr_schedule \
--lr_warmup_steps 200 \
--lr_num_cycles 1 \
--enable_slicing \
--enable_tiling \
--optimizer $optimizer \
--beta1 0.9 \
--beta2 0.95 \
--weight_decay 0.001 \
--max_grad_norm 1.0 \
--allow_tf32 \
--report_to wandb \
--nccl_timeout 1800"
echo "Running command: $cmd"
eval $cmd
echo -ne "-------------------- Finished executing script --------------------\n\n"
done
done
done
done
done
done
requirements.txt
+13
View File
@@ -0,0 +1,13 @@
accelerate
bitsandbytes
diffusers
transformers
huggingface_hub
hf_transfer
peft
decord
wandb
pandas
torch
torchvision
torchao
train_text_to_video_lora.sh
+25 -17
View File
@@ -4,26 +4,34 @@ export WANDB_MODE="offline"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
GPU_IDS="2"
LEARNING_RATES=("1e-4")
LR_SCHEDULES=("cosine_with_restarts")
OPTIMIZERS=("adamw")
MAX_TRAIN_STEPS=("2")
GPU_IDS="0"
DATA_ROOT="dump"
CAPTION_COLUMN="prompts.txt"
# Training Configurations
# Experiment with as many hyperparameters as you want!
LEARNING_RATES=("1e-4" "1e-3")
LR_SCHEDULES=("cosine_with_restarts")
OPTIMIZERS=("adamw", "adam")
MAX_TRAIN_STEPS=("3000")
# Single GPU uncompiled training
ACCELERATE_CONFIG_FILE="accelerate_configs/uncompiled_1.yaml"
# Absolute path to where the data is located. Make sure to have read the README for how to prepare data.
# This example assumes you downloaded an already prepared dataset from HF CLI as follows:
# huggingface-cli download --repo-type dataset Wild-Heart/Disney-VideoGeneration-Dataset --local-dir /path/to/my/datasets/disney-dataset
DATA_ROOT="/path/to/my/datasets/disney-dataset"
CAPTION_COLUMN="prompt.txt"
VIDEO_COLUMN="videos.txt"
# Launch experiments with different hyperparameters
for learning_rate in "${LEARNING_RATES[@]}"; do
for lr_schedule in "${LR_SCHEDULES[@]}"; do
for optimizer in "${OPTIMIZERS[@]}"; do
for steps in "${MAX_TRAIN_STEPS[@]}"; do
cache_dir="/raid/aryan/cogvideox-lora/"
output_dir="/raid/aryan/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
output_dir="/path/to/my/models/cogvideox-lora__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
cmd="accelerate launch --config_file accelerate_configs/uncompiled_1.yaml --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
--pretrained_model_name_or_path THUDM/CogVideoX-2b \
--cache_dir $cache_dir \
cmd="accelerate launch --config_file $ACCELERATE_CONFIG_FILE --gpu_ids $GPU_IDS training/cogvideox_text_to_video_lora.py \
--pretrained_model_name_or_path THUDM/CogVideoX-5b \
--data_root $DATA_ROOT \
--caption_column $CAPTION_COLUMN \
--video_column $VIDEO_COLUMN \
@@ -31,14 +39,14 @@ for learning_rate in "${LEARNING_RATES[@]}"; do
--height_buckets 480 \
--width_buckets 720 \
--frame_buckets 49 \
--validation_prompt \"BW_STYLE A black and white animated scene unfolds with an anthropomorphic goat surrounded by musical notes and symbols, suggesting a playful environment. Mickey Mouse appears, leaning forward in curiosity as the goat remains still. The goat then engages with Mickey, who bends down to converse or react. The dynamics shift as Mickey grabs the goat, potentially in surprise or playfulness, amidst a minimalistic background. The scene captures the evolving relationship between the two characters in a whimsical, animated setting, emphasizing their interactions and emotions::: BW_STYLE A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance\" \
--validation_prompt \"BW_STYLE A black and white animated scene unfolds with an anthropomorphic goat surrounded by musical notes and symbols, suggesting a playful environment. Mickey Mouse appears, leaning forward in curiosity as the goat remains still. The goat then engages with Mickey, who bends down to converse or react. The dynamics shift as Mickey grabs the goat, potentially in surprise or playfulness, amidst a minimalistic background. The scene captures the evolving relationship between the two characters in a whimsical, animated setting, emphasizing their interactions and emotions:::BW_STYLE A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance\" \
--validation_prompt_separator ::: \
--num_validation_videos 1 \
--validation_epochs 10 \
--seed 42 \
--rank 64 \
--lora_alpha 64 \
--mixed_precision fp16 \
--rank 128 \
--lora_alpha 128 \
--mixed_precision bf16 \
--output_dir $output_dir \
--max_num_frames 49 \
--train_batch_size 1 \
@@ -48,7 +56,7 @@ for learning_rate in "${LEARNING_RATES[@]}"; do
--gradient_checkpointing \
--learning_rate $learning_rate \
--lr_scheduler $lr_schedule \
--lr_warmup_steps 200 \
--lr_warmup_steps 400 \
--lr_num_cycles 1 \
--enable_slicing \
--enable_tiling \
train_text_to_video_sft.sh
+21 -14
View File
@@ -1,48 +1,55 @@
# export TORCH_LOGS="+dynamo,recompiles,graph_breaks"
# export TORCHDYNAMO_VERBOSE=1
export TORCH_LOGS="+dynamo,recompiles,graph_breaks"
export TORCHDYNAMO_VERBOSE=1
export WANDB_MODE="offline"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
GPU_IDS="3"
GPU_IDS="0"
# Training Configurations
# Experiment with as many hyperparameters as you want!
LEARNING_RATES=("1e-4")
LR_SCHEDULES=("cosine_with_restarts")
OPTIMIZERS=("adamw")
MAX_TRAIN_STEPS=("20000")
# DATA_ROOT="/raid/aryan/dataset-cogvideox/"
DATA_ROOT="/raid/aryan/openvid-1m"
CAPTION_COLUMN="prompts.txt"
# Single GPU uncompiled training
ACCELERATE_CONFIG_FILE="accelerate_configs/uncompiled_1.yaml"
# Absolute path to where the data is located. Make sure to have read the README for how to prepare data.
# This example assumes you downloaded an already prepared dataset from HF CLI as follows:
# huggingface-cli download --repo-type dataset Wild-Heart/Tom-and-Jerry-VideoGeneration-Dataset --local-dir /path/to/my/datasets/tom-and-jerry-dataset
DATA_ROOT="/path/to/my/datasets/tom-and-jerry-dataset"
CAPTION_COLUMN="captions.txt"
VIDEO_COLUMN="videos.txt"
# Launch experiments with different hyperparameters
for learning_rate in "${LEARNING_RATES[@]}"; do
for lr_schedule in "${LR_SCHEDULES[@]}"; do
for optimizer in "${OPTIMIZERS[@]}"; do
for steps in "${MAX_TRAIN_STEPS[@]}"; do
cache_dir="/raid/aryan/cogvideox-sft/"
output_dir="/raid/aryan/cogvideox-sft__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
output_dir="/path/to/my/models/cogvideox-sft__optimizer_${optimizer}__steps_${steps}__lr-schedule_${lr_schedule}__learning-rate_${learning_rate}/"
cmd="accelerate launch --config_file accelerate_configs/uncompiled_1.yaml --gpu_ids $GPU_IDS training/cogvideox_text_to_video_sft.py \
--pretrained_model_name_or_path THUDM/CogVideoX-2b \
--cache_dir $cache_dir \
cmd="accelerate launch --config_file $ACCELERATE_CONFIG_FILE --gpu_ids $GPU_IDS training/cogvideox_text_to_video_sft.py \
--pretrained_model_name_or_path THUDM/CogVideoX-5b \
--data_root $DATA_ROOT \
--caption_column $CAPTION_COLUMN \
--video_column $VIDEO_COLUMN \
--height_buckets 480 \
--width_buckets 720 \
--frame_buckets 49 \
--validation_prompt \"a man wearing a bicycle helmet, riding a bike through a forested area. The man is wearing a black t-shirt and appears to be in motion, as suggested by the slight blur of the background. The forest is lush and green, with trees and foliage filling the background. The man's helmet is white with a black visor, and he is looking directly at the camera with a slight smile on his face. The style of the video is casual and candid, capturing a moment of outdoor activity:::A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance\" \
--validation_prompt \"Tom, the mischievous gray cat, is sprawled out on a vibrant red pillow, his body relaxed and his eyes half-closed, as if he's just woken up or is about to doze off. His white paws are stretched out in front of him, and his tail is casually draped over the edge of the pillow. The setting appears to be a cozy corner of a room, with a warm yellow wall in the background and a hint of a wooden floor. The scene captures a rare moment of tranquility for Tom, contrasting with his usual energetic and playful demeanor:::A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance\" \
--validation_prompt_separator ::: \
--num_validation_videos 1 \
--validation_epochs 1 \
--seed 42 \
--mixed_precision fp16 \
--mixed_precision bf16 \
--output_dir $output_dir \
--max_num_frames 49 \
--train_batch_size 1 \
--max_train_steps $steps \
--checkpointing_steps 2000 \
--gradient_accumulation_steps 1 \
--gradient_accumulation_steps 4 \
--gradient_checkpointing \
--learning_rate $learning_rate \
--lr_scheduler $lr_schedule \
training/args.py
+20 -2
View File
@@ -301,13 +301,21 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
"--optimizer",
type=lambda s: s.lower(),
default="adam",
choices=["adam", "adamw", "prodigy"],
choices=["adam", "adamw", "prodigy", "came"],
help=("The optimizer type to use."),
)
parser.add_argument(
"--use_8bit",
action="store_true",
help="Whether or not to use 8-bit optimizers from `bitsandbytes`. Ignored if incompatible optimzer selected.",
help="Whether or not to use 8-bit optimizers from `bitsandbytes` or `bitsandbytes`.",
)
parser.add_argument(
"--use_4bit",
action="store_true",
help="Whether or not to use 4-bit optimizers from `torchao`.",
)
parser.add_argument(
"--use_torchao", action="store_true", help="Whether or not to use the `torchao` backend for optimizers."
)
parser.add_argument(
"--beta1",
@@ -355,6 +363,16 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
action="store_true",
help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage.",
)
parser.add_argument(
"--use_cpu_offload_optimizer",
action="store_true",
help="Whether or not to use the CPUOffloadOptimizer from TorchAO to perform optimization step and maintain parameters on the CPU.",
)
parser.add_argument(
"--offload_gradients",
action="store_true",
help="Whether or not to offload the gradients to CPU when using the CPUOffloadOptimizer from TorchAO.",
)
def _get_configuration_args(parser: argparse.ArgumentParser) -> None:
training/cogvideox_text_to_video_lora.py
+38 -31
View File
@@ -25,7 +25,6 @@ from typing import Any, Dict
import diffusers
import torch
import transformers
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import (
@@ -43,11 +42,7 @@ from diffusers import (
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,
is_wandb_available,
)
from diffusers.utils import convert_unet_state_dict_to_peft, export_to_video
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module
from huggingface_hub import create_repo, upload_folder
@@ -56,6 +51,8 @@ from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoTokenizer, T5EncoderModel
import wandb
from args import get_args # isort:skip
from dataset import BucketSampler, VideoDatasetWithResizing # isort:skip
@@ -230,10 +227,6 @@ def main(args):
if torch.backends.mps.is_available():
accelerator.native_amp = False
if args.report_to == "wandb":
if not is_wandb_available():
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -436,6 +429,7 @@ def main(args):
"lr": args.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
use_deepspeed_optimizer = (
accelerator.state.deepspeed_plugin is not None
@@ -459,7 +453,11 @@ def main(args):
prodigy_use_bias_correction=args.prodigy_use_bias_correction,
prodigy_safeguard_warmup=args.prodigy_safeguard_warmup,
use_8bit=args.use_8bit,
use_4bit=args.use_4bit,
use_torchao=args.use_torchao,
use_deepspeed=use_deepspeed_optimizer,
use_cpu_offload_optimizer=args.use_cpu_offload_optimizer,
offload_gradients=args.offload_gradients,
)
# Dataset and DataLoader
@@ -524,24 +522,31 @@ def main(args):
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
if args.use_cpu_offload_optimizer:
lr_scheduler = None
accelerator.print(
"CPU Offload Optimizer cannot be used with DeepSpeed or builtin PyTorch LR Schedulers. If "
"you are training with those settings, they will be ignored."
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
# Prepare everything with our `accelerator`.
transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -567,7 +572,6 @@ def main(args):
# Train!
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
accelerator.print("***** Running training *****")
accelerator.print(f" Num trainable parameters = {num_trainable_parameters}")
@@ -624,6 +628,8 @@ def main(args):
torch.cuda.empty_cache()
torch.cuda.synchronize(accelerator.device)
alphas_cumprod = scheduler.alphas_cumprod.to(accelerator.device, dtype=torch.float32)
for epoch in range(first_epoch, args.num_train_epochs):
transformer.train()
@@ -691,8 +697,7 @@ def main(args):
model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)
alphas_cumprod = scheduler.alphas_cumprod[timesteps]
weights = 1 / (1 - alphas_cumprod)
weights = 1 / (1 - alphas_cumprod[timesteps])
while len(weights.shape) < len(model_pred.shape):
weights = weights.unsqueeze(-1)
@@ -714,7 +719,8 @@ def main(args):
optimizer.step()
optimizer.zero_grad()
lr_scheduler.step()
if not args.use_cpu_offload_optimizer:
lr_scheduler.step()
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
@@ -747,9 +753,10 @@ def main(args):
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
logs = {
"loss": loss.detach().item(),
"lr": lr_scheduler.get_last_lr()[0],
"lr": last_lr,
"gradient_norm_before_clip": gradient_norm_before_clip,
"gradient_norm_after_clip": gradient_norm_after_clip,
}
training/cogvideox_text_to_video_sft.py
+38 -27
View File
@@ -25,7 +25,6 @@ from typing import Any, Dict
import diffusers
import torch
import transformers
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import (
@@ -43,7 +42,7 @@ from diffusers import (
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 export_to_video, is_wandb_available
from diffusers.utils import export_to_video
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module
from huggingface_hub import create_repo, upload_folder
@@ -51,6 +50,8 @@ from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoTokenizer, T5EncoderModel
import wandb
from args import get_args # isort:skip
from dataset import BucketSampler, VideoDatasetWithResizing # isort:skip
@@ -183,10 +184,6 @@ def main(args):
if torch.backends.mps.is_available():
accelerator.native_amp = False
if args.report_to == "wandb":
if not is_wandb_available():
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -363,6 +360,7 @@ def main(args):
"lr": args.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
use_deepspeed_optimizer = (
accelerator.state.deepspeed_plugin is not None
@@ -386,7 +384,11 @@ def main(args):
prodigy_use_bias_correction=args.prodigy_use_bias_correction,
prodigy_safeguard_warmup=args.prodigy_safeguard_warmup,
use_8bit=args.use_8bit,
use_4bit=args.use_4bit,
use_torchao=args.use_torchao,
use_deepspeed=use_deepspeed_optimizer,
use_cpu_offload_optimizer=args.use_cpu_offload_optimizer,
offload_gradients=args.offload_gradients,
)
# Dataset and DataLoader
@@ -451,24 +453,31 @@ def main(args):
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
if args.use_cpu_offload_optimizer:
lr_scheduler = None
accelerator.print(
"CPU Offload Optimizer cannot be used with DeepSpeed or builtin PyTorch LR Schedulers. If "
"you are training with those settings, they will be ignored."
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
# Prepare everything with our `accelerator`.
transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -494,7 +503,6 @@ def main(args):
# Train!
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
accelerator.print("***** Running training *****")
accelerator.print(f" Num trainable parameters = {num_trainable_parameters}")
@@ -551,6 +559,8 @@ def main(args):
torch.cuda.empty_cache()
torch.cuda.synchronize(accelerator.device)
alphas_cumprod = scheduler.alphas_cumprod.to(accelerator.device, dtype=torch.float32)
for epoch in range(first_epoch, args.num_train_epochs):
transformer.train()
@@ -618,8 +628,7 @@ def main(args):
model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)
alphas_cumprod = scheduler.alphas_cumprod[timesteps]
weights = 1 / (1 - alphas_cumprod)
weights = 1 / (1 - alphas_cumprod[timesteps])
while len(weights.shape) < len(model_pred.shape):
weights = weights.unsqueeze(-1)
@@ -641,7 +650,8 @@ def main(args):
optimizer.step()
optimizer.zero_grad()
lr_scheduler.step()
if not args.use_cpu_offload_optimizer:
lr_scheduler.step()
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
@@ -674,9 +684,10 @@ def main(args):
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
logs = {
"loss": loss.detach().item(),
"lr": lr_scheduler.get_last_lr()[0],
"lr": last_lr,
"gradient_norm_before_clip": gradient_norm_before_clip,
"gradient_norm_after_clip": gradient_norm_after_clip,
}
training/utils.py
+82 -29
View File
@@ -1,4 +1,5 @@
import gc
import inspect
from typing import Optional, Tuple, Union
import torch
@@ -22,7 +23,11 @@ def get_optimizer(
prodigy_use_bias_correction: bool = False,
prodigy_safeguard_warmup: bool = False,
use_8bit: bool = False,
use_4bit: bool = False,
use_torchao: bool = False,
use_deepspeed: bool = False,
use_cpu_offload_optimizer: bool = False,
offload_gradients: bool = False,
) -> torch.optim.Optimizer:
optimizer_name = optimizer_name.lower()
@@ -38,18 +43,32 @@ def get_optimizer(
weight_decay=weight_decay,
)
if use_8bit and use_4bit:
raise ValueError("Cannot set both `use_8bit` and `use_4bit` to True.")
if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer:
try:
import torchao
torchao.__version__
except ImportError:
raise ImportError(
"To use optimizers from torchao, please install the torchao library: `USE_CPP=0 pip install torchao`."
)
if not use_torchao and use_4bit:
raise ValueError("4-bit Optimizers are only supported with torchao.")
# Optimizer creation
supported_optimizers = ["adam", "adamw", "prodigy"]
supported_optimizers = ["adam", "adamw", "prodigy", "came"]
if optimizer_name not in supported_optimizers:
logger.warning(
f"Unsupported choice of optimizer: {optimizer_name}. Supported optimizers include {supported_optimizers}. Defaulting to `AdamW`."
)
optimizer_name = "adamw"
if use_8bit and optimizer_name not in ["adam", "adamw"]:
logger.warning(
f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was set to {optimizer_name}."
)
if (use_8bit or use_4bit) and optimizer_name not in ["adam", "adamw"]:
raise ValueError("`use_8bit` and `use_4bit` can only be used with the Adam and AdamW optimizers.")
if use_8bit:
try:
@@ -60,24 +79,32 @@ def get_optimizer(
)
if optimizer_name == "adamw":
optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
if use_torchao:
from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit
optimizer = optimizer_class(
params_to_optimize,
betas=(beta1, beta2),
eps=epsilon,
weight_decay=weight_decay,
)
optimizer_class = AdamW8bit if use_8bit else AdamW4bit if use_4bit else torch.optim.AdamW
else:
optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
init_kwargs = {
"betas": (beta1, beta2),
"eps": epsilon,
"weight_decay": weight_decay,
}
elif optimizer_name == "adam":
optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
if use_torchao:
from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit
optimizer = optimizer_class(
params_to_optimize,
betas=(beta1, beta2),
eps=epsilon,
weight_decay=weight_decay,
)
optimizer_class = Adam8bit if use_8bit else Adam4bit if use_4bit else torch.optim.Adam
else:
optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
init_kwargs = {
"betas": (beta1, beta2),
"eps": epsilon,
"weight_decay": weight_decay,
}
elif optimizer_name == "prodigy":
try:
@@ -92,17 +119,43 @@ def get_optimizer(
"Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
)
optimizer = optimizer_class(
params_to_optimize,
lr=learning_rate,
betas=(beta1, beta2),
beta3=beta3,
weight_decay=weight_decay,
eps=epsilon,
decouple=prodigy_decouple,
use_bias_correction=prodigy_use_bias_correction,
safeguard_warmup=prodigy_safeguard_warmup,
init_kwargs = {
"lr": learning_rate,
"betas": (beta1, beta2),
"beta3": beta3,
"eps": epsilon,
"weight_decay": weight_decay,
"decouple": prodigy_decouple,
"use_bias_correction": prodigy_use_bias_correction,
"safeguard_warmup": prodigy_safeguard_warmup,
}
elif optimizer_name == "came":
try:
import came_pytorch
except ImportError:
raise ImportError("To use CAME, please install the came-pytorch library: `pip install came-pytorch`")
optimizer_class = came_pytorch.CAME
init_kwargs = {
"lr": learning_rate,
"eps": (1e-30, 1e-16),
"betas": (beta1, beta2, beta3),
"weight_decay": weight_decay,
}
if use_cpu_offload_optimizer:
from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
if "fused" in inspect.signature(optimizer_class.__init__).parameters:
init_kwargs.update({"fused": True})
optimizer = CPUOffloadOptimizer(
params_to_optimize, optimizer_class=optimizer_class, offload_gradients=offload_gradients, **init_kwargs
)
else:
optimizer = optimizer_class(params_to_optimize, **init_kwargs)
return optimizer