异步rl：pause/resume

[floatlibai]

异步rl：pause/resume

vLLM文档：https://docs.vllm.ai/en/stable/training/async_rl/

Overview

在标准 RL 训练循环中，生成（generation）和训练（training）是顺序执行的：策略模型先生成 rollout 样本，然后在这些样本上进行训练，如此往复循环。在生成阶段，训练加速器处于空闲状态，反之亦然。

one-off pipelining 方法将生成阶段和训练阶段拆分为两个并行的协程，使模型能够一边生成新样本，一边对先前生成的数据进行训练。这可以带来更高的 GPU 利用率和更大的训练吞吐量。

然而，这种重叠引入了一个复杂性：推理引擎必须在运行过程中更新权重，而此时可能仍有正在处理的请求。

Pause 与 Resume API

为了在推理引擎运行时安全地更新权重，vLLM 提供了 pause_generation 和 resume_generation 方法。这让训练器能够协调出一个干净的窗口来进行权重同步，而不会丢失正在进行的任务。

pause_generation

await engine.pause_generation(mode="keep", clear_cache=True)

mode 参数控制如何处理正在进行的请求：

Mode	Behavior
"abort"	立即中止所有正在进行的请求并返回部分结果（默认）
"wait"	等待所有正在进行的请求完成后再暂停
"keep"	冻结队列中的请求；调用 resume_generation 后它们会恢复执行

clear_cache 参数控制暂停后是否清空 KV 缓存和前缀缓存。

resume_generation

await engine.resume_generation()

暂停后恢复调度器。任何通过 mode="keep" 冻结的请求将继续生成。

HTTP 端点

使用 vLLM HTTP 服务器时，可以通过以下端点使用相同的功能：

• POST /pause?mode=keep - 暂停生成
• POST /resume - 恢复生成

注意：数据并行

当使用 vLLM 的内部负载均衡器（即 data_parallel_backend="ray"）进行数据并行时，pause 和 resume 会在所有 DP rank 上自动处理——一次调用即可。当使用外部负载均衡器（即代理后的多个独立 vLLM 实例）时，必须在权重更新前后分别向每个引擎实例发送 pause 和 resume 请求。

典型异步 RL 流程

典型的带权重同步的异步 RL 循环如下：

