[GOLD] Tool calling support for GOLDTrainer

[Strongich]

What does this PR do?

Adds multi-turn tool-calling support to the GOLDTrainer, for same-family (shared-tokenizer) distillation. This builds on the recently merged [GOLD] VLM support (#5969).

This addition was motivated by these 2 blogs about frontier models' training: first and second.

This PR wires tool calling through both paths: off and on policy, for text LLMs and same-family VLMs, and ports GRPO's environment_factory so environment methods become callable tools.

What's included

• Config: max_tool_calling_iterations (cap on tool-loop turns per rollout); GOLDConfig now overrides remove_unused_columns to default False (like GRPOConfig) so custom dataset columns reach environment.reset(**kwargs).
• Off-policy masking (DataCollatorForChatML / DataCollatorForVisionLanguageChatML, and the pretokenized dataset-prep path): the prompt is masked, assistant turns — including the assistant's tool-call tokens — are supervised, and tool-result tokens are masked to -100. Mirrors sft_trainer.py semantics.
• On-policy tool loop: _tool_call_loop / _get_tool_suffix_ids / _generate_single_turn ported near-verbatim from grpo_trainer.py. A tool_mask tracked during generation excludes tool-result tokens from the JSD/Liger loss.
• environment_factory: ported from GRPO (single-environment form) — pooled per-rollout instances, reset(**example) per rollout, observations injected back into the conversation. get_reward is excluded from tool extraction, matching GRPO's reserved env-owned-reward hook (Environment-owned reward #6238).
• Same-family only, enforced at init: tools + use_uld_loss=True, tools + seq_kd=True, and tools + packing=True raise ValueError; VLM tool rollouts require use_vllm=True; multimodal tool responses (tools returning images) raise NotImplementedError.

Example scripts

Two runnable examples, both pure on-policy (lmbda=1.0): the student generates and actually executes the tools, and the teacher scores the student's own rollouts.

• gold_tool_calling.py (text): Qwen3-8B → Qwen3-1.7B on a browser-agent dataset (DataCreatorAI/tool-calling-browser-agent-tasks), filtered to conversations using six executable browser-form tool stubs passed via tools=. Demonstrates the plain-tools path: shared tokenizer, Qwen3's native tool-calling template, thinking disabled via chat_template_kwargs.
• gold_tool_calling_vlm.py (VLM): Qwen3-VL-8B-Instruct → Qwen3-VL-2B-Instruct on Search-VL VQA tasks (OpenSearch-VL/Search-VL-SFT-36K), with live tools: DuckDuckGo search (smolagents) plus a docling-based layout_parsing environment — this one exercises environment_factory (pooled instances, per-example reset resolving image references). Tool outputs are capped and the tool loop is bounded by max_tool_calling_iterations=10.

Notable details

• Prefix-preserving template: multi-turn tool masking locates message boundaries by incremental prefix rendering, which requires a prefix-preserving chat template. When the tokenizer's own template isn't (e.g. Qwen3 drops historical <think> blocks), all tool-flow renders — off-policy prep, both collators, and the on-policy _get_tool_suffix_ids — use a single training-safe template computed once at init.
• VLM + vLLM image-token expansion (#6294): vLLM is fed the unexpanded (tokenizer-only, single-<image>) prompt IDs; the processor-expanded IDs remain the source of truth for length bookkeeping and the training forward pass.

Testing

• New off-policy masking, on-policy tool-loop smoke, environment_factory, and validation-guard tests in tests/experimental/test_gold_trainer.py.
• Both example scripts trained end-to-end on a single 80GB A100 (transformers 5.x + vLLM colocate + Liger fused JSD).

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue?
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.

---

Note

Medium Risk
Large changes to on-policy generation, loss masking, and VLM vLLM prompting affect core training correctness, but init guards and extensive tests reduce regression risk.

Overview
GOLDTrainer now supports multi-turn tool calling for same-family (shared tokenizer) distillation, aligned with GRPO-style tools and optional environment_factory.

Config adds max_tool_calling_iterations and defaults remove_unused_columns=False so custom dataset columns reach environment.reset(**kwargs).

Training behavior: off-policy paths mask tool-result tokens while supervising assistant (including tool-call) tokens via collators, pretokenized tool_mask, and prompt boundaries at the first assistant turn; on-policy runs a multi-turn execute-and-regenerate loop, logs tool metrics, and applies tool_mask so tool results are excluded from the distillation loss. Init rejects incompatible combos (seq_kd, ULD, packing, VLM without vLLM, multimodal tool responses).

VLM + vLLM: generation uses tokenizer-only unexpanded prompt IDs (#6294) while processor-expanded IDs stay authoritative for training bookkeeping.

Deliverables: runnable gold_tool_calling.py and gold_tool_calling_vlm.py examples plus a large expansion of test_gold_trainer.py (masking, validation, mocked on-policy loop, environment factory).

^{Reviewed by Cursor Bugbot for commit d9b273c. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Strongich]

Hi @kashif! I've added multi-turn tool-calling support to this trainer -- would appreciate a review when you have time.

Scope: same-family LMs and VLMs only. For VLMs, tool rollouts additionally require use_vllm=True.

Two examples added, both pure on-policy (lmbda=1.0) with use_vllm=True:

1. Text: Qwen3-8B -> Qwen3-1.7B (trl/experimental/gold/gold_tool_calling.py)

2. VLM: Qwen3-VL-8B -> Qwen3-VL-2B (trl/experimental/gold/gold_tool_calling_vlm.py)

A few caveats on the examples:

1. The VLM run drifted: as I found while experimenting afterwards, the off-the-shelf teacher itself is not stable with the system prompt and the stand-in tools we use (it often answers from memory instead of calling them), so on-policy distillation pulled the student from ~2 tool calls per rollout initially down to ~0. Producing a high-level example would require fine-tuning the teacher to use these tools first, and then either generating traces to allow $\lambda < 1$ or continuing full on-policy distillation with the fine-tuned teacher -- I don't have enough time & compute for either. Nevertheless, tool distillation is a niche training approach, and the general idea of this PR is to give users who already know what they are doing the tool they can use for it with zero (or almost zero) need to write additional trainer code.
2. Text-only distillation is more stable in terms of training, but the tools used are mocks (the chosen dataset's tools would require real form-submission backends).

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit d9b273c. Configure here.}