Fix Qwen3 thinking mode + increase max_new_tokens: training/train_grpo.py

training/train_grpo.py

@@ -176,7 +176,7 @@ class TrainConfig:
     # Sampling
     sample_temperature: float = 1.0
     sample_top_p: float = 0.95
-    max_new_tokens: int = 96
+    max_new_tokens: int = 384
     max_prompt_tokens: int = 2048
 
     # Tiny behavior warm-start. Smoke showed the raw model sometimes echoes the
@@ -338,11 +338,15 @@ def run_format_warmup(
         for action_dict in _format_warmup_actions(env, task_id):
             user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
             messages.append({"role": "user", "content": user_msg})
-            prompt_text = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-            )
+            try:
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True,
+                    enable_thinking=False,
+                )
+            except TypeError:
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True,
+                )
             completion_text = json.dumps(action_dict, ensure_ascii=False)
             prompt_ids = tokenizer(
                 prompt_text,