feat(eval): add tool-call accuracy evaluator for agent SFT validation - #2338
Merged
HuiyingLi merged 13 commits intoMay 30, 2026
Conversation
khazic
requested review from
HuiyingLi,
ZhiyuLi-Nvidia,
adil-a,
akoumpa,
athitten,
hemildesai,
pthombre and
zyzhou5
as code owners
May 28, 2026 11:10
Loss-only validation cannot catch "model overfits format but emits
wrong tool names or malformed argument JSON" regressions. This adds a
generation-based tool-call accuracy evaluator that runs at every val
step alongside val_loss and surfaces six per-call metrics (has_call,
name_correct, args_json_valid, args_field_recall, args_field_precision,
args_exact_match) through the existing wandb/MLflow/Comet/JSONL loggers.
* New nemo_automodel/components/eval/{tool_call_parser,tool_call_evaluator}.py
with a permissive parser covering Qwen/Hermes/Llama3/Mistral/GPT-OSS
wrappers. The parser uses anchor regexes plus a balanced-brace scanner
so nested JSON arguments survive. The scorer aligns predictions and
ground-truth tool calls positionally so omitting one of N parallel
calls is correctly penalized.
* Extends agent_chat with make_agent_chat_eval_samples() that expands
each dialogue into one eval sample per assistant tool-call position.
* train_ft.py wires the evaluator into _run_validation_epoch and does a
DP all-reduce on (mean*count, count) so corpus means stay correct
under FSDP and any future sharded-eval extensions.
* Example YAMLs gain an opt-in tool_call_eval block (enabled by default
on the full-SFT recipe, commented template on the LoRA recipe).
* 41 new unit tests covering parser formats, nested JSON in arguments,
parallel-call GT alignment, metric aggregation, sample sharding, and
error paths for generate() failures.
Signed-off-by: khazic <khazzz1c@gmail.com>
Calling apply_chat_template with tokenize=True returns a list of token strings (not ids) under some templates / transformers versions. The list-of-str then crashes torch.tensor(..., dtype=long) with "too many dimensions 'str'" inside ToolCallAccuracyEvaluator.evaluate. Render with tokenize=False and run the resulting text through tokenizer.__call__ to get ids. Adds a regression test that asserts the evaluator skips samples whose template returns a non-string. Signed-off-by: khazic <khazzz1c@gmail.com>
Without this, only val_loss / lr / num_label_tokens appear on stdout during validation. Operators monitoring training have to tail the JSONL file or open wandb to verify the new tool-call evaluator is wired up. Appending tool_name_acc / args_json_valid / args_exact_match (plus the scored sample count) to the same INFO line makes presence and direction of the metric obvious from a live tail. Signed-off-by: khazic <khazzz1c@gmail.com>
Two operational issues surfaced on the first end-to-end run: 1. Eval ran but produced n=0 with no diagnostic. The max_prompt_tokens cap was skipping every sample without logging, so it was impossible to tell whether the cause was prompt length, template failure, or tokenization. Promote the skip to a WARNING with the offending length so operators can either raise the cap or trim prompts. 2. The next training step after validation OOMed at cross_entropy on a GPU that was at 22 GiB before val and 73 GiB after. FSDP unshards parameters during generate() and leaves large intermediate buffers cached; without an explicit release, the next training forward has no room for its own logits. Call torch.cuda.empty_cache() at the end of evaluate(). Signed-off-by: khazic <khazzz1c@gmail.com>
When all samples are skipped the user previously saw a single n=0 line with no clue why. Track each skip path (chat_template raised / wrong return type / encode failed / prompt too long / generate failed) and emit a single summary WARNING at the end of evaluate() with the breakdown. Also log the tokenizer class and sample count at the start so unexpected wrappers are obvious in the log. This is purely diagnostic — no behavior change to the metric values. Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
Multi-rank log filtering eats logger.warning output from non-rank-0 workers, so when generate() raises under FSDP2 we currently only see the aggregated skip count without the actual exception. Persist the full traceback of the first failure on each rank to a deterministic file path so the root cause is recoverable post-hoc. This is purely diagnostic; no behavior change. Signed-off-by: khazic <khazzz1c@gmail.com>
…bsent Several Automodel custom model classes (notably Qwen2ForCausalLM) inherit HFCheckpointingMixin + Qwen2PreTrainedModel but not GenerationMixin, so the FSDP-wrapped instance has no .generate() method and every sample fails with AttributeError. This was masking the value of the evaluator under the default fsdp2 training setup used by the agent SFT example. Add _greedy_generate_manual that does a minimal token-by-token greedy decode using only model.forward(). It returns the same shape that .generate() would return (prompt + generated tokens) so the rest of the evaluator (decode, parse, score) is unchanged. No KV cache, so cost is O(L * (P + L)) per sample, but L defaults to 256 tokens which is well within the eval-step budget for 32 - 128 samples. This unblocks the in-loop evaluator under FSDP2 + any custom model without GenerationMixin, which is the common case for Automodel users. Signed-off-by: khazic <khazzz1c@gmail.com>
khazic
force-pushed
the
feat/agent-sft-toolcall-eval
branch
from
May 28, 2026 11:14
1a83f9c to
79803af
Compare
Contributor
|
/ok to test 79803af |
Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/claude review |
Contributor
|
/ok to test 85a59a9 |
Contributor
|
/ok to test 1c89bf6 |
…call-eval Signed-off-by: khazic <khazzz1c@gmail.com> # Conflicts: # tests/unit_tests/datasets/llm/test_agent_chat.py
Contributor
|
/ok to test a6906a6 |
HuiyingLi
approved these changes
May 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Loss-only validation cannot catch the "model overfits format but emits wrong tool names or malformed argument JSON" failure mode during agent SFT —
val_lossis decreasing while the model is silently producing tool calls that no downstream caller can execute. This PR adds a generation-based tool-call accuracy evaluator that runs alongsideval_lossat every val step and surfaces six per-call metrics through the existing wandb / MLflow / Comet / JSONL loggers:tool_call/has_calltool_call/name_correcttool_call/args_json_validtool_call/args_field_recalltool_call/args_field_precisiontool_call/args_exact_matchPlus skip-reason diagnostics (
tool_call/_skip_<reason>) so failures are debuggable without re-running.What's included
nemo_automodel/components/eval/tool_call_parser.py— permissive tool-call parser covering Qwen / Hermes / Llama 3.1 / Mistral / GPT-OSS Harmony wrappers, with a balanced-brace scanner so nested JSON arguments survive parsing. Scorer aligns predictions to ground truth positionally, so a model that emits only one of two parallel tool calls is correctly penalized.nemo_automodel/components/eval/tool_call_evaluator.py—ToolCallAccuracyEvaluatorrenders prompts via the tokenizer's chat template, runs token-by-token generation, parses, aggregates metrics, and releases CUDA cache. Falls back to a manual greedy decode loop when the model class lacks.generate()— needed because several Automodel custom model classes (notablyQwen2ForCausalLM) inherit fromHFCheckpointingMixin + Qwen2PreTrainedModelbut not fromtransformers.generation.GenerationMixin, so the FSDP-wrapped instance has no.generate()method. With the fallback the evaluator works under single-GPU / DDP / FSDP2 uniformly.nemo_automodel/components/datasets/llm/agent_chat.py— addsmake_agent_chat_eval_samples()that expands each dialogue into one eval sample per assistant tool-call position. Reuses the existing_convert_messagesso eval and training see the exact same schema.nemo_automodel/recipes/llm/train_ft.py— wires the evaluator into_run_validation_epochwith DP all-reduce on (mean × count, count) for the regular metrics and SUM all-reduce for the count-like_skip_*diagnostics, so corpus means stay correct under any sharding. Adds tool-call metrics to the existing[val]stdout line so they are visible in a live tail without opening wandb.qwen2_5_3b_function_calling.yamlgains an opt-intool_call_evalblock (enabled by default), companion LoRA yaml has a commented-out template.apply_chat_template, failed encode, prompt too long,generate()OOM,generate()raises).End-to-end verification
8×80GB A100, Qwen2.5-3B + glaive_toolcall_en, FSDP2 strategy:
n=160is 20 samples × 8 ranks summed across the DP group. The base Qwen2.5-3B has not been SFT'd on the tool-call data here; the non-zero scores reflect the format prior it inherits from pretraining and serve as the lower bound the SFT run should improve on.Test plan
pytest tests/unit_tests/eval/ tests/unit_tests/datasets/llm/test_agent_chat.py— 64 passed