Fix training configs

modal_train_diffusion_nft.py

@@ -32,28 +32,28 @@ DEFAULT_GPU_TYPE = "A100-80GB"
 DEFAULT_NUM_GPUS = 4
 DEFAULT_HSDP_REPLICATE_DIM = 1
 DEFAULT_HSDP_SHARD_DIM = DEFAULT_NUM_GPUS
-DEFAULT_MAX_TRAIN_STEPS = 100
-DEFAULT_NUM_SAMPLES_PER_PROMPT = 24
-DEFAULT_NUM_BATCHES_PER_EPOCH = 0
-DEFAULT_COLLECTION_BATCH_SIZE = 6
+DEFAULT_MAX_TRAIN_STEPS = 50
+DEFAULT_NUM_SAMPLES_PER_PROMPT = 4
+DEFAULT_NUM_BATCHES_PER_EPOCH = 1
+DEFAULT_COLLECTION_BATCH_SIZE = 4
 DEFAULT_INNER_EPOCHS = 1
-DEFAULT_TRAIN_BATCH_SIZE = 6
-DEFAULT_GRADIENT_ACCUMULATION_STEPS = 24
+DEFAULT_TRAIN_BATCH_SIZE = 4
+DEFAULT_GRADIENT_ACCUMULATION_STEPS = 1
 DEFAULT_LEARNING_RATE = -1.0
-DEFAULT_SAMPLE_NUM_STEPS = 0
+DEFAULT_SAMPLE_NUM_STEPS = 50
 DEFAULT_SAMPLE_FLOW_SHIFT = -1.0
 DEFAULT_SAMPLE_GUIDANCE_SCALE = -1.0
-DEFAULT_VALIDATION_NUM_STEPS = 0
-DEFAULT_VALIDATION_NUM_PROMPTS = 0
-DEFAULT_VALIDATION_BATCH_SIZE = 0
-DEFAULT_NUM_FRAMES = 77
+DEFAULT_VALIDATION_NUM_STEPS = 50
+DEFAULT_VALIDATION_NUM_PROMPTS = 16
+DEFAULT_VALIDATION_BATCH_SIZE = 4
+DEFAULT_NUM_FRAMES = 17
 DEFAULT_NUM_LATENT_T = 0
-DEFAULT_LOG_SAMPLE_MAX_VIDEOS = 0
+DEFAULT_LOG_SAMPLE_MAX_VIDEOS = 16
 DEFAULT_PREPROCESS_BATCH_SIZE = 128
 DEFAULT_PREPROCESS_NUM_GPUS = 1
-DEFAULT_DATASET = "world-r1-enhanced"
+DEFAULT_DATASET = "world-r1-enhanced-dynamic"
 DEFAULT_REWARD = "videoalign"
-DEFAULT_MAX_PROMPTS = "quarter"
+DEFAULT_MAX_PROMPTS = "512"
 VIDEOALIGN_CKPT_ROOT = f"{MODAL_CACHE_ROOT}/VideoReward"
 MODAL_CONTEXT_IGNORE = [
     ".git/fsmonitor--daemon.ipc",
@@ -131,7 +131,7 @@ image = (modal.Image.from_registry(
 @app.function(
     image=image,
     gpu=f"{DEFAULT_GPU_TYPE}:{DEFAULT_NUM_GPUS}",
-    timeout=12 * 60 * 60,
+    timeout=8 * 60 * 60,
     volumes={
         MODAL_DATA_ROOT: data_vol,
         f"{PROJECT_ROOT}/outputs": runs_vol,
@@ -424,7 +424,7 @@ def main(
     max_prompts: str = DEFAULT_MAX_PROMPTS,
     check_rewards: bool = True,
 ):
-    train.remote(
+    train.spawn(
         max_train_steps=max_train_steps,
         num_samples_per_prompt=num_samples_per_prompt,
         num_batches_per_epoch=num_batches_per_epoch,