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

Module Contents#

Functions#

`_format_refit_key_error`	Format a bounded refit-key diagnostic.
`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.
`_placeholder_range_bounds`
`remap_multimodal_placeholders`	Move vLLM multimodal ranges into the final prompt coordinates.
`_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_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._format_refit_key_error(label: str, keys: set[str]) → str#: Format a bounded refit-key diagnostic.

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_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._placeholder_range_bounds( placeholder_range: Any, ) → tuple[int, int]#

nemo_rl.models.generation.vllm.utils.remap_multimodal_placeholders( *, template_token_ids: list[int], final_token_ids: list[int], mm_placeholders: collections.abc.Mapping[str, list[Any]], splice: Optional[nemo_rl.models.generation.openai_server_utils.PrefixSplice] = None, ) → dict[str, list[Any]]#

Move vLLM multimodal ranges into the final prompt coordinates.

vLLM computes multimodal placeholder offsets while preprocessing the chat-template token sequence. NeMo-RL subsequently replaces the re-tokenized history with the exact model-generated token prefix (splice_prefix_tokens). Every media span that survives the splice must be relocated in the final sequence so its multimodal features stay aligned with the placeholder tokens the engine actually sees.

When splice is given, ranges at or after splice.template_cut_start are shifted arithmetically (the template suffix is copied verbatim, so their new offset is exact) and ranges entirely before the cut are located by search inside the exact prefix final_token_ids[:model_cut_end]. A range straddling the cut cannot be represented and fails closed. Without a splice, every range is located by search in global prompt order.

Ranges are matched in global prompt order because different media items can accumulate different shifts. Coincident ranges resolve to the same offset: Qwen2.5-Omni derives an audio range from its paired video range with an identical (offset, length), so a span is only consumed once.

A span that cannot be located fails closed rather than submitting token IDs with incorrect multimodal positions.

Parameters:

template_token_ids – The chat-template token sequence vLLM used to compute mm_placeholders.
final_token_ids – The exact-token prompt produced by splice_prefix_tokens, which will be submitted to the engine.
mm_placeholders – vLLM’s per-modality placeholder ranges, in template_token_ids coordinates.
splice – The splice that produced final_token_ids; enables the exact arithmetic remap for the template suffix.

Returns:

A new per-modality mapping with the same item ordering, whose ranges are expressed in final_token_ids coordinates.

Raises:

ValueError – If an input range is out of bounds for template_token_ids, straddles the splice boundary, or cannot be relocated in final_token_ids.
TypeError – If a range is neither a mapping nor a dataclass instance.

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`#