distributed validation

README.md

@@ -44,6 +44,8 @@ Scripts are located at scripts/experiment_N.sh
 4. shift 8, euler_steps 50
 5. shift 8, euler_steps 100, adv
 6. pcm_linear_quadratic， euler_steps 50, 0.025, adv
+7. pcm_linear_quadratic， euler_steps 50, 0.05, adv, advweight 0.2
+8. pcm_linear_quadratic， euler_steps 50, 0.05, bs 64
 7. pcm_linear_quadratic， euler_steps 100, 0.05
 8. finetune, 5e-5 
 9. finetune, 1e-5

fastvideo/utils/validation.py

@@ -205,10 +205,11 @@ def log_validation(args, transformer, device, weight_dtype, global_step,  schedu
     # 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_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size
+    num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size
     # pad to multiple of groups
-    validation_prompt_ids += [0] * (num_groups - len(validation_prompt_ids) % num_groups)
-    local_prompt_ids = validation_prompt_ids[nccl_info.group_id::num_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")

scripts/experiment_3.sh

@@ -29,7 +29,7 @@ torchrun --nnodes 2 --nproc_per_node 8\
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
     --checkpointing_steps=500\
-    --validation_steps 250\
+    --validation_steps 125\
     --validation_sampling_steps 8 \
     --checkpoints_total_limit 3\
     --allow_tf32\

scripts/experiment_4.sh

@@ -29,7 +29,7 @@ torchrun --nnodes 2 --nproc_per_node 8\
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
     --checkpointing_steps=500\
-    --validation_steps 250\
+    --validation_steps 125\
     --validation_sampling_steps 8 \
     --checkpoints_total_limit 3\
     --allow_tf32\