code/verl/workers/rollout/vllm_rollout_spmd.py

# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from contextlib import contextmanager
from typing import Any, Optional, Union

import numpy as np
import torch
import torch.distributed
from tensordict import TensorDict
from transformers import PreTrainedTokenizer, ProcessorMixin
from transformers.video_utils import VideoMetadata
from vllm import LLM, RequestOutput, SamplingParams

from ...protocol import DataProto
from ...utils import torch_functional as VF
from ...utils.dataset import process_image
from ...utils.multimodal_contract import load_video_tensors_and_metadata
from ...utils.torch_dtypes import PrecisionType
from .base import BaseRollout
from .config import RolloutConfig


def _repeat_interleave(value: Union[torch.Tensor, np.ndarray, list], repeats: int) -> Union[torch.Tensor, np.ndarray, list]:
    # repeat the elements, supports tensor, numpy array and list
    if isinstance(value, torch.Tensor):
        return value.repeat_interleave(repeats, dim=0)
    elif isinstance(value, np.ndarray):
        return np.repeat(value, repeats, axis=0)
    elif isinstance(value, list):
        out = []
        for v in value:
            out.extend([v] * repeats)
        return out
    else:
        return np.repeat(value, repeats, axis=0)


def _get_logit_bias(processor: Optional[ProcessorMixin]) -> Optional[dict[int, float]]:
    # enforce vllm to not output vision special tokens (image/video placeholders)
    if processor is None:
        return None

    logit_bias = {}
    if hasattr(processor, "image_token"):
        image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token)
        logit_bias[image_token_id] = -100
    if hasattr(processor, "video_token"):
        video_token_id = processor.tokenizer.convert_tokens_to_ids(processor.video_token)
        logit_bias[video_token_id] = -100

    return logit_bias if logit_bias else None


def _process_multi_modal_data(
    multi_modal_data: dict[str, Any],
    image_min_pixels: int,
    image_max_pixels: int,
    video_min_pixels: int,
    video_max_pixels: int,
    video_max_frames: int,
    video_fps: float,
    video_total_pixels: Optional[int],
) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]:
    images, videos = [], []
    mm_kwargs = None

    if "images" in multi_modal_data:
        for image in multi_modal_data["images"]:
            images.append(process_image(image, image_min_pixels, image_max_pixels))

    video_tensors, video_metadatas = load_video_tensors_and_metadata(
        multi_modal_data,
        video_min_pixels=video_min_pixels,
        video_max_pixels=video_max_pixels,
        video_max_frames=video_max_frames,
        video_fps=video_fps,
        video_total_pixels=video_total_pixels,
    )
    if video_tensors:
        if video_metadatas is None or len(video_metadatas) != len(video_tensors):
            metadata_count = 0 if video_metadatas is None else len(video_metadatas)
            raise ValueError(
                "Resolved video data is missing valid metadata entries. "
                f"Got {len(video_tensors)} video tensors and {metadata_count} metadata entries."
            )
        for tensor, metadata in zip(video_tensors, video_metadatas):
            videos.append((tensor, metadata))
        mm_kwargs = {"do_sample_frames": False, "do_resize": False}

    if len(images) != 0:
        return {"image": images}, None

    if len(videos) != 0:
        return {"video": videos}, mm_kwargs

    return None, None


