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

Full Finetuning for LTX possibily extended to other models. (#192)

Browse Source 
* Full Finetuning for LTX possibily extended to other models.

* Change name of the flag

* Used disable grad for component on lora fine tuning enabled

* Suggestions Addressed
Renamed to SFT
Added 2 other models.
Testing required.

* Switching to Full FineTuning

* Run linter.

* parse subfolder when needed.

* tackle saving and loading hooks.

* tackle validation.

* fix subfolder bug.

* remove __class__.

* refactor

* remove unnecessary changes

* handle saving of final model weights correctly

* remove unnecessary changes

* LTX uses a default frame rate of 24 FPS
We need to modify the output validation framerate to match that value.
Add Framerate args.
Add Update video output and inference frame rate

* There was a results_args mapping that needed to be modified.

* update

* update README

* Update README.md

* update docs

* add training configuration in cogvideox

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Aryan <aryan@huggingface.co>
Co-authored-by: Aryan <contact.aryanvs@gmail.com>
This commit is contained in:
ArEnSc
2025-01-13 13:25:39 -05:00
committed by GitHub
parent a3898075d1
commit e5df80cc36
17 changed files with 365 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/trainer.py
+145 -107
View File
@@ -5,7 +5,7 @@ import os
import random
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict
from typing import Any, Dict, List
import diffusers
import torch
@@ -21,6 +21,7 @@ from accelerate.utils import (
gather_object,
set_seed,
)
from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.optimization import get_scheduler
@@ -242,16 +243,7 @@ class Trainer:
condition_components = self.model_config["load_condition_models"](**self._get_load_components_kwargs())
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)
self._disable_grad_for_components([self.text_encoder, self.text_encoder_2, self.text_encoder_3])
if self.args.caption_dropout_p > 0 and self.args.caption_dropout_technique == "empty":
logger.warning(
@@ -305,12 +297,7 @@ class Trainer:
latent_components = self.model_config["load_latent_models"](**self._get_load_components_kwargs())
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)
self._disable_grad_for_components([self.vae])
if self.vae is not None:
if self.args.enable_slicing:
@@ -371,24 +358,22 @@ class Trainer:
diffusion_components = self.model_config["load_diffusion_models"](**self._get_load_components_kwargs())
self._set_components(diffusion_components)
# 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)
components = [self.text_encoder, self.text_encoder_2, self.text_encoder_3, self.vae]
self._disable_grad_for_components(components)
if self.args.training_type == "full-finetune":
logger.info("Finetuning transformer with no additional parameters")
self._enable_grad_for_components([self.transformer])
else:
logger.info("Finetuning transformer with PEFT parameters")
self._disable_grad_for_components([self.transformer])
# 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.
weight_dtype = self._get_training_dtype(accelerator=self.state.accelerator)
if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
# due to pytorch#99272, MPS does not yet support bfloat16.
# Due to pytorch#99272, MPS does not yet support bfloat16.
raise ValueError(
"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
)
@@ -406,13 +391,16 @@ class Trainer:
if self.args.gradient_checkpointing:
self.transformer.enable_gradient_checkpointing()
transformer_lora_config = LoraConfig(
r=self.args.rank,
lora_alpha=self.args.lora_alpha,
init_lora_weights=True,
target_modules=self.args.target_modules,
)
self.transformer.add_adapter(transformer_lora_config)
if self.args.training_type == "lora":
transformer_lora_config = LoraConfig(
r=self.args.rank,
lora_alpha=self.args.lora_alpha,
init_lora_weights=True,
target_modules=self.args.target_modules,
)
self.transformer.add_adapter(transformer_lora_config)
else:
transformer_lora_config = None
# Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
if self.args.allow_tf32 and torch.cuda.is_available():
@@ -432,7 +420,8 @@ class Trainer:
type(unwrap_model(self.state.accelerator, self.transformer)),
):
model = unwrap_model(self.state.accelerator, model)
transformer_lora_layers_to_save = get_peft_model_state_dict(model)
if self.args.training_type == "lora":
transformer_lora_layers_to_save = get_peft_model_state_dict(model)
else:
raise ValueError(f"Unexpected save model: {model.__class__}")
@@ -440,10 +429,18 @@ class Trainer:
if weights:
weights.pop()
self.model_config["pipeline_cls"].save_lora_weights(
output_dir,
transformer_lora_layers=transformer_lora_layers_to_save,
)
if self.args.training_type == "lora":
self.model_config["pipeline_cls"].save_lora_weights(
output_dir,
transformer_lora_layers=transformer_lora_layers_to_save,
)
else:
model.save_pretrained(os.path.join(output_dir, "transformer"))
# In some cases, the scheduler needs to be loaded with specific config (e.g. in CogVideoX). Since we need
# to able to load all diffusion components from a specific checkpoint folder during validation, we need to
# ensure the scheduler config is serialized as well.
self.scheduler.save_pretrained(os.path.join(output_dir, "scheduler"))
def load_model_hook(models, input_dir):
if not self.state.accelerator.distributed_type == DistributedType.DEEPSPEED:
@@ -459,33 +456,39 @@ class Trainer:
f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}"
)
else:
transformer_ = unwrap_model(self.state.accelerator, self.transformer).__class__.from_pretrained(
self.args.pretrained_model_name_or_path, subfolder="transformer"
)
transformer_.add_adapter(transformer_lora_config)
transformer_cls_ = unwrap_model(self.state.accelerator, self.transformer).__class__
lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
transformer_state_dict = {
f'{k.replace("transformer.", "")}': v
for k, v in lora_state_dict.items()
if k.startswith("transformer.")
}
incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
if incompatible_keys is not None:
# check only for unexpected keys
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
if unexpected_keys:
logger.warning(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
f" {unexpected_keys}. "
if self.args.training_type == "lora":
transformer_ = transformer_cls_.from_pretrained(
self.args.pretrained_model_name_or_path, subfolder="transformer"
)
transformer_.add_adapter(transformer_lora_config)
lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
transformer_state_dict = {
f'{k.replace("transformer.", "")}': v
for k, v in lora_state_dict.items()
if k.startswith("transformer.")
}
incompatible_keys = set_peft_model_state_dict(
transformer_, transformer_state_dict, adapter_name="default"
)
if incompatible_keys is not None:
# check only for unexpected keys
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
if unexpected_keys:
logger.warning(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
f" {unexpected_keys}. "
)
# Make sure the trainable params are in float32. This is again needed since the base models
# are in `weight_dtype`. More details:
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
if self.args.mixed_precision == "fp16":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([transformer_])
# Make sure the trainable params are in float32. This is again needed since the base models
# are in `weight_dtype`. More details:
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
if self.args.mixed_precision == "fp16" and self.args.training_type == "lora":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([transformer_], dtype=torch.float32)
else:
transformer_ = transformer_cls_.from_pretrained(os.path.join(input_dir, "transformer"))
self.state.accelerator.register_save_state_pre_hook(save_model_hook)
self.state.accelerator.register_load_state_pre_hook(load_model_hook)
@@ -497,7 +500,7 @@ class Trainer:
self.state.train_steps = self.args.train_steps
# Make sure the trainable params are in float32
if self.args.mixed_precision == "fp16":
if self.args.mixed_precision == "fp16" and self.args.training_type == "lora":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([self.transformer], dtype=torch.float32)
@@ -510,13 +513,13 @@ class Trainer:
* self.state.accelerator.num_processes
)
transformer_lora_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
transformer_parameters_with_lr = {
"params": transformer_lora_parameters,
"params": transformer_trainable_parameters,
"lr": self.state.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
self.state.num_trainable_parameters = sum(p.numel() for p in transformer_lora_parameters)
self.state.num_trainable_parameters = sum(p.numel() for p in transformer_trainable_parameters)
use_deepspeed_opt = (
self.state.accelerator.state.deepspeed_plugin is not None
@@ -608,6 +611,12 @@ class Trainer:
)
self.vae_config = FrozenDict(**vae_config)
# In some cases, the scheduler needs to be loaded with specific config (e.g. in CogVideoX). Since we need
# to able to load all diffusion components from a specific checkpoint folder during validation, we need to
# ensure the scheduler config is serialized as well.
if self.args.training_type == "full-finetune":
self.scheduler.save_pretrained(os.path.join(self.args.output_dir, "scheduler"))
self.state.train_batch_size = (
self.args.batch_size * self.state.accelerator.num_processes * self.args.gradient_accumulation_steps
)
@@ -872,14 +881,17 @@ class Trainer:
accelerator.wait_for_everyone()
if accelerator.is_main_process:
# TODO: consider factoring this out when supporting other types of training algos.
self.transformer = unwrap_model(accelerator, self.transformer)
transformer_lora_layers = get_peft_model_state_dict(self.transformer)
transformer = unwrap_model(accelerator, self.transformer)
self.model_config["pipeline_cls"].save_lora_weights(
save_directory=self.args.output_dir,
transformer_lora_layers=transformer_lora_layers,
)
if self.args.training_type == "lora":
transformer_lora_layers = get_peft_model_state_dict(transformer)
self.model_config["pipeline_cls"].save_lora_weights(
save_directory=self.args.output_dir,
transformer_lora_layers=transformer_lora_layers,
)
else:
transformer.save_pretrained(os.path.join(self.args.output_dir, "transformer"))
self.validate(step=global_step, final_validation=True)
@@ -910,35 +922,7 @@ class Trainer:
memory_statistics = get_memory_statistics()
logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}")
if not final_validation:
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
tokenizer=self.tokenizer,
text_encoder=self.text_encoder,
tokenizer_2=self.tokenizer_2,
text_encoder_2=self.text_encoder_2,
transformer=unwrap_model(accelerator, self.transformer),
vae=self.vae,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
else:
# `torch_dtype` is manually set within `initialize_pipeline()`.
self._delete_components()
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
pipeline.load_lora_weights(self.args.output_dir)
pipeline = self._get_and_prepare_pipeline_for_validation(final_validation=final_validation)
all_processes_artifacts = []
prompts_to_filenames = {}
@@ -953,7 +937,7 @@ class Trainer:
height = self.args.validation_heights[i]
width = self.args.validation_widths[i]
num_frames = self.args.validation_num_frames[i]
frame_rate = self.args.validation_frame_rate
if image is not None:
image = load_image(image)
if video is not None:
@@ -971,6 +955,7 @@ class Trainer:
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_videos_per_prompt=self.args.num_validation_videos_per_prompt,
generator=torch.Generator(device=accelerator.device).manual_seed(
self.args.seed if self.args.seed is not None else 0
@@ -1010,7 +995,7 @@ class Trainer:
elif artifact_type == "video":
logger.debug(f"Saving video to {filename}")
# TODO: this should be configurable here as well as in validation runs where we call the pipeline that has `fps`.
export_to_video(artifact_value, filename, fps=15)
export_to_video(artifact_value, filename, fps=frame_rate)
artifact_value = wandb.Video(filename, caption=prompt)
all_processes_artifacts.append(artifact_value)
@@ -1144,3 +1129,56 @@ class Trainer:
elif self.state.accelerator.mixed_precision == "bf16":
weight_dtype = torch.bfloat16
return weight_dtype
def _get_and_prepare_pipeline_for_validation(self, final_validation: bool = False) -> DiffusionPipeline:
accelerator = self.state.accelerator
if not final_validation:
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
tokenizer=self.tokenizer,
text_encoder=self.text_encoder,
tokenizer_2=self.tokenizer_2,
text_encoder_2=self.text_encoder_2,
transformer=unwrap_model(accelerator, self.transformer),
vae=self.vae,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
else:
self._delete_components()
# Load the transformer weights from the final checkpoint if performing full-finetune
transformer = None
if self.args.training_type == "full-finetune":
transformer = self.model_config["load_diffusion_models"](model_id=self.args.output_dir)["transformer"]
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
transformer=transformer,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
# Load the LoRA weights if performing LoRA finetuning
if self.args.training_type == "lora":
pipeline.load_lora_weights(self.args.output_dir)
return pipeline
def _disable_grad_for_components(self, components: List[torch.nn.Module]):
for component in components:
if component is not None:
component.requires_grad_(False)
def _enable_grad_for_components(self, components: List[torch.nn.Module]):
for component in components:
if component is not None:
component.requires_grad_(True)