From ce95c2df290ef375e3a28c6201bb76da4291f89f Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Fri, 29 Nov 2024 21:46:40 -0800 Subject: [PATCH] [Feat] EMA Distill; Distributed validation (#63) update --- README.md | 2 +- fastvideo/distill.py | 15 ++-- fastvideo/utils/validation.py | 131 +++++++++++++++++----------------- scripts/experiment_13.sh | 11 +-- 4 files changed, 77 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index dfccfbd83..8e377947e 100644 --- a/README.md +++ b/README.md @@ -57,4 +57,4 @@ Scripts are located at scripts/experiment_N.sh 10. pcm_linear_quadratic， euler_steps 50, 0.05, batchsize 32 11. pcm_linear_quadratic， euler_steps 50, learning rate,1e-7 12. shift1, euler_steps 50 -13. pcm_linear_quadratic， euler_steps 50, 0.05, multiphase 250, ema 0.98 \ No newline at end of file +13. pcm_linear_quadratic， euler_steps 50, 0.05, multiphase 250, ema 0.98, decay 0.1 diff --git a/fastvideo/distill.py b/fastvideo/distill.py index 9f28c1630..785fa39be 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -93,7 +93,7 @@ def reshard_fsdp(model): for m in FSDP.fsdp_modules(model): if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD: torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True) -def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg, ema_decay): +def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg): total_loss = 0.0 optimizer.zero_grad() for _ in range(gradient_accumulation_steps): @@ -214,7 +214,7 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti reshard_fsdp(ema_transformer) for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()): with torch.no_grad(): - p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay)) + p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - 0.95)) grad_norm = transformer.clip_grad_norm_(max_grad_norm) optimizer.step() @@ -515,8 +515,7 @@ def main(args): return int(phase) for step in range(init_steps + 1, args.max_train_steps+1): start_time = time.time() - num_phases = get_num_phases(args.multi_phased_distill_schedule, step) if args.multi_phased_distill_schedule is not None else args.validation_sampling_steps - loss, grad_norm= train_one_step_mochi(transformer,teacher_transformer, ema_transformer, optimizer, lr_scheduler, loader, noise_scheduler,solver, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, num_phases, args.not_apply_cfg_solver,args.distill_cfg, args.ema_decay) + loss, grad_norm= train_one_step_mochi(transformer,teacher_transformer, ema_transformer, optimizer, lr_scheduler, loader, noise_scheduler,solver, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, args.validation_sampling_steps, args.not_apply_cfg_solver,args.distill_cfg) step_time = time.time() - start_time step_times.append(step_time) @@ -550,10 +549,10 @@ def main(args): dist.barrier() if args.log_validation and step % args.validation_steps == 0: log_validation(args, transformer, device, - torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,linear_range=args.linear_range, ema=False) + torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False) if args.use_ema: log_validation(args, ema_transformer, device, - torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,linear_range=args.linear_range, ema=True) + torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=True) if args.use_lora: save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps) @@ -583,7 +582,7 @@ if __name__ == "__main__": parser.add_argument('--enable_stable_fp32', action='store_true') # TODO # diffusion setting - parser.add_argument("--ema_decay", type=float, default=0.999) + parser.add_argument("--ema_decay", type=float, default=0.95) parser.add_argument("--ema_start_step", type=int, default=0) parser.add_argument('--cfg', type=float, default=0.1) parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.") @@ -674,10 +673,8 @@ if __name__ == "__main__": # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"] parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.") parser.add_argument("--linear_quadratic_threshold", type=float, default=0.025, help="Threshold for linear quadratic scheduler.") - parser.add_argument("--linear_range", type=float, default=0.5, help="Range for linear quadratic scheduler.") parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.") parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA.") - parser.add_argument("--ema_decay", type=float, default=0.95, help="EMA decay.") parser.add_argument("--multi_phased_distill_schedule", type=str, default=None) args = parser.parse_args() main(args) \ No newline at end of file diff --git a/fastvideo/utils/validation.py b/fastvideo/utils/validation.py index c1dd461b4..4a7dc80eb 100644 --- a/fastvideo/utils/validation.py +++ b/fastvideo/utils/validation.py @@ -200,73 +200,70 @@ def log_validation(args, transformer, device, weight_dtype, global_step, schedu linear_quadraic = True if scheduler_type == "pcm_linear_quadratic" else False scheduler = PCMFMScheduler(1000, shift, num_euler_timesteps, linear_quadraic, linear_quadratic_threshold, linear_range) # args.validation_prompt_dir + videos = [] + # prompt_embed are named embed0 to embedN + # check how many embeds are there + num_embeds = len([f for f in os.listdir(args.validation_prompt_dir) if "embed" in f]) + validation_prompt_ids = list(range(num_embeds)) + num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size + # pad to multiple of groups + validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) + num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups + local_prompt_ids = validation_prompt_ids[nccl_info.group_id * num_embeds_per_group: (nccl_info.group_id + 1) * num_embeds_per_group] - validation_guidance_scale_ls = args.validation_guidance_scale.split(",") - validation_guidance_scale_ls = [float(scale) for scale in validation_guidance_scale_ls] - for validation_guidance_scale in validation_guidance_scale_ls: + for i in local_prompt_ids: + prompt_embed_path = os.path.join(args.validation_prompt_dir, f"embed{i}.pt") + prompt_mask_path = os.path.join(args.validation_prompt_dir, f"mask{i}.pt") + negative_prompt_embed_path = os.path.join(args.uncond_prompt_dir, "embed.pt") + negative_prompt_mask_path = os.path.join(args.uncond_prompt_dir, "mask.pt") + prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) + generator = torch.Generator(device="cuda").manual_seed(12345) + video = sample_validation_video( + transformer, + vae, + scheduler, + scheduler_type=scheduler_type, + num_frames=args.num_frames, + # Peiyuan TODO: remove hardcode + height=480, + width=848, + num_inference_steps=args.validation_sampling_steps, + guidance_scale=args.validation_guidance_scale, + generator=generator, + prompt_embeds = prompt_embeds, + prompt_attention_mask = prompt_attention_mask, + negative_prompt_embeds = negative_prompt_embeds, + negative_prompt_attention_mask = negative_prompt_attention_mask, + )[0] + if nccl_info.rank_within_group == 0: + videos.append(video[0]) + # collect videos from all process to process zero + + gc.collect() + torch.cuda.empty_cache() + # log if main process + torch.distributed.barrier() + all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", '1')))] # remove padded videos + torch.distributed.all_gather_object(all_videos, videos) + if nccl_info.global_rank == 0: + # remove padding + videos = [video for videos in all_videos for video in videos] + videos = videos[:num_embeds] + # linearize all videos + video_filenames = [] + for i, video in enumerate(videos): + filename = os.path.join(args.output_dir, f"validation_step_{global_step}_video_{i}.mp4") + export_to_video(video, filename, fps=30) + video_filenames.append(filename) - videos = [] - # prompt_embed are named embed0 to embedN - # check how many embeds are there - num_embeds = len([f for f in os.listdir(args.validation_prompt_dir) if "embed" in f]) - validation_prompt_ids = list(range(num_embeds)) - num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size - # pad to multiple of groups - validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) - num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups - local_prompt_ids = validation_prompt_ids[nccl_info.group_id * num_embeds_per_group: (nccl_info.group_id + 1) * num_embeds_per_group] - - for i in local_prompt_ids: - prompt_embed_path = os.path.join(args.validation_prompt_dir, f"embed{i}.pt") - prompt_mask_path = os.path.join(args.validation_prompt_dir, f"mask{i}.pt") - prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0) - negative_prompt_embeds = torch.zeros(256, 4096).to(device).to(weight_dtype).unsqueeze(0) - negative_prompt_attention_mask = torch.zeros(256).bool().to(device).unsqueeze(0) - generator = torch.Generator(device="cuda").manual_seed(12345) - video = sample_validation_video( - transformer, - vae, - scheduler, - scheduler_type=scheduler_type, - num_frames=args.num_frames, - # Peiyuan TODO: remove hardcode - height=480, - width=848, - num_inference_steps=args.validation_sampling_steps, - guidance_scale=validation_guidance_scale, - generator=generator, - prompt_embeds = prompt_embeds, - prompt_attention_mask = prompt_attention_mask, - negative_prompt_embeds = negative_prompt_embeds, - negative_prompt_attention_mask = negative_prompt_attention_mask, - )[0] - if nccl_info.rank_within_group == 0: - videos.append(video[0]) - # collect videos from all process to process zero - - gc.collect() - torch.cuda.empty_cache() - # log if main process - torch.distributed.barrier() - all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", '1')))] # remove padded videos - torch.distributed.all_gather_object(all_videos, videos) - if nccl_info.global_rank == 0: - # remove padding - videos = [video for videos in all_videos for video in videos] - videos = videos[:num_embeds] - # linearize all videos - video_filenames = [] - for i, video in enumerate(videos): - filename = os.path.join(args.output_dir, f"validation_step_{global_step}_guidance_{validation_guidance_scale}_video_{i}.mp4") - export_to_video(video, filename, fps=30) - video_filenames.append(filename) - - logs = { - f"{'ema_' if ema else ''}validation_guidance_{validation_guidance_scale}": [ - wandb.Video(filename) - for i, filename in enumerate(video_filenames) - ] - } - wandb.log(logs, step=global_step) + logs = { + f"{'ema_' if ema else ''}validation": [ + wandb.Video(filename) + for i, filename in enumerate(video_filenames) + ] + } + wandb.log(logs, step=global_step) diff --git a/scripts/experiment_13.sh b/scripts/experiment_13.sh index fddd304d6..f43fe37da 100644 --- a/scripts/experiment_13.sh +++ b/scripts/experiment_13.sh @@ -6,11 +6,11 @@ export FI_PROVIDER=efa export FI_EFA_USE_DEVICE_RDMA=1 export NCCL_PROTO=simple -torchrun --nnodes 2 --nproc_per_node 8\ - --node_rank=0 \ +torchrun --nnodes 4 --nproc_per_node 4\ + --node_rank=3 \ --rdzv_id=456 \ --rdzv_backend=c10d \ - --rdzv_endpoint=[MASTER_NODE_IP_ADDRESS]:29500 \ + --rdzv_endpoint=172.23.30.16:29500 \ fastvideo/distill.py\ --seed 42\ --pretrained_model_name_or_path data/mochi\ @@ -34,7 +34,6 @@ torchrun --nnodes 2 --nproc_per_node 8\ --allow_tf32\ --ema_start_step 0\ --cfg 0.0\ - --ema_decay 0.999\ --log_validation\ --output_dir="data/outputs/lq_euler_50_thresh0.05_multiphaseshedule_125-16,250-8,375-4,500-2"\ --tracker_project_name PCM \ @@ -44,5 +43,7 @@ torchrun --nnodes 2 --nproc_per_node 8\ --num_euler_timesteps 50 \ --linear_quadratic_threshold 0.05 \ --multi_phased_distill_schedule "125-16,250-8,375-4,500-2" \ - --use_ema + --use_ema \ + --ema_decay 0.98 \ + --weight_decay 0.1