inference.py

"""
InferenceActor - 使用 vLLM 进行高并发推理采样（AWQ 量化）

该 Actor 运行在 GPU 1 上，负责：
1. 使用 vLLM 引擎进行高效的批量推理（AWQ Int4 量化）
2. 为 GRPO 算法生成 Experience Rollouts（每组 prompt 生成多个候选）
3. 支持动态加载 LoRA 权重（从 TrainingActor 保存的权重中加载）

架构说明：
- Ray 会自动为每个 Actor 分配独立的 GPU（通过 CUDA_VISIBLE_DEVICES）
- 在 Actor 内部，我们统一使用 "cuda:0"，因为 Ray 已经做了设备隔离
- 使用 AWQ (Activation-aware Weight Quantization) Int4 量化：
    * 显存占用减少约 75%（相比 FP16）
    * 推理速度提升 2-3 倍（利用 vLLM 的 AWQ 内核）
    * 精度损失极小（<1%）
- tensor_parallel_size=1: 强制只使用一张 GPU，避免抢占 GPU 0
- gpu_memory_utilization=0.8 预留 20% 显存，避免 OOM 并为 LoRA 权重加载预留空间
- max_model_len=4096 限制上下文长度，适配 RTX 4090 24GB 显存限制
"""

from typing import List, Dict, Optional
import ray
from vllm import LLM, SamplingParams


