update: 500 steps L1 + 300 steps L2, higher lr for 1.5B

Dockerfile

@@ -22,4 +22,4 @@ COPY data/ /app/data/
 COPY train.py .
 COPY evaluate.py .
 
-CMD ["python", "evaluate.py"]
+CMD ["python", "train.py"]

train.py

@@ -169,7 +169,7 @@ train_dataset = Dataset.from_list([
 ])
 
 print("Starting training...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level1")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level1-v2")
 
 trainer = GRPOTrainer(
     model=model,
@@ -179,10 +179,10 @@ trainer = GRPOTrainer(
         output_dir="/tmp/deceit-1.5b",
         bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
         fp16=False,
-        max_steps=150,
+        max_steps=500,
         per_device_train_batch_size=4,
         num_generations=4,
-        learning_rate=5e-6,
+        learning_rate=1e-5,
         warmup_steps=5,
         logging_steps=1,
         save_steps=50,
@@ -267,7 +267,7 @@ def reward_fn_l2(completions, prompts=None, **kwargs):
 
 # Train Level 2
 print("Starting Level 2 training on 1.5B...")
-wandb.init(project=WANDB_PROJECT, name="1.5b-level2")
+wandb.init(project=WANDB_PROJECT, name="1.5b-level2-v2")
 
 trainer_l2 = GRPOTrainer(
     model=model,
@@ -275,7 +275,7 @@ trainer_l2 = GRPOTrainer(
     reward_funcs=[reward_fn_l2],
     args=GRPOConfig(
         output_dir="/tmp/deceit-1.5b-l2",
-        max_steps=80,
+        max_steps=300,
         per_device_train_batch_size=4,
         num_generations=4,
         learning_rate=2e-6,