596 lines
25 KiB
Python
596 lines
25 KiB
Python
import gc
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from functools import partial
|
|
from glob import glob
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import safetensors.torch
|
|
|
|
import deepspeed
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.distributed as dist
|
|
from torch.utils.data import DataLoader
|
|
from torch.distributed.optim import ZeroRedundancyOptimizer
|
|
from torchvision.transforms import functional as TF
|
|
from diffusers.models import AutoencoderKL
|
|
from transformers import BertModel, BertTokenizer, logging as tf_logging
|
|
|
|
from hydit.config import get_args
|
|
from hydit.lr_scheduler import WarmupLR
|
|
from hydit.data_loader.arrow_load_stream import TextImageArrowStream
|
|
from hydit.diffusion import create_diffusion
|
|
from hydit.ds_config import deepspeed_config_from_args
|
|
from hydit.modules.ema import EMA
|
|
from hydit.modules.fp16_layers import Float16Module
|
|
from hydit.modules.models import HUNYUAN_DIT_MODELS
|
|
from hydit.modules.posemb_layers import init_image_posemb
|
|
from hydit.utils.tools import create_logger, set_seeds, create_exp_folder, get_trainable_params
|
|
from IndexKits.index_kits import ResolutionGroup
|
|
from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler
|
|
from peft import LoraConfig, get_peft_model
|
|
|
|
|
|
def deepspeed_initialize(args, logger, model, opt, deepspeed_config):
|
|
logger.info(f"Initialize deepspeed...")
|
|
logger.info(f" Using deepspeed optimizer")
|
|
|
|
def get_learning_rate_scheduler(warmup_min_lr, lr, warmup_num_steps, opt):
|
|
return WarmupLR(opt, warmup_min_lr, lr, warmup_num_steps)
|
|
|
|
logger.info(
|
|
f" Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_num_steps={args.warmup_num_steps}")
|
|
logger.info(
|
|
f" deepspeed_config={deepspeed_config}")
|
|
|
|
model, opt, _, scheduler = deepspeed.initialize(
|
|
model=model,
|
|
model_parameters=get_trainable_params(model),
|
|
config_params=deepspeed_config,
|
|
args=args,
|
|
lr_scheduler=partial(get_learning_rate_scheduler, args.warmup_min_lr,
|
|
args.lr, args.warmup_num_steps) if args.warmup_num_steps > 0 else None,
|
|
)
|
|
return model, opt, scheduler
|
|
|
|
|
|
def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir):
|
|
def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"):
|
|
cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}"
|
|
if rank == 0:
|
|
if args.use_fp16:
|
|
model.module.module.save_pretrained(cur_ckpt_save_dir)
|
|
else:
|
|
model.module.save_pretrained(cur_ckpt_save_dir)
|
|
|
|
checkpoint_path = "[Not rank 0. Disabled output.]"
|
|
|
|
client_state = {
|
|
"steps": train_steps,
|
|
"epoch": epoch,
|
|
"args": args
|
|
}
|
|
if ema is not None:
|
|
client_state['ema'] = ema.state_dict()
|
|
|
|
dst_paths = []
|
|
if train_steps % args.ckpt_every == 0:
|
|
checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
|
|
try:
|
|
if args.training_parts == "lora":
|
|
save_lora_weight(checkpoint_dir, client_state,
|
|
tag=f"{train_steps:07d}.pt")
|
|
else:
|
|
model.save_checkpoint(
|
|
checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt")
|
|
dst_paths.append(checkpoint_path)
|
|
logger.info(f"Saved checkpoint to {checkpoint_path}")
|
|
except:
|
|
logger.error(f"Saved failed to {checkpoint_path}")
|
|
|
|
if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps:
|
|
save_name = "latest.pt"
|
|
checkpoint_path = f"{checkpoint_dir}/{save_name}"
|
|
try:
|
|
if args.training_parts == "lora":
|
|
save_lora_weight(checkpoint_dir, client_state,
|
|
tag=f"{save_name}")
|
|
else:
|
|
model.save_checkpoint(
|
|
checkpoint_dir, client_state=client_state, tag=f"{save_name}")
|
|
dst_paths.append(checkpoint_path)
|
|
logger.info(f"Saved checkpoint to {checkpoint_path}")
|
|
except:
|
|
logger.error(f"Saved failed to {checkpoint_path}")
|
|
|
|
dist.barrier()
|
|
if rank == 0 and len(dst_paths) > 0:
|
|
# Delete optimizer states to avoid occupying too much disk space.
|
|
for dst_path in dst_paths:
|
|
for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"):
|
|
os.remove(opt_state_path)
|
|
|
|
return checkpoint_path
|
|
|
|
|
|
@torch.no_grad()
|
|
def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img):
|
|
try:
|
|
from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
|
|
except:
|
|
from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
|
|
|
|
image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch
|
|
|
|
# clip & mT5 text embedding
|
|
text_embedding = text_embedding.to(device)
|
|
text_embedding_mask = text_embedding_mask.to(device)
|
|
encoder_hidden_states = text_encoder(
|
|
text_embedding.to(device),
|
|
attention_mask=text_embedding_mask.to(device),
|
|
)[0]
|
|
text_embedding_t5 = text_embedding_t5.to(device).squeeze(1)
|
|
text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1)
|
|
with torch.no_grad():
|
|
output_t5 = text_encoder_t5(
|
|
input_ids=text_embedding_t5,
|
|
attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None,
|
|
output_hidden_states=True
|
|
)
|
|
encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach(
|
|
)
|
|
|
|
# additional condition
|
|
image_meta_size = kwargs['image_meta_size'].to(device)
|
|
style = kwargs['style'].to(device)
|
|
|
|
if args.extra_fp16:
|
|
image = image.half()
|
|
image_meta_size = image_meta_size.half() if image_meta_size is not None else None
|
|
|
|
# Map input images to latent space + normalize latents:
|
|
image = image.to(device)
|
|
vae_scaling_factor = vae.config.scaling_factor
|
|
latents = vae.encode(image).latent_dist.sample().mul_(vae_scaling_factor)
|
|
|
|
# positional embedding
|
|
_, _, height, width = image.shape
|
|
reso = f"{height}x{width}"
|
|
cos_cis_img, sin_cis_img = freqs_cis_img[reso]
|
|
|
|
# Model conditions
|
|
model_kwargs = dict(
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
text_embedding_mask=text_embedding_mask,
|
|
encoder_hidden_states_t5=encoder_hidden_states_t5,
|
|
text_embedding_mask_t5=text_embedding_mask_t5,
|
|
image_meta_size=image_meta_size,
|
|
style=style,
|
|
cos_cis_img=cos_cis_img,
|
|
sin_cis_img=sin_cis_img,
|
|
)
|
|
|
|
return latents, model_kwargs
|
|
|
|
|
|
def Core(args):
|
|
if args.training_parts == "lora":
|
|
args.use_ema = False
|
|
|
|
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
|
|
|
|
dist.init_process_group("nccl")
|
|
world_size = dist.get_world_size()
|
|
batch_size = args.batch_size
|
|
grad_accu_steps = args.grad_accu_steps
|
|
global_batch_size = world_size * batch_size * grad_accu_steps
|
|
|
|
rank = dist.get_rank()
|
|
device = rank % torch.cuda.device_count()
|
|
seed = args.global_seed * world_size + rank
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.cuda.set_device(device)
|
|
print(f"Starting rank={rank}, seed={seed}, world_size={world_size}.")
|
|
deepspeed_config = deepspeed_config_from_args(args, global_batch_size)
|
|
|
|
# Setup an experiment folder
|
|
experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank)
|
|
|
|
# Log all the arguments
|
|
logger.info(sys.argv)
|
|
logger.info(str(args))
|
|
# Save to a json file
|
|
args_dict = vars(args)
|
|
args_dict['world_size'] = world_size
|
|
with open(f"{experiment_dir}/args.json", 'w') as f:
|
|
json.dump(args_dict, f, indent=4)
|
|
|
|
# Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel."
|
|
# If needed, just comment the following line.
|
|
tf_logging.set_verbosity_error()
|
|
|
|
# ===========================================================================
|
|
# Building HYDIT
|
|
# ===========================================================================
|
|
|
|
logger.info("Building HYDIT Model.")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Training sample base size, such as 256/512/1024. Notice that this size is
|
|
# just a base size, not necessary the actual size of training samples. Actual
|
|
# size of the training samples are correlated with `resolutions` when enabling
|
|
# multi-resolution training.
|
|
# ---------------------------------------------------------------------------
|
|
image_size = args.image_size
|
|
if len(image_size) == 1:
|
|
image_size = [image_size[0], image_size[0]]
|
|
if len(image_size) != 2:
|
|
raise ValueError(f"Invalid image size: {args.image_size}")
|
|
assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \
|
|
f"got {image_size}"
|
|
latent_size = [image_size[0] // 8, image_size[1] // 8]
|
|
|
|
# initialize model by deepspeed
|
|
assert args.deepspeed, f"Must enable deepspeed in this script: train_deepspeed.py"
|
|
with deepspeed.zero.Init(data_parallel_group=torch.distributed.group.WORLD,
|
|
remote_device=None if args.remote_device == 'none' else args.remote_device,
|
|
config_dict_or_path=deepspeed_config,
|
|
mpu=None,
|
|
enabled=args.zero_stage == 3):
|
|
model = HUNYUAN_DIT_MODELS[args.model](args,
|
|
input_size=latent_size,
|
|
log_fn=logger.info,
|
|
)
|
|
# Multi-resolution / Single-resolution training.
|
|
if args.multireso:
|
|
resolutions = ResolutionGroup(image_size[0],
|
|
align=16,
|
|
step=args.reso_step,
|
|
target_ratios=args.target_ratios).data
|
|
else:
|
|
resolutions = ResolutionGroup(image_size[0],
|
|
align=16,
|
|
target_ratios=['1:1']).data
|
|
|
|
freqs_cis_img = init_image_posemb(args.rope_img,
|
|
resolutions=resolutions,
|
|
patch_size=model.patch_size,
|
|
hidden_size=model.hidden_size,
|
|
num_heads=model.num_heads,
|
|
log_fn=logger.info,
|
|
rope_real=args.rope_real,
|
|
)
|
|
|
|
# Create EMA model and convert to fp16 if needed.
|
|
ema = None
|
|
if args.use_ema:
|
|
ema = EMA(args, model, device, logger)
|
|
|
|
# Setup FP16 main model:
|
|
if args.use_fp16:
|
|
model = Float16Module(model, args)
|
|
logger.info(
|
|
f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}")
|
|
|
|
diffusion = create_diffusion(
|
|
noise_schedule=args.noise_schedule,
|
|
predict_type=args.predict_type,
|
|
learn_sigma=args.learn_sigma,
|
|
mse_loss_weight_type=args.mse_loss_weight_type,
|
|
beta_start=args.beta_start,
|
|
beta_end=args.beta_end,
|
|
noise_offset=args.noise_offset,
|
|
)
|
|
|
|
try:
|
|
from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
|
|
except:
|
|
from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
|
|
|
|
# Setup VAE
|
|
logger.info(f" Loading vae from {VAE_EMA_PATH}")
|
|
vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH)
|
|
# Setup BERT text encoder
|
|
logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}")
|
|
text_encoder = BertModel.from_pretrained(
|
|
TEXT_ENCODER, False, revision=None)
|
|
# Setup BERT tokenizer:
|
|
logger.info(f" Loading Bert tokenizer from {TOKENIZER}")
|
|
tokenizer = BertTokenizer.from_pretrained(TOKENIZER)
|
|
# Setup T5 text encoder
|
|
from hydit.modules.text_encoder import MT5Embedder
|
|
mt5_path = T5_ENCODER['MT5']
|
|
if mt5_path is None:
|
|
embedder_t5 = CustomizeEmbeds()
|
|
else:
|
|
embedder_t5 = MT5Embedder(
|
|
mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5)
|
|
tokenizer_t5 = embedder_t5.tokenizer
|
|
text_encoder_t5 = embedder_t5.model
|
|
|
|
if args.extra_fp16:
|
|
logger.info(f" Using fp16 for extra modules: vae, text_encoder")
|
|
vae = vae.half().to(device)
|
|
text_encoder = text_encoder.half().to(device)
|
|
text_encoder_t5 = text_encoder_t5.half().to(device)
|
|
else:
|
|
vae = vae.to(device)
|
|
text_encoder = text_encoder.to(device)
|
|
text_encoder_t5 = text_encoder_t5.to(device)
|
|
|
|
logger.info(
|
|
f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}")
|
|
logger.info(" Using deepspeed optimizer")
|
|
opt = None
|
|
|
|
# ===========================================================================
|
|
# Building Dataset
|
|
# ===========================================================================
|
|
|
|
logger.info(f"Building Streaming Dataset.")
|
|
logger.info(f" Loading index file {args.index_file} (v2)")
|
|
|
|
dataset = TextImageArrowStream(args=args,
|
|
resolution=image_size[0],
|
|
random_flip=args.random_flip,
|
|
log_fn=logger.info,
|
|
index_file=args.index_file,
|
|
multireso=args.multireso,
|
|
batch_size=batch_size,
|
|
world_size=world_size,
|
|
random_shrink_size_cond=args.random_shrink_size_cond,
|
|
merge_src_cond=args.merge_src_cond,
|
|
uncond_p=args.uncond_p,
|
|
text_ctx_len=args.text_len,
|
|
tokenizer=tokenizer,
|
|
uncond_p_t5=args.uncond_p_t5,
|
|
text_ctx_len_t5=args.text_len_t5,
|
|
tokenizer_t5=tokenizer_t5,
|
|
)
|
|
if args.multireso:
|
|
sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
|
|
shuffle=False, drop_last=True, batch_size=batch_size)
|
|
else:
|
|
sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
|
|
shuffle=False, drop_last=True)
|
|
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler,
|
|
num_workers=args.num_workers, pin_memory=True, drop_last=True)
|
|
logger.info(f" Dataset contains {len(dataset):,} images.")
|
|
logger.info(f" Index file: {args.index_file}.")
|
|
if args.multireso:
|
|
logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} '
|
|
f'and base size {dataset.index_manager.base_size}')
|
|
logger.info(f'\n {dataset.index_manager.resolutions}')
|
|
|
|
# ===========================================================================
|
|
# Loading parameter
|
|
# ===========================================================================
|
|
|
|
logger.info(f"Loading parameter")
|
|
start_epoch = 0
|
|
start_epoch_step = 0
|
|
train_steps = 0
|
|
# Resume checkpoint if needed
|
|
# if args.resume is not None or len(args.resume) > 0:
|
|
if True:
|
|
model, ema, start_epoch, start_epoch_step, train_steps = model_resume(
|
|
args, model, ema, logger)
|
|
|
|
if args.training_parts == "lora":
|
|
lora_ckpt = args.lora_ckpt
|
|
if lora_ckpt is not None:
|
|
lastest_checkpoint = lora_ckpt
|
|
from peft.peft_model import PeftModel
|
|
print(f"Loading lora model from {lastest_checkpoint}")
|
|
if args.use_fp16:
|
|
model.module = PeftModel.from_pretrained(
|
|
model.module, lastest_checkpoint, is_trainable=True)
|
|
else:
|
|
model = PeftModel.from_pretrained(
|
|
model, lastest_checkpoint, is_trainable=True)
|
|
else:
|
|
loraconfig = LoraConfig(
|
|
r=args.rank,
|
|
lora_alpha=args.rank,
|
|
target_modules=args.target_modules
|
|
)
|
|
|
|
if args.use_fp16:
|
|
model.module = get_peft_model(model.module, loraconfig)
|
|
else:
|
|
model = get_peft_model(model, loraconfig)
|
|
|
|
logger.info(f" Training parts: {args.training_parts}")
|
|
|
|
model, opt, scheduler = deepspeed_initialize(
|
|
args, logger, model, opt, deepspeed_config)
|
|
|
|
# ===========================================================================
|
|
# Training
|
|
# ===========================================================================
|
|
|
|
# print model structure
|
|
# for name, param in model.named_parameters():
|
|
# print(name, param.size())
|
|
# raise Exception("stop")
|
|
|
|
model.train()
|
|
if args.use_ema:
|
|
ema.eval()
|
|
|
|
print(f" Worker {rank} ready.")
|
|
dist.barrier()
|
|
|
|
iters_per_epoch = len(loader)
|
|
logger.info(
|
|
" ****************************** Running training ******************************")
|
|
logger.info(f" Number GPUs: {world_size}")
|
|
logger.info(f" Number training samples: {len(dataset):,}")
|
|
logger.info(
|
|
f" Number parameters: {sum(p.numel() for p in model.parameters()):,}")
|
|
logger.info(
|
|
f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}")
|
|
logger.info(
|
|
" ------------------------------------------------------------------------------")
|
|
logger.info(f" Iters per epoch: {iters_per_epoch:,}")
|
|
logger.info(f" Batch size per device: {batch_size}")
|
|
logger.info(
|
|
f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)")
|
|
logger.info(f" Gradient Accu steps: {args.grad_accu_steps}")
|
|
logger.info(
|
|
f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}")
|
|
|
|
logger.info(
|
|
f" Training epochs: {start_epoch}/{args.epochs}")
|
|
logger.info(
|
|
f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}")
|
|
logger.info(
|
|
f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}")
|
|
logger.info(
|
|
" ------------------------------------------------------------------------------")
|
|
logger.info(f" Noise schedule: {args.noise_schedule}")
|
|
logger.info(
|
|
f" Beta limits: ({args.beta_start}, {args.beta_end})")
|
|
logger.info(f" Learn sigma: {args.learn_sigma}")
|
|
logger.info(f" Prediction type: {args.predict_type}")
|
|
logger.info(f" Noise offset: {args.noise_offset}")
|
|
|
|
logger.info(
|
|
" ------------------------------------------------------------------------------")
|
|
logger.info(
|
|
f" Using EMA model: {args.use_ema} ({args.ema_dtype})")
|
|
if args.use_ema:
|
|
logger.info(
|
|
f" Using EMA decay: {ema.max_value if args.use_ema else None}")
|
|
logger.info(
|
|
f" Using EMA warmup power: {ema.power if args.use_ema else None}")
|
|
logger.info(f" Using main model fp16: {args.use_fp16}")
|
|
logger.info(f" Using extra modules fp16: {args.extra_fp16}")
|
|
logger.info(
|
|
" ------------------------------------------------------------------------------")
|
|
logger.info(f" Experiment directory: {experiment_dir}")
|
|
logger.info(
|
|
" *******************************************************************************")
|
|
|
|
if args.gc_interval > 0:
|
|
gc.disable()
|
|
gc.collect()
|
|
|
|
# Variables for monitoring/logging purposes:
|
|
log_steps = 0
|
|
running_loss = 0
|
|
start_time = time.time()
|
|
|
|
if args.async_ema:
|
|
ema_stream = torch.cuda.Stream()
|
|
|
|
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
|
|
target_height=768, target_width=1280, train_steps=0)
|
|
pbar = PBar(args.epochs * len(loader))
|
|
|
|
# Training loop
|
|
for epoch in range(start_epoch, args.epochs):
|
|
logger.info(f" Start random shuffle with seed={seed}")
|
|
# Makesure all processors use the same seed to shuffle dataset.
|
|
dataset.shuffle(seed=args.global_seed + epoch, fast=True)
|
|
logger.info(f" End of random shuffle")
|
|
|
|
# Move sampler to start_index
|
|
if not args.multireso:
|
|
start_index = start_epoch_step * world_size * batch_size
|
|
if start_index != sampler.start_index:
|
|
sampler.start_index = start_index
|
|
# Reset start_epoch_step to zero, to ensure next epoch will start from the beginning.
|
|
start_epoch_step = 0
|
|
logger.info(f" Iters left this epoch: {len(loader):,}")
|
|
|
|
logger.info(f" Beginning epoch {epoch}...")
|
|
step = 0
|
|
for batch in loader:
|
|
step += 1
|
|
|
|
latents, model_kwargs = prepare_model_inputs(
|
|
args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img)
|
|
|
|
# training model by deepspeed while use fp16
|
|
if args.use_fp16:
|
|
if args.use_ema and args.async_ema:
|
|
with torch.cuda.stream(ema_stream):
|
|
ema.update(model.module.module, step=step)
|
|
torch.cuda.current_stream().wait_stream(ema_stream)
|
|
|
|
loss_dict = diffusion.training_losses(
|
|
model=model, x_start=latents, model_kwargs=model_kwargs)
|
|
loss = loss_dict["loss"].mean()
|
|
model.backward(loss)
|
|
last_batch_iteration = (
|
|
train_steps + 1) // (global_batch_size // (batch_size * world_size))
|
|
model.step(
|
|
lr_kwargs={'last_batch_iteration': last_batch_iteration})
|
|
|
|
if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1):
|
|
if args.use_fp16:
|
|
ema.update(model.module.module, step=step)
|
|
else:
|
|
ema.update(model.module, step=step)
|
|
|
|
# ===========================================================================
|
|
# Log loss values:
|
|
# ===========================================================================
|
|
running_loss += loss.item()
|
|
log_steps += 1
|
|
train_steps += 1
|
|
if train_steps % args.log_every == 0:
|
|
# Measure training speed:
|
|
torch.cuda.synchronize()
|
|
end_time = time.time()
|
|
steps_per_sec = log_steps / (end_time - start_time)
|
|
# Reduce loss history over all processes:
|
|
avg_loss = torch.tensor(
|
|
running_loss / log_steps, device=device)
|
|
dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
|
|
avg_loss = avg_loss.item() / world_size
|
|
# get lr from deepspeed fused optimizer
|
|
logger.info(f"(step={train_steps:07d}) " +
|
|
(f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") +
|
|
f"Train Loss: {avg_loss:.4f}, "
|
|
f"Lr: {opt.param_groups[0]['lr']:.6g}, "
|
|
f"Steps/Sec: {steps_per_sec:.2f}, "
|
|
f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}")
|
|
# Reset monitoring variables:
|
|
running_loss = 0
|
|
log_steps = 0
|
|
start_time = time.time()
|
|
|
|
# collect gc:
|
|
if args.gc_interval > 0 and (step % args.gc_interval == 0):
|
|
gc.collect()
|
|
|
|
pbar.step(
|
|
f"Epoch {epoch}, step {step}, loss {loss.item():.4f}", args.epochs * len(loader), train_steps)
|
|
|
|
if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps
|
|
) and train_steps > 0:
|
|
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
|
|
target_height=768, target_width=1280, train_steps=train_steps)
|
|
save_checkpoint(args, rank, logger, model, ema,
|
|
epoch, train_steps, checkpoint_dir)
|
|
|
|
if train_steps >= args.max_training_steps:
|
|
logger.info(f"Breaking step loop at {train_steps}.")
|
|
break
|
|
|
|
if train_steps >= args.max_training_steps:
|
|
logger.info(f"Breaking epoch loop at {epoch}.")
|
|
break
|
|
|
|
dist.destroy_process_group()
|