• 从当前策略开始生成 rollout 样本
• 当训练器有新权重需要更新时，使用 mode="keep" 暂停生成
• 将更新后的权重从训练器同步到推理引擎（参见 [Weight Transfer](https://docs.vllm.ai/en/stable/training/weight_transfer/)）
• 恢复生成——正在进行的请求将使用新权重继续执行
• 重复此过程

关键洞察在于：使用 mode="keep" 暂停的请求，在暂停前产生的 token 来自旧权重，而在恢复后产生的 token 则来自新权重。clear_cache 参数控制在暂停期间是否使 KV 缓存失效：

• 当 clear_cache=True 时，之前缓存的 key-value 条目将被丢弃，因此恢复后生成的所有 token 都将完全使用新权重计算。
• 当 clear_cache=False 时，现有的 KV 缓存条目将被保留，这意味着上下文中的某些 token 可能仍然反映旧权重的状态（即过期的 KV 缓存）。

Partial rollout 就是clear_cache=False吗？

---

源码实现分析

调用链路

HTTP API (api_router.py)
  → AsyncLLM (async_llm.py)
    → EngineCoreClient (core_client.py)  [跨进程 ZMQ 通信]
      → EngineCoreProc/DPEngineCoreProc (core.py)
        → Scheduler.set_pause_state() (scheduler.py)
        → Scheduler.schedule() / _process_input_queue() / idle callbacks

核心数据结构

PauseState 枚举 (vllm/v1/core/sched/interface.py:22-32)

class PauseState(enum.IntEnum):
    UNPAUSED = 0     # 正常运行
    PAUSED_NEW = 1   # 不调度新请求，但已有 running 请求继续执行
    PAUSED_ALL = 2   # 不调度任何请求

PauseMode 类型 (vllm/v1/engine/__init__.py:26)

PauseMode = Literal["abort", "wait", "keep"]

---

pause_generation

pause_generation(mode="keep", clear_cache=True) 逐层实现

第一层：AsyncLLM (vllm/v1/engine/async_llm.py:750-793)

async def pause_generation(self, *, mode="abort", clear_cache=True):
    if clear_cache:
        await self.renderer.clear_mm_cache_async()  # ① 清理多模态缓存
    await self.engine_core.pause_scheduler_async(
        mode=mode, clear_cache=clear_cache
    )  # ② 跨进程 RPC
    await asyncio.sleep(0.02)  # ③ 等待 in-flight 输出的最终返回

三步详解:

1. 清理多模态缓存: 如果 clear_cache=True，清空 renderer 中缓存的 image/audio 等多模态编码特征
2. 跨进程 RPC: 通过 ZMQ 将 ("pause_scheduler", mode, clear_cache) 发送到 EngineCore 进程
3. 20ms 延迟: 确保 in-flight 请求的最终输出在 pause_generation() 返回前被 output_handler 取出并推送给调用方。这不是正确性要求，只是为了让调用方的视角更直观（看到所有输出后才返回）

第二层：EngineCoreClient (vllm/v1/engine/core_client.py:1130-1133)

async def pause_scheduler_async(self, mode, clear_cache):
    await self.call_utility_async("pause_scheduler", mode, clear_cache)

call_utility_async 将方法名和参数序列化为 EngineCoreRequestType.UTILITY 类型的消息，通过 ZMQ 发送到 EngineCore 进程。如果 engine core 返回了一个 Future，则等待该 Future 完成。

async def call_utility_async(self, method: str, *args) -> Any:
    return await self._call_utility_async(method, *args, engine=self.core_engine)

async def _call_utility_async(
    self, method: str, *args, engine: EngineIdentity
) -> Any:
    call_id = uuid.uuid1().int >> 64
    future = asyncio.get_running_loop().create_future()
    self.utility_results[call_id] = future
    message = (
        EngineCoreRequestType.UTILITY.value,
        *self.encoder.encode((self.client_index, call_id, method, args)),
    )
    await self._send_input_message(message, engine, args)
    self._ensure_output_queue_task()
    return await future

第三层：EngineCoreProc — 核心逻辑 (vllm/v1/engine/core.py:1665-1705)

def pause_scheduler(self, mode, clear_cache):
    if mode == "abort":
        # 立即终止所有请求，发送 ABORT 完成信号
        aborted_reqs = self.scheduler.finish_requests(
            None, RequestStatus.FINISHED_ABORTED
        )
        self._send_abort_outputs(aborted_reqs)

    # mode="keep" → PAUSED_ALL（冻结一切）
    # mode="abort"/"wait" → PAUSED_NEW（冻结新请求，running 继续）
    pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
    self.scheduler.set_pause_state(pause_state)

    if self._pause_complete():  # 引擎已空闲 → 同步返回
        if clear_cache:
            self._reset_caches()
        return None  # Future 已完成，call_utility_async 直接返回

    # 引擎还有工作在进行中 → 异步等待
    future = Future()
    self._idle_state_callbacks.append(
        partial(engine_idle_callback, future=future)
    )
    return future  # call_utility_async 会等待这个 Future

engine_idle_callback 的注册与触发

  run_busy_loop()
    │
    ├─ _process_input_queue()
    │    │
    │    ├─ has_work() → PAUSED_ALL 时返回 False
    │    │    │
    │    │    └─ 进入 while not has_work() 死循环
    │    │         │
    │    │         ├─ _notify_idle_state_callbacks()  ← engine_idle_callback 在这里被调用
    │    │         │    future.set_result(None)        ← pause 的 RPC 返回
    │    │         │
    │    │         └─ input_queue.get(block=True)      ← 阻塞等待新消息
    │    │              │
    │    │              └─ 收到 resume → _handle_client_request()
    │    │                   → resume_scheduler() → set_pause_state(UNPAUSED)
    │    │                   → has_work() 重新返回 True，退出内层循环
    │    │
    │    └─ 退出 _process_input_queue
    │
    └─ _process_engine_step()  ← 恢复正常的调度+推理

当引擎进入空闲循环 (_process_input_queue → has_work() == False) 时触发:

def _process_input_queue(self):
    while not self.has_work() and self.is_running():
        self._notify_idle_state_callbacks()  # ← 触发所有注册的回调
        # ... 等待新输入 ...

def _notify_idle_state_callbacks(self):
    while self._idle_state_callbacks:
        callback = self._idle_state_callbacks.pop()
        callback(self)  # 调用 engine_idle_callback

def engine_idle_callback(engine, future):
    if clear_cache:
        engine._reset_caches()  # 清空所有缓存
    future.set_result(None)     # Future 完成 → pause 调用返回

has_work() 的判断逻辑

def has_work(self):
    return (
        self.engines_running or
        self.scheduler.has_requests() or
        bool(self.batch_queue)
    )

当 PAUSED_ALL 时，scheduler 的 get_num_unfinished_requests() 返回 0，所以 has_requests() 返回 False。如果也没有 batch_queue 工作，has_work() 返回 False，引擎进入空闲循环，触发 idle callbacks。

第四层：Scheduler — 暂停状态如何影响调度

schedule() 方法中的 token_budget 控制 (vllm/v1/core/sched/scheduler.py:409-411):

if self._pause_state == PauseState.PAUSED_ALL:
    # 不调度任何请求
    token_budget = 0

get_num_unfinished_requests() 对暂停状态的响应 (scheduler.py:2136-2145):

def get_num_unfinished_requests(self):
    if self._pause_state == PauseState.PAUSED_ALL:
        return 0   # 对外宣称没有未完成的请求 → has_work() 返回 False
    if self._pause_state == PauseState.PAUSED_NEW:
        return len(self.running)  # 只算 running，不算 waiting
    return len(self.waiting) + len(self.skipped_waiting) - ... + len(self.running)

schedule() 中新请求的处理 (scheduler.py:629):

# 只有在 UNPAUSED 状态才会从 waiting 队列取新请求
if not preempted_reqs and self._pause_state == PauseState.UNPAUSED:
    step_skipped_waiting = create_request_queue(self.policy)
    while (self.waiting or self.skipped_waiting) and token_budget > 0:
        # ... 调度 waiting 中的请求 ...

这意味着在 PAUSED_NEW 或 PAUSED_ALL 状态下，新添加的请求会被放入 waiting 队列，但不会被调度执行。

_reset_caches() (core.py:706-719)

def _reset_caches(self):
    self.reset_prefix_cache(
        reset_running_requests=True, reset_connector=True
    )
    self.reset_mm_cache()
    self.reset_encoder_cache()

清空三类缓存:

• Prefix Cache: 前缀匹配缓存（最关键的 KV 缓存）
• MM Cache: 多模态特征缓存
• Encoder Cache: 编码器输出缓存

resume_generation

第一层：AsyncLLM (vllm/v1/engine/async_llm.py:795-797)

async def resume_generation(self):
    await self.engine_core.resume_scheduler_async()

第二层：EngineCoreClient (vllm/v1/engine/core_client.py:1135-1136)

async def resume_scheduler_async(self):
    await self.call_utility_async("resume_scheduler")

第三层：EngineCore — 同步返回 (vllm/v1/engine/core.py:753-755)

def resume_scheduler(self):
    self.scheduler.set_pause_state(PauseState.UNPAUSED)

只需将调度器状态切回 UNPAUSED。下一次 schedule() 调用时:

• token_budget 恢复正常计算
• waiting 队列中的请求重新被调度
• 之前被 mode="keep" 冻结的请求继续执行

DPEngineCoreProc — Data Parallel 场景 (vllm/v1/engine/core.py:1852-1876)

def resume_scheduler(self):
    # 安全检查：确保 pause 已完成
    if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave):
        raise RuntimeError(
            "resume_scheduler called while pause is still in flight..."
        )
    if self.engines_running:
        return  # 已经在运行

    super().resume_scheduler()      # 设为 UNPAUSED
    self.ignore_start_dp_wave = False

    # Barrier: 通过 all-reduce 等待所有 DP rank 都完成 resume
    has_global_unfinished = ParallelConfig.has_unfinished_dp(
        self.dp_group, self.scheduler.has_unfinished_requests()
    )
    if has_global_unfinished:
        self.engines_running = True  # 启动 step 循环

