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

fix: resuming from a checkpoint when using deepspeed. (#38)

Browse Source 
* fix: resuming from a checkpoint when using deepspeed.

* remove changes to prepare_dataset.py

* propagate to others.

* tackle gradnorm.

---------

Co-authored-by: Aryan <aryan@huggingface.co>
This commit is contained in:
Sayak Paul
2024-10-20 02:56:32 +05:30
committed by GitHub
parent f0d9908bbd
commit 6c00cf094b
4 changed files with 88 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -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
training/cogvideox_image_to_video_lora.py
+27 -16
View File
@@ -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)
training/cogvideox_text_to_video_lora.py
+29 -18
View File
@@ -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)
training/cogvideox_text_to_video_sft.py
+31 -18
View File
@@ -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)