From 989dc28dcf71237339497e2a5f30b6b1263b6a44 Mon Sep 17 00:00:00 2001 From: wailovet Date: Fri, 28 Jun 2024 03:12:26 +0800 Subject: [PATCH] 'update' --- hook_HYDiT_main.py | 31 +- hook_HYDiT_main_train_deepspeed.py | 570 +++++++++++++++++++++++++++++ hook_HYDiT_run.py | 63 ++-- hook_HYDiT_utils.py | 278 ++++++++++++++ mz_train_tools_core_HYDiT.py | 30 +- 5 files changed, 893 insertions(+), 79 deletions(-) create mode 100644 hook_HYDiT_main_train_deepspeed.py create mode 100644 hook_HYDiT_utils.py diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py index b65064b..1d7ff01 100644 --- a/hook_HYDiT_main.py +++ b/hook_HYDiT_main.py @@ -39,37 +39,21 @@ from tqdm import tqdm UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" - - def set_unet_path(path): global UNET_PATH UNET_PATH = path - - VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" - - def set_vae_ema_path(path): global VAE_EMA_PATH VAE_EMA_PATH = path - - TOKENIZER = "ckpts/t2i/tokenizer" - - def set_tokenizer_path(path): global TOKENIZER TOKENIZER = path - - TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' - - def set_text_encoder_path(path): global TEXT_ENCODER TEXT_ENCODER = path - - T5_ENCODER = { 'MT5': 'ckpts/t2i/mt5', 'attention_mask': True, @@ -78,20 +62,13 @@ T5_ENCODER = { 'torch_dtype': torch.float16, 'learnable_replace': True } - - -global TRAIN_CONFIG - - -def set_train_config(train_config): - global TRAIN_CONFIG - TRAIN_CONFIG = train_config - - def set_t5_encoder_path(path): global T5_ENCODER T5_ENCODER['MT5'] = path - +global TRAIN_CONFIG +def set_train_config(train_config): + global TRAIN_CONFIG + TRAIN_CONFIG = train_config def create_logger(log_dir=None): logger = logging.getLogger(__name__) diff --git a/hook_HYDiT_main_train_deepspeed.py b/hook_HYDiT_main_train_deepspeed.py new file mode 100644 index 0000000..e494616 --- /dev/null +++ b/hook_HYDiT_main_train_deepspeed.py @@ -0,0 +1,570 @@ +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 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}") + 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 + except: + from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar + + # 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'] + 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": + 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 + # =========================================================================== + + 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}, mean_loss {running_loss / step:.4f}", epoch * 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() diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py index 3aa8eed..11ed62c 100644 --- a/hook_HYDiT_run.py +++ b/hook_HYDiT_run.py @@ -23,34 +23,26 @@ class SimpleNamespaceCNWarrper(SimpleNamespace): return self.__dict__.get(name, None) -import requests - -master_port = 0 - - -def LOG(log): - if master_port == 0: - raise Exception("master_port is 0") - # 发送http - resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={ - "Content-Type": "application/json"}) - if resp.status_code != 200: - raise Exception(f"LOG failed: {resp.text}") - - if __name__ == "__main__": parser = argparse.ArgumentParser( conflict_handler='resolve', ) parser.add_argument("--sys_path", type=str, default="") parser.add_argument("--train_config_file", type=str, default="") - parser.add_argument("--master_port", type=int, default=0) + parser.add_argument("--mz_master_port", type=int, default=0) args = parser.parse_args() - master_port = args.master_port + master_port = args.mz_master_port print(f"master_port = {master_port}") + try: + from . import hook_HYDiT_utils + except Exception as e: + import hook_HYDiT_utils + + hook_HYDiT_utils.set_master_port(master_port) + sys_path = args.sys_path if sys_path != "": sys.path.append(sys_path) @@ -62,11 +54,11 @@ if __name__ == "__main__": except Exception as e: import hook_HYDiT_main try: - from . import hook_HYDiT_config + from . import hook_HYDiT_main_train_deepspeed except Exception as e: - import hook_HYDiT_config + import hook_HYDiT_main_train_deepspeed - + import hydit.config def _handle_conflict_error(self, *args, **kwargs): pass @@ -87,7 +79,7 @@ if __name__ == "__main__": argparse.ArgumentParser._handle_conflict_error = _handle_conflict_error argparse.ArgumentParser._handle_conflict_resolve = _handle_conflict_error argparse.ArgumentParser.parse_args = parse_args - margs = hook_HYDiT_config.get_args() + margs = hydit.config.get_args() margs.model = train_config.get("model", "DiT-g/2") margs.task_flag = train_config.get( @@ -111,7 +103,9 @@ if __name__ == "__main__": margs.target_ratios = train_config.get( "target_ratios", ['1:1', '3:4', '4:3', '16:9', '9:16']) margs.rope_img = train_config.get("rope_img", "base1024") + margs.image_size = train_config.get("image_size", 1024) + margs.rope_real = train_config.get("rope_real", True) margs.index_file = train_config.get("index_file", None) @@ -120,7 +114,9 @@ if __name__ == "__main__": margs.noise_offset = train_config.get("noise_offset", 0.1) - margs.log_every = train_config.get("log_every", 10) + margs.log_every = train_config.get("log_every", 99999999999999) + margs.use_zero_stage = 2 + margs.deepspeed = True margs.results_dir = train_config.get("results_dir") margs.mse_loss_weight_type = train_config.get( @@ -130,21 +126,30 @@ if __name__ == "__main__": if hasattr(margs, k): setattr(margs, k, v) - hook_HYDiT_main.set_unet_path( + hook_HYDiT_utils.set_unet_path( train_config.get("unet_path")) - hook_HYDiT_main.set_vae_ema_path( + hook_HYDiT_utils.set_vae_ema_path( train_config.get("vae_ema_path")) - hook_HYDiT_main.set_text_encoder_path( + hook_HYDiT_utils.set_text_encoder_path( train_config.get("text_encoder_path")) - hook_HYDiT_main.set_tokenizer_path( + hook_HYDiT_utils.set_tokenizer_path( train_config.get("tokenizer_path")) - hook_HYDiT_main.set_t5_encoder_path( + hook_HYDiT_utils.set_t5_encoder_path( train_config.get("t5_encoder_path")) - hook_HYDiT_main.set_train_config(train_config) + hook_HYDiT_utils.set_train_config(train_config) - hook_HYDiT_main.Core(margs, LOG) + try: + # deepspeed/runtime/engine + import deepspeed.runtime.engine + deepspeed.runtime.engine.DeepSpeedEngine._do_sanity_check = lambda x: None + except Exception as e: + pass + + if type(margs.image_size) == int: + margs.image_size = [margs.image_size, margs.image_size] + hook_HYDiT_main_train_deepspeed.Core(margs) diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py new file mode 100644 index 0000000..2c33682 --- /dev/null +++ b/hook_HYDiT_utils.py @@ -0,0 +1,278 @@ + + +import json +import os +import torch + + +UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" + + +def set_unet_path(path): + global UNET_PATH + UNET_PATH = path + + +VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" + + +def set_vae_ema_path(path): + global VAE_EMA_PATH + VAE_EMA_PATH = path + + +TOKENIZER = "ckpts/t2i/tokenizer" + + +def set_tokenizer_path(path): + global TOKENIZER + TOKENIZER = path + + +TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' + + +def set_text_encoder_path(path): + global TEXT_ENCODER + TEXT_ENCODER = path + + +T5_ENCODER = { + 'MT5': 'ckpts/t2i/mt5', + 'attention_mask': True, + 'layer_index': -1, + 'attention_pool': True, + 'torch_dtype': torch.float16, + 'learnable_replace': True +} + + +def set_t5_encoder_path(path): + global T5_ENCODER + T5_ENCODER['MT5'] = path + + +global TRAIN_CONFIG + + +def set_train_config(train_config): + global TRAIN_CONFIG + TRAIN_CONFIG = train_config + + +def easy_sample_images( + args, + vae=None, + text_encoder=None, + tokenizer=None, + model=None, + embedder_t5=None, + target_height=768, + target_width=1280, + prompt="A photo of a girl with a hat on a sunny day", + negative_prompt="", + batch_size=1, + guidance_scale=5.0, + infer_steps=20, + sampler='ddpm', + train_steps=0, +): + from hydit.diffusion.pipeline import StableDiffusionPipeline + from diffusers import schedulers + from hydit.constants import SAMPLER_FACTORY + from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed + from hydit.modules.models import HUNYUAN_DIT_CONFIG + with torch.no_grad(): + + workspace_dir = TRAIN_CONFIG.get("workspace_dir") + sample_config_file = TRAIN_CONFIG.get("sample_config_file", None) + if sample_config_file is None: + print("sample_config_file is not set.") + return + try: + sample_config = json.load(open(sample_config_file, "r")) + except Exception as e: + print(f"Failed to load sample_config_file: {sample_config_file}") + return + sample_images_dir = os.path.join(workspace_dir, "sample_images") + os.makedirs(sample_images_dir, exist_ok=True) + + # Load sampler from factory + kwargs = SAMPLER_FACTORY[sampler]['kwargs'] + scheduler = SAMPLER_FACTORY[sampler]['scheduler'] + + # Build scheduler according to the sampler. + scheduler_class = getattr(schedulers, scheduler) + scheduler = scheduler_class(**kwargs) + + # Set timesteps for inference steps. + scheduler.set_timesteps(infer_steps, "cuda") + + def calc_rope(height, width): + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + patch_size = model_config['patch_size'] + head_size = model_config['hidden_size'] // model_config['num_heads'] + th = height // 8 // patch_size + tw = width // 8 // patch_size + base_size = 512 // 8 // patch_size + start, stop = get_fill_resize_and_crop((th, tw), base_size) + sub_args = [start, stop, (th, tw)] + rope = get_2d_rotary_pos_embed(head_size, *sub_args) + return rope + + pipeline = StableDiffusionPipeline(vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=model.module, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + embedder_t5=embedder_t5, + ).to("cuda") + + style = torch.as_tensor([0, 0] * batch_size, device="cuda") + + src_size_cond = (target_width, target_height) + size_cond = list(src_size_cond) + [target_width, target_height, 0, 0] + image_meta_size = torch.as_tensor( + [size_cond] * 2 * batch_size, device="cuda",) + + if type(sample_config) != list: + sample_config = [sample_config] + + for i, sample in enumerate(sample_config): + prompt = sample.get("prompt", "") + negative_prompt = sample.get("negative_prompt", "") + guidance_scale = sample.get("cfg", guidance_scale) + infer_steps = sample.get("steps", infer_steps) + width = sample.get("width", target_width) + height = sample.get("height", target_height) + + freqs_cis_img = calc_rope(height, width) + + samples = pipeline( + height=height, + width=width, + prompt=prompt, + negative_prompt=negative_prompt, + num_images_per_prompt=batch_size, + guidance_scale=guidance_scale, + num_inference_steps=infer_steps, + style=style, + return_dict=False, + use_fp16=True, + learn_sigma=args.learn_sigma, + freqs_cis_img=freqs_cis_img, + image_meta_size=image_meta_size, + )[0] + + try: + pass + except Exception as e: + print(f"Failed to sample images: {e}") + continue + + # print("samples:",type(samples),) + # input("Press Enter to continue...") + # print("samples:",samples,) + + if type(samples) == list: + pil_image = samples[0] + else: + pil_image = samples + + sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png" + sample_filename_path = os.path.join( + sample_images_dir, sample_filename) + pil_image.save(sample_filename_path) + + return None + + +def model_resume(args, model, ema, logger): + """ + Load pretrained weights. + """ + start_epoch = 0 + start_epoch_step = 0 + train_steps = 0 + resume_path = UNET_PATH + + logger.info(f"Resume from checkpoint {resume_path}") + + if args.resume_split: + # Resume main model + + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + model.load_state_dict(resume_ckpt_module, strict=False) + + # Resume ema model + if args.use_ema: + if args.module_to_ema: + if "resume_ckpt_module" in locals(): + logger.info(f" Resume ema model from main states.") + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + logger.info(f" Resume ema model from module states.") + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + if "resume_ckpt_ema" in locals(): + logger.info(f" Resume ema model from EMA states.") + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + logger.info(f" Resume ema model from EMA states.") + resume_ckpt_ema = torch.load(resume_path, + map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + raise ValueError( + " “If `resume` is True, then either `resume_split` must be true.”") + + return model, ema, start_epoch, start_epoch_step, train_steps + + +import tqdm + +import requests + + +def set_master_port(port): + global master_port + master_port = port + + +master_port = 0 + + +def LOG(log): + if master_port == 0: + raise Exception("master_port is 0") + # 发送http + resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={ + "Content-Type": "application/json"}) + if resp.status_code != 200: + raise Exception(f"LOG failed: {resp.text}") + + +# with tqdm(total=total_steps, initial=train_steps) as pbar: +# pbar.update(1) +# pbar.set_description( +# f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") +class PBar: + def __init__(self, total): + self.pbar = tqdm.tqdm(total=total) + + def step(self, desc, total_steps, train_steps): + self.pbar.update(1) + self.pbar.set_description(desc) + + LOG({ + "type": "sample_images", + "global_step": train_steps, + "total_steps": total_steps, + # "latent": noise_pred_latent_path, + }) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index bc0b9bd..07b64d1 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -297,27 +297,6 @@ def check_model_auto_download(args): def check_required(): - # timm==0.9.5 - # diffusers==0.21.2 - # peft==0.10.0 - # protobuf==3.19.0 - # torchvision==0.14.1 - # transformers==4.39.1 - # peft==0.10.0 - # accelerate==0.29.3 - # loguru==0.7.2 - # einops==0.7.0 - # sentencepiece==0.1.99 - # cuda-python==11.7.1 - # onnxruntime==1.12.1 - # onnx==1.12.0 - # nvidia-pyindex==1.0.9 - # onnx-graphsurgeon==0.3.27 - # polygraphy==0.47.1 - # pandas==2.0.3 - # gradio==3.50.2 - # deepspeed==0.6.3 - # pyarrow==16.1.0 packages = [ "pandas", "pyarrow", "diffusers", "transformers", "timm", "peft", "accelerate", "loguru", "einops", "sentencepiece", @@ -330,6 +309,11 @@ def check_required(): subprocess.run([sys.executable, "-m", "pip", "install", package], check=True) + try: + import deepspeed + except ImportError: + raise Exception("请手动安装合适的deepspeed版本") + def MZ_HYDiTTrain_call(args={}): check_required() @@ -584,8 +568,8 @@ def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir): try: subprocess.run( - [sys.executable, exec_pyfile, "--sys_path", HYDiT_tool_dir, - "--master_port", str(port), "--train_config_file", train_config_path], + [sys.executable, "-m", "torch.distributed.launch", "--use_env", exec_pyfile, "--sys_path", HYDiT_tool_dir, + "--mz_master_port", str(port), "--train_config_file", train_config_path], check=True, ) stop_server()