Repository navigation
[DistillationTrainer refactor] Support tool calling - #6723
Conversation
…rompts Mirror GRPO verbatim (minus the tools/environments/chat_template it lacks): add self._is_vlm (canonical if/elif/else on the processing_class), import prepare_multimodal_messages + transformers, and rewrite _tokenize_prompts to normalize multimodal messages, extract images, and return the processor's multimodal fields. Steps 2-4 (batch threading, loss threading, vLLM+tests+docs) follow.
…batch _generate now returns multimodal_fields; _generate_and_score_completions stores the processor's multimodal keys (pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, token_type_ids, mm_token_type_ids, image_position_ids) in the output batch, mirroring GRPO's block verbatim. Text path unaffected (multimodal_fields is empty).
…repare_inputs Mirror GRPO: split_pixel_values_by_grid before splitting the generation batch into micro-batches, and unsplit_pixel_values_by_grid on each buffered micro-batch, so each training step gets the pixel_values for its own prompts. No-op for the text path. _compute_loss threading (step 3b) still to come.
…essor auto-build - Pass the multimodal fields to the student/teacher hidden-state forward in _compute_loss. - Auto-build a processor (AutoProcessor) instead of a tokenizer, and resolve pad_token via self._tokenizer, so VLM processors work — mirrors GRPO.
Mirror GRPO's TestGRPOTrainerVLM: parametrize over the tiny VLM checkpoints and train on zen-image. Uses self-distillation (no tiny+small VLM fixture pair) and asserts the loss stayed finite rather than params-changed.
Mirror GRPO's skipped vLLM+images placeholder (server mode, pending a vLLM server mock). The colocate vLLM path already threads images identically to GRPO.
Add a 'Training Vision Language Models' section: pass a VLM as student and teacher, use a prompt-only dataset with an image/images column, tested architectures + the compatibility tip.
…vice-prep, token_type_ids Mirror GRPO's multimodal handling in _generate_and_score_completions: - Read the dataset image/images column and inject it into prompts via prepare_multimodal_messages (previously the image column was dropped, so VLM training ran text-only). - Rebuild forward_kwargs from the images with the processor (batched, device-placed tensors), emit num_images / num_tiles so split_pixel_values_by_grid regroups packed pixel_values under gradient accumulation, and extend token_type_ids / mm_token_type_ids over the completion. - Thread token_type_ids / mm_token_type_ids into _get_last_hidden_state (required by Qwen-VL M-RoPE); restore GRPO's pre-5.0 VLM backbone fallback. - Add a gradient-accumulation VLM test exercising the num_images split.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c0fae09aca
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
c0fae09 to
d3bb76e
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
…riable-size images)
Replace the What's New section (GRPO multi-environment RL + KTO stable) with the DistillationTrainer graduation to the stable API, in README.md and docs/source/index.md.
Add the tools ctor param + guards (transformers>=5.0, jmespath/response_template, supports_tool_calling), build the tool dict, wire the response schema + training-safe chat template, and set max_tool_calling_iterations. Async tools rejected (deferred to a future PR); environments dropped (no reward path). Mirrors GRPO.
… masking - _tokenize_prompts: pass tools + chat_template. - _generate: parse_response decode + _tool_call_loop; return tool_mask/images + tool metrics. - _get_tool_suffix_ids + _tool_call_loop ported from GRPO, dropping async + logprobs (justified; every other line diff-verified against GRPO). - _generate_and_score_completions: pad tool_mask, loss_mask = completion_mask * tool_mask, store tool_mask in the batch. _compute_loss masks the JSD by loss_mask (tool-result tokens stay attended but are not distilled).
Mirror GRPO's tool test — mock generation to emit a valid call, an invalid call, and a non-call, and assert tools/call_frequency==2/3 and tools/failure_frequency==1/2 (async variant dropped).
d9a1e9e to
a2ec0f8
Compare
d3bb76e to
083be80
Compare
…ling # Conflicts: # tests/test_distillation_trainer.py # trl/trainer/distillation_trainer.py
…h GRPO - Document tool calling: add an Agent Training section (Tools + Multimodal Tool Responses) mirroring grpo_trainer.md, and list tools/call_frequency and tools/failure_frequency in Logged metrics (plus the non-tool token qualifier GRPO carries on num_tokens and completions/*_length). - Move max_tool_calling_iterations to the training section of DistillationConfig, matching GRPOConfig. - Port test_train_with_tools_multimodal_response and test_train_with_tools_text_response_multimodal_prompt from test_grpo_trainer.py; both fail without the VLM x tools fixes in 8f11ec3. - Restore the is-not-None tool metric assertions in test_train_with_tools.
| break | ||
| tid = self._tokenizer.convert_tokens_to_ids("<|video_pad|>") | ||
| if tid != self._tokenizer.unk_token_id: | ||
| self._video_pad_token_id = tid |
There was a problem hiding this comment.
Tool-image token IDs miss models
Medium Severity
Vision type IDs for tool-returned images are rebuilt by matching only <|image_pad|> and <|image|>. Advertised VLMs such as Gemma 3 (<start_of_image>) and LLaVA never match, so mm_ids stay zeros and image placeholders in the completion are not marked for generate or the JSD forward.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 49747ae. Configure here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 740ef37. Configure here.
|
Deep review done before merging: cursor's vllm_mode catch was real (AttributeError with use_vllm + tools), fixed in 8ed5196 along with gating the decode on tools and doc alignment; the other bot findings were fixed earlier or exist verbatim in GRPO. |
What
Important
For reviewers: it's mostly copy-pasted from GRPO, nothing really new here!
Adds tool-calling support to
DistillationTrainer, mirroring GRPO. The student generates multi-turn tool-calling trajectories on-policy; tool-result tokens are masked out of the JSD loss.max_tool_calling_iterationsconfig field (caps tool turns;None= unlimited)toolsconstructor param +_tool_dict,supports_tool_calling/ response-schema / training-chat-template setup_tool_call_loop,parse_response,_get_tool_suffix_idsin the generation pathtool_maskloss masking:loss_mask = completion_mask * tool_masktest_train_with_tools(mockedgenerate, assertstools/call_frequencyandtools/failure_frequency)All mirrored blocks are reproduced verbatim from GRPO.
Deferred (future PR)
Async tools, environments, and per-token logprob capture — none exist on a stable base yet.
Stacking
Stacked on #6650 (VLM). Re-base to
mainonce #6650 lands.Part of #6449.
Note
Medium Risk
Touches core on-policy generation, loss masking, and VLM forward kwargs in
DistillationTrainer; behavior is largely ported from GRPO but multimodal tool loops add integration risk.Overview
Adds on-policy tool-calling to
DistillationTrainer, aligned with GRPO: pass atoolslist, run a multi-turn generate→execute→regenerate loop, and distill only on student-generated tokens by masking tool-result tokens out of the JSD loss (loss_mask = completion_mask * tool_mask).DistillationConfiggainsmax_tool_calling_iterationsto cap tool turns. Init wires response parsing, prefix-preserving chat templates, transformers ≥5.0 checks, and rejects async tools. Logging addstools/call_frequency/tools/failure_frequency, and token/completion length metrics note they exclude tool tokens when tools are enabled.VLM paths merge tool-returned images into generation and the forward pass (including
mm_token_type_ids/ truncation safety). Docs add an Agent Training section for distillation; GRPO docs document the same iteration cap.Tests cover mocked multi-turn tool calls (success/failure metrics) and VLM tool responses (image and text).
Reviewed by Cursor Bugbot for commit 8ed5196. Bugbot is set up for automated code reviews on this repo. Configure here.