Skip to content

[TOP 2/N] Add batch-invariant scorer integration coverage - #1964

Closed
kaixih wants to merge 4 commits into
radixark:mainfrom
kaixih:dsv4-top-2-logprob-scoring
Closed

[TOP 2/N] Add batch-invariant scorer integration coverage#1964
kaixih wants to merge 4 commits into
radixark:mainfrom
kaixih:dsv4-top-2-logprob-scoring

Conversation

@kaixih

@kaixih kaixih commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Add one GPU integration test for the existing Miles true-on-policy scorer.

With Megatron batch-invariant mode enabled, the scorer keeps its existing path:

TP full-vocab gather
-> padded-vocab truncation
-> torch.log_softmax
-> selected-token gather

The test uses BF16 logits with the DSV4 129280-token vocabulary. It verifies
that Miles reaches Megatron's registered BI log_softmax implementation and
that the same row scored alone or in a 17-row batch produces a bitwise-identical
selected-token logprob.

Scope

  • No runtime code or CLI changes.
  • No duplicate BI kernel or custom op.
  • DSV4 prefill-vs-trainer E2E parity remains a follow-up.

This is split 2/N from #1788 and is orthogonal to #1912.

Validation

  • New GPU integration test: passed on GB200 and in stage-b-2-gpu-h200 CI.
  • Focused local scorer/config tests: 59 passed.
  • GitHub pre-commit and CPU matrix: passed.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@kaixih kaixih changed the title [TOP 2/N] Add opt-in FP32 batch-invariant log-probability scoring [TOP 2/N] Add opt-in BF16/FP32 batch-invariant log-probability scoring Jul 30, 2026
@kaixih kaixih changed the title [TOP 2/N] Add opt-in BF16/FP32 batch-invariant log-probability scoring [TOP 2/N] Add batch-invariant scorer integration coverage Jul 30, 2026
@kaixih

kaixih commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Closing this test-only split because no additional Miles runtime integration is required. Megatron-LM's enable_batch_invariant_mode() globally registers the accelerator implementation of aten::_log_softmax; therefore Miles' existing torch.log_softmax scorer automatically uses the batch-invariant kernel under --batch-invariant-mode. This conclusion is also recorded in #1788.

@kaixih kaixih closed this Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant