clean code

train_deepspeed.py

@@ -32,25 +32,12 @@ from diffusion import create_diffusion
 from diffusers.models import AutoencoderKL
 from download import find_model
 
-
 import deepspeed
 
 #################################################################################
 #                             Training Helper Functions                         #
 #################################################################################
 
-@torch.no_grad()
-def update_ema(ema_model, model, decay=0.9999):
-    """
-    Step the EMA model towards the current model.
-    """
-    ema_params = OrderedDict(ema_model.named_parameters())
-    model_params = OrderedDict(model.named_parameters())
-
-    for name, param in model_params.items():
-        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
-
 
 def requires_grad(model, flag=True):
     """
@@ -112,15 +99,12 @@ def center_crop_arr(pil_image, image_size):
 
 def main(args):
     """
-    Trains a new DiT model.
+    Trains a new DiT-MoE model.
     """
     assert torch.cuda.is_available(), "Training currently requires at least one GPU."
     
     deepspeed.init_distributed()
 
-    # Setup DDP:
-    #dist.init_process_group("nccl")
-    #assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
     rank = args.local_rank
     device = rank % torch.cuda.device_count()
     seed = args.global_seed * dist.get_world_size() + rank
@@ -160,19 +144,10 @@ def main(args):
         model.load_state_dict(state_dict)
 
 
-    # Note that parameter initialization is done within the DiT constructor
-    # ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
-    # requires_grad(ema, False)
-    # model = DDP(model.to(device), device_ids=[rank])
-    # model = DDP(model.to(device), device_ids=[device])
     diffusion = create_diffusion(timestep_respacing="")  # default: 1000 steps, linear noise schedule 
     vae = AutoencoderKL.from_pretrained(args.vae_path).to(device)
     logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
 
-    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
-    # opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0)
-
-    # Setup data:
     transform = transforms.Compose([
         transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
         transforms.RandomHorizontalFlip(),
@@ -198,11 +173,6 @@ def main(args):
     )
     logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
 
-    # Prepare models for training:
-    # update_ema(ema, model.module, decay=0)  # Ensure EMA is initialized with synced weights
-    # model.train()  # important! This enables embedding dropout for classifier-free guidance
-    # ema.eval()  # EMA model should always be in eval mode
-
     model_engine, opt, _, __ = deepspeed.initialize(
         args=args, model=model, model_parameters=model.parameters())
 
@@ -228,14 +198,9 @@ def main(args):
             model_kwargs = dict(y=y)
             with torch.autocast(device_type='cuda'):
                 loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
-            loss = loss_dict["loss"].mean()
-            #if (data_iter_step + 1) % args.accum_iter == 0:
-            #    opt.zero_grad()
-            #loss.backward()
+            loss = loss_dict["loss"].mean() 
             model_engine.backward(loss) 
             model_engine.step()
-            # opt.step()
-            # update_ema(ema, model.module)
             
             data_iter_step += 1
             # Log loss values:
@@ -270,15 +235,12 @@ def main(args):
                     logger.info(f"Saved checkpoint to {checkpoint_path}")
                 dist.barrier()
 
-    model.eval()  # important! This disables randomized embedding dropout
-    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
-
+    # model.eval()  # important! This disables randomized embedding dropout
     logger.info("Done!")
     cleanup()
 
 
-if __name__ == "__main__":
-    # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
+if __name__ == "__main__": 
     parser = argparse.ArgumentParser()
     parser.add_argument("--data-path", type=str, required=True)
     parser.add_argument("--results-dir", type=str, default="results")