'update'

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__)

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()

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)

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,
+        })

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()