`nemo_rl.models.generation.vllm.utils`#

Module Contents#

Functions#

`assert_reload_refit_config_supported`	Reject pure-config combinations unsupported by vLLM reload refit.
`is_grouped_moe_expert_weight_name`	Return whether a checkpoint key is a grouped MoE expert slab.
`assert_refit_unsupported_grouped_moe_params`	Reject grouped MoE MXFP8 state-dict params before refit starts.
`_as_routed_experts_tensor`	Convert backend routed-expert ids to the resolved carry dtype.
`validate_rollout_prompt`	Reject rollouts generated from a different prompt than the learner’s.
`format_prompt_for_vllm_generation`	Format a list of prompts for vllm generation (which requires a specific format for its own `generate` method).
`pad_and_align_routed_expert_indices`	Return full-sequence-aligned routed experts as `[S, L, topk]` in `routed_experts_dtype`.
`attach_routed_experts_to_chat_response_choices`	Attach aligned routed experts to OpenAI chat response choices.
`attach_token_information_to_chat_response_choices`	Attach engine-native token information to OpenAI chat response choices.
`model_dump_chat_response_with_dynamic_message_fields`	Dump a vLLM OpenAI chat response while preserving dynamic message fields.
`encode_counter_key`	Encode a histogram component or a label into a flat counter key.
`_is_kept`	Whether a snapshot key belongs to a family the driver consumes.
`aggregate_spec_decode_counters`	Aggregate vLLM engine counters from multiple workers.
`_first_delta`	Return the delta for the first candidate name present, else `None`.
`_mean_from_histogram`	Mean of a histogram over the step, from its `sum`/`count` delta.
`compute_engine_step_metrics`	Compute per-step vLLM engine token, sequence-length and outcome metrics.
`compute_spec_decode_metrics`	Compute delta and derived metrics for speculative decoding.
`resolve_generation_worker_cls`	Return the quantized vLLM generation worker FQN if `quant_cfg` is set, else `default_cls`.

Data#

`R3_MISSING_ROUTE_SENTINEL`
`VLLM_LOGPROB_FLOOR`
`G_ROUTED_EXPERTS_RANGE_CHECKED`
`GROUPED_MOE_MXFP8_REFIT_ERROR`
`_GROUPED_MOE_EXPERT_WEIGHT_SUFFIXES`
`COUNTER_KEY_SEP`
`HISTOGRAM_SUM_PART`
`HISTOGRAM_COUNT_PART`
`_KEPT_COUNTER_NAMES`
`GENERATION_WORKER_OVERRIDES`

API#

nemo_rl.models.generation.vllm.utils.R3_MISSING_ROUTE_SENTINEL#: None

nemo_rl.models.generation.vllm.utils.VLLM_LOGPROB_FLOOR#: None

nemo_rl.models.generation.vllm.utils.G_ROUTED_EXPERTS_RANGE_CHECKED#: False

nemo_rl.models.generation.vllm.utils.GROUPED_MOE_MXFP8_REFIT_ERROR#: ‘MXFP8 refit does not support grouped MoE expert weights.’

nemo_rl.models.generation.vllm.utils._GROUPED_MOE_EXPERT_WEIGHT_SUFFIXES#: (‘mlp.experts.gate_up_proj’, ‘mlp.experts.down_proj’)

nemo_rl.models.generation.vllm.utils.assert_reload_refit_config_supported( config: nemo_rl.models.generation.vllm.config.VllmConfig, ) → None#: Reject pure-config combinations unsupported by vLLM reload refit.

nemo_rl.models.generation.vllm.utils.is_grouped_moe_expert_weight_name(name: str) → bool#: Return whether a checkpoint key is a grouped MoE expert slab.

nemo_rl.models.generation.vllm.utils.assert_refit_unsupported_grouped_moe_params( config: nemo_rl.models.generation.vllm.config.VllmConfig, state_dict_info: dict[str, Any], ) → None#: Reject grouped MoE MXFP8 state-dict params before refit starts.

nemo_rl.models.generation.vllm.utils._as_routed_experts_tensor( value: Any, *, device: torch.device, dtype: torch.dtype, ) → torch.Tensor#

Convert backend routed-expert ids to the resolved carry dtype.

Guards against expert ids overflowing dtype before the narrowing cast, which would otherwise wrap silently (e.g. if the expert count was mis-detected when resolving the dtype).

nemo_rl.models.generation.vllm.utils.validate_rollout_prompt( expected: list[int], actual: list[int] | None, ) → None#: Reject rollouts generated from a different prompt than the learner’s.

nemo_rl.models.generation.vllm.utils.format_prompt_for_vllm_generation( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], sample_idx: Optional[int] = None, ) → list[dict[str, Any]]#

Format a list of prompts for vllm generation (which requires a specific format for its own generate method).

See https://docs.vllm.ai/en/v0.9.1/features/multimodal_inputs.html for prompt format for multimodal inputs.

nemo_rl.models.generation.vllm.utils.pad_and_align_routed_expert_indices( request_output: Any, completion_output: Any, *, valid_length: int, padded_length: int, device: torch.device, require_complete_routed_experts: bool = False, allow_missing_routed_experts_fallback: bool = True, return_stats: bool = False, routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE, ) → Optional[torch.Tensor] | tuple[Optional[torch.Tensor], dict[str, int]]#: Return full-sequence-aligned routed experts as [S, L, topk] in routed_experts_dtype.

