Llama-3.1-8B-Instruct, GRPO-trained on search_plan_act

Full-weight checkpoint (LoRA merged) of meta-llama/Llama-3.1-8B-Instruct, trained with GRPO (via TRL's multi-turn environment_factory tool-calling API) on search_plan_act: a procedurally-generated, closed-book multi-turn task requiring the model to search for facts, read_record/ update_record/link_records on retrieved entities, and finish, looping through tool calls until the goal is achieved or it gives up. Reward decomposes into outcome / grounding / stop-behavior components, with an explicit adversarial-pilot check that no reward-hacking strategy (reckless guessing, spamming searches, never finishing) scores above genuine goal-directed play.

Training required a real, non-obvious fix: Llama-3.1's auto-derived tool-call response_template (via transformers/TRL's chat-template-schema derivation) required a tool call's closing brace to be immediately followed by the end-of-turn token -- too strict, since Llama doesn't reliably stop right after one call. Fixed with a lookbehind regex ((?<=\})\}\s*) that correctly extracts tool-call content regardless of what follows, verified against single-call, multi-param, and pathological multi-call completions.

Training: LoRA (r=32, alpha=64) on all attention+MLP projections, GRPO via TRL 1.9.2 + vLLM colocate generation, 250 steps, effective batch size 32, 4x GH200 GPUs.

Evaluation

Evaluated pre-RL (base meta-llama/Llama-3.1-8B-Instruct) vs post-RL (this checkpoint) on held-out real-world and synthetic benchmarks, all using vLLM-accelerated generation with the model's own chat template (apples-to-apples generation path in both rows):

Benchmark Pre-RL Post-RL (this checkpoint)
MuSiQue (EM / F1, n=2417) 0.321 / 0.434 0.318 / 0.430
2WikiMultihopQA (EM / F1, n=12576) 0.375 / 0.474 0.374 / 0.473
search_plan_ood (mean reward, n=100) 0.024 0.064
BALROG BabyAI (mean episode return, n=50) 0.262 0.286

Honest finding: this checkpoint shows real, if modest, gains on the two benchmarks closest to the training distribution (search_plan_ood, BALROG's closed-environment agentic loop), but is essentially flat on the two real-world QA benchmarks (MuSiQue, 2WikiMultihopQA). Compare against the sibling bmonikraj/qwen3-8b-search-plan-act checkpoint, trained with the identical recipe, which showed much stronger and more consistent transfer (including a large MuSiQue gain) -- RL transfer here is genuinely model-specific, not just a function of the training recipe.

Files

  • model.safetensors, config.json, tokenizer.json, tokenizer_config.json, generation_config.json, chat_template.jinja -- the full merged checkpoint.
  • adapter/ -- the raw LoRA adapter (pre-merge), kept for reference.
Downloads last month
8
Safetensors
Model size
8B params
Tensor type
BF16
·
Video Preview
loading

Model tree for bmonikraj/llama3-8b-search-plan-act

Finetuned
(2921)
this model