Skip to content

[data, test] feat: add unit tests for HuggingFace dataset processors - #3680

Closed
lonexreb wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
lonexreb:training/test-hf-processors
Closed

[data, test] feat: add unit tests for HuggingFace dataset processors#3680
lonexreb wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
lonexreb:training/test-hf-processors

Conversation

@lonexreb

@lonexreb lonexreb commented May 5, 2026

Copy link
Copy Markdown
Contributor

Summary

src/megatron/bridge/data/hf_processors/ exports three pure-function dataset processors used by default_squad_config, default_gsm8k_config, and default_openmathinstruct2_config in finetune_utils.py:

  • process_squad_example
  • process_gsm8k_example (+ private _extract_final_answer helper)
  • process_openmathinstruct2_example

Each is a dict-in / dict-out function with no I/O and no tokenizer dependency.

They had functional tests under tests/functional_tests/test_groups/data/hf_processors/ but zero unit tests — meaning every regression in input/output formatting required a GPU container CI slot to catch.

This PR adds 22 fast unit tests across 4 classes (+287 LoC). Tests-only — no production changes.

What's covered

Class Tests Coverage
TestProcessSquadExample 5 Context: ... Question: ... Answer: formatting, first-answer-as-output rule, original_answers preservation, single-answer case, bare Answer: suffix, tokenizer-arg-is-no-op contract, missing-field KeyError
TestExtractFinalAnswer 5 extraction after ####, whitespace stripping, no-delimiter fallback, empty-after-delimiter case, multiple-delimiter LAST-split rule
TestProcessGsm8kExample 4 Question: ... Answer: formatting, full-answer output, extracted-final-answer in original_answers, no-#### flow, tokenizer-no-op
TestProcessOpenMathInstruct2Example 5 Problem: ... Solution: formatting, generated-solution as output, expected_answer verbatim preservation, missing-field KeyError, tokenizer-no-op
TestProcessorOutputContract 3 (parametrized) Cross-processor invariant: every processor returns {input: str, output: str, original_answers: list[str] with ≥1 element}

Why this matters