DP 场景的 resume 需要额外的同步:

1. 确保 pending_pause 已清除、ignore_start_dp_wave 已重置
2. 调用父类的 resume_scheduler() 将状态设为 UNPAUSED
3. 通过 all-reduce barrier 等待所有 DP rank 都完成 resume
4. 如果全局有未完成的请求，启动 step 循环

---

完整时序：mode="keep" 的 pause → resume 流程

  调用方                          AsyncLLM                   EngineCore                    Scheduler
    │                                │                           │                             │
    │  pause_generation(mode="keep")  │                           │                             │
    │───────────────────────────────>│                           │                             │
    │                                │  clear_mm_cache_async()   │                             │
    │                                │──────────────────────────>│                             │
    │                                │                           │                             │
    │                                │  pause_scheduler_async    │                             │
    │                                │  ("keep", clear_cache)    │                             │
    │                                │──────────────────────────>│                             │
    │                                │                           │  set_pause_state(PAUSED_ALL) │
    │                                │                           │────────────────────────────>│
    │                                │                           │                             │
    │                                │                           │  [schedule() 中 token_budget=0]│
    │                                │                           │  [get_num_unfinished 返回 0] │
    │                                │                           │                             │
    │                                │                           │  [_process_input_queue 空闲循环]│
    │                                │                           │──┐                          │
    │                                │                           │  │ _notify_idle_callbacks   │
    │                                │                           │<─┘                          │
    │                                │                           │  engine_idle_callback:      │
    │                                │                           │    _reset_caches()          │
    │                                │                           │    future.set_result(None)  │
    │                                │                           │                             │
    │                                │  <── RPC 返回 ───────────│                             │
    │                                │                           │                             │
    │                                │  await asyncio.sleep(0.02)│                             │
    │  <──── pause 完成 ─────────────│                           │                             │
    │                                │                           │                             │
    │  [权重更新期间]                  │                           │                             │
    │                                │                           │                             │
    │  resume_generation()           │                           │                             │
    │───────────────────────────────>│                           │                             │
    │                                │  resume_scheduler_async   │                             │
    │                                │──────────────────────────>│                             │
    │                                │                           │  set_pause_state(UNPAUSED)  │
    │                                │                           │────────────────────────────>│
    │                                │  <── RPC 返回 ────────────│                             │
    │  <──── resume 完成 ────────────│                           │                             │
    │                                │                           │                             │
    │  [下一个 generate() 调用]        │                           │  [schedule() 恢复正常]       │
    │                                │                           │  [waiting 中请求恢复执行]     │

---

mode 对比总结

方面	mode="abort"	mode="wait"	mode="keep"
调度器状态	PAUSED_NEW	PAUSED_NEW	PAUSED_ALL
正在运行的请求	立即终止，发送 ABORT	继续执行直到完成	冻结保留（不终止不执行）
新请求	排队但不调度	排队但不调度	排队但不调度
KV Cache	空闲后由 callback 清空	空闲后由 callback 清空	空闲后由 callback 清空
pause 返回时机	等 aborted 输出发送完毕	等所有 running 请求完成	等引擎进入空闲循环
resume 后	被 abort 的请求不可恢复	所有请求已完成	冻结的请求继续执行
适用场景	快速权重切换，丢弃旧结果	优雅暂停，保留所有结果	异步 RL 训练，保留请求上下文

[floatlibai]

RL侧调用示例

examples/rl/rlhf_async_new_apis.py

以下对示例源码进行逐段注释，展示在异步 RL 训练中如何通过 pause/resume + weight transfer 在推理运行期间切换权重。

"""
Demonstrates async reinforcement learning using vLLM and Ray,
with native weight syncing APIs and batch-invariant generation.

The script separates training and inference workloads onto distinct GPUs
so that Ray can manage process placement and inter-process communication.
A Hugging Face Transformer model occupies one GPU for training, and a
vLLM AsyncLLMEngine occupies another GPU for inference.

Batch invariance is enabled so that generation output is deterministic
regardless of how many requests are batched together.
"""
# 概述：Trainer 占一个 GPU，vLLM 引擎占另一个 GPU。
#       通过 NCCL 通道在 GPU 间直接传输权重。
#       开启 batch_invariant 保证输出确定性，用于后续验证。

import asyncio
import uuid
from dataclasses import asdict

import ray
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import vllm
from vllm import SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.base import (
    WeightTransferInitRequest,   # 初始化权重传输通道的请求
    WeightTransferUpdateRequest, # 触发一次权重广播的请求
)
from vllm.distributed.weight_transfer.nccl_engine import (
    NCCLTrainerSendWeightsArgs,
    NCCLWeightTransferEngine,
    NCCLWeightTransferInitInfo,    # NCCL 通道元信息 (地址/端口/rank)
    NCCLWeightTransferUpdateInfo,  # 要传输的权重元信息 (名称/dtype/shape)
)
from vllm.platforms import current_platform
from vllm.utils.network_utils import get_ip, get_open_port
from vllm.v1.executor import Executor

# ============================
# 全局配置
# ============================

MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"  # 引擎初始加载的旧模型 (旧权重)
MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"        # Trainer 持有的新模型 (新权重)
PAUSE_TOKEN_THRESHOLD = 10               # 任一请求生成 ≥10 个 token 时触发 pause
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"

# ================================================================
# MyLLM: 继承 AsyncLLMEngine，增加并发 pause 协调逻辑
# ================================================================

class MyLLM(vllm.AsyncLLMEngine):
    """Configure the vLLM worker for Ray placement group execution."""

    def __init__(self, **kwargs):
        engine_args = vllm.AsyncEngineArgs(**kwargs)
        vllm_config = engine_args.create_engine_config()
        executor_class = Executor.get_class(vllm_config)
        super().__init__(
            vllm_config=vllm_config,
            executor_class=executor_class,
            log_requests=engine_args.enable_log_requests,
            log_stats=not engine_args.disable_log_stats,
        )
        # ────────────────────────────────────────────────────
        # 2 个共享标志位用于并发协程之间的协调：
        #   _generation_paused:  pause_generation() 调用已返回，> 引擎处于冻结状态
        #   _request_pause_flag: 任意请求已达到 PAUSE_TOKEN_THRESHOLD，> 可以触发 pause 了
        # ────────────────────────────────────────────────────
        self._generation_paused = False
        self._request_pause_flag = False

    async def do_generate(
        self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
    ) -> tuple[vllm.RequestOutput, int]:
        """生成单个请求的完整输出。

        多个 do_generate 协程可以同时运行，它们共享 self 上的两个标志位。
        返回 (output, pause_token_index)，其中 pause_token_index
        是权重切换前生成的 token 数（即新旧权重的分界点），-1 表示未暂停。
        """
        pause_token_index = -1
        prev_token_count = 0

        # ──────────────────────────────────────────────
        # self.generate() 返回 AsyncGenerator。
        # 引擎每次产出一个新 token，协程就 yield 一次。
        # 在 pause 期间（PAUSED_ALL），引擎停止产出，
        # 协程在 async for 上阻塞但不报错，resume 后继续产出。
        # ──────────────────────────────────────────────
        async for request_output in self.generate(
            {"prompt_token_ids": prompt_token_ids},
            sampling_params,
            request_id=str(uuid.uuid4()),
        ):
            output = request_output
            cur_token_count = len(output.outputs[0].token_ids)

            # ① 任意一个请求达到阈值 → 设置暂停标志，通知 pause_after_n_tokens
            if (
                cur_token_count >= PAUSE_TOKEN_THRESHOLD
                and not self._request_pause_flag
            ):
                self._request_pause_flag = True

            # ② pause 已完成 → 记录此刻的 token 数作为新旧权重的分界点
            #    （每个协程独立记录自己的分界点）
            if self._generation_paused and pause_token_index == -1:
                pause_token_index = prev_token_count

            prev_token_count = cur_token_count
        return output, pause_token_index

    async def pause_after_n_tokens(self):
        """等待任一请求触发暂停信号，然后执行 pause。

        与 do_generate 协程在同一事件循环中并发运行。
        """
        # ────────────────────────────────────────
        # ③ 忙等：每次迭代把控制权交还给事件循环 (sleep(0))，
        #    让 do_generate 协程有机会推进并设置 _request_pause_flag。
        #    这是无锁协调机制 —— 不阻塞事件循环的"自旋"。
        # ────────────────────────────────────────
        while not self._request_pause_flag:
            await asyncio.sleep(0)

        # ④ 调用父类的 pause_generation，mode="keep" = 冻结请求不丢失
        await super().pause_generation(mode="keep")

        # ⑤ 等待输出排空：pause 本身有 20ms 延迟，这里额外等 5s 确保
        #    pause 前产出的所有 token 都已通过 output_handler 推送给生成协程
        await asyncio.sleep(5)

        # ⑥ 通知所有生成协程：引擎已冻结，可以标记分界点了
        self._generation_paused = True

