Enhance logging in AjetDataParallelPPOActor and increase max_parallel in AIMESwarmTrainer

Author: binary-husky

ajet/backbone/verl/dp_actor.py

@@ -198,7 +198,8 @@ def update_policy(self, data: DataProto):
 
                 self.actor_optimizer.zero_grad()
 
-                for micro_batch in micro_batches:
+                num_micro_batches = len(micro_batches)
+                for micro_batch_idx, micro_batch in enumerate(micro_batches, 1):
                     micro_batch = micro_batch.to(get_device_id())
                     micro_batch_metrics = {}
                     model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch, "pad_token_id": pad_token_id}
@@ -207,7 +208,7 @@ def update_policy(self, data: DataProto):
                     advantages = model_inputs["advantages"]
                     # [AJET] Debug logging for tensor shapes
                     input_ids = model_inputs["input_ids"]
-                    print(f'[Update Policy] -> Micro batch shape, input_ids {input_ids.shape}, response {response_mask.shape}')
+                    print(f'[Update Policy] -> Micro batch shape, input_ids {input_ids.shape}, response {response_mask.shape} @{micro_batch_idx}/{num_micro_batches}')
 
                     entropy_coeff = self.config.entropy_coeff
                     loss_agg_mode = self.config.loss_agg_mode

tutorial/opencode_build_aime/agent_roll.py

@@ -186,7 +186,7 @@ def train(self):
         task_count = 0
         executor = PeriodicDrainThreadPoolExecutor(
             workers=self.grpo_n * self.remote_batch_size,
-            max_parallel=64,
+            max_parallel=256,
             auto_retry=True
         )
 

tutorial/opencode_build_aime/agent_run_v3.py

@@ -20,6 +20,7 @@
 import tempfile
 import time
 from dataclasses import dataclass
+from textwrap import dedent
 from uuid import uuid4
 
 from openai import OpenAI
@@ -96,39 +97,39 @@ def __init__(self, timeout: int = 30, memory_limit_mb: int = 512):
         self.memory_limit_mb = memory_limit_mb
 
     def execute(self, code: str, stdin: str = "") -> ProcessExecuteResult:
-        pre_template = f"""
-import signal
-import resource
-import os
-import sys
+        pre_template = dedent(f"""\
+            import signal
+            import resource
+            import os
+            import sys
 
-os.environ['OPENBLAS_NUM_THREADS'] = '1'
+            os.environ['OPENBLAS_NUM_THREADS'] = '1'
 
-def _exec_set_alarm_timeout(timeout):
-    signal.signal(signal.SIGALRM, _exec_time_exceeded)
-    signal.alarm(timeout)
+            def _exec_set_alarm_timeout(timeout):
+                signal.signal(signal.SIGALRM, _exec_time_exceeded)
+                signal.alarm(timeout)
 
-def _exec_time_exceeded(*_):
-    print('Suicide from timeout.', flush=True)
-    try:
-        os.killpg(0, 9)
-    except Exception:
-        pass
-    os._exit({TIMEOUT_EXIT_CODE})
+            def _exec_time_exceeded(*_):
+                print('Suicide from timeout.', flush=True)
+                try:
+                    os.killpg(0, 9)
+                except Exception:
+                    pass
+                os._exit({TIMEOUT_EXIT_CODE})
 
-def _exec_set_max_runtime(seconds):
-    soft, hard = resource.getrlimit(resource.RLIMIT_CPU)
-    resource.setrlimit(resource.RLIMIT_CPU, (seconds, hard))
+            def _exec_set_max_runtime(seconds):
+                soft, hard = resource.getrlimit(resource.RLIMIT_CPU)
+                resource.setrlimit(resource.RLIMIT_CPU, (seconds, hard))
 
-_exec_set_alarm_timeout({self.timeout})
-_exec_set_max_runtime({self.timeout})
+            _exec_set_alarm_timeout({self.timeout})
+            _exec_set_max_runtime({self.timeout})
 
-_exec_time_start = time.perf_counter()
-"""
-        post_template = f"""
-_exec_time_end = time.perf_counter()
-_exec_duration = _exec_time_end - _exec_time_start
-"""
+            _exec_time_start = time.perf_counter()
+            """)
+        post_template = dedent("""\
+            _exec_time_end = time.perf_counter()
+            _exec_duration = _exec_time_end - _exec_time_start
+            """)
 
         with tempfile.TemporaryDirectory() as tmp_path:
             source_path = f"{tmp_path}/source.py"
@@ -307,23 +308,24 @@ async def run(self, messages: list[dict], sampling_params: dict) -> tuple[str, l
         user_turns, assistant_turns = 0, 0
         all_response_content = []
 
-        system_prompt = """You are an expert mathematician specialized in solving challenging math competition problems.
+        system_prompt = dedent("""\
+            You are an expert mathematician specialized in solving challenging math competition problems.
 
-You have access to a Python code execution tool. Use it to:
-1. Perform calculations and verify your answers
-2. Run code when you need precise computation
-3. Test your hypotheses before giving final answers
+            You have access to a Python code execution tool. Use it to:
+            1. Perform calculations and verify your answers
+            2. Run code when you need precise computation
+            3. Test your hypotheses before giving final answers
 
-Instructions:
-1. Think through the problem step by step
-2. Use the python_code_with_standard_io tool when you need to execute code
-3. Show your reasoning clearly
-4. Put your final numerical answer inside \\boxed{} at the end
+            Instructions:
+            1. Think through the problem step by step
+            2. Use the python_code_with_standard_io tool when you need to execute code
+            3. Show your reasoning clearly
+            4. Put your final numerical answer inside \\boxed{} at the end
 
-For each function call, return a json object within <tool_call></tool_call> XML tags:
-<tool_call>
-{"name": "python_code_with_standard_io", "arguments": {"code": "your python code", "input": "stdin input if needed"}}
-</tool_call>"""
+            For each function call, return a json object within <tool_call></tool_call> XML tags:
+            <tool_call>
+            {"name": "python_code_with_standard_io", "arguments": {"code": "your python code", "input": "stdin input if needed"}}
+            </tool_call>""")
 
         formatted_messages = [msg for msg in messages if msg.get("role") != "system"]
         if not any(msg.get("role") == "system" for msg in messages):