Recipes built via _sft_common / _peft_common rely on these processors implicitly through the default_*_config helpers. A formatting regression (e.g. dropping the trailing Answer:, breaking _extract_final_answer's #### parsing) silently changes every fine-tuning recipe's input templates without the recipe-level tests noticing.

Test plan

  • python3 -m ast parse clean
  • ruff check clean
  • ruff format applied
  • CI: cicd-unit-tests-core picks up the new module under tests/unit_tests/data/

Risk

Zero — tests only.

Self-verification (lesson from #3648)

Before writing, ran git ls-files | grep hf_processors. Confirmed: only tests/functional_tests/test_groups/data/hf_processors/ exists; no tests/unit_tests/data/hf_processors/ directory present on main. This PR creates that directory.

`src/megatron/bridge/data/hf_processors/` exports three pure-function
dataset processors (squad, gsm8k, openmathinstruct2) used by the
`default_*_config` helpers in finetune_utils.py. Each one is a
dict-in / dict-out function with no I/O and no tokenizer dependency.

They had **functional tests** under
`tests/functional_tests/test_groups/data/hf_processors/` but **zero
unit tests** — meaning every regression in the input/output formatting
required a GPU container CI slot to catch.

This PR adds unit-test coverage so processor-format regressions get
caught at L0-unit-test time. 22 tests across 4 classes:

`TestProcessSquadExample` (5):
  - basic example produces documented `Context: ... Question: ...
    Answer:` formatting; output is the FIRST answer in the answers
    list; original_answers preserves all alternatives
  - single-answer example produces a valid output
  - input strictly ends with bare `Answer:`
  - tokenizer arg is a no-op (None and a sentinel produce identical
    output) — locks in the documented contract
  - missing required field surfaces a clear KeyError

`TestExtractFinalAnswer` (5):
  - extracts value after `####`, strips whitespace, returns full
    answer stripped when no delimiter, returns empty string when
    `####` is followed by whitespace, uses LAST split when `####`
    appears multiple times

`TestProcessGsm8kExample` (4):
  - basic example: `Question: ... Answer:` formatting; output is the
    full chain-of-thought; original_answers contains ONLY the
    extracted final numerical answer
  - no-`####` answer flows through `_extract_final_answer` correctly
  - input strictly ends with `Answer:`
  - tokenizer arg is a no-op

`TestProcessOpenMathInstruct2Example` (5):
  - basic example: `Problem: ... Solution:` formatting; output is the
    generated solution; original_answers wraps expected_answer in a
    one-element list
  - input strictly ends with `Solution:`
  - expected_answer is preserved verbatim (no stripping)
  - tokenizer arg is a no-op
  - missing required field surfaces KeyError

`TestProcessorOutputContract` (3, parametrized):
  - cross-processor invariant: every processor returns a dict with
    `input` (str), `output` (str), and `original_answers` (list[str]
    with at least one element)

Tests-only — no production code changes. Locks in the formatting
contract every recipe that uses these processors implicitly depends on.

Signed-off-by: lonexreb <reach2shubhankar@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented May 5, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@cuichenx cuichenx left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these tests already exist in tests/functional_tests/, but they actually don't use GPUs. it's better to move them instead of creating duplicate tests.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label May 6, 2026
@yaoyu-33 yaoyu-33 added area:data Dataset builders, preprocessing, and samplers ci CI, automation, test queue, or workflow infrastructure work labels May 7, 2026
Per reviewer feedback (cuichenx) on NVIDIA-NeMo#3680: the tests for the squad,
gsm8k, and openmathinstruct2 processors do not require GPUs, so they
belong under tests/unit_tests/, not tests/functional_tests/. Moves the
three existing test files instead of creating duplicates.

- git mv tests/functional_tests/test_groups/data/hf_processors/{test_squad,test_gsm8k,test_openmathinstruct2}.py
  → tests/unit_tests/data/hf_processors/
- remove now-empty functional test directory (and __init__.py)
- delete the duplicate consolidated test_processors.py added in the
  original PR
- restore tests/unit_tests/training/utils/test_flop_utils.py that was
  accidentally dropped in the initial commit

No test logic changes — only relocation.

Signed-off-by: lonexreb <reach2shubhankar@gmail.com>
@lonexreb

Copy link
Copy Markdown
Contributor Author

Thanks @cuichenx — addressed in 3f0b6e7:

  • git mv'd the three existing test files from tests/functional_tests/test_groups/data/hf_processors/tests/unit_tests/data/hf_processors/ (preserves git history)
  • removed the duplicate consolidated test_processors.py I'd originally added
  • also restored tests/unit_tests/training/utils/test_flop_utils.py that was accidentally dropped from the first commit during a rebase

Net change: ~30 existing pure-pytest tests now run under L0 unit tests instead of L0 functional tests, no GPU slot needed.

@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label May 12, 2026
@yaoyu-33 yaoyu-33 added the needs-review PR is ready for code review and waiting on a reviewer label May 12, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label May 14, 2026
@yaoyu-33 yaoyu-33 removed the needs-review PR is ready for code review and waiting on a reviewer label May 14, 2026
@lonexreb

Copy link
Copy Markdown
Contributor Author

Closing as superseded by upstream work.

After fetching the latest main, the contributions in this PR are now obsolete:

  1. HF processor unit tests (tests/unit_tests/data/hf_processors/{test_gsm8k,test_openmathinstruct2,test_squad}.py) — the move from functional → unit was already landed upstream, and the upstream versions have additional coverage. For test_openmathinstruct2.py specifically, upstream is 305 lines vs this branch's 159 lines (adds TestStripIntermediateBoxed cases that aren't in this branch).
  2. test_flop_utils.py — upstream landed a more comprehensive version via [training, test] test: add MLA, MTP, and provider-override coverage for FLOPs calculator #3695 ([training, test] test: add MLA, MTP, and provider-override coverage for FLOPs calculator), now 1664 lines vs this branch's 1393.

Rebasing this branch onto current main would delete 419 lines of better tests, so closing is the right call. Thanks @cuichenx for the original review feedback — the move to tests/unit_tests/ is reflected in upstream now.

@lonexreb lonexreb closed this May 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:data Dataset builders, preprocessing, and samplers ci CI, automation, test queue, or workflow infrastructure work community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants