diff --git a/Makefile b/Makefile index 78644f9..be64981 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ check_dirs := training tests quality: ruff check $(check_dirs) - ruff format --check $(check_dirs) setup.py + ruff format --check $(check_dirs) style: ruff check $(check_dirs) --fix diff --git a/training/cogvideox_image_to_video_lora.py b/training/cogvideox_image_to_video_lora.py index da1973d..305aac4 100644 --- a/training/cogvideox_image_to_video_lora.py +++ b/training/cogvideox_image_to_video_lora.py @@ -26,6 +26,7 @@ from typing import Any, Dict import diffusers import torch import transformers +import wandb from accelerate import Accelerator, DistributedType from accelerate.logging import get_logger from accelerate.utils import ( @@ -52,8 +53,6 @@ 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, VideoDatasetWithResizeAndRectangleCrop # isort:skip @@ -385,13 +384,15 @@ def main(args): transformer_lora_layers_to_save = None for model in models: - if isinstance(model, type(unwrap_model(transformer))): + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + model = unwrap_model(model) transformer_lora_layers_to_save = get_peft_model_state_dict(model) else: raise ValueError(f"unexpected save model: {model.__class__}") # make sure to pop weight so that corresponding model is not saved again - weights.pop() + if weights: + weights.pop() CogVideoXImageToVideoPipeline.save_lora_weights( output_dir, @@ -401,13 +402,20 @@ def main(args): def load_model_hook(models, input_dir): transformer_ = None - while len(models) > 0: - model = models.pop() + # This is a bit of a hack but I don't know any other solution. + if not accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() - if isinstance(model, type(unwrap_model(transformer))): - transformer_ = model - else: - raise ValueError(f"Unexpected save model: {model.__class__}") + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}") + else: + transformer_ = CogVideoXTransformer3DModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="transformer" + ) + transformer_.add_adapter(transformer_lora_config) lora_state_dict = CogVideoXImageToVideoPipeline.lora_state_dict(input_dir) @@ -795,12 +803,15 @@ def main(args): 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": last_lr, - "gradient_norm_before_clip": gradient_norm_before_clip, - "gradient_norm_after_clip": gradient_norm_after_clip, - } + logs = {"loss": loss.detach().item(), "lr": last_lr} + # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555 + if accelerator.distributed_type != DistributedType.DEEPSPEED: + logs.update( + { + "gradient_norm_before_clip": gradient_norm_before_clip, + "gradient_norm_after_clip": gradient_norm_after_clip, + } + ) progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) diff --git a/training/cogvideox_text_to_video_lora.py b/training/cogvideox_text_to_video_lora.py index 1c4d31d..cb94d76 100644 --- a/training/cogvideox_text_to_video_lora.py +++ b/training/cogvideox_text_to_video_lora.py @@ -25,6 +25,7 @@ from typing import Any, Dict import diffusers import torch import transformers +import wandb from accelerate import Accelerator, DistributedType from accelerate.logging import get_logger from accelerate.utils import ( @@ -51,8 +52,6 @@ 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, VideoDatasetWithResizeAndRectangleCrop # isort:skip @@ -358,13 +357,15 @@ def main(args): transformer_lora_layers_to_save = None for model in models: - if isinstance(model, type(unwrap_model(transformer))): + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + model = unwrap_model(model) transformer_lora_layers_to_save = get_peft_model_state_dict(model) else: raise ValueError(f"unexpected save model: {model.__class__}") # make sure to pop weight so that corresponding model is not saved again - weights.pop() + if weights: + weights.pop() CogVideoXPipeline.save_lora_weights( output_dir, @@ -374,13 +375,20 @@ def main(args): def load_model_hook(models, input_dir): transformer_ = None - while len(models) > 0: - model = models.pop() + # This is a bit of a hack but I don't know any other solution. + if not accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() - if isinstance(model, type(unwrap_model(transformer))): - transformer_ = model - else: - raise ValueError(f"Unexpected save model: {model.__class__}") + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}") + else: + transformer_ = CogVideoXTransformer3DModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="transformer" + ) + transformer_.add_adapter(transformer_lora_config) lora_state_dict = CogVideoXPipeline.lora_state_dict(input_dir) @@ -553,7 +561,7 @@ def main(args): # We need to initialize the trackers we use, and also store our configuration. # The trackers initializes automatically on the main process. - if accelerator.is_main_process: + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: tracker_name = args.tracker_name or "cogvideox-lora" accelerator.init_trackers(tracker_name, config=vars(args)) @@ -731,7 +739,7 @@ def main(args): progress_bar.update(1) global_step += 1 - if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED: + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: if global_step % args.checkpointing_steps == 0: # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` if args.checkpoints_total_limit is not None: @@ -758,12 +766,15 @@ def main(args): 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": last_lr, - "gradient_norm_before_clip": gradient_norm_before_clip, - "gradient_norm_after_clip": gradient_norm_after_clip, - } + logs = {"loss": loss.detach().item(), "lr": last_lr} + # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555 + if accelerator.distributed_type != DistributedType.DEEPSPEED: + logs.update( + { + "gradient_norm_before_clip": gradient_norm_before_clip, + "gradient_norm_after_clip": gradient_norm_after_clip, + } + ) progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) diff --git a/training/cogvideox_text_to_video_sft.py b/training/cogvideox_text_to_video_sft.py index 033dd36..72592d6 100644 --- a/training/cogvideox_text_to_video_sft.py +++ b/training/cogvideox_text_to_video_sft.py @@ -25,7 +25,8 @@ from typing import Any, Dict import diffusers import torch import transformers -from accelerate import Accelerator, DistributedType +import wandb +from accelerate import Accelerator, DistributedType, init_empty_weights from accelerate.logging import get_logger from accelerate.utils import ( DistributedDataParallelKwargs, @@ -50,8 +51,6 @@ 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, VideoDatasetWithResizeAndRectangleCrop # isort:skip @@ -336,8 +335,9 @@ def main(args): def save_model_hook(models, weights, output_dir): if accelerator.is_main_process: for model in models: - if isinstance(model, type(unwrap_model(transformer))): + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): model: CogVideoXTransformer3DModel + model = unwrap_model(model) model.save_pretrained( os.path.join(output_dir, "transformer"), safe_serialization=True, max_shard_size="5GB" ) @@ -345,22 +345,32 @@ def main(args): raise ValueError(f"Unexpected save model: {model.__class__}") # make sure to pop weight so that corresponding model is not saved again - weights.pop() + if weights: + weights.pop() def load_model_hook(models, input_dir): transformer_ = None + init_under_meta = False - while len(models) > 0: - model = models.pop() + # This is a bit of a hack but I don't know any other solution. + if not accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() - if isinstance(model, type(unwrap_model(transformer))): - transformer_: CogVideoXTransformer3DModel = model - else: - raise ValueError(f"Unexpected save model: {model.__class__.__name__}") + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}") + else: + with init_empty_weights(): + transformer_ = CogVideoXTransformer3DModel.from_config( + args.pretrained_model_name_or_path, subfolder="transformer" + ) + init_under_meta = True load_model = CogVideoXTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer")) transformer_.register_to_config(**load_model.config) - transformer_.load_state_dict(load_model.state_dict()) + transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta) del load_model # Make sure the trainable params are in float32. This is again needed since the base models @@ -722,12 +732,15 @@ def main(args): 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": last_lr, - "gradient_norm_before_clip": gradient_norm_before_clip, - "gradient_norm_after_clip": gradient_norm_after_clip, - } + logs = {"loss": loss.detach().item(), "lr": last_lr} + # gradnorm + deepspeed: https://github.com/microsoft/DeepSpeed/issues/4555 + if accelerator.distributed_type != DistributedType.DEEPSPEED: + logs.update( + { + "gradient_norm_before_clip": gradient_norm_before_clip, + "gradient_norm_after_clip": gradient_norm_after_clip, + } + ) progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step)