Skip to content

fix(misc): normalize zero-temperature sampling - #5759

Open
yaoyu-33 wants to merge 1 commit into
mainfrom
yuya/verified-bug-hunter-temperature-zero-20260824
Open

fix(misc): normalize zero-temperature sampling#5759
yaoyu-33 wants to merge 1 commit into
mainfrom
yuya/verified-bug-hunter-temperature-zero-20260824

Conversation

@yaoyu-33

Copy link
Copy Markdown
Contributor

Summary

The maintained standalone sync and async text-generation CLIs accept --temperature 0 --top_p 0.9. Their shared sampling builder previously resolved that to active nucleus sampling with a zero temperature. The pinned MCore torch sampler then divided finite logits by zero and aborted generation with a non-finite probability error after model and distributed initialization.

Normalize exact zero-temperature requests to greedy decoding (top_k=1, top_p=0). This matches the pinned MCore OpenAI completion frontends and preserves every nonzero-temperature path.

Root cause and fix

add_sampling_args accepts zero temperature, and build_sampling_params selected top_k=0 whenever positive top-p was supplied. Both offline entrypoints forwarded that tuple unchanged to MCore. The minimal fix handles exact temperature zero at the shared Bridge ownership boundary before the existing omitted-top-k and incompatibility logic.

The regression parses the supported Bridge CLI, exercises the production builder, and calls the exact pinned CPU-pure torch_sampling.py implementation. Only the sampler's abstract base import is stubbed so this focused unit test does not initialize unrelated CUDA/MCore package dependencies.

Regression evidence

Before the production fix, the new test failed with:

temperature = 0.0, top_k = 0, top_p = 0.9
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

After the unchanged production fix:

tests/unit_tests/scripts/test_text_generation.py::test_zero_temperature_top_p_sampling_is_greedy
1 passed

Validation

uv run --no-sync python -m pytest --confcutdir=tests/unit_tests/scripts \
  tests/unit_tests/scripts/test_text_generation.py -q
# 21 passed

uv run --no-sync python -m pytest --confcutdir=tests/unit_tests/scripts \
  tests/unit_tests/scripts/test_text_generation_entrypoint.py \
  tests/unit_tests/scripts/test_async_text_generation_entrypoint.py -q
# 7 passed

uv run --no-sync pre-commit run --all-files
# all hooks passed

git diff --check
# passed

Scope

No dependencies, lockfiles, workflows, public signatures, MCore source, server request handling, model code, or nonzero-temperature sampling behavior changed. No GPU, distributed, convergence, performance, or full-suite validation is claimed.

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 24, 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.

@yaoyu-33 yaoyu-33 added area:misc Cross-cutting utilities, logging, helpers, and other changes bug Something isn't working needs-review PR is ready for code review and waiting on a reviewer labels Aug 24, 2026
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 34dc1cd

@claude

claude Bot commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Light review — LGTM

The change correctly fixes the temperature == 0 + --top_p > 0 case. Previously that combination set top_k = 0, routing MCore down the top-p branch where logits / temperature divides by zero. Forcing top_k = 1 (and top_p = 0.0) makes it greedy, which short-circuits the temperature division in TorchSampling.sample_from_logits. Correct and well-scoped.

Minor observations (non-blocking):

  • The new branch silently overrides any explicit --top_k/--top_p when temperature == 0, so no ValueError is raised even if the user passed both. That is reasonable (temperature 0 = deterministic greedy), just worth confirming it is intended.
  • The unit test loads torch_sampling.py directly from the 3rdparty/Megatron-LM submodule via spec_from_file_location, coupling it to the MCore greedy (top_k == 1 -> argmax) implementation. Acceptable, but note the coupling.

Suggested test cases

  • test_zero_temperature_top_p_sampling_is_greedy (added — verifies greedy selection and top_k==1, top_p==0.0)
  • test_top_p_sampling_is_compatible_with_default_cli_values (existing — regression guard for the non-zero-temperature top-p path)
  • test_default_sampling_remains_greedy (existing — default CLI still greedy)
  • Consider adding a case with --temperature 0 --top_k 50 to confirm the zero-temperature branch overrides an explicit top_k without raising the both-positive error.
  • No perf tests impacted.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:misc Cross-cutting utilities, logging, helpers, and other changes bug Something isn't working needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant