diff --git a/docs/transferqueue_integration_plan.md b/docs/transferqueue_integration_plan.md new file mode 100644 index 00000000..56b7c3f8 --- /dev/null +++ b/docs/transferqueue_integration_plan.md @@ -0,0 +1,143 @@ +# verl-SpeCo TransferQueue 落地方案 + +> 目标:在**不修改上游 verl**的前提下,把 SpeCo online 训练里的逐样本特征流 +> 从「`SpecoRayPPOTrainer` driver 中转 + Ray object store」改为「TransferQueue +> 直传」,干掉 driver 这个数据瓶颈,并解锁流式消费与跨副本负载均衡。 +> +> 约束:仅 hook,与 SpeCo 现有 hook 模式一致;TQ 作为独立库使用,**不复用** verl +> 的 `main_ppo_sync` TQ 集成。 + +--- + +## 0. 现状:SpeCo online 特征流的 controller 瓶颈 + +SpeCo online 路径以 `SpecoRayPPOTrainer` 为 hub,所有跨进程大张量都被 driver 串行 +中转,介质是 Ray object store(`ray.put`/`ray.get`/`parallel_put`)。这与 verl 引入 +TQ 想解决的痛点 1:1 对应,只是 verl 干掉的是 `RayPPOTrainer`,我们要干掉的是 +SpeCo 在它之上加的 drafter 管线中转。 + +| # | 流向 | 当前机制 | 是否逐样本 | hook 位置(SpeCo 侧) | +|---|---|---|---|---| +| **a1** | target hidden states(SGLang 采集)-> drafter | `drafter_sample` 塞进 `DataProto.non_tensor_batch` -> driver pop/bucket -> `parallel_put` -> drafter `ray.get` | ✅ | `speco_ray_trainer.py` `generate_sequences_with_speco`;`sglang_adapter.py` `pop_drafter_samples`/`bucket_drafter_samples_by_replica`;`sglang_runtime.py` 组装 `drafter_sample` | +| **a2** | target hidden states(old-logprob hook)-> drafter | actor 前向 hook 截行 -> `ray.put` chunk -> driver 重打包 -> 分发 | ✅ | `oldlogprob_runtime.py` `_install_oldlogprob_hidden_hooks`/`_put_oldlogprob_hidden_refs`;`speco_ray_trainer.py` `_speco_collect_oldlogprob_features` | +| **b2** | target top-logprobs -> drafter(`use_logits=true`) | 随 a1 同一 side-channel | ✅ | `sglang_runtime.py` `target_logprobs`/`hidden_raw_target_logprobs` | +| **d** | rollout tokens -> drafter 训练集 | 随 a1 同一 side-channel(online)/`torch.save` 分片(offline) | ✅ | `sglang_runtime.py`;`speco_worker.py` `_store_rollout_sample` | +| b1 | target **lm_head 权重**(行)-> drafter `TargetHead` | ONE_TO_ALL Ray 分发 | ❌ 参数广播 | `rollout_publish.py` `export_actor_lm_head_weight`/`get_actor_lm_head_weight`;`speco_ray_trainer.py` `_speco_sync_target_lm_head_weight` | +| c | drafter 权重 -> rollout 引擎 | `ray.put` -> driver -> actor;vLLM 末段 ZMQ+SHM,SGLang 进程内 | ❌ 参数广播 | `speco_worker.py` `maybe_publish`;`rollout_publish.py` `update_draft_weights`;`vllm_runtime.py` `BucketedWeightSender` | + +**关键事实**:hidden states 跨进程前一律 CPU 物化(`oldlogprob_runtime.py`、 +`sglang_runtime.py`、`feature_store.py` 均 `.cpu()`),a1 路径下 driver 进程的 +host memory 会真正承载整批 hidden states 并做一次 Ray store 往返。这正是 TQ 要 +消除的往返。 + +--- + +## 1. 为什么不把"替换 feature_store"作为第一刀 + +`TorchShardFeatureStore`(`feature_store.py`)是 `torch.save` 分片 + JSONL manifest, +**只服务于 `collect_only`/`offline`**,不参与 online 热路径。替换它能统一离线存储 +抽象、换更快的分布式后端,但**不解决 controller 瓶颈**,性能收益有限。降级为 +可选尾项(见 §5 P3)。 + +--- + +## 2. 目标方案:TQ 直传逐样本特征流(a1 / a2 / b2 / d) + +### 2.1 角色映射 + +| TQ 角色 | SpeCo 对应 | +|---|---| +| Producer(写) | rollout worker(SGLang 路径,a1/b2/d)/ actor worker(old-logprob 路径,a2)——均在 SpeCo 既有 hook 内 | +| Consumer(读) | drafter worker `collect_rollout_features`(SpeCo 侧) | +| TransferQueueController(control plane) | SpeCo launcher 启动一个 Ray actor;drafter 经 `Sampler`/`StreamingDataLoader` 拉取 | +| Storage backend | `SimpleStorage`(ZMQ,跨节点 CPU 内存);进阶可切 `MooncakeStore`(RDMA,GPU-DRAM) | + +### 2.2 partition / key / 字段设计 + +- `partition_id`:`speco_train`(验证集用 `speco_val`)。 +- `key`:`{uid}_{session_id}_{index}`,与 verl TQ 一致;`uid` SpeCo 已有。 +- `tags`:`global_steps`、`source`∈{`rollout`,`oldlogprob`}、`replica_rank`/`owner_rank`、`status`、`prompt_len`/`response_len`/`seq_len`。ReplayBuffer/负载均衡按 tag 匹配。 +- `fields`(列):`input_ids`、`loss_mask`、`position_ids`、`hidden_states`、`last_hidden_states`/`target`、`target_logprobs`、`hidden_positions`、`prompts`、`responses`。与 `DraftFeatureSample`(`feature_store.py`)字段对齐,便于 online/offline 复用。 + +### 2.3 数据流(目标) + +``` +rollout/actor worker (SpeCo hook) + │ 生成/截取 hidden states 后,就地 tq.kv_batch_put(samples) + ▼ +TransferQueue (SimpleStorage, 跨节点 CPU 内存;可选 MooncakeStore RDMA) + │ control plane 按 sample 粒度追踪 ready 状态,Sampler 跨 drafter 副本均衡 + ▼ +drafter worker + │ tq.kv_batch_get / StreamingDataLoader 消费 → 喂入既有 DataBuffer / collect_online_data + ▼ +drafter 训练 (不变) +``` + +driver 只下发触发与轻量 key/meta,**不再承载 hidden states**。 + +--- + +## 3. 落地改动点(全部在 SpeCo 侧,hook-only) + +### 3.1 启动与配置 +- `draft_train_launcher.py` / `main.py`:`tq.init(config.transfer_queue)`;起 `TransferQueueController.remote(Sampler)`。 +- `config/speco_base.yaml`:新增 `drafter.transfer_queue` 块(backend、partition、enable 开关)。参考 verl `ppo_trainer.yaml` 的 `transfer_queue:` 结构,但**独立配置**,不复用 verl 的。 + +### 3.2 Producer 侧 +- **a1/b2/d(SGLang)**:`sglang_runtime.py` 组装 `drafter_sample` 处(~1594-1648),增加 `tq.kv_batch_put`;返回给 driver 的 `drafter_sample` 只保留 key/meta(或整段不再走 DataProto side-channel,driver 仅触发)。 +- **a2(old-logprob)**:`oldlogprob_runtime.py` `_put_oldlogprob_hidden_refs`(~216),把 `ray.put(hidden_chunk)` 换成 `tq.kv_batch_put`;`OLD_LOGPROB_HIDDEN_CHUNK_REFS_KEY` 改为 TQ key 列表。 + +### 3.3 Consumer 侧 +- `speco_worker.py` `collect_rollout_features`(~665):把 `_resolve_ray_object_ref`/`_resolve_hidden_state_chunks`(`ray.get`)换成 `tq.kv_batch_get`;`_dispatch_nd_compute`(~159)的 `parallel_put` 退化为只传 key(或 drafter 直接从 TQ Sampler 拉,driver 不参与分发)。 +- drafter 内部 `DataBuffer`/`collect_online_data`(`base_trainer.py`)保持不变,只是数据来源由 `ray.get` 改为 TQ get。 + +### 3.4 Driver 侧 +- `speco_ray_trainer.py`:`_speco_collect_rollout_features_rpc`/`speco_collect_rollout_features`(~351)、`_speco_collect_oldlogprob_features`(~1114)不再搬数据,只做触发/传 key;`bucket_drafter_samples_by_replica` 可由 TQ `Sampler` 替代(逐步迁移,先保留作回退)。 + +### 3.5 不改动 +- **b1(lm_head 权重)、c(drafter 权重)**:保持现状。与 verl 上游一致(权重不走 TQ),且 c 的 vLLM 末段已有专用 ZMQ+SHM 通道。 +- verl 本体:零改动。 + +--- + +## 4. 收益与边界(诚实评估) + +### 收益 +1. **去掉 driver 对 hidden states 的 host-memory 中转 + Ray store 往返**:producer 直存 TQ,consumer 直取,driver 不再承载整批特征。 +2. **流式消费**:drafter 在样本 ready 时即可消费,不必等整批 `generate_sequences` 返回,采集与训练可重叠。 +3. **跨 drafter 副本负载均衡**:TQ `Sampler`/`RankAwareSampler` 替代手写 `bucket_drafter_samples_by_replica`/`owner_rank` 分配。 +4. **(若采纳 P3)统一 online/collect_only/offline 存储**:同一 TQ partition,`collect_only` 写、`offline` 读,消掉 on-disk 分片层。 + +### 边界 / 不解决的事 +- 只优化**特征采集**这一子阶段,**不加速** rollout 本身、actor update、reward;e2e 增益取决于该子阶段在 step 中的占比。 SpeCo README 的 20% rollout / 11% e2e 提升来自 acceptance length,与本方案是不同机制,不要混为一谈。 +- **权重同步(b1/c)不放进 TQ**,与 verl 上游保持一致。 +- hidden states 跨进程前**仍需 CPU 物化**(现状如此);要避免物化需切 `MooncakeStore` RDMA,属进阶项。 +- 引入 TQ 依赖与一个 control-plane Ray actor,增加少量运维面。 + +### 风险 +- TQ 与 SpeCo 现有 `owner_rank`/`replica_rank` 路由语义需对齐(Sampler 要复刻「按 owner 分桶」语义,否则样本会错配 drafter 副本)。 +- old-logprob 的 chunk 拆分(`hidden_states_ref_chunks`)映射到 TQ 列式存储时,需保证 chunk meta 与 key 的一致性。 +- 回退路径:保留 `enable_transfer_queue=False` 时走原 Ray 路径,渐进切换。 + +--- + +## 5. 分阶段实施 + +| 阶段 | 范围 | 产出 | +|---|---|---| +| **P0** | a1(SGLang hidden states)走 TQ 直传;drafter `kv_batch_get` 消费;driver 仅触发 | 验证 controller-bypass 闭环 + 正确性 | +| **P1** | a2(old-logprob hidden states)走 TQ;chunk 拆分映射 TQ 列 | 覆盖第二条采集路径 | +| **P2** | b2(top-logprobs)+ d(tokens)随 a1 同 partition 传输;Sampler 替代手写 bucket | 完整特征流 + 跨副本均衡 | +| **P3(可选)** | `TorchShardFeatureStore` → TQ partition,统一 online/collect_only/offline | 离线工作流统一 | + +每个阶段保留 `enable_transfer_queue` 开关与原 Ray 路径回退。 + +--- + +## 6. 待确认决策 + +1. **TQ backend**:`SimpleStorage`(CPU 内存,默认)起步,还是直接上 `MooncakeStore`(RDMA,省 CPU 物化)?后者依赖 RDMA 网络,建议 P0 用 SimpleStorage。 +2. **drafter 消费模式**:`kv_batch_get`(主动拉,改动小)还是 `StreamingDataLoader`(全自动流式,改动大、收益高)?建议 P0 用前者,P2 再考虑后者。 +3. **driver 角色**:P0 先保留 driver 传 key(最小改动),还是直接让 drafter 从 TQ Sampler 自取(driver 彻底退出数据路径)?前者风险低,建议 P0 用前者。 +4. **是否做 P3**:离线统一是否在本次范围内,还是单独立项。 diff --git a/verl_speco/config/speco_base.yaml b/verl_speco/config/speco_base.yaml index f8eeb435..a01b8290 100644 --- a/verl_speco/config/speco_base.yaml +++ b/verl_speco/config/speco_base.yaml @@ -179,3 +179,16 @@ actor_rollout_ref: repeat: true prefetch_depth: 2 strict_schema: true + # TransferQueue transport for drafter features (P0: hidden states only). + # When enabled and the `transfer_queue` package is installed, the SGLang + # rollout server writes hidden states directly to TQ and the drafter + # worker reads them by key, bypassing the RayPPOTrainer driver and the + # Ray object store for the dominant tensor. Default off -> unchanged + # Ray path. Requires `pip install TransferQueue==0.1.7`. + transfer_queue: + enable: false + backend: + storage_backend: SimpleStorage + SimpleStorage: + total_storage_size: 100000 + num_data_storage_units: 8 diff --git a/verl_speco/integration/oldlogprob_runtime.py b/verl_speco/integration/oldlogprob_runtime.py index 51ae0b1c..0130a3b3 100644 --- a/verl_speco/integration/oldlogprob_runtime.py +++ b/verl_speco/integration/oldlogprob_runtime.py @@ -48,6 +48,9 @@ OLD_LOGPROB_HIDDEN_LAYOUT_KEY = "speco_oldlogprob_hidden_layout" OLD_LOGPROB_TIMING_KEY = "speco_oldlogprob_timing" OLD_LOGPROB_SELECTED_BATCH_INDICES_KEY = "speco_oldlogprob_selected_batch_indices" +# Stamped on the micro-batch by the collect plan so the actor-worker producer +# can build step-unique TransferQueue keys for old-logprob hidden chunks (P1). +OLD_LOGPROB_GLOBAL_STEP_KEY = "speco_oldlogprob_global_step" _TIMING_SELECT_US = 0 _TIMING_SP_MERGE_US = 1 @@ -204,6 +207,45 @@ def _oldlogprob_hidden_object_ref_enabled(micro_batch: Any) -> bool: return bool(value) +def _oldlogprob_global_step(micro_batch: Any) -> int: + """Read the global step stamped on the micro-batch by the collect plan.""" + + value = 0 + try: + from verl.utils import tensordict_utils as tu + + value = tu.get_non_tensor_data(data=micro_batch, key=OLD_LOGPROB_GLOBAL_STEP_KEY, default=0) + except Exception: # noqa: BLE001 + try: + value = micro_batch.get(OLD_LOGPROB_GLOBAL_STEP_KEY, 0) + except Exception: # noqa: BLE001 + if _tensor_key_present(micro_batch, OLD_LOGPROB_GLOBAL_STEP_KEY): + value = micro_batch[OLD_LOGPROB_GLOBAL_STEP_KEY] + value = getattr(value, "data", value) + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _speco_tq_enabled_for_oldlogprob() -> bool: + """Configure (idempotently) and report whether TQ transport is usable here. + + The actor worker reaches its drafter training config through the same env + serialization the SGLang server uses (``SPECO_SGLANG_DRAFTER_CONFIG_ENV``). + """ + + from verl_speco.integration.transferqueue_bridge import ( + configure_transfer_queue, + is_transfer_queue_enabled, + ) + + drafter = _load_drafter_env() + training = _get_nested(drafter, ("training",), None) + configure_transfer_queue(training) + return is_transfer_queue_enabled() + + def _is_sparse_sp_non_source_context(context: dict[str, Any]) -> bool: return ( bool(context.get("sparse_sp_merge")) @@ -491,7 +533,27 @@ def _put_oldlogprob_hidden_refs( else tensors[0].contiguous() ) ray_put_started = time.perf_counter() - chunk_ref = ray.put(hidden_chunk) + # P1: when TQ is enabled, store the owner's concatenated hidden chunk + # in TransferQueue and carry the key in place of the Ray ObjectRef. + # The driver treats the token as opaque; the drafter consumer + # resolves "speco:" keys via TQ instead of ray.get. Falls back to + # ray.put otherwise (unchanged behavior). + if _speco_tq_enabled_for_oldlogprob(): + from verl_speco.integration.transferqueue_bridge import make_sample_key, put_sample + + tq_key = make_sample_key( + _oldlogprob_global_step(micro_batch), + int(owner), + f"chunk{len(chunk_refs)}", + ) + put_sample( + tq_key, + {"hidden": hidden_chunk}, + tag={"global_step": _oldlogprob_global_step(micro_batch), "owner": int(owner)}, + ) + chunk_ref = tq_key + else: + chunk_ref = ray.put(hidden_chunk) ray_put_us += (time.perf_counter() - ray_put_started) * 1_000_000.0 chunk_index = len(chunk_refs) chunk_refs.append(chunk_ref) diff --git a/verl_speco/integration/sglang_runtime.py b/verl_speco/integration/sglang_runtime.py index d13bf59d..8bb5ebfb 100644 --- a/verl_speco/integration/sglang_runtime.py +++ b/verl_speco/integration/sglang_runtime.py @@ -2020,6 +2020,44 @@ async def generate( "global_step": collection_global_steps, "replica_rank": self.replica_rank, } + # P0: offload the dominant hidden_states tensor to + # TransferQueue so it bypasses the RayPPOTrainer driver and + # the Ray object store. The key rides with the sample dict; + # the drafter worker fetches by key. No-op when TQ disabled. + from verl_speco.integration.transferqueue_bridge import ( + configure_transfer_queue, + is_transfer_queue_enabled, + make_sample_key, + put_sample, + ) + configure_transfer_queue(training_cfg) + if is_transfer_queue_enabled(): + tq_key = make_sample_key(collection_global_steps, self.replica_rank, request_id) + tq_payload = {"hidden_states": hidden_states.unsqueeze(0).cpu()} + # P2: also offload the other large tensors so they bypass + # the driver too. The drafter consumer restores them from + # this same TQ payload. + if target_logprobs is not None: + tq_payload["target_logprobs"] = target_logprobs.unsqueeze(0).cpu() + if torch.is_tensor(hidden_raw_target_logprobs): + tq_payload["hidden_raw_target_logprobs"] = hidden_raw_target_logprobs.unsqueeze(0).cpu() + if torch.is_tensor(hidden_raw_target_logprobs_positions): + tq_payload["hidden_raw_target_logprobs_positions"] = ( + hidden_raw_target_logprobs_positions.unsqueeze(0).cpu() + ) + put_sample( + tq_key, + tq_payload, + tag={"global_step": collection_global_steps, "replica_rank": self.replica_rank}, + ) + drafter_sample["hidden_states_tq_key"] = tq_key + drafter_sample["hidden_states"] = None + if target_logprobs is not None: + drafter_sample["target_logprobs"] = None + if torch.is_tensor(hidden_raw_target_logprobs): + drafter_sample["hidden_raw_target_logprobs"] = None + if torch.is_tensor(hidden_raw_target_logprobs_positions): + drafter_sample["hidden_raw_target_logprobs_positions"] = None else: self._speco_log_missing_hidden_states_once( collection_global_steps=collection_global_steps, diff --git a/verl_speco/integration/task_runner.py b/verl_speco/integration/task_runner.py index a0498ceb..2ab07c9b 100644 --- a/verl_speco/integration/task_runner.py +++ b/verl_speco/integration/task_runner.py @@ -309,5 +309,17 @@ def _run_with_speco_trainer(self, config): speco_worker_cls=speco_worker_cls, ) - trainer.init_workers() - trainer.fit() + # Bootstrap TransferQueue before spawning Ray actors so worker processes + # inherit the TQ environment (mirrors verl main_ppo_sync tq.init). The + # TaskRunner owns shutdown so a failed fit cannot leak the named + # controller/storage. No-op when transfer_queue.enable=false or the + # package is not installed. + from verl_speco.integration.transferqueue_bridge import close_transfer_queue, init_transfer_queue + + transfer_queue_started = init_transfer_queue(config) + try: + trainer.init_workers() + trainer.fit() + finally: + if transfer_queue_started: + close_transfer_queue() diff --git a/verl_speco/integration/transferqueue_bridge.py b/verl_speco/integration/transferqueue_bridge.py new file mode 100644 index 00000000..7062b7f4 --- /dev/null +++ b/verl_speco/integration/transferqueue_bridge.py @@ -0,0 +1,305 @@ +"""TransferQueue bridge for SPECO drafter feature transport. + +This module lets SPECO route large per-sample drafter-training tensors (hidden +states, target logprobs) through TransferQueue (TQ) instead of funneling them +through the ``SpecoRayPPOTrainer`` driver process and the Ray object store. It +is the SpeCo-side analog of verl's ``transferqueue_utils.py``, but used as a +standalone transport library -- it does **not** depend on verl's +``main_ppo_sync`` TQ integration and does **not** modify upstream verl. + +Design: +- P0: offload SGLang-collected ``hidden_states`` (a1 path). +- P1: offload old-logprob-collected ``hidden_states`` (a2 path) -- replaces the + ``ray.put`` chunk + driver relay. +- P2: also offload ``target_logprobs`` / ``hidden_raw_target_logprobs``. +- Default ``enable: false`` -> behavior is bit-identical to the current Ray + path. TQ is only touched when explicitly enabled and the ``transfer_queue`` + package is importable. + +A sample remains in TQ until task cleanup: a SPECO drafter replica contains +multiple TP/SP ranks and each rank reads the same sample (the owner-route +dispatch duplicates a DP bucket to all SP ranks of one replica; only the SP +leader is ``is_collect``). Deleting a sample after the first read would race the +remaining ranks. Garbage collection is therefore deferred to task teardown; a +finer-grained leader-clears-after-barrier is future work. + +Note: the TQ call sites follow the documented KV API (``kv_put`` / +``kv_batch_get`` / ``kv_close``) of TransferQueue 0.1.7. The exact signatures +(keyword names, return shapes) must be verified against the installed TQ +version on first run; the bridge fails loud, never silently. +""" + +from __future__ import annotations + +import logging +import os +import threading +from typing import Any, Optional + +import torch + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +# Partition under which SPECO drafter samples are stored. +_SPECO_TQ_PARTITION = "speco_drafter_features" + +try: + import transfer_queue as tq # type: ignore + from transfer_queue import KVBatchMeta # noqa: F401 (re-exported for symmetry) + + _TQ_IMPORTABLE = True +except ImportError: + + _TQ_IMPORTABLE = False + + class KVBatchMeta: # type: ignore[no-redef] + """Stand-in used only when TransferQueue is not installed.""" + + class _MockTQ: + """Mock that raises on any use; only hit if enabled without TQ installed.""" + + def __getattr__(self, name: str) -> Any: + def _raise(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError( + f"transfer_queue is not installed. Cannot call tq.{name}(). " + "Install with `pip install TransferQueue==0.1.7` or disable " + "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable." + ) + + return _raise + + tq = _MockTQ() # type: ignore[assignment] + + +# --------------------------------------------------------------------------- +# State +# --------------------------------------------------------------------------- + +_state_lock = threading.Lock() +_state = { + "enabled": False, # config says enable=true + "configured": False, # configure_transfer_queue has run + "initialized": False, # tq.init() has run in this process + "config": None, # the transfer_queue sub-config (plain dict) + "owner": False, # this process created the task-level TQ system +} + + +def configure_transfer_queue(training_cfg: Any) -> bool: + """Read the ``transfer_queue`` sub-config from the drafter training config. + + Called from both the SGLang server process (via env-serialized drafter + config) and the drafter worker process (via full hydra config). Idempotent. + Returns whether TQ transport is usable in this process. + """ + + global _state + tq_cfg = _extract_tq_config(training_cfg) + with _state_lock: + _state["config"] = tq_cfg + _state["enabled"] = bool(tq_cfg.get("enable")) if tq_cfg else False + _state["configured"] = True + usable = _state["enabled"] and _TQ_IMPORTABLE + if _state["enabled"] and not _TQ_IMPORTABLE: + logger.warning( + "[SpeCo TQ] transfer_queue.enable=true but transfer_queue package is " + "not installed; falling back to inline Ray transport." + ) + return usable + + +def _extract_tq_config(training_cfg: Any) -> Optional[dict]: + if training_cfg is None: + return None + transfer_queue_cfg = None + if hasattr(training_cfg, "get"): + transfer_queue_cfg = training_cfg.get("transfer_queue", None) + elif isinstance(training_cfg, dict): + transfer_queue_cfg = training_cfg.get("transfer_queue", None) + if transfer_queue_cfg is None: + return None + return _to_plain_dict(transfer_queue_cfg) + + +def _to_plain_dict(value: Any) -> dict: + """Convert OmegaConf DictConfig / nested mapping to a plain dict.""" + + if hasattr(value, "to_container"): + try: + import omegaconf + + return omegaconf.OmegaConf.to_container(value, resolve=True) # type: ignore[arg-type] + except Exception: # noqa: BLE001 + pass + if isinstance(value, dict): + return {k: _to_plain_dict(v) for k, v in value.items()} + return value + + +def is_transfer_queue_enabled() -> bool: + """True only if configured enabled AND the TQ package is importable.""" + + return bool(_state["enabled"]) and _TQ_IMPORTABLE + + +def init_transfer_queue(config: Any) -> bool: + """Cluster-wide TQ bootstrap, called once from the SpecoTaskRunner. + + Mirrors verl ``main_ppo_sync`` calling ``tq.init(config.transfer_queue)`` + in the TaskRunner before workers spawn. Other Ray processes lazily call + ``tq.init()`` and connect to the named TransferQueue controller. Returns + whether TQ is usable; no-op (returns False) when disabled or not installed. + """ + + tq_cfg = _extract_tq_config(_drafter_training_cfg(config)) + if tq_cfg is None or not bool(tq_cfg.get("enable")) or not _TQ_IMPORTABLE: + return False + tq.init(_to_plain_dict(tq_cfg)) + with _state_lock: + _state["config"] = _to_plain_dict(tq_cfg) + _state["enabled"] = True + _state["initialized"] = True + _state["owner"] = True + logger.info("[SpeCo TQ] TransferQueue bootstrapped in task runner (partition=%s)", _SPECO_TQ_PARTITION) + return True + + +def _drafter_training_cfg(config: Any) -> Any: + try: + return config.actor_rollout_ref.rollout.drafter.training + except AttributeError: + return None + + +def _ensure_initialized() -> None: + """Lazily ``tq.init()`` once per worker process (mirrors verl TQ_INITIALIZED). + + A no-argument initialization discovers the named TransferQueue controller + on the connected Ray cluster. It deliberately does not create a separate + per-worker configuration. + """ + + if _state["initialized"]: + return + with _state_lock: + if _state["initialized"]: + return + tq.init() + _state["initialized"] = True + + +# --------------------------------------------------------------------------- +# Key / put / get / close +# --------------------------------------------------------------------------- + +def make_sample_key(global_step: Any, replica_rank: Any, request_id: Any) -> str: + """Build a deterministic, cluster-unique key for one drafter sample. + + Uniqueness space: (global_step, replica_rank, request_id). Each rollout + request produces exactly one drafter_sample, so this is unique per sample. + """ + + return f"speco:{global_step}:{replica_rank}:{request_id}" + + +def put_sample( + key: str, + tensor_dict: dict, + *, + tag: Optional[dict] = None, +) -> None: + """Store a dict of CPU tensors under ``key`` in the SPECO TQ partition. + + ``tensor_dict`` values must be CPU ``torch.Tensor`` (or None). None values + are dropped. Raises if TQ is enabled but the call fails -- never silently + degrades, so a transport failure surfaces immediately rather than dropping + a sample. + """ + + if not is_transfer_queue_enabled(): + raise RuntimeError("put_sample called while TransferQueue is not enabled.") + payload = {k: v for k, v in tensor_dict.items() if torch.is_tensor(v)} + if not payload: + return + _ensure_initialized() + # Pass a plain single-sample dict of columns. TQ's kv_put adds its required + # batch dimension internally; constructing a scalar TensorDict here would be + # incorrect. (Exact kwarg names verified against TQ 0.1.7 on first run.) + tq.kv_put( + key=key, + partition_id=_SPECO_TQ_PARTITION, + fields=payload, + tag=tag or {}, + ) + + +def get_sample(key: str) -> dict: + """Retrieve one tensor dict without deleting it. + + A drafter replica may execute this method on multiple TP/SP ranks; each + rank reads the same key. TQ storage is released once at task shutdown by + the process that initialized it, after every consumer has finished. + """ + + if not is_transfer_queue_enabled(): + raise RuntimeError("get_sample called while TransferQueue is not enabled.") + _ensure_initialized() + # TQ returns the stored sample (TensorDict-like). Return shape is version + # dependent, so handle both a direct value and a {key: value} mapping. + result = tq.kv_batch_get(keys=[key], partition_id=_SPECO_TQ_PARTITION) + value = _extract_value(result, key) + if value is None: + return {} + return _tensordict_to_dict(value) + + +def close_transfer_queue() -> None: + """Close task-level TQ resources if this process initialized them. + + Only the TaskRunner owns controller/storage teardown. Worker-side lazy + clients must not call this because they may still be serving other ranks. + """ + + with _state_lock: + if not _state["owner"]: + return + _state["owner"] = False + _state["initialized"] = False + try: + tq.close() + except AttributeError: + # tq.close() is not part of the public TQ API on some versions; nothing + # to tear down explicitly. The named controller/storage actors are + # reaped when the Ray job exits. + logger.debug("[SpeCo TQ] tq.close() unavailable; skipping explicit teardown") + except Exception: # noqa: BLE001 + logger.debug("[SpeCo TQ] tq.close() raised; ignoring shutdown error") + + +def _extract_value(result: Any, key: str) -> Any: + if result is None: + return None + if isinstance(result, dict): + return result.get(key) + if isinstance(result, (list, tuple)): + return result[0] if len(result) > 0 else None + return result + + +def _tensordict_to_dict(value: Any) -> dict: + if hasattr(value, "items"): + return {k: v for k, v in value.items()} + return dict(value) + + +__all__ = [ + "KVBatchMeta", + "configure_transfer_queue", + "close_transfer_queue", + "init_transfer_queue", + "is_transfer_queue_enabled", + "make_sample_key", + "put_sample", + "get_sample", +] diff --git a/verl_speco/trainer/speco_ray_trainer.py b/verl_speco/trainer/speco_ray_trainer.py index 018d8f81..f76feb14 100644 --- a/verl_speco/trainer/speco_ray_trainer.py +++ b/verl_speco/trainer/speco_ray_trainer.py @@ -39,6 +39,7 @@ from verl_speco.integration.oldlogprob_runtime import ( OLD_LOGPROB_AUX_LAYER_IDS_KEY, OLD_LOGPROB_COLLECT_MASK_KEY, + OLD_LOGPROB_GLOBAL_STEP_KEY, OLD_LOGPROB_HIDDEN_CAPTURE_IMPL_KEY, OLD_LOGPROB_HIDDEN_CHUNK_META_KEY, OLD_LOGPROB_HIDDEN_CHUNK_REFS_KEY, @@ -2223,6 +2224,9 @@ def compute_old_log_prob_without_collection(): self._speco_oldlogprob_hidden_layout(), ) tu.assign_non_tensor_data(batch_td, OLD_LOGPROB_HIDDEN_OBJECT_REF_KEY, True) + # Stamp the global step so the actor-worker producer can build + # step-unique TransferQueue keys for old-logprob hidden chunks (P1). + tu.assign_non_tensor_data(batch_td, OLD_LOGPROB_GLOBAL_STEP_KEY, self.global_steps) self._speco_last_oldlogprob_prepare_elapsed_sec = ( time.perf_counter() - prepare_started diff --git a/verl_speco/workers/speco_worker.py b/verl_speco/workers/speco_worker.py index 9ca71d5b..e2bc018b 100644 --- a/verl_speco/workers/speco_worker.py +++ b/verl_speco/workers/speco_worker.py @@ -69,10 +69,54 @@ def _resolve_ray_object_ref(value): return value -def _resolve_hidden_state_chunks(chunks, expected_rows: int | None = None): +def _densify_tq_tensor(tensor): + """Unwrap a tensor returned by TransferQueue into a plain dense tensor. + + TQ stores each ``put_sample`` payload inside a TensorDict and + ``kv_batch_get`` returns it with an added batch dimension, as a NestedTensor + (jagged on dim 0). The old-logprob chunk resolver slices + ``tensor[start:start + length]`` on dim 0, which NestedTensor does not + support (``slice(): not supported for NestedTensor on dim=0``). The producer + put a dense ``[rows, hidden]`` tensor, so flatten the NestedTensor back to + that 2-D form. Mirrors ``_speco_tensor_rows`` which uses ``tensor.unbind()`` + for the same nested-tensor case. + """ + if not torch.is_tensor(tensor): + return tensor + if tensor.is_nested: + parts = [p for p in tensor.unbind() if p.numel() > 0] + if not parts: + return None + tensor = torch.cat(parts, dim=0) + if tensor.dim() == 3: + tensor = tensor.squeeze(0) + elif tensor.dim() == 1: + tensor = tensor.unsqueeze(0) + return tensor.contiguous() + + +def _resolve_tq_or_ray_ref(ref): + # P1: old-logprob chunk refs may be TransferQueue keys ("speco:" prefix) + # instead of Ray ObjectRefs. Fetch the stored chunk via TQ and fall back to + # ray.get otherwise. Multiple SP ranks of one replica read the same key; + # TQ samples are not cleared on read (see transferqueue_bridge.get_sample). + if isinstance(ref, str) and ref.startswith("speco:"): + from verl_speco.integration.transferqueue_bridge import get_sample + + return _densify_tq_tensor(get_sample(ref).get("hidden")) + return _resolve_ray_object_ref(ref) + + +def _resolve_hidden_state_chunks(chunks, expected_rows: int | None = None, cache=None): if not chunks: return None - resolved_cache = {} + # Cross-sample cache: callers pass a per-step cache so multiple samples + # sharing the same owner chunk (same TQ key / ObjectRef) are fetched only + # once. Without this, ~16 samples in one owner each trigger a full TQ + # kv_batch_get (or ray.get) on the same ~400MB chunk. ray.get dedups at + # the object store; TQ get_sample does NOT, so the cache is essential. + if cache is None: + cache = {} pieces = [] full_rows = int(expected_rows or 0) hidden_size = None @@ -83,10 +127,10 @@ def _resolve_hidden_state_chunks(chunks, expected_rows: int | None = None): ref = chunk.get("ref") if ref is None: continue - cache_key = id(ref) - if cache_key not in resolved_cache: - resolved_cache[cache_key] = _resolve_ray_object_ref(ref) - tensor = resolved_cache[cache_key] + cache_key = ref if isinstance(ref, str) else id(ref) + if cache_key not in cache: + cache[cache_key] = _resolve_tq_or_ray_ref(ref) + tensor = cache[cache_key] if not torch.is_tensor(tensor): continue tensor = cast(torch.Tensor, tensor) @@ -365,6 +409,14 @@ def __init__( self.config.rollout.drafter.training.get("step", 100) ) + # Configure TransferQueue transport for drafter features. No-op when + # disabled or when the transfer_queue package is not installed; the + # existing inline Ray path is used otherwise. Cached on the instance so + # collect_rollout_features can branch without re-reading config. + from verl_speco.integration.transferqueue_bridge import configure_transfer_queue + + self._speco_tq_enabled = configure_transfer_queue(self.config.rollout.drafter.training) + def _ensure_process_group_initialized(self): if not dist.is_initialized(): initialize_global_process_group_ray( @@ -796,9 +848,39 @@ def _flush_rollout_features_for_step(self) -> None: def collect_rollout_features(self, samples: list[dict]): if not samples: return + # Per-step cross-sample cache for chunk fetches. Reset every collect + # call so keys (which carry global_step) never stale and the cache + # cannot grow unbounded. Shared across all samples in this step. + self._tq_chunk_cache = {} for sample in samples: if not sample: continue + # P2: restore TransferQueue-offloaded tensors into the sample before + # building the batch. One fetch restores hidden_states plus the other + # large tensors (target_logprobs, raw target logprobs) so they all + # bypass the driver. a2 old-logprob samples (which carry + # hidden_states_ref_chunks, not a tq key) skip this and resolve below. + tq_key = sample.get("hidden_states_tq_key") + if tq_key is not None and self._speco_tq_enabled: + from verl_speco.integration.transferqueue_bridge import get_sample + + payload = get_sample(tq_key) + for _field in ( + "hidden_states", + "target_logprobs", + "hidden_raw_target_logprobs", + "hidden_raw_target_logprobs_positions", + ): + if payload.get(_field) is not None: + sample[_field] = payload[_field] + if sample.get("hidden_states") is None: + # Fail loud: a TQ key was produced but the payload is + # missing -> transport is broken. Do NOT silently drop the + # sample (that would corrupt training data). + raise RuntimeError( + f"[SpeCo TQ] drafter worker got empty hidden_states for " + f"key={tq_key}; TQ enabled but producer payload missing." + ) batch = { "input_ids": sample["input_ids"], "prompts": sample["prompts"], @@ -830,6 +912,8 @@ def collect_rollout_features(self, samples: list[dict]): batch[key] = sample[key] hidden = sample.get("hidden_states") if hidden is None: + # a2 old-logprob path: resolve Ray ObjectRefs (or TQ keys, P1) + # from per-owner chunks / single ref. hidden_chunks = sample.get("hidden_states_ref_chunks") if hidden_chunks: expected_rows = None @@ -838,7 +922,9 @@ def collect_rollout_features(self, samples: list[dict]): hidden_positions = cast(torch.Tensor, hidden_positions) expected_rows = int(hidden_positions.numel()) hidden = _resolve_hidden_state_chunks( - hidden_chunks, expected_rows=expected_rows + hidden_chunks, + expected_rows=expected_rows, + cache=self._tq_chunk_cache, ) else: hidden = _resolve_ray_object_ref(sample.get("hidden_states_ref"))