[Feat] EMA Distill; Distributed validation (#63)

update

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
+13. pcm_linear_quadratic， euler_steps 50, 0.05, multiphase 250, ema 0.98, decay 0.1

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)

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)
 

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