# ================================================================
# TrainModel: 训练侧 Ray Actor，持有训练后的新权重
# ================================================================

@ray.remote(num_gpus=1)
class TrainModel:
    """Ray actor that wraps the training model on a dedicated GPU."""

    def __init__(self, model_name: str):
        from vllm.model_executor.layers.batch_invariant import (
            init_batch_invariance,
        )
        init_batch_invariance()
        # 把模型加载到 cuda:0（该 actor 独占的 GPU）
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, dtype=torch.bfloat16
        ).to("cuda:0")
        self.port = get_open_port()
        self.master_address = get_ip()

    def get_master_address_and_port(self):
        """暴露 NCCL master 地址，供 engine 侧连接同一通信组。"""
        return self.master_address, self.port

    def get_weight_metadata(self):
        """收集所有参数的名称、dtype、shape，供 engine 侧验证和准备缓冲区。"""
        names = []
        dtype_names = []
        shapes = []
        for name, p in self.model.named_parameters():
            names.append(name)
            dtype_names.append(str(p.dtype).split(".")[-1])
            shapes.append(list(p.shape))
        return names, dtype_names, shapes

    def init_weight_transfer_group(self, world_size):
        """初始化 NCCL 通信组 (world_size=2: trainer + engine worker)。"""
        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
            dict(
                master_address=self.master_address,
                master_port=self.port,
                world_size=world_size,
            ),
        )

    def broadcast_weights(self, packed: bool = True):
        """通过 NCCL 一次性广播所有权重到 engine。"""
        trainer_args = NCCLTrainerSendWeightsArgs(
            group=self.model_update_group,
            packed=packed,
        )
        NCCLWeightTransferEngine.trainer_send_weights(
            iterator=self.model.named_parameters(),
            trainer_args=trainer_args,
        )

    @torch.inference_mode()
    def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
        """Greedy 解码（验证阶段用）。"""
        input_ids = torch.tensor([token_ids], device="cuda:0")
        output = self.model.generate(
            input_ids, max_new_tokens=max_new_tokens, do_sample=False,
        )
        new_token_ids = output[0, len(token_ids):].tolist()
        return new_token_ids

# ================================================================
# 主流程 Phase 1: 并发生成 + 权重切换
# ================================================================

# ── 初始化 Ray 环境和 env vars ──
# batch_invariant 保证同一批输入在任何 batching 下输出一致（NVIDIA 专属）
ray_env_vars = {}
if current_platform.is_rocm():
    ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1"
    ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
else:
    ray_env_vars["VLLM_BATCH_INVARIANT"] = "1"

ray.init(runtime_env={"env_vars": ray_env_vars})

# ── 启动 Trainer Actor（GPU 0）──
train_model = TrainModel.remote(MODEL_NAME_V2)

# ── 启动 vLLM Engine Actor（GPU 1）──
llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote(
    model=MODEL_NAME_V1,               # 加载旧模型
    enforce_eager=True,
    max_model_len=8192,
    distributed_executor_backend="ray",
    gpu_memory_utilization=0.75,
    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # 启用 NCCL 权重传输
    ...
)

# ── 建立 NCCL 通信通道 ──
master_address, master_port = ray.get(
    train_model.get_master_address_and_port.remote()
)

world_size = 2  # 1 trainer + 1 inference worker

# engine 侧：注册 rank=1，等待 trainer 连接
inference_handle = llm.init_weight_transfer_engine.remote(
    WeightTransferInitRequest(
        init_info=asdict(
            NCCLWeightTransferInitInfo(
                master_address=master_address,
                master_port=master_port,
                rank_offset=1,          # engine worker 的 NCCL rank 从 1 开始
                world_size=world_size,
            )
        )
    )
)

# trainer 侧：创建 NCCL group (trainer rank=0)
train_handle = train_model.init_weight_transfer_group.remote(world_size)
ray.get([train_handle, inference_handle])  # 双向握手完成

# 预先收集权重元信息（每次 update 都相同，避免重复调用）
names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())

# ── 核心流程：触发生成 → 暂停 → 切换权重 → 恢复 ──

N_NEW_TOKENS = 100
sampling_params = SamplingParams(
    temperature=0,
    max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS,  # 10 + 100 = 110
)

# ① 同时触发所有生成请求（Ray remote call，不等待结果）
#    返回的 gen_futures 只是 Future 句柄，协程在引擎内部并发运行
gen_futures = [
    llm.do_generate.remote(ptids, sampling_params)
    for ptids in batch_prompt_token_ids
]
# ↑ 这些 Future 此时还未完成，生成协程正在远程引擎中逐渐产出 token

# ② 启动暂停监控器（与 gen_futures 并发运行）
#    ray.get 阻塞直到 pause 完成
ray.get(llm.pause_after_n_tokens.remote())
# ↑ 当这块返回时：
#    - 所有 do_generate 协程都已在 async for 中阻塞
#    - 引擎处于 PAUSED_ALL 状态
#    - 请求数据和 KV cache 被保留

# ③ 执行权重更新（pause/resume 之间的事务区间）
#    start → update → broadcast → finish 四步原子协议
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))

inference_handle = llm.update_weights.remote(
    WeightTransferUpdateRequest(
        update_info=asdict(
            NCCLWeightTransferUpdateInfo(
                names=names,            # 参数名列表
                dtype_names=dtype_names,# 参数 dtype
                shapes=shapes,          # 参数 shape
                packed=True,            # 一次传输所有参数
            )
        )
    )
)
train_handle = train_model.broadcast_weights.remote(packed=True)
# engine 接受权重元信息 + trainer 广播权重数据 —— 两条路径并行
ray.get([train_handle, inference_handle])

ray.get(llm.finish_weight_update.remote())
# ↑ 权重更新完成。引擎当前持有的模型权重 = MODEL_NAME_V2

# ④ 恢复生成
ray.get(llm.resume_generation.remote())
# ↑ 引擎解除 PAUSED_ALL → UNPAUSED，scheduler 恢复调度
#    所有冻结的 do_generate 协程继续产出 token

# ⑤ 收集所有结果
results = ray.get(gen_futures)
# ↑ 此时所有生成协程已完成，返回 (output, pause_idx)

# ── 打印结果 ──

for i, (output, pause_idx) in enumerate(results):
    all_token_ids = list(output.outputs[0].token_ids)
    # pause_idx 之前的 token 来自旧权重
    before_text = tokenizer.decode(all_token_ids[:pause_idx])
    # pause_idx 及之后的 token 来自新权重
    after_text = tokenizer.decode(all_token_ids[pause_idx:])
    print(f"  Old weights ({pause_idx} tokens): {before_text!r}")
    print(f"  New weights ({len(all_token_ids) - pause_idx} tokens): {after_text!r}")

# ================================================================
# Phase 2: 正确性验证
# ================================================================

# 验证逻辑：
#   旧权重生成的 before tokens + 新权重 = weight-synced 引擎的 after tokens
#                                        = 独立 V2 引擎用相同前缀生成的结果
#   如果两者一致，说明 weight transfer 正确。

# 关闭 Phase 1 的引擎，释放 GPU
ray.get(llm.shutdown.remote())
ray.kill(llm)
ray.kill(train_model)

# 启动一个全新的 V2 引擎，直接加载 MODEL_NAME_V2
llm_v2 = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote(
    model=MODEL_NAME_V2,  # 直接加载训练后的模型
    ...
)

# 用 V2 引擎以 before tokens 为前缀续写，对比 after tokens
val_futures = [
    llm_v2.do_generate.remote(
        # prompt = 原始 prompt + pause 前旧权重生成的 token
        list(output.prompt_token_ids)
        + list(output.outputs[0].token_ids)[:pause_idx],
        SamplingParams(
            temperature=0,
            max_tokens=len(output.outputs[0].token_ids) - pause_idx,
        ),
    )
    for output, pause_idx in results
]
val_results = ray.get(val_futures)

# 逐 token 比对
for i, ((output, pause_idx), (val_output, _)) in enumerate(
    zip(results, val_results)
):
    expected = list(output.outputs[0].token_ids)[pause_idx:]  # synced 引擎的 after
    actual = list(val_output.outputs[0].token_ids)             # V2 引擎的结果
    if actual == expected:
        print(f"  [PASS]")
    else:
        print(f"  [FAIL] first divergence at token {j}: ...")

---

两个协程类型并发协作的时序总结