@ray.remote(num_gpus=1)
class InferenceActor:
    """
    推理 Actor，使用 vLLM 进行高并发采样
    
    该 Actor 被 Ray 调度到独立的 GPU 上（通常是 GPU 0），
    负责为训练提供大量的推理样本。
    """
    
    def __init__(
        self,
        model_path: str,
        lora_path: Optional[str] = None,
    ):
        """
        初始化 vLLM 推理引擎
        
        Args:
            model_path: HuggingFace 模型路径或本地路径（Qwen3-8B-AWQ）
            lora_path: LoRA 权重路径（可选），如果提供则加载 LoRA 适配器
        
        注意：
        - vLLM 会自动从模型路径加载 tokenizer
        - Qwen3-8B-AWQ 和 Qwen3-8B-Base 共享相同的 tokenizer
        - 这确保了推理和训练时的 tokenization 一致性
        """
        # vLLM 配置参数（AWQ 量化）
        # quantization="awq": 使用 AWQ (Activation-aware Weight Quantization) Int4 量化
        #   AWQ 是一种感知激活的权重量化方法，相比 GPTQ 等量化方法，精度更高
        #   显存占用减少约 75%，推理速度提升 2-3 倍
        #
        # dtype="float16": AWQ 量化模型通常配合 FP16 使用
        #   虽然 AWQ 权重本身是 Int4，但激活值和计算仍使用 FP16
        #
        # tensor_parallel_size=1: 强制只使用一张 GPU，确保不会抢占 GPU 0
        #   这对于双卡环境至关重要，确保 GPU 0 专门用于训练，GPU 1 专门用于推理
        #
        # gpu_memory_utilization=0.8: 只使用 80% 的 GPU 显存
        # 原因：
        #   1. 避免显存碎片化导致的 OOM
        #   2. 为动态加载 LoRA 权重预留空间（LoRA 权重需要额外的显存）
        #   3. 为 KV Cache 的扩展预留缓冲空间
        #
        # max_model_len=4096: 限制最大上下文长度
        #   适配 RTX 4090 24GB 显存限制，防止 OOM
        #   4096 tokens 对于大多数任务已经足够
        #
        # enable_lora=True: 启用 LoRA 支持，允许在运行时动态加载/切换 LoRA 适配器
        # 这对于 RLHF/GRPO 训练至关重要，因为我们需要在推理时使用最新的训练权重
        #
        # trust_remote_code=True: 允许执行模型仓库中的自定义代码（Qwen3 需要）
        # 
        # Tokenizer 一致性说明：
        # - vLLM 会自动从 model_path 加载 tokenizer
        # - Qwen3-8B-AWQ 和 Qwen3-8B-Base 使用相同的 tokenizer
        # - 这确保了推理（InferenceActor）和训练（TrainingActor）的 tokenization 一致
        llm_kwargs = {
            "model": model_path,  # Qwen3-8B-AWQ 模型路径
            "quantization": "awq",  # 使用 AWQ Int4 量化
            "dtype": "float16",  # AWQ 配合 FP16 使用
            "tensor_parallel_size": 1,  # 强制只使用一张 GPU，避免抢占 GPU 0
            "gpu_memory_utilization": 0.8,
            "max_model_len": 4096,  # 限制上下文长度，适配 24GB 显存
            "enable_lora": True,
            "trust_remote_code": True,  # Qwen3 需要信任远程代码
        }
        
        # 如果提供了 LoRA 路径，则在初始化时加载
        if lora_path is not None:
            llm_kwargs["enable_lora"] = True
            # 注意：vLLM 的 LoRA 加载方式可能因版本而异
            # 这里假设使用 LoRA 适配器路径
            # 实际使用时可能需要根据 vLLM 版本调整
        
        self.llm = LLM(**llm_kwargs)
        
        # 如果初始化时没有提供 LoRA，后续可以通过 update_lora 方法动态加载
        self.current_lora_path = lora_path
    
    def generate_rollouts(
        self,
        prompts: List[str],
        n: int = 8,
        temperature: float = 0.8,
        max_tokens: int = 1024,
        lora_path: Optional[str] = None,
    ) -> List[Dict[str, str]]:
        """
        为给定的 prompts 生成多个候选回复（Rollouts）
        
        这是 GRPO 算法的核心：对于每个 prompt，我们需要生成多个候选回复，
        然后由训练 Actor 计算这些候选的奖励并更新策略。
        
        Args:
            prompts: 输入提示列表
            n: 每个 prompt 生成的候选数量（GRPO Group Size）
               在 GRPO 中，通常设置为 8，表示每组生成 8 个候选
            temperature: 采样温度，控制生成的随机性
               0.8 是一个平衡值，既保证多样性又不过于随机
            max_tokens: 最大生成 token 数
            lora_path: 可选的 LoRA 权重路径。如果提供且与当前路径不同，
               则会在生成前更新 LoRA 权重
        
        Returns:
            List[Dict[str, str]]: 每个元素包含 'prompt' 和 'generated_text'
                格式示例: [{"prompt": "...", "generated_text": "..."}, ...]
        """
        # 如果提供了新的 LoRA 路径且与当前路径不同，则更新 LoRA 权重
        if lora_path is not None and lora_path != self.current_lora_path:
            self.update_lora(lora_path)
        # 配置采样参数
        # n=8: GRPO Group Size，每个 prompt 生成 8 个不同的候选
        # 这些候选将用于后续的奖励计算和策略更新
        sampling_params = SamplingParams(
            n=n,  # 每个 prompt 生成 n 个候选
            temperature=temperature,
            max_tokens=max_tokens,
            # 不使用 top_p/top_k，让 temperature 控制多样性
        )
        
        # 执行批量推理
        # vLLM 会自动进行批处理优化，提高吞吐量
        outputs = self.llm.generate(prompts, sampling_params)
        
        # 将 vLLM 的输出格式转换为标准格式
        results = []
        for output in outputs:
            prompt = output.prompt
            # output.outputs 是一个列表，包含 n 个生成的候选
            for generated_output in output.outputs:
                results.append({
                    "prompt": prompt,
                    "generated_text": generated_output.text,
                })
        
        return results
    
    def update_lora(self, lora_path: str) -> bool:
        """
        动态更新 LoRA 权重
        
        在训练过程中，TrainingActor 会定期保存更新后的 LoRA 权重。
        该方法允许 InferenceActor 加载最新的权重，实现推理和训练的同步。
        
        Args:
            lora_path: 新的 LoRA 权重路径
        
        Returns:
            bool: 是否成功更新
        """
        # TODO: 根据 vLLM 版本实现 LoRA 权重更新
        # vLLM 的 LoRA 更新 API 可能因版本而异
        # 通常需要调用类似 self.llm.add_lora() 或 self.llm.update_lora() 的方法
        try:
            # 示例代码（需要根据实际 vLLM API 调整）:
            # self.llm.add_lora(lora_path)
            self.current_lora_path = lora_path
            return True
        except Exception as e:
            print(f"Failed to update LoRA weights: {e}")
            return False