nemo_rl.models.generation.score_centering#

vLLM request preparation and sample extraction for score centering.

The current strategy collects the top-k head. Future strategies must keep request preparation, validation and extraction consistent, and return the same pair of sampler token IDs and their sampler logprobs. Staging consumes this tensor contract regardless of how the candidates were selected.

Module Contents#

Classes#

ScoreCenteringRequestOverrides

HTTP request fields used by the current top-k collection strategy.

Functions#

build_score_centering_request_overrides

Build request overrides for collecting score-centering candidates.

validate_score_centering_request

Check that the HTTP request advertises the configured collection strategy.

extract_score_centering_samples

Extract sampler candidate IDs and logprobs from one vLLM completion.

extract_flat_score_centering_samples

Vectorized extract_score_centering_samples for vLLM FlatLogprobs.

Data#

VLLM_LOGPROB_FLOOR

API#

nemo_rl.models.generation.score_centering.VLLM_LOGPROB_FLOOR#

None

class nemo_rl.models.generation.score_centering.ScoreCenteringRequestOverrides#

Bases: typing.TypedDict

HTTP request fields used by the current top-k collection strategy.

Initialization

Initialize self. See help(type(self)) for accurate signature.

logprobs: bool#

None

top_logprobs: int#

None

nemo_rl.models.generation.score_centering.build_score_centering_request_overrides(
*,
num_candidates: int,
) → nemo_rl.models.generation.score_centering.ScoreCenteringRequestOverrides#

Build request overrides for collecting score-centering candidates.

Gym applies these fields to each outgoing policy request. Extend this function together with validation and extraction for another selection strategy, keeping the HTTP request consistent with engine sampling.

Parameters:

num_candidates – Number of sampler candidates to collect per token.

Returns:

HTTP request overrides. Callers merge them with other request settings.

nemo_rl.models.generation.score_centering.validate_score_centering_request(
request: Any,
*,
num_candidates: int,
) → None#

Check that the HTTP request advertises the configured collection strategy.

Parameters:

  • request – Resolved OpenAI-compatible chat request on the vLLM worker.

  • num_candidates – Candidate count configured for capture.

nemo_rl.models.generation.score_centering.extract_score_centering_samples(
output: Any,
*,
num_candidates: int,
) → tuple[torch.Tensor, torch.Tensor]#

Extract sampler candidate IDs and logprobs from one vLLM completion.

The current strategy takes the lowest-ranked candidates from each engine logprob mapping, excluding an extra sampled token outside the top-k head. Candidate logprobs retain their original mass rather than being normalized over the selected set. Candidates the sampler filtered out (top-k, top-p) keep their -inf processed logprob.

Parameters:

  • output – One engine CompletionOutput, including per-position logprobs.

  • num_candidates – Number of sampler candidates to retain per position.

Returns:

CPU (sampler_ids, sampler_logprobs) tensors, both shaped [generated_len, num_candidates], with int32 IDs and float32 logprobs. Callers retain the existing sampler_topk_token_ids and sampler_topk_logprobs transport fields.

nemo_rl.models.generation.score_centering.extract_flat_score_centering_samples(
output: Any,
*,
num_candidates: int,
) → tuple[torch.Tensor, torch.Tensor]#

Vectorized extract_score_centering_samples for vLLM FlatLogprobs.

With flat logprobs vLLM stores num_candidates + 1 entries per generated position: the sampled token first, then the top-k head by rank (a sampled token inside the head appears twice). Reshaping the flat lists replaces the per-position dict walk and applies the same checks to whole tensors.

Parameters:

  • output – One engine CompletionOutput whose logprobs is a vllm.logprobs.FlatLogprobs.

  • num_candidates – Number of sampler candidates to retain per position.

Returns:

CPU (sampler_ids, sampler_logprobs) tensors, both shaped [generated_len, num_candidates], with int32 IDs and float32 logprobs, matching extract_score_centering_samples.