class vLLMRollout(BaseRollout):
    def __init__(
        self,
        model_path: str,
        config: RolloutConfig,
        tokenizer: PreTrainedTokenizer,
        processor: Optional[ProcessorMixin],
    ):
        """A vLLM rollout. It requires the module is supported by the vllm.

        Args:
            module: module here follows huggingface APIs
            config: DictConfig
            tokenizer: the task/model tokenizer
        """
        super().__init__()
        self.rank = int(os.getenv("RANK", "0"))
        self.config = config
        self.tokenizer = tokenizer
        self.pad_token_id = tokenizer.pad_token_id
        self.use_tqdm = (self.rank == 0) and (not config.disable_tqdm)

        # Bypass mode: have vLLM emit the logprob of every sampled token so the
        # trainer can reuse it as old_log_probs, skipping one FSDP recompute and
        # keeping the first update's ratio away from a constant 1.
        self.calculate_log_probs = bool(getattr(config, "calculate_log_probs", False))
        if self.calculate_log_probs and self.rank == 0:
            print(
                "[rollout] calculate_log_probs=True: vLLM rollout logprobs will be reused as old_log_probs "
                "(bypass mode, equivalent to verl's algorithm.rollout_correction.bypass_mode=True)."
            )

        # Filter-only path: derive a per-sequence mean logprob from the vLLM
        # token logprobs. It never becomes old_log_probs, so it stays compatible
        # with oracle rows.
        self.collect_seq_logprob_for_filter = bool(
            getattr(config, "collect_seq_logprob_for_filter", False)
        )
        if self.collect_seq_logprob_for_filter and self.rank == 0:
            print(
                "[rollout] collect_seq_logprob_for_filter=True: per-seq mean logprob will be "
                "written into non_tensor_batch['seq_logprob_for_filter']."
            )
        if config.tensor_parallel_size > torch.distributed.get_world_size():
            raise ValueError("Tensor parallelism size should be less than world size.")

        if config.max_num_batched_tokens < config.prompt_length + config.response_length:
            raise ValueError("max_num_batched_tokens should be greater than prompt_length + response_length.")

        engine_kwargs = {}
        if processor is not None:  # only VLMs have processor
            if config.limit_images:
                engine_kwargs["limit_mm_per_prompt"] = {"image": config.limit_images, "video": 1}

        kv_cache_dtype = str(getattr(config, "kv_cache_dtype", "auto") or "auto").lower()
        if kv_cache_dtype not in {"auto", "fp8", "fp8_e5m2", "fp8_e4m3"}:
            raise ValueError(
                f"rollout.kv_cache_dtype must be one of 'auto'/'fp8'/'fp8_e5m2'/'fp8_e4m3', "
                f"got {kv_cache_dtype!r}."
            )
        if kv_cache_dtype != "auto" and self.rank == 0:
            print(
                f"[rollout] kv_cache_dtype={kv_cache_dtype}: KV cache stored in 1 byte/elt "
                f"(vs bf16 2 bytes), expecting ~2x decode speed on Hopper."
            )

        self.inference_engine = LLM(
            model=model_path,
            skip_tokenizer_init=False,
            trust_remote_code=config.trust_remote_code,
            load_format="dummy",
            dtype=PrecisionType.to_str(PrecisionType.to_dtype(config.dtype)),
            seed=config.seed,
            max_model_len=config.max_model_len or config.prompt_length + config.response_length,
            distributed_executor_backend="external_launcher",
            tensor_parallel_size=config.tensor_parallel_size,
            gpu_memory_utilization=config.gpu_memory_utilization,
            max_num_batched_tokens=config.max_num_batched_tokens,
            disable_log_stats=config.disable_log_stats,
            enforce_eager=config.enforce_eager,
            disable_custom_all_reduce=True,
            enable_chunked_prefill=config.enable_chunked_prefill,
            enable_sleep_mode=True,
            kv_cache_dtype=kv_cache_dtype,
            # Disable the vLLM v1 multimodal IPC preprocessing cache, which keeps
            # one LRU on the P0 sender and another on the P1 receiver. Repeated
            # sleep(level=1)/wake_up cycles drift the two eviction orders apart,
            # so P0 sends only None believing the receiver still holds mm_hash
            # while P1 has already evicted it, raising
            # "AssertionError: Expected a cached item for mm_hash=...", most
            # often around _validate() / save_checkpoint. RL rollout rarely
            # repeats a video, so disabling this cache costs no throughput.
            mm_processor_cache_gb=0,
            **engine_kwargs,
        )

        # Offload vllm model to reduce peak memory usage
        self.inference_engine.sleep(level=1)

        sampling_kwargs = {
            "max_tokens": config.response_length,
            "detokenize": False,
            "logit_bias": _get_logit_bias(processor),
        }
        default_sampling_params = SamplingParams()
        for key in config.to_dict().keys():
            # ``config.seed`` initializes the vLLM engine RNG above. Passing the
            # same value again as SamplingParams.seed resets every request to an
            # identical per-request RNG stream, unlike Transformers generation
            # where each device's RNG advances continuously.
            if key == "seed":
                continue
            if hasattr(default_sampling_params, key):
                sampling_kwargs[key] = getattr(config, key)

        # logprobs=1 returns the top-1 logprob plus the chosen token's logprob
        # even when the chosen token is not the top-1. Do not use logprobs=0:
        # vLLM v1 engines treat it inconsistently across versions and some return
        # output.logprobs=None, which would silently fall back to 0.0 and degrade
        # training into probability-weighted asymmetric REINFORCE (visible as
        # ppo_kl close to entropy) instead of a PPO bypass.
        if self.calculate_log_probs or self.collect_seq_logprob_for_filter:
            sampling_kwargs["logprobs"] = 1

        print(f"Sampling params: {sampling_kwargs}.")
        self.sampling_params = SamplingParams(**sampling_kwargs)

    @contextmanager
    def update_sampling_params(self, **kwargs):
        # update sampling params
        old_sampling_params = self.sampling_params
        if kwargs:
            # vLLM 0.17+: SamplingParams is a msgspec.Struct with read-only properties
            # Get the valid constructor fields
            struct_fields = getattr(self.sampling_params, "__struct_fields__", None)
            if struct_fields is not None:
                # msgspec.Struct: rebuild with current values + overrides
                current_args = {}
                for field in struct_fields:
                    if not field.startswith("_"):
                        current_args[field] = getattr(self.sampling_params, field)
                for key, value in kwargs.items():
                    if key in current_args:
                        current_args[key] = value
                self.sampling_params = SamplingParams(**current_args)
            else:
                # Older vLLM: use setattr directly
                for key, value in kwargs.items():
                    if hasattr(self.sampling_params, key):
                        setattr(self.sampling_params, key, value)

        yield
        # roll back to previous sampling params
        self.sampling_params = old_sampling_params

    @torch.no_grad()
    def generate_sequences(self, prompts: DataProto) -> DataProto:
        # left-padded attention_mask
        input_ids: torch.Tensor = prompts.batch["input_ids"]  # (bs, prompt_length)
        attention_mask: torch.Tensor = prompts.batch["attention_mask"]
        position_ids: torch.Tensor = prompts.batch["position_ids"]
        eos_token_id: int = prompts.meta_info["eos_token_id"]
        batch_size = input_ids.size(0)

        non_tensor_batch = prompts.non_tensor_batch
        batch_raw_prompt_ids = non_tensor_batch.pop("raw_prompt_ids")
        batch_multi_modal_data = non_tensor_batch.pop("multi_modal_data", None)

        if batch_size != len(batch_raw_prompt_ids):
            raise RuntimeError("vllm sharding manager is not work properly.")

        if batch_multi_modal_data is not None:
            vllm_inputs = []
            for raw_prompt_ids, multi_modal_data in zip(batch_raw_prompt_ids, batch_multi_modal_data):
                item = {"prompt_token_ids": list(raw_prompt_ids)}
                mm_data, mm_kwargs = _process_multi_modal_data(
                    multi_modal_data,
                    prompts.meta_info["image_min_pixels"],
                    prompts.meta_info["image_max_pixels"],
                    prompts.meta_info["video_min_pixels"],
                    prompts.meta_info["video_max_pixels"],
                    prompts.meta_info["video_max_frames"],
                    prompts.meta_info["video_fps"],
                    prompts.meta_info.get("video_total_pixels"),
                )
                if mm_data is not None:
                    if "video" in mm_data:
                        videos = []
                        for tensor, metadata in mm_data["video"]:
                            if isinstance(metadata, dict):
                                metadata_obj = VideoMetadata(
                                    total_num_frames=metadata.get("total_num_frames", tensor.shape[0] if hasattr(tensor, "shape") else len(tensor)),
                                    fps=metadata.get("fps"),
                                    frames_indices=metadata.get("frames_indices"),
                                    video_backend=metadata.get("video_backend"),
                                    width=metadata.get("width"),
                                    height=metadata.get("height"),
                                    duration=metadata.get("duration"),
                                )
                            else:
                                metadata_obj = metadata
                            videos.append((tensor, metadata_obj))
                        item["multi_modal_data"] = {"video": videos}
                    else:
                        item["multi_modal_data"] = mm_data
                    if mm_kwargs is not None:
                        item["mm_processor_kwargs"] = mm_kwargs
                vllm_inputs.append(item)
        else:
            vllm_inputs = [{"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in batch_raw_prompt_ids]

        # users can customize different sampling_params at different run
        with self.update_sampling_params(**prompts.meta_info):
            completions: list[RequestOutput] = self.inference_engine.generate(
                prompts=vllm_inputs, sampling_params=self.sampling_params, use_tqdm=self.use_tqdm
            )
            response_ids = [output.token_ids for completion in completions for output in completion.outputs]
            response_ids = VF.pad_2d_list_to_length(
                response_ids, self.pad_token_id, max_length=self.config.response_length
            ).to(input_ids.device)

            # vLLM only attaches logprobs when SamplingParams.logprobs is set.
            # Validation goes through update_sampling_params, which leaves that
            # field alone, so training runs always see them and other runs skip
            # this whole block.
            collect_log_probs = (
                (self.calculate_log_probs or self.collect_seq_logprob_for_filter)
                and self.sampling_params.logprobs is not None
            )
            if collect_log_probs:
                response_log_probs: list[list[float]] = []
                missing_seqs = 0
                missing_chosen = 0
                total_tokens = 0
                for completion in completions:
                    for out in completion.outputs:
                        token_ids_seq = out.token_ids
                        lp_dicts = out.logprobs  # list[dict[token_id, Logprob]]
                        if lp_dicts is None or len(lp_dicts) == 0:
                            missing_seqs += 1
                            response_log_probs.append([])
                            continue
                        seq_lps: list[float] = []
                        for tok_id, lp_dict in zip(token_ids_seq, lp_dicts):
                            total_tokens += 1
                            entry = lp_dict.get(tok_id) if lp_dict else None
                            if entry is None:
                                missing_chosen += 1
                                if lp_dict:
                                    entry = next(iter(lp_dict.values()))
                            seq_lps.append(float(entry.logprob) if entry is not None else 0.0)
                        response_log_probs.append(seq_lps)

                # Fail loudly when vLLM returned no logprobs, usually because a
                # vLLM v1 engine silently ignored SamplingParams.logprobs. A
                # silent fallback to old_log_probs == 0 would degrade training
                # into probability-weighted asymmetric REINFORCE, visible only as
                # actor/ppo_kl close to actor/entropy.
                if missing_seqs > 0:
                    raise RuntimeError(
                        f"[rollout] calculate_log_probs=True but vLLM returned no logprobs for "
                        f"{missing_seqs} sequence(s). Check that "
                        f"SamplingParams.logprobs={self.sampling_params.logprobs} is supported by "
                        f"this vLLM version (use logprobs>=1, never 0)."
                    )
                if total_tokens > 0 and missing_chosen / total_tokens > 0.01:
                    raise RuntimeError(
                        f"[rollout] {missing_chosen}/{total_tokens} chosen tokens are absent from "
                        f"the vLLM logprob dicts (>1%); the values are not trustworthy."
                    )

                # Pad to response_length with 0.0; response_mask hides the
                # padding, so it never reaches the ratio.
                response_log_probs = VF.pad_2d_list_to_length(
                    response_log_probs, 0.0, max_length=self.config.response_length
                ).to(input_ids.device, dtype=torch.float32)

                # One-shot self check: print the mean/std of the vLLM logprobs on
                # the first rollout so the log shows real values rather than
                # silent zeros. A working bypass reports mean between -2.0 and
                # -0.5 and std above 0.5, depending on temperature and entropy.
                if self.rank == 0 and not getattr(self, "_logprob_self_check_done", False):
                    flat = response_log_probs.flatten()
                    nonzero = flat[flat != 0.0]
                    print(
                        f"[rollout][bypass self-check] vLLM rollout_log_probs "
                        f"shape={tuple(response_log_probs.shape)} "
                        f"all_mean={flat.mean().item():.4f} all_std={flat.std().item():.4f} "
                        f"nonzero_mean={nonzero.mean().item() if nonzero.numel() > 0 else float('nan'):.4f} "
                        f"nonzero_std={nonzero.std().item() if nonzero.numel() > 1 else float('nan'):.4f} "
                        f"nonzero_min={nonzero.min().item() if nonzero.numel() > 0 else float('nan'):.4f} "
                        f"nonzero_max={nonzero.max().item() if nonzero.numel() > 0 else float('nan'):.4f} "
                        f"({nonzero.numel()}/{flat.numel()} non-zero)"
                    )
                    self._logprob_self_check_done = True

            if self.sampling_params.n > 1:
                batch_size = batch_size * self.sampling_params.n
                input_ids = _repeat_interleave(input_ids, self.sampling_params.n)
                attention_mask = _repeat_interleave(attention_mask, self.sampling_params.n)
                position_ids = _repeat_interleave(position_ids, self.sampling_params.n)
                if batch_multi_modal_data is not None:
                    batch_multi_modal_data = _repeat_interleave(batch_multi_modal_data, self.sampling_params.n)

        sequence_ids = torch.cat([input_ids, response_ids], dim=-1)
        response_length = response_ids.size(1)
        delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device)
        delta_position_id = delta_position_id.view(1, -1).expand(batch_size, -1)
        if position_ids.ndim == 3:  # qwen2vl mrope: (batch_size, 4, seq_length)
            delta_position_id = delta_position_id.view(batch_size, 1, -1).expand(batch_size, position_ids.size(1), -1)

        # prompt: left pad + response: right pad
        # attention_mask: [0,0,0,0,1,1,1,1 | 1,1,1,0,0,0,0,0]
        # position_ids:   [0,0,0,0,0,1,2,3 | 4,5,6,7,8,9,10,11]
        response_position_ids = position_ids[..., -1:] + delta_position_id
        position_ids = torch.cat([position_ids, response_position_ids], dim=-1)
        response_mask = VF.get_response_mask(
            response_ids=response_ids, eos_token_id=eos_token_id, dtype=attention_mask.dtype
        )
        attention_mask = torch.cat((attention_mask, response_mask), dim=-1)

        # all the tp ranks should contain the same data here. data in all ranks are valid
        batch_tensors = {
            "prompts": input_ids,
            "responses": response_ids,
            "input_ids": sequence_ids,  # here input_ids become the whole sentences
            "attention_mask": attention_mask,
            "response_mask": response_mask,
            "position_ids": position_ids,
        }
        # Only bypass mode puts token logprobs into batch.rollout_log_probs. The
        # filter-only path needs the per-sequence mean alone, which avoids both
        # the oracle-row token mismatch and one (B, T) tensor of dispatch traffic.
        if collect_log_probs and self.calculate_log_probs:
            batch_tensors["rollout_log_probs"] = response_log_probs
        batch = TensorDict(batch_tensors, batch_size=batch_size)

        non_tensor_batch: dict = {}
        if batch_multi_modal_data is not None:
            non_tensor_batch["multi_modal_data"] = batch_multi_modal_data

        # Per-sequence mean logprob, never used in the gradient path. Masking
        # with response_mask is more robust than treating != 0 as padding.
        if collect_log_probs and self.collect_seq_logprob_for_filter:
            valid_lengths = response_mask.sum(dim=-1).clamp_min(1).to(response_log_probs.dtype)
            seq_logprob_sum = (response_log_probs * response_mask.to(response_log_probs.dtype)).sum(dim=-1)
            seq_logprob_mean = (seq_logprob_sum / valid_lengths).detach().cpu().numpy().astype(np.float32)
            non_tensor_batch["seq_logprob_for_filter"] = seq_logprob_mean

        return DataProto(batch=batch, non_tensor_batch=non_tensor_batch, meta_info=prompts.meta_info)