nemo_rl.models.generation.vllm.utils.attach_routed_experts_to_chat_response_choices( response: Any, final_request_output: Any, *, device: torch.device, logger: Any = None, routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE, ) → Any#: Attach aligned routed experts to OpenAI chat response choices.

nemo_rl.models.generation.vllm.utils.attach_token_information_to_chat_response_choices( response: Any, final_request_output: Any, ) → Any#: Attach engine-native token information to OpenAI chat response choices.

nemo_rl.models.generation.vllm.utils.model_dump_chat_response_with_dynamic_message_fields( response: Any, ) → dict[str, Any]#: Dump a vLLM OpenAI chat response while preserving dynamic message fields.

nemo_rl.models.generation.vllm.utils.COUNTER_KEY_SEP#: ‘|’

nemo_rl.models.generation.vllm.utils.HISTOGRAM_SUM_PART#: ‘sum’

nemo_rl.models.generation.vllm.utils.HISTOGRAM_COUNT_PART#: ‘count’

nemo_rl.models.generation.vllm.utils._KEPT_COUNTER_NAMES: tuple[str, ...]#: None

nemo_rl.models.generation.vllm.utils.encode_counter_key(name: str, part: str) → str#: Encode a histogram component or a label into a flat counter key.

nemo_rl.models.generation.vllm.utils._is_kept(metric_name: str) → bool#: Whether a snapshot key belongs to a family the driver consumes.

nemo_rl.models.generation.vllm.utils.aggregate_spec_decode_counters( worker_metrics: list[dict[str, float | list[float]]], ) → dict[str | tuple[str, int], float]#

Aggregate vLLM engine counters from multiple workers.

Combines the metrics collected from DP leader workers into a single aggregated counter dictionary. Retains the spec-decode family and the engine series listed above; everything else in the snapshot is dropped.

Parameters:: worker_metrics – List of metric dictionaries from each worker. Each dict maps metric names to float values or lists of floats (for per-position metrics).
Returns:: Dictionary mapping metric names to their aggregated float values. Per-position metrics use (name, position) tuples as keys.

.. rubric:: Example

metrics_from_workers = policy_generation.get_metrics() counters = aggregate_spec_decode_counters(metrics_from_workers) print(counters.get(“vllm:spec_decode_num_drafts”, 0)) 1234.0

nemo_rl.models.generation.vllm.utils._first_delta( delta: dict[str | tuple[str, int], float], candidates: tuple[str, ...], part: Optional[str] = None, ) → Optional[float]#: Return the delta for the first candidate name present, else None.

nemo_rl.models.generation.vllm.utils._mean_from_histogram( delta: dict[str | tuple[str, int], float], candidates: tuple[str, ...], ) → Optional[float]#

Mean of a histogram over the step, from its sum/count delta.

Exact rather than bucket-interpolated: vLLM tracks both, so the mean needs no approximation. None when the engine served no request in the step, which is a real state (a step spent entirely in training) and not a zero-length sequence.

nemo_rl.models.generation.vllm.utils.compute_engine_step_metrics( start_counters: dict[str | tuple[str, int], float], end_counters: dict[str | tuple[str, int], float], ) → dict[str, float]#

Compute per-step vLLM engine token, sequence-length and outcome metrics.

These are the engine’s own accounting, which is why they are worth carrying even though the driver already derives token counts from the tensors a generate() call returns: sequence-length distributions and aborted requests leave no trace in the returned tensors at all.

Parameters:

start_counters – Counter snapshot taken before generation.
end_counters – Counter snapshot taken after generation.

Returns:

Metrics for logging, keyed with a vllm/ prefix. Absent series are omitted rather than reported as zero, so a vLLM release that renames one leaves a gap in the dashboard instead of a plausible-looking zero.

nemo_rl.models.generation.vllm.utils.compute_spec_decode_metrics( start_counters: dict[str | tuple[str, int], float], end_counters: dict[str | tuple[str, int], float], ) → dict[str, float]#

Compute delta and derived metrics for speculative decoding.

Calculates the difference between two counter snapshots and derives acceptance rate and acceptance length metrics for logging.

Parameters:

start_counters – Counter snapshot taken before generation.
end_counters – Counter snapshot taken after generation.

Returns:

Dictionary of metrics suitable for logging to wandb/tensorboard. Keys are prefixed with “vllm/” for namespace consistency. Includes: - vllm/spec_num_drafts: Total number of draft batches - vllm/spec_num_draft_tokens: Total draft tokens generated - vllm/spec_num_accepted_tokens: Total tokens accepted - vllm/spec_acceptance_length: Average accepted tokens per draft + 1 - vllm/spec_acceptance_rate: Ratio of accepted to draft tokens - vllm/{metric}-{position}: Per-position acceptance counts - vllm/spec_acceptance_rate-pos-{position}: Per-position acceptance rates

nemo_rl.models.generation.vllm.utils.GENERATION_WORKER_OVERRIDES#: None

nemo_rl.models.generation.vllm.utils.resolve_generation_worker_cls(default_cls: str, config: dict) → str#

Return the quantized vLLM generation worker FQN if quant_cfg is set, else default_cls.

Safe to call even when ModelOpt is not installed — returns default_cls unchanged whenever quant_cfg is None, so the core generation path stays import-free of ModelOpt.

nemo_rl.models.generation.vllm.utils#

Module Contents#

Functions#

Data#

API#

`nemo_rl.models.generation.vllm.utils`#