[Bugfix][Rollout] Reject missing or misaligned sampled-token log probabilities

[0z5a]

Problem

The non-streaming vLLM rollout parser can replace missing sampled-token log probabilities with zero. Two token IDs paired with only one probability can therefore become apparently valid training data.

Changes

• Validate token IDs, aligned log-probability entries and finite numeric values before appending generated data or calculating rewards.
• Preserve real zero probabilities, finite vLLM sentinels, empty terminal responses and non-trainable tool tokens.
• Preserve multimodal canonical token prefixes and prepared training inputs on the existing early-return paths.
• Include a locally generated request ID in errors without including prompts or token contents.

The patch contains the rollout parser and its tests, based on upstream ce92eff12ecdc81396bf41a2f94e62dd5b0aca32.

Validation

Check	Result
Fresh publication checkout: rollout tests	99 passed
Full repository pre-commit hooks	All 9 passed
Earlier native rollout and Sample tests, local and Linux	111 passed
Earlier Qwen3-0.6B integration	2 complete RL rounds; 16 samples verified the submitted parser hash
Checkpoint-to-serving weight reload	Both saved checkpoints matched all 226 served parameter tensors byte-for-byte

The integration used real vLLM generation, response-dependent reward, full-model backpropagation, Megatron distributed AdamW, checkpoint persistence and weight reload.

Scope and remaining validation

This remains a draft. The small-model integration required separately recorded local-attention/runtime compatibility changes outside this patch. The 111-test and GPU results are archived validation; the fresh publication run covered the 99 rollout tests and repository hooks.

This is a data-validation correctness fix. No performance or convergence improvement is claimed.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34526977 | 📁 Comparing 327eb5b against latest (ce92eff)

  🔍 Preview build  

45 files changed · ± 45 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces strict validation for vLLM generation metadata (such as token IDs and log probabilities) to prevent silent fallbacks, ensures prompt preparation changes remain local until validation succeeds, and adds unique request IDs to headers and error messages. It also significantly expands the test suite. The reviewer suggests simplifying the validation logic in _inference_generate_tokens_and_logprobs by checking for empty token_ids early to avoid redundant checks and improve clarity.

[gemini-code-assist]

The validation logic for empty token_ids can be significantly simplified. By checking if not token_ids: early, we can immediately return [], [] and avoid redundant checks for logprobs and content being None when token_ids is empty. This also prevents potentially misleading error messages (e.g., raising 'logprobs must be an object for nonempty trainable token_ids' when token_ids is actually empty but logprobs is of an invalid type).

    if not isinstance(choice, dict):
        raise ValueError("choice must be an object")
    token_ids = choice.get("token_ids")
    if not isinstance(token_ids, list) or any(type(token_id) is not int or token_id < 0 for token_id in token_ids):
        raise ValueError("token_ids must be a list of non-negative integers")

    if not token_ids:
        return [], []

    logprobs = choice.get("logprobs")
    if not isinstance(logprobs, dict):
        raise ValueError("logprobs must be an object for nonempty trainable token_ids")
    content = logprobs.get("content")
    if not isinstance(content, list):
        raise ValueError("logprobs.content must be a list")
    if len(content) != len(token_ids):
        raise ValueError(f"token/logprob length mismatch: {len(token_ids)} tokens, {len(content)} entries")

    log_probs = []
    for index, entry in enumerate(content):
        if not isinstance(entry, dict) or "logprob" not in entry:
            raise ValueError(f"missing logprob at token index {index}")
        value = entry["logprob"]
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            raise ValueError(f"non-numeric logprob at token index {index}")
        try:
            value = float(value)
        except OverflowError as exc:
            raise ValueError(f"non-finite logprob at token index {index}") from exc
        if not math.isfinite(value):
            raise ValueError(f"non-finite logprob at token index {index}")
        log_probs.append(value)
    return token_ids, log_probs