do_generate(0)     do_generate(1)     pause_after_n_tokens        Trainer
    │                   │                    │                        │
    │  生成 token...     │                    │                        │
    │                   │  生成 token...      │                        │
    │  token=11          │                    │                        │
    │  → _request_pause_flag = True ─────────>│  检测到 flag            │
    │                   │                    │                        │
    │  [在 async for    │  [在 async for     │  pause_generation("keep")│
    │   迭代器阻塞]      │   迭代器阻塞]       │  引擎 PAUSED_ALL        │
    │                   │                    │                        │
    │                   │                    │  _generation_paused=True │
    │  → 记录 pause_idx │  → 记录 pause_idx  │                        │
    │                   │                    │                        │
    │  [继续阻塞等待      │  [继续阻塞等待      │                        │
    │   resume]         │   resume]         │      NCCL 广播权重       │
    │                   │                    │◄───────────────────────│
    │                   │                    │                        │
    │                   │                    │  resume_generation()    │
    │  async for 恢复   │  async for 恢复    │                        │
    │  yield(after)     │  yield(after)      │                        │
    │                   │                    │                        │
    │  return (output,  │  return (output,   │                        │
    │   pause_idx=11)   │   pause_idx=5)     │                        │

• do_generate 协程: 在 async for 上阻塞，不抛异常不退出，resume 后继续产出
• pause_after_n_tokens 协程: 忙等 + 调用 pause + sleep + 通知生成协程
• 二者共享 _request_pause_flag 和 _generation_paused 两个 bool，无需锁

[floatlibai]

vime框架调用示例

vime 将 vLLM 作为 HTTP 子进程运行 (vllm serve)，通过 VLLMEngine Ray Actor 管理。pause/resume 全部走 HTTP 端点 (POST /pause, POST /resume)。

VLLMEngine: HTTP 控制面 (vllm_engine.py)

class VLLMEngine(RayActor):
    """Ray actor for vLLM OpenAI HTTP rollout (connect or spawn local ``vllm serve``)."""

    def __init__(self, args, rank, ...):
        # ...
        self._weight_version: str | None = None  # 跟踪当前权重版本号

    # ── pause_generation ──────────────────────────────────────────
    # 通过 HTTP POST /pause?mode=keep&clear_cache=false 发送到 vLLM 进程。
    # clear_cache=false = 保留 KV cache，请求在 resume 后继续产出 token。
    # 只有 node_rank==0 的 HTTP-serving 节点执行；worker 节点返回 None。
    # ────────────────────────────────────────────────────────────────
    def pause_generation(self):
        if self.node_rank != 0:
            return None
        response = requests.post(
            f"{self._http_base()}/pause",
            params={"mode": "keep", "clear_cache": "false"},
            json={},
            timeout=120,
        )
        response.raise_for_status()
        return response

    # ── continue_generation ───────────────────────────────────────
    # 就是 POST /resume。vime 内部命名为 continue_generation 以更好
    # 地描述语义：被暂停的请求将"继续"生成。
    # ────────────────────────────────────────────────────────────────
    def continue_generation(self):
        if self.node_rank != 0:
            return None
        response = requests.post(
            f"{self._http_base()}/resume", json={}, timeout=120
        )
        response.raise_for_status()
        return response

    # ── 权重更新相关的控制面方法 ────────────────────────────────────
    # 所有这些方法都遵循同一种模式：node_rank 0 发起 HTTP POST，其他 rank 直接返回 None

    def start_weight_update(self, is_checkpoint_format=False):
        return self._make_request("start_weight_update", {"is_checkpoint_format": is_checkpoint_format})

    def finish_weight_update(self):
        return self._make_request("finish_weight_update", {})

    def flush_cache(self):
        """POST /reset_prefix_cache — 清空前缀缓存"""
        requests.post(f"{self._http_base()}/reset_prefix_cache", params={"reset_running_requests": False})

_make_request 是所有控制面 POST 的统一入口:

    def _make_request(self, endpoint, payload=None):
        """单节点 HTTP POST，headless worker 直接返回 None。"""
        if self.node_rank != 0:
            return None
        url = f"{self._http_base()}/{endpoint.lstrip('/')}"
        return _response_json(requests.post(url, json=payload or {}))

---

引擎启动方式: vLLM 作为子进程 (vllm_engine.py)

vime 不以库形式调用 vLLM，而是通过 vllm serve CLI 拉起子进程：

def build_vllm_cmd_and_env(server_args):
    # 构造命令: vllm serve <model> --tensor-parallel-size N --port P ...
    cmd = ["vllm", "serve", str(server_args["model_path"]),
           "--tensor-parallel-size", str(server_args["tp_size"]),
           "--port", str(server_args["port"]), ...]

    # 关键配置:
    # - weight_transfer_config: colocate 场景用 IPC 后端，否则用 NCCL
    # - worker_extension_cls: 注册 vime 的 Worker 扩展，用于处理权重更新
    if colocate:
        cmd += ["--weight-transfer-config", '{"backend":"ipc"}']
    else:
        cmd += ["--weight-transfer-config", '{"backend":"nccl"}']
    cmd += ["--worker-extension-cls", vime_worker_extension_cls]
    return cmd, env

