Cleanups

Author: SurbhiJainUSC

src/maxtext/configs/post_train/rl.yml

@@ -191,8 +191,8 @@ penalty_incorrect_format: 0.0
 penalty_incorrect_answer: 0.0
 
 # ====== Configuration for math_verify Pool ======
-# Per-item wall-clock timeout (seconds) for math_verify
-math_verify_timeout: 15
+# Timeout (seconds) for math_verify
+math_verify_timeout: 300
 # Max worker processes for the math_verify pool. null ⇒ min(batch_size, cpu_count())
 math_verify_num_procs: null
 

src/maxtext/configs/types.py

@@ -1813,7 +1813,7 @@ class Reward(BaseModel):
   )
   penalty_incorrect_format: float = Field(-0.5, description="Penalty for an incorrect format.")
   penalty_incorrect_answer: float = Field(-1.0, description="Penalty for an incorrect answer.")
-  math_verify_timeout: int = Field(15, description="Per-item wall-clock timeout (seconds) for math_verify.")
+  math_verify_timeout: int = Field(300, description="Timeout (seconds) for math_verify call per batch.")
   math_verify_num_procs: int | None = Field(
       None,
       description=("Max worker processes for the math_verify pool. None ⇒ " "min(batch_size, cpu_count())."),