1500 steps working

Author: wukaixingxp

apps/julia-grpo/llama3_8b_julia.yaml

@@ -2,12 +2,12 @@
 # >>> python -m apps.julia-grpo.main --config apps/julia-grpo/llama3_8b_julia.yaml
 
 # Global configuration
-group_size: 8  # num_generations from unsloth.py
+group_size: 16  # num_generations from unsloth.py
 batch_size: 4   # per_device_train_batch_size from unsloth.py
-max_req_tokens: 2048  # max_prompt_length from unsloth.py
-max_res_tokens: 1024  # max_completion_length from unsloth.py
+max_req_tokens: 1024  # max_prompt_length from unsloth.py
+max_res_tokens: 2048  # max_completion_length from unsloth.py
 model: "meta-llama/Meta-Llama-3.1-8B-Instruct"
-off_by_n: 0 # Off by one by default
+off_by_n: 2 # Off by one by default
 
 # Main loop configuration
 rollout_threads: 1   # Single thread for Julia code generation
@@ -38,7 +38,9 @@ openenv_config:
   env_vars: {}  # Additional environment variables if needed
   container_timeout_s: 180.0  # Timeout for container operations
   request_timeout_s: 120.0  # Timeout for code execution requests
-  container_memory_gb: 4  # Memory limit for containers
+  container_memory_gb: 128  # Memory limit for containers
+  port: 8000          # port for container communication
+  num_worker: 8    # number of workers
 
 # Policy configuration
 policy:
@@ -66,16 +68,16 @@ trainer:
     hf_assets_path: hf://${model}
   optimizer:
     name: AdamW
-    lr: 1e-5  # learning_rate from unsloth.py
+    lr: 5e-6  # learning_rate from unsloth.py
     eps: 1e-8
     weight_decay: 0.01  # weight_decay from unsloth.py
   lr_scheduler:
-    warmup_steps: 50  # warmup_ratio=0.1 * max_steps=500 from unsloth.py
+    warmup_steps: 0  # warmup_ratio=0.1 * max_steps=500 from unsloth.py
   training:
     local_batch_size: ${batch_size}
     seq_len: ${sum:${max_req_tokens},${max_res_tokens}}  # seq_len >= max_req_tokens + max_res_tokens
     max_norm: 1.0
-    steps: 500  # max_steps from unsloth.py
+    steps: 3000  # max_steps from unsloth.py
     dtype: bfloat16
     gc_freq: 1
   compile:
@@ -90,11 +92,11 @@ trainer:
     disable_loss_parallel: true
   checkpoint:
     enable: true
-    folder: "checkpoint_llama3_8b_julia1107"
+    folder: "checkpoint_llama3_8b_julia1109"
     initial_load_path: hf://${model}
     initial_load_in_hf: true
     last_save_in_hf: true
-    interval: 100  # save_steps from unsloth.py
+    interval: 150  # save_steps from unsloth.py
     async_mode: "disabled"
   activation_checkpoint:
     mode: selective

apps/julia-grpo/main.py

@@ -281,6 +281,7 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
 
             # Extract reward from result
             reward = result.reward if result.reward is not None else 0.0
+            record_metric("reward/julia/reward", reward, Reduce.MEAN)
             obs = result.observation
 
             passed = obs.tests_passed
@@ -293,19 +294,20 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
             print(f"  Tests Passed: {passed}")
             print(f"  Tests Failed: {failed}")
             print(f"  Total Tests: {total}")
+            print(f"  Exit Code: {obs.exit_code}")
+            print(f"  Code Compiles: {obs.code_compiles}")
 
             if obs.stderr:
-                print(f"  Stderr: {obs.stderr[:200]}")
+                print(f"  Stderr: {obs.stderr[:500]}")
                 record_metric("reward/julia/has_errors", 1, Reduce.SUM)
 
-            if obs.error_message:
-                print(f"  Error Message: {obs.error_message[:200]}")
+            if obs.stdout:
+                print(f"  Stdout (first 200 chars): {obs.stdout[:200]}")
 
             # Log metrics
-            record_metric("reward/julia/tests_passed", passed, Reduce.SUM)
-            record_metric("reward/julia/tests_failed", failed, Reduce.SUM)
-            record_metric("reward/julia/tests_total", total, Reduce.SUM)
-            record_metric("reward/julia/pass_rate", reward, Reduce.MEAN)
+            pass_rate = passed / total if total > 0 else 0.0
+
+            record_metric("reward/julia/pass_rate", pass_rate, Reduce.MEAN)
 
             print(f"Final Reward: {reward:.3f}")
             print("=" * 80)
@@ -337,7 +339,7 @@ def _extract_code(self, response: str) -> str:
 class ComputeAdvantages(ForgeActor):
     @endpoint
     async def compute(self, group: Group) -> list[float]:
-        rewards = torch.tensor([[e.reward for e in group]])
+        rewards = torch.tensor([[e.reward for e in group]], dtype=torch.float32)
         mean = rewards.mean(1, keepdim=True)
         std = rewards.std(1, keepdim=True)
         advantages = (rewards - mean) / (std + 1e-4)
@@ -517,6 +519,14 @@ async def main(cfg: DictConfig):
     request_timeout_s = openenv_config.get("request_timeout_s", 120.0)
     container_memory_gb = openenv_config.get("container_memory_gb", 4)
 
+    # Set PORT and NUM_WORKER environment variables for the Julia server
+    # These match the Dockerfile defaults
+    if "PORT" not in env_vars:
+        env_vars["PORT"] = str(openenv_config.get("port", 8000))
+    if "NUM_WORKER" not in env_vars:
+        env_vars["NUM_WORKER"] = str(openenv_config.get("num_worker", 4))
+    if "JULIA_MAX_WORKERS" not in env_vars:
+        env_vars["JULIA_MAX_WORKERS"] = str(openenv_config.get("julia_max_workers", 16))
     julia_env_actor = await GenericOpenEnvActor.options(
         **cfg.actors.julia_env
     ).as_actor(
@@ -587,12 +597,14 @@ async def continuous_rollouts():
             responses: list[Completion] = await policy.generate.route(prompt)
             t.step("policy_generation")
 
-            # Construct episodes and calculate rewards
+            # Construct episodes and calculate rewards in parallel
             episodes = []
             input_ids = torch.ones(
                 (group_size, max_req_tokens + max_res_tokens),
                 dtype=torch.long,
             )
+
+            # Create episodes first
             for i, response in enumerate(responses):
                 episode = Episode(
                     episode_id=str(uuid.uuid4()),
@@ -602,12 +614,20 @@ async def continuous_rollouts():
                     target=target,
                     completion=response,
                 )
-                episode.reward = await reward_actor.evaluate_response.route(
+                episodes.append(episode)
+
+            # Evaluate all rewards in parallel
+            reward_tasks = [
+                reward_actor.evaluate_response.route(
                     prompt=prompt, response=response.text, target=target
                 )
-                episodes.append(episode)
+                for response in responses
+            ]
+            rewards = await asyncio.gather(*reward_tasks)
 
-                # Build input_ids for reference logprobs
+            # Assign rewards and build input_ids
+            for i, (episode, reward) in enumerate(zip(episodes, rewards)):
+                episode.reward = reward
                 input_ids[i, :max_req_tokens] = episode.request_tensor
                 input_ids[i, max_req_tokens:] = episode.response_tensor