def launch_server_process(server_args):
    cmd, env = build_vllm_cmd_and_env(server_args)
    p = multiprocessing.Process(target=os.execvpe, args=(cmd[0], cmd, env))
    p.start()
    return p

VLLMEngine.init() 中：

def init(self, dist_init_addr, port, ...):
    if self.args.rollout_external:
        self._init_external()    # 外部引擎：仅连 HTTP 检测 /health
    else:
        self._init_normal()      # 启动子进程 → /health 等待启动完毕

---

UpdateWeightFromTensor: Colocate (IPC) 权重更新路径

Colocate 模式下，训练和推理在同一台机器上，用 CUDA IPC 共享内存传输权重。

# update_weight_from_tensor.py (lines 266-319)

@torch.no_grad()
def update_weights(self):
    """
    version++, flush caches, process buckets. Progress on rank 0.
    """
    self.weight_version += 1
    rank = dist.get_rank()

    # ① 暂停所有 rollout 引擎（仅 rank 0 通过 Ray 调度）
    if rank == 0:
        ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
        # ② 刷新前缀缓存
        ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())  # 所有训练 rank 同步

    # ③ 进入权重更新模式（每个引擎的 leader rank 调用 start_weight_update）
    if self._ipc_engine is not None and rank == self._ipc_gather_src:
        ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True))
    dist.barrier(group=get_gloo_group())

    # ④ 迭代 HF 权重分片，逐个 chunk 发送到 vLLM 引擎
    megatron_local_weights = self.weights_getter()
    for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(...):
        refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
        ray.get(refs)          # 等待引擎侧接收完成
        torch.cuda.ipc_collect()  # 释放 CUDA IPC 缓冲区

    dist.barrier(group=get_gloo_group())

    # ⑤ 退出权重更新模式
    if self._ipc_engine is not None and rank == self._ipc_gather_src:
        ray.get(self._ipc_engine.finish_weight_update.remote())
    dist.barrier(group=get_gloo_group())

    # ⑥ 恢复所有引擎的生成
    if rank == 0:
        ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())

---

UpdateWeightFromDistributed: NCCL 权重更新路径

非 Colocate 模式通过 NCCL 广播权重，支持多节点。PP 的每个 stage 有独立的 NCCL group。

# update_weight_from_distributed.py (lines 125-164)

@torch.no_grad()
def update_weights(self):
    """
    Pause → flush → _send_weights → continue. Progress on PP source.
    """
    self.weight_version += 1

    # ① 暂停所有引擎（仅 rank 0）
    if dist.get_rank() == 0:
        ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
        # ② 刷新缓存在 pause 之后执行（保证一致性）
        ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())

    # ③ 权重更新会话: start → 发送权重 → finish
    _begin_vllm_weight_update_session(self.rollout_engines)
    try:
        self._send_weights(pbar)
        torch.cuda.synchronize()
    finally:
        _end_vllm_weight_update_session(self.rollout_engines)

    dist.barrier(group=get_gloo_group())

    # ④ 恢复生成
    if dist.get_rank() == 0:
        ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
    dist.barrier(group=get_gloo_group())

_begin_vllm_weight_update_session 和 _end_vllm_weight_update_session:

def _begin_vllm_weight_update_session(rollout_engines):
    if dist.get_rank() == 0:
        ray.get([engine.start_weight_update.remote(is_checkpoint_format=True)
                 for engine in rollout_engines])
    dist.barrier(group=get_gloo_group())

def _end_vllm_weight_update_session(rollout_engines):
    if dist.get_rank() == 0:
        ray.get([engine.finish_weight_update.remote()
                 for engine in rollout_engines])
    dist.barrier(group=get_gloo_group())

connect_rollout_engines_from_distributed 中 NCCL 通信组初始化 (vllm_engine.py):

def init_weights_update_group(self, master_address, master_port,
                               rank_offset, world_size, group_name, backend):
    """调用 POST /init_weight_transfer_engine 建立 NCCL 通道。"""
    payload = {
        "init_info": {
            "master_address": master_address,
            "master_port": master_port,
            "rank_offset": rank_offset,   # 各引擎在 NCCL group 中的起始 rank
            "world_size": world_size,      # = sum(engine_gpus) + 1 (trainer)
        }
    }
    return self._make_request("init_weight_transfer_engine", payload)

---

关键设计选择:

1. clear_cache=False + 手动 flush_cache: vime 在 pause(mode="keep", clear_cache=False) 后立即调用单独的 flush_cache()。这给 pause 的 idle callback 和 flush_cache 的同步提供了间隔，避免了 clear_cache=True 时 callback 内部自调 _reset_caches 的竞态。
2. ray.get([engine.pause_generation.remote() for engine in engines]): 多引擎并行的 pause — vime 可能有多个 rollout 引擎（例如 prefill engine + decode engine），对所有引擎同时发出 pause，ray.get 阻塞直到所有完成。
3. dist.barrier(group=gloo_group): 除了引擎侧的同步，vime 还在训练 rank 之间做了 barrier。这确保所有训练 rank 在 pause 完成后才进入权重发送阶段，避免某些 rank 过早开始 NCCL 广播。