diff --git a/README.md b/README.md index 044aacf8..7d5304d7 100644 --- a/README.md +++ b/README.md @@ -300,19 +300,23 @@ actor_rollout_ref.rollout.drafter.speculative_algorithm=EAGLE3 ## Separate Draft Model Training -verl-SpeCo also supports a separate draft model training workflow. In this -mode, rollout workers collect drafter training features into a feature store, -and the draft model can be trained separately after feature collection. +verl-SpeCo also supports standalone DSpark draft-model training from a finite +verl-style prompt Parquet or prompt/response JSONL/Parquet file. For prompt-only +rows, a producer asks the target vLLM service to generate the response while +extracting prompt/output hidden states. It transfers each global batch through +TransferQueue, and a consumer trains the drafter independently of PPO. Quickstart: ```bash -bash examples/run_qwen3-8b_drafter_separate_training.sh +bash examples/run_qwen3-8b_drafter_dspark_separate_training.sh ``` -Replace the model, drafter, dataset, feature-store, and checkpoint paths in -the script before running it. The script uses `collect_only` mode for rollout -feature collection and `offline` mode for standalone drafter training. +Set the same model, dataset, drafter, checkpoint, GPU, and optimization values +used by ordinary standalone training near the top of the script. Transport +identity, Ray/TQ connection settings, and the Producer/Consumer lifecycle are +derived and managed internally. The target hidden-state vLLM service uses the +local port 8000 convention. The main mode values are: diff --git a/examples/run_qwen3-8b_drafter_dspark_separate_training.sh b/examples/run_qwen3-8b_drafter_dspark_separate_training.sh new file mode 100644 index 00000000..1a08a43b --- /dev/null +++ b/examples/run_qwen3-8b_drafter_dspark_separate_training.sh @@ -0,0 +1,165 @@ +#!/usr/bin/env bash +# Copyright 2026 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. +set -euo pipefail +set -x + +script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +repo_root=$(cd -- "${script_dir}/.." && pwd) +cd "${repo_root}" + +# Standalone DSpark draft-model training using an already-running hidden-state +# vLLM. Start tools/run_qwen3-8b_drafter_hidden_state_vllm.sh in another terminal +# first. This process owns Ray/TQ, Producer and Consumer, but it must not own +# the target vLLM so inference and training can use different accelerators. + +project_name=${PROJECT_NAME:-verl_dspark_drafter} +exp_name=${EXP_NAME:-qwen3_8b_dspark_separate_training} + +draft_train_gpus_per_node=${TRAIN_GPUS:-2} + +MODEL_PATH=${MODEL_PATH:-/path/to/Qwen3-8B} +# Ordinary verl prompt Parquet is supported; target vLLM generates responses. +TRAIN_FILE=${TRAIN_FILE:-/path/to/train_file.parquet} +# Optional. Leave empty to initialize DSpark from the target-model/config +# fallback; set it only when loading or resuming an existing drafter. +DRAFTER_PATH=${DRAFTER_PATH:-} +DRAFT_CKPTS_DIR=${DRAFT_CKPTS_DIR:-/path/to/dspark_draft_checkpoints} + +PYTHON_BIN=${PYTHON_BIN:-python3} +DEVICE_ENV=${DEVICE_ENV:-ASCEND_RT_VISIBLE_DEVICES} +TRAIN_DEVICES=${TRAIN_DEVICES:-2,3} +SPECO_VLLM_ENDPOINTS=${SPECO_VLLM_ENDPOINTS:-'[http://127.0.0.1:8000/v1,http://127.0.0.1:8001/v1]'} +VLLM_READY_TIMEOUT_SECONDS=${VLLM_READY_TIMEOUT_SECONDS:-120} + +# Producer -> vLLM concurrency and bounded queues. MAX_INFLIGHT_REQUESTS is the +# process-wide request limit; PER_ENDPOINT_CONCURRENCY applies independently to +# every URL in SPECO_VLLM_ENDPOINTS. +VLLM_REQUEST_TIMEOUT=${VLLM_REQUEST_TIMEOUT:-120} +VLLM_MAX_INFLIGHT_REQUESTS=${VLLM_MAX_INFLIGHT_REQUESTS:-16} +VLLM_PER_ENDPOINT_CONCURRENCY=${VLLM_PER_ENDPOINT_CONCURRENCY:-4} +PRODUCER_INPUT_QUEUE_SIZE=${PRODUCER_INPUT_QUEUE_SIZE:-32} +PRODUCER_PUBLISH_QUEUE_SIZE=${PRODUCER_PUBLISH_QUEUE_SIZE:-16} +PRODUCER_MAX_PENDING_SAMPLES=${PRODUCER_MAX_PENDING_SAMPLES:-1024} +PRODUCER_PENDING_POLL_INTERVAL=${PRODUCER_PENDING_POLL_INTERVAL:-0.5} +PRODUCER_MAX_SEQUENCE_LENGTH=${PRODUCER_MAX_SEQUENCE_LENGTH:-8192} +PRODUCER_MAX_FEATURE_LENGTH=${PRODUCER_MAX_FEATURE_LENGTH:-512} +PRODUCER_GENERATION_MAX_TOKENS=${PRODUCER_GENERATION_MAX_TOKENS:-512} + +# Standalone trainer. +MAX_STEPS=${MAX_STEPS:-10} +SAVE_INTERVAL_STEPS=${SAVE_INTERVAL_STEPS:-5} +SAVE_FINAL_CHECKPOINT=${SAVE_FINAL_CHECKPOINT:-true} +BATCH_SIZE_PER_GPU=${BATCH_SIZE_PER_GPU:-2} +LEARNING_RATE=${LEARNING_RATE:-1e-6} +LR_WARMUP_STEPS=${LR_WARMUP_STEPS:-0} +LR_SCHEDULER_TYPE=${LR_SCHEDULER_TYPE:-constant} +LR_DECAY_STEPS=${LR_DECAY_STEPS:-100} +MIN_LR_RATIO=${MIN_LR_RATIO:-0.1} +PARAM_OFFLOAD=${PARAM_OFFLOAD:-true} +OPTIMIZER_OFFLOAD=${OPTIMIZER_OFFLOAD:-true} + +# DSpark architecture, sampling and losses. TARGET_LAYER_IDS must match the +# auxiliary layers exposed by both hidden-state vLLM services. +DSPARK_BLOCK_SIZE=${DSPARK_BLOCK_SIZE:-7} +DSPARK_NUM_ANCHORS=${DSPARK_NUM_ANCHORS:-32} +DSPARK_MAX_WINDOW=${DSPARK_MAX_WINDOW:-512} +DSPARK_LOSS_MODE=${DSPARK_LOSS_MODE:-full_vocab} +DSPARK_SAMPLED_CE_NEGATIVES=${DSPARK_SAMPLED_CE_NEGATIVES:-0} +DSPARK_LOSS_DECAY_GAMMA=${DSPARK_LOSS_DECAY_GAMMA:-7} +DSPARK_NUM_TARGET_LAYERS=${DSPARK_NUM_TARGET_LAYERS:-5} +DSPARK_NUM_HIDDEN_LAYERS=${DSPARK_NUM_HIDDEN_LAYERS:-5} +DSPARK_TARGET_LAYER_IDS=${DSPARK_TARGET_LAYER_IDS:-'[1,9,17,25,33]'} +DSPARK_MARKOV_RANK=${DSPARK_MARKOV_RANK:-256} +DSPARK_MARKOV_HEAD_TYPE=${DSPARK_MARKOV_HEAD_TYPE:-vanilla} +DSPARK_CE_LOSS_ALPHA=${DSPARK_CE_LOSS_ALPHA:-0.1} +DSPARK_L1_LOSS_ALPHA=${DSPARK_L1_LOSS_ALPHA:-0.45} +DSPARK_L1_CHUNK_SIZE=${DSPARK_L1_CHUNK_SIZE:-0} +# The current DSpark trainer rejects nonzero confidence loss because target +# acceptance labels are not part of the standalone feature protocol yet. +DSPARK_CONFIDENCE_LOSS_ALPHA=${DSPARK_CONFIDENCE_LOSS_ALPHA:-0.0} +DSPARK_DEBUG_LOG=${DSPARK_DEBUG_LOG:-false} +DSPARK_DEBUG_LOG_FIRST_N=${DSPARK_DEBUG_LOG_FIRST_N:-2} +DSPARK_DEBUG_LOG_INTERVAL=${DSPARK_DEBUG_LOG_INTERVAL:-100} + +export "${DEVICE_ENV}=${TRAIN_DEVICES}" +export SPECO_VLLM_ENDPOINTS + +# Fail before entering the unified launcher when the separately managed vLLM +# is absent. Otherwise a localhost endpoint would make the launcher start its +# fallback vLLM inside the training process and on the training devices. +if ! "${PYTHON_BIN}" tools/wait_for_vllm_endpoints.py \ + --endpoints "${SPECO_VLLM_ENDPOINTS}" \ + --timeout-seconds "${VLLM_READY_TIMEOUT_SECONDS}"; then + echo "Start tools/run_qwen3-8b_drafter_hidden_state_vllm.sh first" >&2 + exit 1 +fi + +PYTHONUNBUFFERED=1 "${PYTHON_BIN}" -m verl_speco.standalone_tq_training_launcher \ + speco.draft_training.num_gpus_per_node=${draft_train_gpus_per_node} \ + speco.draft_training.nnodes=1 \ + speco.draft_training.standalone=True \ + data.train_files=${TRAIN_FILE} \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.fsdp_config.param_offload=${PARAM_OFFLOAD} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${OPTIMIZER_OFFLOAD} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.drafter.enable=True \ + actor_rollout_ref.rollout.drafter.enable_drafter_training=True \ + actor_rollout_ref.rollout.drafter.model_path=${DRAFTER_PATH} \ + actor_rollout_ref.rollout.drafter.checkpoint_path=${DRAFT_CKPTS_DIR} \ + actor_rollout_ref.rollout.drafter.speculative_algorithm=DSPARK \ + actor_rollout_ref.rollout.drafter.training.mode=offline \ + actor_rollout_ref.rollout.drafter.training.max_steps=${MAX_STEPS} \ + actor_rollout_ref.rollout.drafter.training.save_interval_steps=${SAVE_INTERVAL_STEPS} \ + actor_rollout_ref.rollout.drafter.training.save_final_checkpoint=${SAVE_FINAL_CHECKPOINT} \ + actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${BATCH_SIZE_PER_GPU} \ + actor_rollout_ref.rollout.drafter.training.lr=${LEARNING_RATE} \ + actor_rollout_ref.rollout.drafter.training.lr_warmup_steps=${LR_WARMUP_STEPS} \ + actor_rollout_ref.rollout.drafter.training.lr_scheduler_type=${LR_SCHEDULER_TYPE} \ + actor_rollout_ref.rollout.drafter.training.lr_decay_steps=${LR_DECAY_STEPS} \ + actor_rollout_ref.rollout.drafter.training.min_lr_ratio=${MIN_LR_RATIO} \ + actor_rollout_ref.rollout.drafter.training.use_logits=False \ + actor_rollout_ref.rollout.drafter.training.dspark_block_size=${DSPARK_BLOCK_SIZE} \ + actor_rollout_ref.rollout.drafter.training.dspark_num_anchors=${DSPARK_NUM_ANCHORS} \ + actor_rollout_ref.rollout.drafter.training.dspark_max_window=${DSPARK_MAX_WINDOW} \ + actor_rollout_ref.rollout.drafter.training.dspark_loss_mode=${DSPARK_LOSS_MODE} \ + actor_rollout_ref.rollout.drafter.training.dspark_sampled_ce_negatives=${DSPARK_SAMPLED_CE_NEGATIVES} \ + actor_rollout_ref.rollout.drafter.training.dspark_loss_decay_gamma=${DSPARK_LOSS_DECAY_GAMMA} \ + actor_rollout_ref.rollout.drafter.training.dspark_num_target_layers=${DSPARK_NUM_TARGET_LAYERS} \ + actor_rollout_ref.rollout.drafter.training.dspark_num_hidden_layers=${DSPARK_NUM_HIDDEN_LAYERS} \ + actor_rollout_ref.rollout.drafter.training.dspark_target_layer_ids=${DSPARK_TARGET_LAYER_IDS} \ + actor_rollout_ref.rollout.drafter.training.dspark_markov_rank=${DSPARK_MARKOV_RANK} \ + actor_rollout_ref.rollout.drafter.training.dspark_markov_head_type=${DSPARK_MARKOV_HEAD_TYPE} \ + actor_rollout_ref.rollout.drafter.training.dspark_ce_loss_alpha=${DSPARK_CE_LOSS_ALPHA} \ + actor_rollout_ref.rollout.drafter.training.dspark_l1_loss_alpha=${DSPARK_L1_LOSS_ALPHA} \ + actor_rollout_ref.rollout.drafter.training.dspark_l1_chunk_size=${DSPARK_L1_CHUNK_SIZE} \ + actor_rollout_ref.rollout.drafter.training.dspark_confidence_loss_alpha=${DSPARK_CONFIDENCE_LOSS_ALPHA} \ + actor_rollout_ref.rollout.drafter.training.dspark_debug_log=${DSPARK_DEBUG_LOG} \ + actor_rollout_ref.rollout.drafter.training.dspark_debug_log_first_n=${DSPARK_DEBUG_LOG_FIRST_N} \ + actor_rollout_ref.rollout.drafter.training.dspark_debug_log_interval=${DSPARK_DEBUG_LOG_INTERVAL} \ + speco.standalone_tq_producer.request_timeout=${VLLM_REQUEST_TIMEOUT} \ + speco.standalone_tq_producer.max_inflight_requests=${VLLM_MAX_INFLIGHT_REQUESTS} \ + speco.standalone_tq_producer.per_endpoint_concurrency=${VLLM_PER_ENDPOINT_CONCURRENCY} \ + speco.standalone_tq_producer.input_queue_size=${PRODUCER_INPUT_QUEUE_SIZE} \ + speco.standalone_tq_producer.publish_queue_size=${PRODUCER_PUBLISH_QUEUE_SIZE} \ + speco.standalone_tq_producer.max_pending_samples=${PRODUCER_MAX_PENDING_SAMPLES} \ + speco.standalone_tq_producer.pending_poll_interval_seconds=${PRODUCER_PENDING_POLL_INTERVAL} \ + speco.standalone_tq_producer.max_sequence_length=${PRODUCER_MAX_SEQUENCE_LENGTH} \ + speco.standalone_tq_producer.max_feature_length=${PRODUCER_MAX_FEATURE_LENGTH} \ + speco.standalone_tq_producer.generation_max_tokens=${PRODUCER_GENERATION_MAX_TOKENS} \ + trainer.project_name=${project_name} \ + trainer.experiment_name=${exp_name} \ + "$@" diff --git a/examples/run_qwen3-8b_drafter_eagle3_separate_training.sh b/examples/run_qwen3-8b_drafter_eagle3_separate_training.sh new file mode 100644 index 00000000..83c70fac --- /dev/null +++ b/examples/run_qwen3-8b_drafter_eagle3_separate_training.sh @@ -0,0 +1,164 @@ +#!/usr/bin/env bash +# Copyright 2026 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. +set -euo pipefail +set -x + +# Standalone EAGLE3 drafter training. Start +# examples/run_qwen3-8b_drafter_hidden_state_vllm.sh in another terminal first. +# The target-model vLLM and the drafter trainer can therefore use disjoint GPUs. +# +# EAGLE3 can initialize its drafter structure from the target model config, so +# no pre-initialized drafter directory is required. Set model_path only when +# loading an existing drafter checkpoint/config is desired. +# +# The vLLM hidden-state layer IDs must be EAGLE3_TARGET_LAYER_IDS followed by +# the target model's final layer. For Qwen3-8B the default is: +# [1,9,17,25,33,36] +# The EAGLE3 drafter config must have the same number (five) of aux states. + +project_name=${PROJECT_NAME:-verl_eagle3_drafter} +exp_name=${EXP_NAME:-qwen3_8b_eagle3_separate_training} + +draft_train_gpus_per_node=${TRAIN_GPUS:-2} +MODEL_PATH=${MODEL_PATH:-/path/to/Qwen3-4B} +TRAIN_FILE=${TRAIN_FILE:-/path/to/data} +DRAFT_CKPTS_DIR=${DRAFT_CKPTS_DIR:-/path/to/ckpt} + +PYTHON_BIN=${PYTHON_BIN:-python3} +DEVICE_ENV=${DEVICE_ENV:-CUDA_VISIBLE_DEVICES} +TRAIN_DEVICES=${TRAIN_DEVICES:-6,7} +SPECO_VLLM_ENDPOINTS=${SPECO_VLLM_ENDPOINTS:-'[http://127.0.0.1:8000/v1]'} +VLLM_READY_TIMEOUT_SECONDS=${VLLM_READY_TIMEOUT_SECONDS:-120} + +# These IDs must equal the auxiliary prefix of VLLM_HIDDEN_STATE_LAYER_IDS in +# run_qwen3-8b_drafter_hidden_state_vllm.sh. Do not include the final layer. +EAGLE3_TARGET_LAYER_IDS=${EAGLE3_TARGET_LAYER_IDS:-'[1,9,17,25,33]'} + +# Producer throughput and bounded queues. +VLLM_REQUEST_TIMEOUT=${VLLM_REQUEST_TIMEOUT:-120} +VLLM_MAX_INFLIGHT_REQUESTS=${VLLM_MAX_INFLIGHT_REQUESTS:-16} +VLLM_PER_ENDPOINT_CONCURRENCY=${VLLM_PER_ENDPOINT_CONCURRENCY:-4} +PRODUCER_INPUT_QUEUE_SIZE=${PRODUCER_INPUT_QUEUE_SIZE:-32} +PRODUCER_PUBLISH_QUEUE_SIZE=${PRODUCER_PUBLISH_QUEUE_SIZE:-16} +PRODUCER_MAX_PENDING_SAMPLES=${PRODUCER_MAX_PENDING_SAMPLES:-1024} +PRODUCER_PENDING_POLL_INTERVAL=${PRODUCER_PENDING_POLL_INTERVAL:-0.5} +PRODUCER_MAX_SEQUENCE_LENGTH=${PRODUCER_MAX_SEQUENCE_LENGTH:-8192} +PRODUCER_MAX_FEATURE_LENGTH=${PRODUCER_MAX_FEATURE_LENGTH:-512} +PRODUCER_GENERATION_MAX_TOKENS=${PRODUCER_GENERATION_MAX_TOKENS:-511} + +# Standalone EAGLE3 trainer settings. +MAX_STEPS=${MAX_STEPS:-1000} +SAVE_INTERVAL_STEPS=${SAVE_INTERVAL_STEPS:-100} +SAVE_FINAL_CHECKPOINT=${SAVE_FINAL_CHECKPOINT:-true} +BATCH_SIZE_PER_GPU=${BATCH_SIZE_PER_GPU:-2} +LEARNING_RATE=${LEARNING_RATE:-1e-5} +LR_WARMUP_STEPS=${LR_WARMUP_STEPS:-0} +LR_SCHEDULER_TYPE=${LR_SCHEDULER_TYPE:-constant} +LR_DECAY_STEPS=${LR_DECAY_STEPS:-1000} +MIN_LR_RATIO=${MIN_LR_RATIO:-0.1} +PARAM_OFFLOAD=${PARAM_OFFLOAD:-true} +OPTIMIZER_OFFLOAD=${OPTIMIZER_OFFLOAD:-true} + +if [[ "${MODEL_PATH}" == /path/to/* || "${TRAIN_FILE}" == /path/to/* ]]; then + echo "Set MODEL_PATH and TRAIN_FILE before starting training." >&2 + exit 2 +fi + +export "${DEVICE_ENV}=${TRAIN_DEVICES}" +export SPECO_VLLM_ENDPOINTS + +# Avoid the launcher's localhost fallback vLLM: this job must consume the +# separately managed hidden-state services, which keep target inference off the +# training devices. +"${PYTHON_BIN}" - "${SPECO_VLLM_ENDPOINTS}" "${VLLM_READY_TIMEOUT_SECONDS}" <<'PY' +import sys +import time +from urllib.error import URLError +from urllib.request import urlopen + +raw_endpoints = sys.argv[1].strip() +if not (raw_endpoints.startswith("[") and raw_endpoints.endswith("]")): + raise SystemExit("SPECO_VLLM_ENDPOINTS must use [url0,url1] syntax") +endpoints = [ + item.strip().strip("'\"").rstrip("/") + for item in raw_endpoints[1:-1].split(",") + if item.strip() +] +if not endpoints: + raise SystemExit("SPECO_VLLM_ENDPOINTS must contain at least one URL") +deadline = time.monotonic() + float(sys.argv[2]) +pending = set(endpoints) +while pending: + for endpoint in list(pending): + try: + with urlopen(f"{endpoint}/models", timeout=2) as response: + if 200 <= response.status < 300: + print(f"EXTERNAL_VLLM_READY endpoint={endpoint}", flush=True) + pending.remove(endpoint) + except (OSError, URLError): + pass + if pending and time.monotonic() >= deadline: + raise SystemExit( + "external hidden-state vLLM is not ready at: " + + ", ".join(sorted(pending)) + + "; start examples/run_qwen3-8b_drafter_hidden_state_vllm.sh first" + ) + if pending: + time.sleep(1) +PY + +PYTHONUNBUFFERED=1 "${PYTHON_BIN}" -m verl_speco.standalone_tq_training_launcher \ + speco.draft_training.num_gpus_per_node=${draft_train_gpus_per_node} \ + speco.draft_training.nnodes=1 \ + speco.draft_training.standalone=True \ + data.train_files=${TRAIN_FILE} \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.fsdp_config.param_offload=${PARAM_OFFLOAD} \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=${OPTIMIZER_OFFLOAD} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.drafter.enable=true \ + actor_rollout_ref.rollout.drafter.enable_drafter_training=true \ + actor_rollout_ref.rollout.drafter.checkpoint_path=${DRAFT_CKPTS_DIR} \ + actor_rollout_ref.rollout.drafter.speculative_algorithm=EAGLE3 \ + actor_rollout_ref.rollout.drafter.rollout.spec_steps=3 \ + actor_rollout_ref.rollout.drafter.rollout.spec_topk=1 \ + actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=4 \ + actor_rollout_ref.rollout.drafter.training.mode=offline \ + actor_rollout_ref.rollout.drafter.training.max_steps=${MAX_STEPS} \ + actor_rollout_ref.rollout.drafter.training.save_interval_steps=${SAVE_INTERVAL_STEPS} \ + actor_rollout_ref.rollout.drafter.training.save_final_checkpoint=${SAVE_FINAL_CHECKPOINT} \ + actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${BATCH_SIZE_PER_GPU} \ + actor_rollout_ref.rollout.drafter.training.lr=${LEARNING_RATE} \ + actor_rollout_ref.rollout.drafter.training.lr_warmup_steps=${LR_WARMUP_STEPS} \ + actor_rollout_ref.rollout.drafter.training.lr_scheduler_type=${LR_SCHEDULER_TYPE} \ + actor_rollout_ref.rollout.drafter.training.lr_decay_steps=${LR_DECAY_STEPS} \ + actor_rollout_ref.rollout.drafter.training.min_lr_ratio=${MIN_LR_RATIO} \ + actor_rollout_ref.rollout.drafter.training.use_logits=false \ + actor_rollout_ref.rollout.drafter.training.eagle3_target_layer_ids=${EAGLE3_TARGET_LAYER_IDS} \ + speco.standalone_tq_producer.target_layer_ids=${EAGLE3_TARGET_LAYER_IDS} \ + speco.standalone_tq_producer.request_timeout=${VLLM_REQUEST_TIMEOUT} \ + speco.standalone_tq_producer.max_inflight_requests=${VLLM_MAX_INFLIGHT_REQUESTS} \ + speco.standalone_tq_producer.per_endpoint_concurrency=${VLLM_PER_ENDPOINT_CONCURRENCY} \ + speco.standalone_tq_producer.input_queue_size=${PRODUCER_INPUT_QUEUE_SIZE} \ + speco.standalone_tq_producer.publish_queue_size=${PRODUCER_PUBLISH_QUEUE_SIZE} \ + speco.standalone_tq_producer.max_pending_samples=${PRODUCER_MAX_PENDING_SAMPLES} \ + speco.standalone_tq_producer.pending_poll_interval_seconds=${PRODUCER_PENDING_POLL_INTERVAL} \ + speco.standalone_tq_producer.max_sequence_length=${PRODUCER_MAX_SEQUENCE_LENGTH} \ + speco.standalone_tq_producer.max_feature_length=${PRODUCER_MAX_FEATURE_LENGTH} \ + speco.standalone_tq_producer.generation_max_tokens=${PRODUCER_GENERATION_MAX_TOKENS} \ + trainer.project_name=${project_name} \ + trainer.experiment_name=${exp_name} \ + "$@" diff --git a/pyproject.toml b/pyproject.toml index ca8a9c7b..72672830 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,9 @@ dependencies = [ "packaging>=24", ] +[project.optional-dependencies] +transfer-queue = ["TransferQueue==0.1.10"] + [project.urls] Repository = "https://github.com/verl-project/verl-SpeCo" @@ -23,6 +26,8 @@ Repository = "https://github.com/verl-project/verl-SpeCo" verl-speco = "verl_speco.main:main" verl-speco-draft-train = "verl_speco.draft_train_launcher:main" verl-speco-inspect-features = "verl_speco.inspect_feature_store:main" +verl-speco-tq-owner = "verl_speco.tq_owner:main" +verl-speco-tq-producer = "verl_speco.standalone_tq_producer:main" [tool.setuptools.dynamic] version = { attr = "verl_speco.__version__" } diff --git a/tests/config/test_speco_config_overlay.py b/tests/config/test_speco_config_overlay.py index 749b62fc..2b29e464 100644 --- a/tests/config/test_speco_config_overlay.py +++ b/tests/config/test_speco_config_overlay.py @@ -60,6 +60,8 @@ def _copy_overlay_configs( def test_overlay_has_expected_default_drafter_shape() -> None: raw = OmegaConf.load(CONFIG_DIR / "speco_base.yaml") drafter = raw.actor_rollout_ref.rollout.drafter + standalone = OmegaConf.load(CONFIG_DIR / "draft_trainer.yaml") + standalone_training = standalone.actor_rollout_ref.rollout.drafter.training assert raw.speco.verl_base.version == "0.8.0" assert raw.speco.verl_base.branch == "release/v0.8.0" @@ -76,6 +78,17 @@ def test_overlay_has_expected_default_drafter_shape() -> None: assert drafter.training.warmup_style is None assert drafter.training.resume_trainer_state_from_checkpoint is True assert drafter.training.eagle1_num_hidden_layers == 1 + assert drafter.training.mode == "online" + assert drafter.training.feature_store.type == "torch_shard" + assert "target_feature_replay" not in drafter.training + assert standalone_training.target_feature_replay.cache.enabled is False + assert standalone_training.target_feature_replay.cache.max_size_gb == 0 + assert standalone_training.target_feature_replay.vllm_endpoints is None + assert standalone_training.target_feature_replay.endpoint_cooldown == 5 + assert standalone_training.target_feature_pipeline.enabled is False + assert standalone_training.target_feature_pipeline.concurrency == 16 + assert standalone_training.target_feature_pipeline.producer_prefetch_depth == 4 + assert standalone_training.target_feature_pipeline.prefetch_depth == 2 def test_overlay_composes_with_release_upstream_verl(tmp_path: Path) -> None: @@ -119,6 +132,13 @@ def test_draft_trainer_composes_as_primary_config(tmp_path: Path) -> None: config = compose(config_name="draft_trainer") assert config.actor_rollout_ref.rollout.drafter.training.mode == "offline" + assert config.actor_rollout_ref.rollout.drafter.training.feature_store.type == ( + "torch_shard" + ) + assert ( + config.actor_rollout_ref.rollout.drafter.training.target_feature_replay.cache.enabled + is False + ) assert config.speco.draft_training.enable is True assert "trainer" in config assert "algorithm" in config diff --git a/tests/examples/test_example_scripts.py b/tests/examples/test_example_scripts.py index dbb70a06..2b23b5e0 100644 --- a/tests/examples/test_example_scripts.py +++ b/tests/examples/test_example_scripts.py @@ -22,6 +22,11 @@ ROOT = Path(__file__).resolve().parents[2] EXAMPLES = sorted((ROOT / "examples").glob("*.sh")) +PPO_EXAMPLES = [ + script + for script in EXAMPLES + if not script.name.endswith("_separate_training.sh") +] def _require_working_bash() -> str: @@ -40,7 +45,7 @@ def test_example_shell_syntax_is_valid(script: Path) -> None: subprocess.run([bash, "-n", str(script)], check=True) -@pytest.mark.parametrize("script", EXAMPLES, ids=lambda path: path.name) +@pytest.mark.parametrize("script", PPO_EXAMPLES, ids=lambda path: path.name) def test_example_keeps_speco_entrypoint_and_required_drafter_switches( script: Path, ) -> None: @@ -62,6 +67,37 @@ def test_example_keeps_speco_entrypoint_and_required_drafter_switches( assert "actor_rollout_ref.rollout.drafter.training.publish_async=" in source +def test_standalone_tq_training_example_uses_unified_launcher() -> None: + source = ( + ROOT / "examples" / "run_qwen3-8b_drafter_dspark_separate_training.sh" + ).read_text(encoding="utf-8") + + assert "-m verl_speco.standalone_tq_training_launcher" in source + assert "data.train_files=${TRAIN_FILE}" in source + assert "actor_rollout_ref.rollout.drafter.enable=True" in source + assert "actor_rollout_ref.rollout.drafter.enable_drafter_training=True" in source + assert "actor_rollout_ref.rollout.drafter.model_path=${DRAFTER_PATH}" in source + assert "actor_rollout_ref.rollout.drafter.speculative_algorithm=DSPARK" in source + assert "speco.standalone_tq_producer.max_inflight_requests=" in source + assert "speco.standalone_tq_producer.per_endpoint_concurrency=" in source + assert "actor_rollout_ref.rollout.drafter.training.dspark_ce_loss_alpha=" in source + assert "actor_rollout_ref.rollout.drafter.training.dspark_l1_loss_alpha=" in source + + +def test_standalone_tq_hidden_state_vllm_uses_separate_devices() -> None: + source = ( + ROOT / "tools" / "run_qwen3-8b_drafter_hidden_state_vllm.sh" + ).read_text(encoding="utf-8") + + assert 'VLLM_DEVICES=${VLLM_DEVICES:-0,1,2,3,4,5}' in source + assert "service_count=$((device_count / VLLM_TP))" in source + assert 'env "${DEVICE_ENV}=${devices}" vllm serve "${MODEL_PATH}"' in source + assert '--tensor-parallel-size "${VLLM_TP}"' in source + assert '--max-num-seqs "${VLLM_MAX_NUM_SEQS}"' in source + assert '"${VLLM_HIDDEN_STATE_LAYER_IDS}"' in source + assert '"kv_connector":"ExampleHiddenStatesConnector"' in source + + def test_vllm_eagle3_example_keeps_runtime_agnostic_training_switches() -> None: source = (ROOT / "examples" / "run_qwen3-8b_drafter_eagle3_vllm.sh").read_text( encoding="utf-8" diff --git a/tests/integration/test_drafter_lr_scheduler.py b/tests/integration/test_drafter_lr_scheduler.py index d5272e1a..b7d7a83f 100644 --- a/tests/integration/test_drafter_lr_scheduler.py +++ b/tests/integration/test_drafter_lr_scheduler.py @@ -21,6 +21,7 @@ from verl_speco.backends.lr_scheduler import ( # noqa: E402 ClampedGlobalCosineLR, + LinearWarmupDecayLR, build_drafter_lr_scheduler, ) @@ -92,6 +93,55 @@ def test_scheduler_builder_uses_configured_global_cosine_values() -> None: assert optimizer.param_groups[0]["lr"] == pytest.approx(5e-6) +def test_linear_warmup_decay_reaches_zero_after_decay_steps() -> None: + optimizer = _optimizer() + scheduler = LinearWarmupDecayLR( + optimizer, + decay_steps=100, + min_lr_ratio=0.0, + warmup_steps=10, + ) + + assert optimizer.param_groups[0]["lr"] == pytest.approx(0.0) + + _step(optimizer, scheduler, 5) + assert optimizer.param_groups[0]["lr"] == pytest.approx(5e-6) + + _step(optimizer, scheduler, 5) + assert optimizer.param_groups[0]["lr"] == pytest.approx(1e-5) + + _step(optimizer, scheduler, 45) + assert optimizer.param_groups[0]["lr"] == pytest.approx(5e-6) + + _step(optimizer, scheduler, 45) + assert optimizer.param_groups[0]["lr"] == pytest.approx(0.0) + + _step(optimizer, scheduler, 10) + assert optimizer.param_groups[0]["lr"] == pytest.approx(0.0) + + +def test_scheduler_builder_uses_linear_warmup_decay_values() -> None: + optimizer = _optimizer(lr=2e-5) + scheduler = build_drafter_lr_scheduler( + optimizer, + { + "lr_scheduler_type": "linear", + "lr_decay_steps": 100, + "min_lr_ratio": 0.1, + "lr_warmup_steps": 10, + }, + ) + + _step(optimizer, scheduler, 10) + assert optimizer.param_groups[0]["lr"] == pytest.approx(2e-5) + + _step(optimizer, scheduler, 45) + assert optimizer.param_groups[0]["lr"] == pytest.approx(1.1e-5) + + _step(optimizer, scheduler, 45) + assert optimizer.param_groups[0]["lr"] == pytest.approx(2e-6) + + def test_scheduler_builder_resumes_from_successful_optimizer_steps() -> None: optimizer = _optimizer() scheduler = build_drafter_lr_scheduler( @@ -114,6 +164,27 @@ def test_scheduler_builder_resumes_from_successful_optimizer_steps() -> None: assert optimizer.param_groups[0]["lr"] == pytest.approx(1e-5 * expected_ratio) +def test_linear_scheduler_builder_resumes_from_successful_optimizer_steps() -> None: + optimizer = _optimizer() + scheduler = build_drafter_lr_scheduler( + optimizer, + { + "lr_scheduler_type": "linear", + "lr_decay_steps": 100, + "min_lr_ratio": 0.0, + "lr_warmup_steps": 10, + "_resume_optimizer_steps": 55, + }, + ) + + assert scheduler.last_epoch == 55 + assert optimizer.param_groups[0]["lr"] == pytest.approx(5e-6) + + _step(optimizer, scheduler, 1) + assert scheduler.last_epoch == 56 + assert optimizer.param_groups[0]["lr"] == pytest.approx(4.888888888888889e-6) + + def test_scheduler_builder_does_not_replace_explicit_invalid_decay() -> None: with pytest.raises(ValueError, match="lr_decay_steps"): build_drafter_lr_scheduler( @@ -136,3 +207,16 @@ def test_scheduler_builder_does_not_replace_explicit_invalid_decay() -> None: def test_clamped_global_cosine_rejects_invalid_config(kwargs, message) -> None: with pytest.raises(ValueError, match=message): ClampedGlobalCosineLR(_optimizer(), **kwargs) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"decay_steps": 0}, "lr_decay_steps"), + ({"decay_steps": 10, "warmup_steps": 10}, "lr_warmup_steps"), + ({"min_lr_ratio": -0.1}, "min_lr_ratio"), + ], +) +def test_linear_warmup_decay_rejects_invalid_config(kwargs, message) -> None: + with pytest.raises(ValueError, match=message): + LinearWarmupDecayLR(_optimizer(), **kwargs) diff --git a/tests/integration/test_eagle3_aux_hidden_contract.py b/tests/integration/test_eagle3_aux_hidden_contract.py index 3a3f3b18..f6c8474c 100644 --- a/tests/integration/test_eagle3_aux_hidden_contract.py +++ b/tests/integration/test_eagle3_aux_hidden_contract.py @@ -84,3 +84,21 @@ def test_eagle3_model_uses_dynamic_aux_hidden_count() -> None: with pytest.raises(ValueError, match="num_aux_hidden_states=5"): model.project_hidden_states(torch.randn(2, 3, 12)) + + +def test_eagle3_model_defaults_missing_pretraining_tp() -> None: + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.auto import AutoDraftModelConfig + from verl_speco.models.eagle.llama_eagle import LlamaMLP + + raw_config = _minimal_eagle3_config() + raw_config.pop("pretraining_tp") + config = AutoDraftModelConfig._config_mapping["LlamaForCausalLMEagle3"].from_dict( + raw_config + ) + mlp = LlamaMLP(config) + + output = mlp(torch.randn(2, 3, config.hidden_size)) + + assert output.shape == (2, 3, config.hidden_size) diff --git a/tests/special_sanity/test_check_example_naming.py b/tests/special_sanity/test_check_example_naming.py index cac73db5..81708416 100644 --- a/tests/special_sanity/test_check_example_naming.py +++ b/tests/special_sanity/test_check_example_naming.py @@ -69,7 +69,7 @@ def test_missing_actor_backend_rejected(): def test_separate_training_entrypoint_passes(): - assert _violations("run_qwen3-8b_drafter_separate_training.sh") == [] + assert _violations("run_qwen3-8b_drafter_dspark_separate_training.sh") == [] def test_separate_training_may_name_its_drafter_backends(): diff --git a/tests/unit/test_draft_feature_store.py b/tests/unit/test_draft_feature_store.py index d75549f9..272780ca 100644 --- a/tests/unit/test_draft_feature_store.py +++ b/tests/unit/test_draft_feature_store.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import importlib +import json import pytest @@ -22,7 +23,12 @@ DraftFeatureDataLoader = draft_dataset.DraftFeatureDataLoader DraftFeatureDataLoaderConfig = draft_dataset.DraftFeatureDataLoaderConfig DraftFeatureSample = feature_store.DraftFeatureSample +DraftReplaySample = feature_store.DraftReplaySample +JsonlTokenReplayFeatureStore = feature_store.JsonlTokenReplayFeatureStore +TokenReplayFeatureStore = feature_store.TokenReplayFeatureStore TorchShardFeatureStore = feature_store.TorchShardFeatureStore +VllmSafetensorsFeatureStore = feature_store.VllmSafetensorsFeatureStore +build_feature_store_from_config = feature_store.build_feature_store_from_config def _sample(index: int = 0): @@ -68,6 +74,272 @@ def test_torch_shard_feature_store_roundtrip(tmp_path): assert reader.get_metadata()["num_samples"] == 2 +def test_token_replay_feature_store_roundtrip(tmp_path): + sample = DraftReplaySample( + algorithm="DSPARK", + input_ids=torch.arange(12, dtype=torch.long), + loss_mask=torch.ones(12, dtype=torch.float32), + attention_mask=torch.ones(12, dtype=torch.bool), + position_ids=torch.arange(12, dtype=torch.long), + feature_positions=torch.arange(4, 10, dtype=torch.long), + draft_position_ids=torch.arange(5, 11, dtype=torch.long), + metadata={"target_model_path": "/target", "global_step": 3}, + ) + store = TokenReplayFeatureStore(tmp_path, max_samples_per_shard=1) + store.write_many([sample]) + store.close() + + reader = TokenReplayFeatureStore(tmp_path, read_only=True) + loaded = reader.read(next(reader.iter_keys())) + + assert loaded.algorithm == "DSPARK" + assert loaded.input_ids.dtype == torch.int32 + assert loaded.attention_mask.dtype == torch.bool + assert torch.equal(loaded.feature_positions, torch.arange(4, 10, dtype=torch.int32)) + assert "hidden_states" not in loaded.to_dict() + assert reader.get_metadata()["format"] == "token_replay" + + +def test_token_replay_rejects_non_contiguous_feature_positions(): + sample = DraftReplaySample( + input_ids=torch.arange(8), + loss_mask=torch.ones(8), + attention_mask=torch.ones(8, dtype=torch.bool), + position_ids=torch.arange(8), + feature_positions=torch.tensor([2, 4]), + draft_position_ids=torch.tensor([3, 5]), + ) + + with pytest.raises(ValueError, match="contiguous"): + sample.validate(strict=True) + + +def test_jsonl_token_replay_reads_input_ids_and_loss_mask(tmp_path): + path = tmp_path / "samples.jsonl" + row = { + "id": "sample-0", + "input_ids": list(range(10)), + "loss_mask": [0, 0, 0, 0, 1, 1, 1, 0, 0, 0], + "text": "ignored for replay", + "metadata": {"source": "unit"}, + } + path.write_text(json.dumps(row) + "\n", encoding="utf-8") + + store = JsonlTokenReplayFeatureStore(path, read_only=True, max_seq_len=4) + keys = list(store.iter_keys(shuffle=False)) + loaded = store.read(keys[0]) + + assert keys == ["samples.jsonl:0"] + assert loaded.algorithm == "EAGLE3" + assert torch.equal(loaded.input_ids, torch.arange(10)) + assert torch.equal(loaded.loss_mask, torch.tensor(row["loss_mask"], dtype=torch.float32)) + assert torch.equal(loaded.attention_mask, torch.ones(10, dtype=torch.bool)) + assert torch.equal(loaded.position_ids, torch.arange(10)) + assert torch.equal(loaded.feature_positions, torch.arange(3, 7)) + assert torch.equal(loaded.draft_position_ids, torch.arange(4, 8)) + assert loaded.metadata["source"] == "jsonl_token_replay" + assert loaded.metadata["id"] == "sample-0" + assert store.get_metadata()["format"] == "jsonl_token_replay" + assert store.get_metadata()["num_samples"] == 1 + + +def test_build_feature_store_from_config_supports_jsonl_token_replay(tmp_path): + path = tmp_path / "samples.jsonl" + path.write_text( + json.dumps({"input_ids": [1, 2, 3], "loss_mask": [0, 1, 1]}) + "\n", + encoding="utf-8", + ) + + store = build_feature_store_from_config( + { + "type": "jsonl_token_replay", + "path": path, + "max_seq_len": 8, + }, + read_only=True, + ) + loaded = store.read(next(store.iter_keys(shuffle=False))) + + assert isinstance(loaded, DraftReplaySample) + assert torch.equal(loaded.feature_positions, torch.arange(0, 3)) + + +def test_jsonl_token_replay_reads_conversations_with_chat_template(tmp_path): + class FakeTokenizer: + def apply_chat_template( + self, messages, *, tokenize, add_generation_prompt + ): + assert tokenize is True + token_ids = [] + for message in messages: + role_id = {"user": 10, "assistant": 20, "system": 30}[message["role"]] + token_ids.extend([role_id, len(message["content"])]) + if add_generation_prompt: + token_ids.append(20) + return token_ids + + path = tmp_path / "samples.jsonl" + row = { + "id": "conv-0", + "conversations": [ + {"from": "human", "value": "question"}, + {"from": "assistant", "value": "answer"}, + ], + "algorithm": "DSPARK", + } + path.write_text(json.dumps(row) + "\n", encoding="utf-8") + + store = JsonlTokenReplayFeatureStore( + path, + read_only=True, + max_seq_len=8, + tokenizer_path="/target", + ) + store._tokenizer = FakeTokenizer() + loaded = store.read(next(store.iter_keys(shuffle=False))) + + assert loaded.algorithm == "DSPARK" + assert torch.equal(loaded.input_ids, torch.tensor([10, 8, 20, 6])) + assert torch.equal(loaded.loss_mask, torch.tensor([0.0, 0.0, 0.0, 1.0])) + assert torch.equal(loaded.feature_positions, torch.arange(2, 4)) + assert torch.equal(loaded.draft_position_ids, torch.arange(3, 5)) + assert loaded.metadata["source"] == "jsonl_conversations" + assert loaded.metadata["id"] == "conv-0" + + +def test_jsonl_token_replay_accepts_tensor_chat_template_output(tmp_path): + class TensorTokenizer: + def apply_chat_template( + self, messages, *, tokenize, add_generation_prompt + ): + token_ids = [] + for message in messages: + token_ids.extend([1 if message["role"] == "user" else 2, 3]) + if add_generation_prompt: + token_ids.append(2) + return torch.tensor([token_ids], dtype=torch.long) + + path = tmp_path / "samples.jsonl" + row = { + "conversations": [ + {"from": "human", "value": "question"}, + {"from": "assistant", "value": "answer"}, + ], + } + path.write_text(json.dumps(row) + "\n", encoding="utf-8") + + store = JsonlTokenReplayFeatureStore( + path, + read_only=True, + tokenizer_path="/target", + ) + store._tokenizer = TensorTokenizer() + loaded = store.read(next(store.iter_keys(shuffle=False))) + + assert torch.equal(loaded.input_ids, torch.tensor([1, 3, 2, 3])) + assert torch.equal(loaded.loss_mask, torch.tensor([0.0, 0.0, 0.0, 1.0])) + + +def test_jsonl_token_replay_accepts_batch_encoding_chat_template_output(tmp_path): + class FakeBatchEncoding: + def __init__(self, input_ids): + self.data = {"input_ids": input_ids} + + class BatchEncodingTokenizer: + def apply_chat_template( + self, messages, *, tokenize, add_generation_prompt + ): + token_ids = [] + for message in messages: + token_ids.extend([1 if message["role"] == "user" else 2, 3]) + if add_generation_prompt: + token_ids.append(2) + return FakeBatchEncoding([token_ids]) + + path = tmp_path / "samples.jsonl" + row = { + "conversations": [ + {"from": "human", "value": "question"}, + {"from": "assistant", "value": "answer"}, + ], + } + path.write_text(json.dumps(row) + "\n", encoding="utf-8") + + store = JsonlTokenReplayFeatureStore( + path, + read_only=True, + tokenizer_path="/target", + ) + store._tokenizer = BatchEncodingTokenizer() + loaded = store.read(next(store.iter_keys(shuffle=False))) + + assert torch.equal(loaded.input_ids, torch.tensor([1, 3, 2, 3])) + assert torch.equal(loaded.loss_mask, torch.tensor([0.0, 0.0, 0.0, 1.0])) + + +def test_vllm_safetensors_feature_store_records_manifest_path_and_roundtrips(tmp_path): + pytest.importorskip("safetensors.torch") + hidden_positions = torch.arange(160, dtype=torch.long) + sample = DraftFeatureSample( + algorithm="DSPARK", + input_ids=torch.arange(160, dtype=torch.long), + loss_mask=torch.ones(160, dtype=torch.float32), + hidden_states=torch.randn(160, 16, dtype=torch.float32), + position_ids=torch.arange(10, 170, dtype=torch.long), + metadata={ + "source": "token_replay_vllm_file", + "global_step": 7, + "hidden_states_layout": "dflash_aux_plus_last", + "hidden_positions": hidden_positions, + "feature_start": 10, + "feature_end": 170, + }, + ) + store = VllmSafetensorsFeatureStore(tmp_path) + + keys = store.write_many([sample]) + store.close() + + manifest_lines = (tmp_path / "manifest.jsonl").read_text(encoding="utf-8").splitlines() + assert len(manifest_lines) == 1 + entry = json.loads(manifest_lines[0]) + assert entry["path"].endswith(".safetensors") + assert (tmp_path / entry["path"]).exists() + assert entry["sample"]["metadata"]["hidden_positions"] == { + "__tensor__": True, + "dtype": "torch.int64", + "shape": [160], + } + + reader = VllmSafetensorsFeatureStore(tmp_path, read_only=True) + loaded = reader.read(keys[0]) + + assert loaded.algorithm == "DSPARK" + assert torch.equal(loaded.input_ids, sample.input_ids) + assert torch.equal(loaded.position_ids, sample.position_ids) + assert torch.equal(loaded.hidden_states, sample.hidden_states) + assert torch.equal(loaded.metadata["hidden_positions"], hidden_positions) + assert reader.get_metadata()["format"] == "vllm_safetensors" + assert reader.get_metadata()["num_samples"] == 1 + + +def test_build_feature_store_from_config_supports_vllm_safetensors(tmp_path): + pytest.importorskip("safetensors.torch") + writer = build_feature_store_from_config( + {"type": "vllm_safetensors", "path": tmp_path} + ) + writer.write_many([_sample(0)]) + writer.close() + + reader = build_feature_store_from_config( + {"type": "vllm_safetensors", "path": tmp_path}, read_only=True + ) + loaded = reader.read(next(reader.iter_keys(shuffle=False))) + + assert isinstance(loaded, DraftFeatureSample) + assert torch.equal(loaded.input_ids, torch.tensor([1, 2, 3, 4])) + + def test_feature_sample_normalizes_singleton_position_ids(): sample = DraftFeatureSample( input_ids=torch.tensor([1, 2, 3, 4], dtype=torch.long), diff --git a/tests/unit/test_draft_train_launcher.py b/tests/unit/test_draft_train_launcher.py index 654a0c1f..6894d165 100644 --- a/tests/unit/test_draft_train_launcher.py +++ b/tests/unit/test_draft_train_launcher.py @@ -19,6 +19,7 @@ build_torch_distributed_command, normalize_training_args, resolve_launch_config, + validate_tq_launch_config, ) @@ -110,3 +111,46 @@ def test_launcher_rejects_standalone_multinode() -> None: "speco.draft_training.standalone=true", ] ) + + +def test_launcher_accepts_complete_tq_consumer_config() -> None: + validate_tq_launch_config( + [ + "actor_rollout_ref.rollout.drafter.training.feature_store.type=tq", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable=true", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.ray.address=ray:6379", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.run_id=run-a", + ] + ) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ([], "enable=true"), + ( + [ + "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable=true", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.run_id=run-a", + ], + "ray.address", + ), + ( + [ + "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable=true", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.ray.address=ray:6379", + ], + "run_id", + ), + ], +) +def test_launcher_tq_validation_requires_canonical_connection_overrides( + overrides, message +) -> None: + with pytest.raises(ValueError, match=message): + validate_tq_launch_config( + [ + "actor_rollout_ref.rollout.drafter.training.feature_store.type=tq", + *overrides, + ] + ) diff --git a/tests/unit/test_draft_training_loop.py b/tests/unit/test_draft_training_loop.py index 86d14b14..986ac8d3 100644 --- a/tests/unit/test_draft_training_loop.py +++ b/tests/unit/test_draft_training_loop.py @@ -19,15 +19,22 @@ import pytest -pytest.importorskip("torch") +torch = pytest.importorskip("torch") -from omegaconf import OmegaConf +from omegaconf import OmegaConf # noqa: E402 -from verl_speco.trainer.draft_training_loop import ( +from verl_speco.trainer.draft_training_loop import ( # noqa: E402 _build_backend, + _clear_tq_batch_across_ranks, + _connect_tq_store_across_ranks, + _contains_replay_samples, + _is_out_of_memory_error, + _next_batch_across_ranks, _rewrite_standalone_block_runtime_config, _save_standalone_checkpoint, + _should_log_batch_progress, ) +from verl_speco.trainer.feature_store import DraftReplaySample # noqa: E402 class _FakeTrainer: @@ -44,11 +51,69 @@ def _save_checkpoint_async(self, step: int): return self.future +class _FakeTQLoader: + def __init__(self, error: BaseException | None = None): + self.error = error + self.clear_calls: list[list[str] | None] = [] + + def clear_completed_batch(self, keys): + self.clear_calls.append(keys) + if self.error is not None: + raise self.error + + +class _FakeTQStore: + def __init__(self, error: BaseException | None = None): + self.error = error + self.connect_calls = 0 + + def connect(self): + self.connect_calls += 1 + if self.error is not None: + raise self.error + + +def test_tq_completed_batch_is_cleared_once_on_rank_zero() -> None: + loader = _FakeTQLoader() + _clear_tq_batch_across_ranks( + loader, + ["k0", "k1"], + rank=0, + device=torch.device("cpu"), + ) + assert loader.clear_calls == [["k0", "k1"]] + + +def test_tq_clear_failure_is_reported_and_not_retried() -> None: + loader = _FakeTQLoader(RuntimeError("clear failed")) + with pytest.raises(RuntimeError, match="failed to clear"): + _clear_tq_batch_across_ranks( + loader, + ["k0", "k1"], + rank=0, + device=torch.device("cpu"), + ) + assert loader.clear_calls == [["k0", "k1"]] + + +def test_tq_store_connection_failure_is_reported() -> None: + store = _FakeTQStore(RuntimeError("connect failed")) + with pytest.raises(RuntimeError, match="failed to connect"): + _connect_tq_store_across_ranks( + store, + rank=0, + device=torch.device("cpu"), + ) + assert store.connect_calls == 1 + + def _export_trainer(model_type: str, model_path=None): """Minimal trainer stand-in for the standalone checkpoint export helpers.""" return SimpleNamespace( backend=SimpleNamespace(model_type=model_type), - config=SimpleNamespace(rollout=SimpleNamespace(drafter=SimpleNamespace(model_path=model_path))), + config=SimpleNamespace( + rollout=SimpleNamespace(drafter=SimpleNamespace(model_path=model_path)) + ), ) @@ -56,11 +121,50 @@ def _standalone_config(algorithm: str): return OmegaConf.create( { "model": {"path": "/does/not/exist"}, - "rollout": {"drafter": {"speculative_algorithm": algorithm, "training": {}}}, + "rollout": { + "drafter": {"speculative_algorithm": algorithm, "training": {}} + }, } ) +@pytest.mark.parametrize( + ("attempted_batches", "expected"), + [ + (1, True), + (2, True), + (3, True), + (4, False), + (99, False), + (100, True), + (101, False), + ], +) +def test_should_log_standalone_batch_progress(attempted_batches, expected): + assert _should_log_batch_progress(attempted_batches) is expected + + +def test_is_out_of_memory_error_matches_npu_oom_message(): + error = RuntimeError("NPU out of memory. Tried to allocate 258.00 MiB") + + assert _is_out_of_memory_error(error) + assert not _is_out_of_memory_error(RuntimeError("bad batch")) + + +def test_contains_replay_samples_detects_draft_replay_sample(): + sample = DraftReplaySample( + input_ids=torch.arange(4), + loss_mask=torch.ones(4), + attention_mask=torch.ones(4, dtype=torch.bool), + position_ids=torch.arange(4), + feature_positions=torch.arange(1, 3), + draft_position_ids=torch.arange(2, 4), + ) + + assert _contains_replay_samples([sample]) + assert not _contains_replay_samples([{"input_ids": [1, 2]}]) + + @pytest.mark.parametrize( ("algorithm", "expected_backend", "expected_model_type"), [ @@ -245,6 +349,56 @@ def test_standalone_dspark_checkpoint_preserves_source_runtime_config(tmp_path): assert saved_training_config == training_config +def test_standalone_dspark_checkpoint_rewrites_generic_qwen3_architecture(tmp_path): + checkpoint_dir = tmp_path / "draft_step_5" + checkpoint_dir.mkdir() + source_dir = tmp_path / "source_dspark" + source_dir.mkdir() + target_dir = tmp_path / "target_qwen3" + target_dir.mkdir() + (source_dir / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3", + "architectures": ["DSparkDraftModel"], + "markov_head_type": "vanilla", + } + ), + encoding="utf-8", + ) + (target_dir / "config.json").write_text( + json.dumps({"model_type": "qwen3"}), encoding="utf-8" + ) + (checkpoint_dir / "config.json").write_text( + json.dumps( + { + "model_type": "dspark", + "architectures": ["DSparkDraftModel"], + "markov_head_type": "vanilla", + } + ), + encoding="utf-8", + ) + trainer = SimpleNamespace( + backend=SimpleNamespace(model_type="dspark"), + config=SimpleNamespace( + model=SimpleNamespace(path=str(target_dir)), + rollout=SimpleNamespace( + drafter=SimpleNamespace(model_path=str(source_dir)) + ), + ), + ) + + _rewrite_standalone_block_runtime_config(trainer, str(checkpoint_dir)) + + runtime_config = json.loads( + (checkpoint_dir / "config.json").read_text(encoding="utf-8") + ) + assert runtime_config["model_type"] == "qwen3" + assert runtime_config["architectures"] == ["DSparkDraftModel"] + assert runtime_config["speco_training_model_type"] == "dspark" + + def test_standalone_domino_checkpoint_exports_dflash_projector_config(tmp_path): checkpoint_dir = tmp_path / "draft_step_5" checkpoint_dir.mkdir() @@ -337,3 +491,163 @@ def test_standalone_dflash_checkpoint_preserves_source_runtime_config(tmp_path): assert runtime_config["dflash_config"]["target_layer_ids"] == [2, 10, 18] assert runtime_config["eagle_aux_hidden_state_layer_ids"] == [3, 11, 19] assert saved_training_config == training_config + + +def test_standalone_block_checkpoint_uses_target_model_type_without_source_config( + tmp_path, +): + checkpoint_dir = tmp_path / "draft_step_5" + checkpoint_dir.mkdir() + target_dir = tmp_path / "target_qwen3" + target_dir.mkdir() + missing_source_dir = tmp_path / "missing_source_dspark" + (target_dir / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3", + "head_dim": 128, + "rope_theta": 1000000.0, + "max_position_embeddings": 40960, + } + ), + encoding="utf-8", + ) + training_config = { + "model_type": "dspark", + "architectures": ["DSparkDraftModel"], + "target_layer_ids": [1, 9, 17], + "markov_head_type": "vanilla", + "head_dim": 80, + "rope_theta": 10000.0, + } + (checkpoint_dir / "config.json").write_text( + json.dumps(training_config), encoding="utf-8" + ) + trainer = SimpleNamespace( + backend=SimpleNamespace(model_type="dspark"), + config=SimpleNamespace( + model=SimpleNamespace(path=str(target_dir)), + rollout=SimpleNamespace( + drafter=SimpleNamespace(model_path=str(missing_source_dir)) + ), + ), + ) + + _rewrite_standalone_block_runtime_config(trainer, str(checkpoint_dir)) + + runtime_config = json.loads( + (checkpoint_dir / "config.json").read_text(encoding="utf-8") + ) + saved_training_config = json.loads( + (checkpoint_dir / "speco_training_config.json").read_text(encoding="utf-8") + ) + assert runtime_config["model_type"] == "dspark" + assert runtime_config["architectures"] == ["DSparkDraftModel"] + assert runtime_config["speco_training_model_type"] == "dspark" + assert runtime_config["dspark_config"]["markov_head_type"] == "vanilla" + assert runtime_config["head_dim"] == 80 + assert runtime_config["rope_theta"] == 10000.0 + assert saved_training_config == training_config + + +def test_standalone_eagle3_checkpoint_exports_vllm_llama_runtime_config(tmp_path): + checkpoint_dir = tmp_path / "draft_step_5" + checkpoint_dir.mkdir() + target_dir = tmp_path / "target_qwen3" + target_dir.mkdir() + missing_source_dir = tmp_path / "missing_source_eagle3" + (target_dir / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3", + "hidden_size": 4096, + "head_dim": 128, + "rope_theta": 1000000, + "max_position_embeddings": 40960, + } + ), + encoding="utf-8", + ) + training_config = { + "model_type": "qwen3", + "architectures": ["LlamaForCausalLMEagle3"], + "num_hidden_layers": 1, + "hidden_size": 4096, + "vocab_size": 151936, + "tie_word_embeddings": False, + } + (checkpoint_dir / "config.json").write_text( + json.dumps(training_config), encoding="utf-8" + ) + trainer = SimpleNamespace( + backend=SimpleNamespace(model_type="eagle3"), + config=SimpleNamespace( + model=SimpleNamespace(path=str(target_dir)), + rollout=SimpleNamespace( + drafter=SimpleNamespace(model_path=str(missing_source_dir)) + ), + ), + ) + + _rewrite_standalone_block_runtime_config(trainer, str(checkpoint_dir)) + + runtime_config = json.loads( + (checkpoint_dir / "config.json").read_text(encoding="utf-8") + ) + assert runtime_config["model_type"] == "qwen3" + assert runtime_config["architectures"] == ["LlamaForCausalLMEagle3"] + assert runtime_config["num_hidden_layers"] == 1 + assert runtime_config["tie_word_embeddings"] is False + assert not (checkpoint_dir / "speco_training_config.json").exists() + + +def test_next_batch_across_ranks_returns_local_batch_without_distributed(): + batch = [object()] + + assert _next_batch_across_ranks( + iter([batch]), rank=0, device=torch.device("cpu") + ) is batch + + +def test_next_batch_across_ranks_returns_none_when_source_is_exhausted(): + assert ( + _next_batch_across_ranks( + iter(()), rank=0, device=torch.device("cpu") + ) + is None + ) + + +def test_next_batch_across_ranks_preserves_local_producer_error(): + def broken_source(): + raise ValueError("producer failed") + yield [] + + with pytest.raises(RuntimeError, match="failed on rank=0") as exc_info: + _next_batch_across_ranks( + iter(broken_source()), rank=0, device=torch.device("cpu") + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + + +def test_next_batch_across_ranks_stops_for_remote_rank_failure(monkeypatch): + monkeypatch.setattr( + "verl_speco.trainer.draft_training_loop.dist.is_initialized", lambda: True + ) + monkeypatch.setattr( + "verl_speco.trainer.draft_training_loop.dist.get_world_size", lambda: 2 + ) + + def fake_all_reduce(state, op): + del op + state[0] = 1 + + monkeypatch.setattr( + "verl_speco.trainer.draft_training_loop.dist.all_reduce", fake_all_reduce + ) + + with pytest.raises(RuntimeError, match="failed on another rank"): + _next_batch_across_ranks( + iter([[object()]]), rank=1, device=torch.device("cpu") + ) diff --git a/tests/unit/test_drafter_sample_protocol.py b/tests/unit/test_drafter_sample_protocol.py new file mode 100644 index 00000000..3a93e018 --- /dev/null +++ b/tests/unit/test_drafter_sample_protocol.py @@ -0,0 +1,144 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from verl_speco.trainer.feature_store import DraftFeatureSample +from verl_speco.transport.drafter_sample_protocol import ( + PROTOCOL_SCHEMA_VERSION, + ExpectedFeatureConfig, + SampleMetadata, + decode_sample, + encode_sample, + is_ready_sample_tag, + make_eos_record, + make_ready_tag, + make_sample_key, + parse_ready_tag, +) + + +def _metadata() -> SampleMetadata: + return SampleMetadata( + schema_version=PROTOCOL_SCHEMA_VERSION, + run_id="run-a", + sample_id="train-000017", + sequence_no=17, + ) + + +def _sample(*, algorithm: str = "DSPARK") -> DraftFeatureSample: + return DraftFeatureSample( + algorithm=algorithm, + input_ids=torch.tensor([10, 11, 12, 13]), + loss_mask=torch.tensor([0.0, 1.0, 1.0, 1.0]), + position_ids=torch.tensor([6, 7, 8, 9]), + hidden_states=torch.arange(64, dtype=torch.bfloat16).reshape(4, 16), + last_hidden_states=torch.arange(32, dtype=torch.float32).reshape(4, 8), + metadata={ + "hidden_states_layout": "dflash_aux_plus_last", + "target_layer_ids": [2, 8, 14], + "hidden_positions": torch.tensor([6, 7, 8, 9]), + "nested": {"pair": ("a", 2)}, + }, + ) + + +def test_sample_round_trip_preserves_complete_sample() -> None: + meta = _metadata() + key = make_sample_key(meta) + fields = encode_sample(_sample(), meta) + restored = decode_sample( + key, + make_ready_tag(meta), + fields, + ExpectedFeatureConfig(run_id=meta.run_id), + ) + + assert key == "drafter:v2:run-a:000000000017:train-000017" + assert restored.algorithm == "DSPARK" + assert torch.equal(restored.input_ids, _sample().input_ids) + assert torch.equal(restored.hidden_states, _sample().hidden_states) + assert torch.equal(restored.last_hidden_states, _sample().last_hidden_states) + assert torch.equal( + restored.metadata["hidden_positions"], _sample().metadata["hidden_positions"] + ) + assert restored.metadata["nested"]["pair"] == ("a", 2) + assert fields["sample__manifest_json"].dtype == torch.uint8 + + +def test_hidden_state_tensor_list_round_trip() -> None: + sample = replace( + _sample(algorithm="EAGLE3"), + hidden_states=[torch.ones(4, 3), torch.zeros(4, 5)], + ) + meta = _metadata() + restored = decode_sample( + make_sample_key(meta), + make_ready_tag(meta), + encode_sample(sample, meta), + ExpectedFeatureConfig(run_id="run-a"), + ) + assert restored.algorithm == "EAGLE3" + assert isinstance(restored.hidden_states, list) + assert [tuple(value.shape) for value in restored.hidden_states] == [(4, 3), (4, 5)] + + +def test_ready_parser_is_the_shared_discovery_contract() -> None: + meta = _metadata() + tag = make_ready_tag(meta) + assert parse_ready_tag(tag, run_id="run-a") == meta + assert is_ready_sample_tag(tag, run_id="run-a") + assert not is_ready_sample_tag(tag, run_id="another-run") + assert parse_ready_tag({**tag, "sequence_no": "bad"}, run_id="run-a") is None + assert parse_ready_tag({**tag, "schema_version": 1}, run_id="run-a") is None + + +def test_decode_rejects_identity_mismatch() -> None: + meta = _metadata() + fields = encode_sample(_sample(), meta) + bad_tag = {**make_ready_tag(meta), "sample_id": "wrong"} + with pytest.raises(ValueError, match="key mismatch"): + decode_sample( + make_sample_key(meta), + bad_tag, + fields, + ExpectedFeatureConfig(run_id="run-a"), + ) + + +def test_metadata_codec_rejects_lossy_unknown_values() -> None: + sample = replace(_sample(), metadata={"unsupported": object()}) + with pytest.raises(TypeError, match="metadata.unsupported"): + encode_sample(sample, _metadata()) + + +@pytest.mark.parametrize("algorithm", ["EAGLE3", "DFLASH", "DSPARK", "DOMINO"]) +def test_protocol_algorithm_is_not_hardcoded(algorithm: str) -> None: + meta = _metadata() + sample = _sample(algorithm=algorithm) + restored = decode_sample( + make_sample_key(meta), + make_ready_tag(meta), + encode_sample(sample, meta), + ExpectedFeatureConfig(run_id="run-a"), + ) + assert restored.algorithm == algorithm + assert "algorithm" not in make_ready_tag(meta) + + +def test_eos_record_is_control_only() -> None: + key, fields, tag = make_eos_record("run-a", 18) + assert key == "control:v2:run-a:eos" + assert fields["marker"].tolist() == [1] + assert tag == { + "record_type": "control", + "status": "eos", + "schema_version": 2, + "run_id": "run-a", + "total_samples": 18, + } diff --git a/tests/unit/test_producer_input_reader.py b/tests/unit/test_producer_input_reader.py new file mode 100644 index 00000000..2cb976c2 --- /dev/null +++ b/tests/unit/test_producer_input_reader.py @@ -0,0 +1,313 @@ +# Copyright 2026 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. + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from verl_speco.producer import input_reader + + +def test_iter_input_records_reads_jsonl(tmp_path: Path) -> None: + input_path = tmp_path / "train.jsonl" + input_path.write_text( + json.dumps({"prompt": "Q: ", "response": "A", "source": "test"}) + "\n", + encoding="utf-8", + ) + + records = list(input_reader.iter_input_records(input_path)) + + assert len(records) == 1 + assert records[0].sample_id == "train-000000" + assert records[0].prompt == "Q: " + assert records[0].response == "A" + assert records[0].source_metadata == {"source": "test"} + + +def test_iter_input_records_reads_parquet_rows(monkeypatch, tmp_path: Path) -> None: + input_path = tmp_path / "train.parquet" + input_path.write_bytes(b"PAR1-test-fixture") + rows = [ + {"sample_id": "first", "prompt": "Q1: ", "response": "A1"}, + {"prompt": "Q2: ", "response": "A2", "split": "train"}, + ] + + class FakeParquetFile: + def __init__(self, path: Path) -> None: + assert path == input_path + + def iter_batches(self): + yield SimpleNamespace(to_pylist=lambda: rows) + + monkeypatch.setattr( + input_reader.importlib, + "import_module", + lambda name: SimpleNamespace(ParquetFile=FakeParquetFile), + ) + + records = list(input_reader.iter_input_records(input_path)) + + assert [record.sample_id for record in records] == ["first", "train-000001"] + assert records[1].source_metadata == {"split": "train"} + + +def test_iter_input_records_reads_real_parquet_when_available(tmp_path: Path) -> None: + pyarrow = pytest.importorskip("pyarrow") + parquet = pytest.importorskip("pyarrow.parquet") + input_path = tmp_path / "train.parquet" + parquet.write_table( + pyarrow.Table.from_pylist( + [{"prompt": "real Q: ", "response": "real A", "split": "train"}] + ), + input_path, + ) + + records = list(input_reader.iter_input_records(input_path)) + + assert len(records) == 1 + assert records[0].prompt == "real Q: " + assert records[0].response == "real A" + assert records[0].source_metadata == {"split": "train"} + + +def test_iter_input_records_reads_real_dapo_style_parquet_when_available( + tmp_path: Path, +) -> None: + pyarrow = pytest.importorskip("pyarrow") + parquet = pytest.importorskip("pyarrow.parquet") + input_path = tmp_path / "dapo.parquet" + parquet.write_table( + pyarrow.Table.from_pylist( + [ + { + "data_source": "math_dapo", + "prompt": [{"role": "user", "content": "Solve Q"}], + "reward_model": { + "ground_truth": "42", + "style": "rule-lighteval/MATH_v2", + }, + "extra_info": {"index": "dapo-real-row"}, + } + ] + ), + input_path, + ) + + record = next(input_reader.iter_input_records(input_path)) + + assert record.prompt == ({"role": "user", "content": "Solve Q"},) + assert record.response is None + assert record.sample_id == "dapo-real-row" + + +def test_dapo_parquet_prompt_is_prepared_for_target_generation( + monkeypatch, tmp_path: Path +) -> None: + input_path = tmp_path / "train.parquet" + input_path.write_bytes(b"PAR1-test-fixture") + + class FakeParquetFile: + def __init__(self, path: Path) -> None: + assert path == input_path + + def iter_batches(self): + yield SimpleNamespace( + to_pylist=lambda: [ + { + "prompt": [{"role": "user", "content": "Solve Q"}], + "reward_model": {"ground_truth": "42"}, + "extra_info": {"index": "dapo-row-id"}, + } + ] + ) + + monkeypatch.setattr( + input_reader.importlib, + "import_module", + lambda name: SimpleNamespace(ParquetFile=FakeParquetFile), + ) + + record = next(input_reader.iter_input_records(input_path)) + + class ChatTokenizer: + def apply_chat_template(self, messages, *, tokenize, add_generation_prompt): + assert messages == [{"role": "user", "content": "Solve Q"}] + assert tokenize is True + assert add_generation_prompt is True + return [10, 11, 12] + + generation = input_reader.prepare_generation_request( + record, + ChatTokenizer(), + {"max_sequence_length": 16, "generation_max_tokens": 4}, + ) + finalized = input_reader.finalize_generated_request( + generation, + [10, 11, 12, 20, 21], + {"max_sequence_length": 16, "max_feature_length": 8}, + ) + + assert record.sample_id == "dapo-row-id" + assert record.response is None + assert generation.prompt_token_ids == (10, 11, 12) + assert generation.max_tokens == 4 + assert finalized.input_ids.tolist() == [10, 11, 12, 20, 21] + assert finalized.loss_mask.tolist() == [0, 0, 0, 1, 1] + assert finalized.prompt_token_ids == [10, 11, 12, 20, 21] + + +def test_finalize_generated_request_aligns_connector_excluding_final_token() -> None: + request = input_reader.GenerationRequest( + sequence_no=0, + sample_id="generated-row", + prompt_token_ids=(10, 11, 12), + max_tokens=4, + source_metadata={}, + ) + + finalized = input_reader.finalize_generated_request( + request, + [10, 11, 12, 20], + {"max_sequence_length": 16, "max_feature_length": 8}, + expected_response_token_ids=[20, 21], + ) + + assert finalized.input_ids.tolist() == [10, 11, 12, 20, 21] + assert finalized.loss_mask.tolist() == [0, 0, 0, 1, 1] + assert finalized.prompt_token_ids == [10, 11, 12, 20] + assert finalized.feature_positions.tolist() == [2, 3] + assert finalized.draft_position_ids.tolist() == [3, 4] + + +def test_prefilled_response_limits_vllm_prefix_after_selecting_training_window() -> None: + prepared = input_reader._build_tokenized_request( + sequence_no=0, + sample_id="long-response", + prompt_length=3, + full_ids=list(range(20)), + source_metadata={}, + config={"max_sequence_length": 8, "max_feature_length": 4}, + ) + + assert prepared.input_ids.numel() == 20 + assert prepared.feature_positions.tolist() == [2, 3, 4, 5] + assert prepared.prompt_token_ids == [0, 1, 2, 3, 4, 5] + + +def test_prefilled_response_rejects_prompt_prefix_beyond_vllm_limit() -> None: + with pytest.raises( + ValueError, + match=r"vLLM prefill of 13 tokens.*max_sequence_length=8", + ): + input_reader._build_tokenized_request( + sequence_no=0, + sample_id="long-prompt", + prompt_length=10, + full_ids=list(range(20)), + source_metadata={}, + config={"max_sequence_length": 8, "max_feature_length": 4}, + ) + + +def test_iter_jsonl_conversation_splits_final_assistant_response(tmp_path: Path) -> None: + input_path = tmp_path / "conversation.jsonl" + input_path.write_text( + json.dumps( + { + "conversation": [ + {"role": "human", "content": "Question"}, + {"role": "assistant", "content": "Answer"}, + ] + } + ) + + "\n", + encoding="utf-8", + ) + + record = next(input_reader.iter_input_records(input_path)) + + assert record.prompt == ({"role": "user", "content": "Question"},) + assert record.response == "Answer" + + +def test_iter_jsonl_legacy_conversations_normalizes_from_value(tmp_path: Path) -> None: + input_path = tmp_path / "conversations.jsonl" + input_path.write_text( + json.dumps( + { + "conversations": [ + {"from": "human", "value": "Question"}, + {"from": "gpt", "value": "Answer"}, + ] + } + ) + + "\n", + encoding="utf-8", + ) + + record = next(input_reader.iter_input_records(input_path)) + + assert record.prompt == ({"role": "user", "content": "Question"},) + assert record.response == "Answer" + + +def test_prepare_generated_prefill_request_uses_full_sequence_without_final_token() -> None: + request = input_reader.GenerationRequest( + sequence_no=0, + sample_id="generated-row", + prompt_token_ids=(10, 11, 12), + max_tokens=4, + source_metadata={}, + ) + + prepared = input_reader.prepare_generated_prefill_request( + request, + [20, 21], + {"max_sequence_length": 16, "max_feature_length": 8}, + ) + + assert prepared.input_ids.tolist() == [10, 11, 12, 20, 21] + assert prepared.loss_mask.tolist() == [0, 0, 0, 1, 1] + assert prepared.prompt_token_ids == [10, 11, 12, 20] + assert prepared.feature_positions.tolist() == [2, 3] + + +def test_finalize_generated_request_rejects_misaligned_connector_tokens() -> None: + request = input_reader.GenerationRequest( + sequence_no=0, + sample_id="generated-row", + prompt_token_ids=(10, 11), + max_tokens=4, + source_metadata={}, + ) + + with pytest.raises(ValueError, match="excluding its final token"): + input_reader.finalize_generated_request( + request, + [10, 11, 99], + {"max_sequence_length": 16, "max_feature_length": 8}, + expected_response_token_ids=[20, 21], + ) + + +def test_non_utf8_non_parquet_input_has_actionable_error(tmp_path: Path) -> None: + input_path = tmp_path / "train.data" + input_path.write_bytes(b"plain-prefix\xc0binary") + + with pytest.raises(ValueError, match="not UTF-8 JSONL or a Parquet file"): + list(input_reader.iter_input_records(input_path)) diff --git a/tests/unit/test_standalone_resume.py b/tests/unit/test_standalone_resume.py new file mode 100644 index 00000000..9708c0f1 --- /dev/null +++ b/tests/unit/test_standalone_resume.py @@ -0,0 +1,55 @@ +# Copyright 2026 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. + +from pathlib import Path + +import pytest + +from verl_speco.trainer.standalone_resume import ( + load_standalone_resume, + save_standalone_resume, +) + + +def test_standalone_resume_round_trip_and_input_validation(tmp_path: Path) -> None: + input_path = tmp_path / "train.jsonl" + input_path.write_text('{"prompt":"q","response":"a"}\n', encoding="utf-8") + checkpoint_path = tmp_path / "draft_step_7" + + save_standalone_resume( + checkpoint_path, + [5, 1, 5, 3], + optimizer_step=7, + input_path=input_path, + ) + + consumed, metadata = load_standalone_resume( + checkpoint_path, input_path=input_path + ) + assert consumed == {1, 3, 5} + assert metadata is not None + assert metadata["optimizer_step"] == 7 + assert metadata["consumed_count"] == 3 + + input_path.write_text('{"prompt":"changed","response":"a"}\n', encoding="utf-8") + with pytest.raises(ValueError, match="input file changed"): + load_standalone_resume(checkpoint_path, input_path=input_path) + + +def test_missing_standalone_resume_is_not_a_resume_checkpoint( + tmp_path: Path, +) -> None: + consumed, metadata = load_standalone_resume(tmp_path / "pretrained") + assert consumed == set() + assert metadata is None diff --git a/tests/unit/test_standalone_tq_training_launcher.py b/tests/unit/test_standalone_tq_training_launcher.py new file mode 100644 index 00000000..a8992740 --- /dev/null +++ b/tests/unit/test_standalone_tq_training_launcher.py @@ -0,0 +1,405 @@ +# Copyright 2026 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. +from __future__ import annotations + +from pathlib import Path +import json +import threading + +from omegaconf import OmegaConf +import pytest + +from verl_speco.standalone_tq_training_launcher import ( + _preflight_input_file, + _producer_max_samples, + _target_final_layer_id, + build_pipeline_commands, + resolve_pipeline_config, + run_pipeline, + start_ray_session, +) +import verl_speco.tq_owner as tq_owner + + +def _training_args() -> list[str]: + return [ + "data.train_files=/data/train.jsonl", + "actor_rollout_ref.model.path=/models/Qwen3-8B", + "actor_rollout_ref.rollout.drafter.model_path=/models/dspark", + "actor_rollout_ref.rollout.drafter.speculative_algorithm=DSPARK", + "actor_rollout_ref.rollout.drafter.training.max_steps=10", + ] + + +def test_pipeline_config_derives_transport_identity_from_training_args() -> None: + config = resolve_pipeline_config(_training_args(), environ={}) + + assert config.input_path == "/data/train.jsonl" + assert config.model_path == "/models/Qwen3-8B" + assert config.tokenizer_path == "/models/Qwen3-8B" + assert config.algorithm == "DSPARK" + assert config.target_layer_ids == (1, 9, 17, 25, 33) + assert config.vllm_endpoints == ("http://127.0.0.1:8000/v1",) + assert config.run_id.startswith("dspark-") + + +def test_producer_max_samples_uses_remaining_total_steps() -> None: + args = [ + *_training_args(), + "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=2", + "speco.draft_training.nproc_per_node=4", + "speco.draft_training.nnodes=1", + ] + + assert _producer_max_samples(args, resumed_optimizer_step=6) == 32 + + +def test_pipeline_config_reads_non_dspark_algorithm_from_training_args() -> None: + args = [ + item.replace("speculative_algorithm=DSPARK", "speculative_algorithm=DFLASH") + for item in _training_args() + ] + args.append( + "actor_rollout_ref.rollout.drafter.training.dflash_target_layer_ids=[2,10,20]" + ) + + config = resolve_pipeline_config(args, environ={}) + + assert config.algorithm == "DFLASH" + assert config.target_layer_ids == (2, 10, 20) + assert config.run_id.startswith("dflash-") + + +def test_pipeline_config_prefers_generic_producer_layer_ids() -> None: + args = [ + *_training_args(), + "speco.standalone_tq_producer.target_layer_ids=[3,11,21]", + "actor_rollout_ref.rollout.drafter.training.dspark_target_layer_ids=[1,9,17]", + ] + + config = resolve_pipeline_config(args, environ={}) + + assert config.target_layer_ids == (3, 11, 21) + + +def test_pipeline_config_accepts_one_hydra_list_train_file() -> None: + args = _training_args() + args[0] = "data.train_files=['/data/train.jsonl']" + + config = resolve_pipeline_config(args, environ={}) + + assert config.input_path == "/data/train.jsonl" + + +def test_pipeline_config_allows_missing_drafter_path_for_fresh_training() -> None: + args = [ + item + for item in _training_args() + if not item.startswith("actor_rollout_ref.rollout.drafter.model_path=") + ] + + config = resolve_pipeline_config(args, environ={}) + + assert config.model_path == "/models/Qwen3-8B" + + +def test_pipeline_config_accepts_multiple_vllm_endpoints() -> None: + config = resolve_pipeline_config( + _training_args(), + environ={ + "SPECO_VLLM_ENDPOINTS": ( + "[http://127.0.0.1:8000/v1,http://127.0.0.1:8001/v1]" + ) + }, + ) + + assert config.vllm_endpoints == ( + "http://127.0.0.1:8000/v1", + "http://127.0.0.1:8001/v1", + ) + + +def test_target_final_layer_id_uses_local_model_config(tmp_path) -> None: + (tmp_path / "config.json").write_text( + json.dumps({"text_config": {"num_hidden_layers": 48}}), + encoding="utf-8", + ) + + assert _target_final_layer_id(str(tmp_path), (2, 10, 20)) == 48 + + +def test_pipeline_config_rejects_multiple_train_files() -> None: + args = _training_args() + args[0] = "data.train_files=[a.jsonl,b.jsonl]" + + with pytest.raises(ValueError, match="exactly one train file"): + resolve_pipeline_config(args, environ={}) + + +def test_preflight_accepts_verl_prompt_parquet(monkeypatch, tmp_path) -> None: + input_path = tmp_path / "train.parquet" + input_path.write_bytes(b"PAR1-test-fixture") + + class FakeParquetFile: + def __init__(self, path: Path) -> None: + assert path == input_path + + def iter_batches(self): + class Batch: + @staticmethod + def to_pylist(): + return [ + { + "prompt": [{"role": "user", "content": "Solve Q"}], + "reward_model": {"ground_truth": "42"}, + } + ] + + yield Batch() + + from verl_speco.producer import input_reader + + monkeypatch.setattr( + input_reader.importlib, + "import_module", + lambda name: type("ParquetModule", (), {"ParquetFile": FakeParquetFile}), + ) + + _preflight_input_file(str(input_path)) + + +def test_pipeline_commands_hide_and_replace_tq_overrides() -> None: + args = [ + *_training_args(), + "actor_rollout_ref.rollout.drafter.training.feature_store.type=torch_shard", + "actor_rollout_ref.rollout.drafter.training.transfer_queue.run_id=user-value", + ] + config = resolve_pipeline_config(args, environ={}) + + commands = build_pipeline_commands( + config, + args, + ray_address="10.0.0.1:6379", + python_executable="python", + ) + + assert commands.owner[:3] == ["python", "-m", "verl_speco.tq_owner"] + assert commands.vllm is not None + assert commands.vllm[:3] == ["vllm", "serve", "/models/Qwen3-8B"] + assert "ExampleHiddenStatesConnector" in " ".join(commands.vllm) + assert commands.vllm_endpoints == ("http://127.0.0.1:8000/v1",) + assert commands.producer[:3] == [ + "python", + "-m", + "verl_speco.standalone_tq_producer", + ] + assert commands.consumer[:3] == [ + "python", + "-m", + "verl_speco.draft_train_launcher", + ] + assert any(item.endswith("feature_store.type=tq") for item in commands.consumer) + assert not any("run_id=user-value" in item for item in commands.consumer) + assert any(f"run_id={config.run_id}" in item for item in commands.consumer) + assert any( + "vllm_endpoints=[http://127.0.0.1:8000/v1]" in item + for item in commands.producer + ) + assert any("max_samples=40" in item for item in commands.producer) + + +def test_pipeline_commands_pass_all_external_vllm_endpoints_to_producer() -> None: + endpoints = "[http://127.0.0.1:8000/v1,http://127.0.0.1:8001/v1]" + config = resolve_pipeline_config( + _training_args(), environ={"SPECO_VLLM_ENDPOINTS": endpoints} + ) + + commands = build_pipeline_commands( + config, + _training_args(), + ray_address="127.0.0.1:6379", + python_executable="python", + ) + + # Multiple services are started by the dedicated shell script. The unified + # launcher only verifies them and passes both URLs into the Producer pool. + assert commands.vllm is None + assert any(f"vllm_endpoints={endpoints}" in item for item in commands.producer) + + +def test_pipeline_commands_forward_producer_tuning_overrides() -> None: + args = [ + *_training_args(), + "speco.standalone_tq_producer.max_inflight_requests=32", + "speco.standalone_tq_producer.per_endpoint_concurrency=8", + "speco.standalone_tq_producer.max_feature_length=384", + ] + config = resolve_pipeline_config(args, environ={}) + + commands = build_pipeline_commands( + config, + args, + ray_address="127.0.0.1:6379", + python_executable="python", + ) + + assert ( + "speco.standalone_tq_producer.max_inflight_requests=32" + in commands.producer + ) + assert ( + "speco.standalone_tq_producer.per_endpoint_concurrency=8" + in commands.producer + ) + assert "speco.standalone_tq_producer.max_feature_length=384" in commands.producer + + +class _FakeRuntimeContext: + gcs_address = "127.0.0.1:61234" + + +class _FakeRay: + def __init__(self) -> None: + self.init_kwargs = None + self.shutdown_called = False + + def init(self, **kwargs) -> None: + self.init_kwargs = kwargs + + def get_runtime_context(self) -> _FakeRuntimeContext: + return _FakeRuntimeContext() + + def shutdown(self) -> None: + self.shutdown_called = True + + +def test_ray_session_starts_local_control_plane_without_exposed_address() -> None: + ray = _FakeRay() + + session = start_ray_session( + environ={"RAY_ADDRESS": "172.51.9.253:35195"}, ray_module=ray + ) + + assert session.address == "127.0.0.1:61234" + assert ray.init_kwargs == { + "address": "local", + "namespace": "speco-drafter", + "include_dashboard": False, + } + session.close() + assert ray.shutdown_called + + +class _FakeProcess: + def __init__(self, role: str) -> None: + self.role = role + self.returncode = 0 if role == "consumer" else None + self.terminated = False + + def poll(self): + return self.returncode + + def terminate(self) -> None: + self.terminated = True + self.returncode = -15 + + def wait(self, timeout=None): + return self.returncode + + def kill(self) -> None: + self.returncode = -9 + + +def test_pipeline_starts_owner_then_consumer_then_producer() -> None: + started: list[str] = [] + processes: list[_FakeProcess] = [] + child_environments: list[dict[str, str]] = [] + vllm_started = False + + def fake_popen(command, *, env): + nonlocal vllm_started + if command[0] == "vllm": + role = "vllm" + vllm_started = True + else: + module = command[2] + role = { + "verl_speco.tq_owner": "owner", + "verl_speco.draft_train_launcher": "consumer", + "verl_speco.standalone_tq_producer": "producer", + }[module] + process = _FakeProcess(role) + started.append(role) + processes.append(process) + child_environments.append(dict(env)) + if role == "owner": + Path(env["SPECO_TQ_OWNER_READY_FILE"]).touch() + return process + + config = resolve_pipeline_config(_training_args(), environ={}) + commands = build_pipeline_commands( + config, + _training_args(), + ray_address="127.0.0.1:61234", + ) + + assert ( + run_pipeline( + commands, + ray_address="127.0.0.1:61234", + environ={"RAY_ADDRESS": "172.51.9.253:35195"}, + popen=fake_popen, + endpoint_ready=lambda _: vllm_started, + ) + == 0 + ) + assert started == ["vllm", "owner", "consumer", "producer"] + assert all(env["RAY_ADDRESS"] == "127.0.0.1:61234" for env in child_environments) + assert all(process.poll() is not None for process in processes) + + +def test_owner_writes_internal_ready_file(monkeypatch, tmp_path) -> None: + ready_file = tmp_path / "owner.ready" + monkeypatch.setenv("SPECO_TQ_OWNER_READY_FILE", str(ready_file)) + monkeypatch.setattr(tq_owner, "configure_transfer_queue", lambda config: True) + monkeypatch.setattr(tq_owner, "connect_ray_cluster", lambda *args: None) + monkeypatch.setattr(tq_owner, "start_transfer_queue_owner", lambda config: None) + monkeypatch.setattr(tq_owner, "close_transfer_queue_owner", lambda: None) + monkeypatch.setattr(tq_owner, "publish_owner_ready", lambda *args: "ready-key") + stop_event = threading.Event() + stop_event.set() + config = OmegaConf.create( + { + "actor_rollout_ref": { + "rollout": { + "drafter": { + "training": { + "transfer_queue": { + "enable": True, + "run_id": "test-run", + "schema_version": 1, + "ray": { + "address": "127.0.0.1:6379", + "namespace": "speco-drafter", + }, + } + } + } + } + } + } + ) + + assert tq_owner.run_owner(config, stop_event=stop_event) == 0 + assert ready_file.is_file() diff --git a/tests/unit/test_target_feature_pipeline.py b/tests/unit/test_target_feature_pipeline.py new file mode 100644 index 00000000..6063ca45 --- /dev/null +++ b/tests/unit/test_target_feature_pipeline.py @@ -0,0 +1,78 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); + +import time + +import pytest + +torch = pytest.importorskip("torch") + +from verl_speco.trainer.feature_store import DraftFeatureSample # noqa: E402 +from verl_speco.trainer.target_feature_pipeline import ( # noqa: E402 + TargetFeatureProducer, +) + + +def _feature(value: int) -> DraftFeatureSample: + return DraftFeatureSample( + input_ids=torch.tensor([value]), + loss_mask=torch.ones(1), + hidden_states=torch.zeros(1, 4), + position_ids=torch.ones(1, dtype=torch.long), + ) + + +class _Replayer: + backend = "vllm_file" + + def materialize(self, samples): + time.sleep(0.01) + return [_feature(int(samples[0]))] + + +def test_target_feature_producer_preserves_batch_order_and_prefetches(): + producer = TargetFeatureProducer( + [[0, 1], [2, 3]], + _Replayer(), + rank=0, + concurrency=2, + producer_prefetch_depth=2, + prefetch_depth=2, + queue_timeout=2, + ) + try: + batches = list(producer) + finally: + producer.close() + + assert [[int(x.input_ids[0]) for x in batch] for batch in batches] == [ + [0, 1], + [2, 3], + ] + assert producer.metrics()["producer/samples_total"] == 4 + + +class _FailingReplayer: + backend = "vllm_file" + + def materialize(self, samples): + raise ValueError("broken replay") + + +def test_target_feature_producer_propagates_background_failure(): + producer = TargetFeatureProducer( + [[0]], + _FailingReplayer(), + rank=0, + concurrency=1, + producer_prefetch_depth=1, + prefetch_depth=1, + queue_timeout=2, + ) + try: + with pytest.raises(RuntimeError, match="producer failed") as exc_info: + next(producer) + assert isinstance(exc_info.value.__cause__, ValueError) + finally: + producer.close() diff --git a/tests/unit/test_target_feature_replay.py b/tests/unit/test_target_feature_replay.py new file mode 100644 index 00000000..8b7730c0 --- /dev/null +++ b/tests/unit/test_target_feature_replay.py @@ -0,0 +1,410 @@ +# Copyright 2026 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 threading +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from verl_speco.trainer.feature_store import DraftFeatureSample, DraftReplaySample # noqa: E402 +from verl_speco.trainer.target_feature_replay import ( # noqa: E402 + BoundedReplayCache, + FeatureContract, + TargetFeatureReplayer, + _VllmEndpointState, + _hidden_capture_target, + _normalize_vllm_endpoints, + feature_from_vllm_payload, + load_vllm_final_norm, +) + + +def _feature_sample() -> DraftFeatureSample: + return DraftFeatureSample( + input_ids=torch.arange(8), + loss_mask=torch.ones(8), + hidden_states=torch.zeros(8, 16), + position_ids=torch.arange(1, 9), + ) + + +def test_hidden_capture_target_matches_transformers_hidden_state_indices(): + assert _hidden_capture_target(0, 36) == ("layer", 0) + assert _hidden_capture_target(34, 36) == ("layer", 34) + assert _hidden_capture_target(35, 36) == ("final", None) + + +def test_normalize_vllm_endpoints_prefers_pool_and_deduplicates(): + assert _normalize_vllm_endpoints( + { + "vllm_endpoint": "http://legacy:8000/v1", + "vllm_endpoints": [ + "http://host1:8000/v1/", + "http://host2:8000/v1", + "http://host1:8000/v1", + ], + } + ) == ["http://host1:8000/v1", "http://host2:8000/v1"] + + +def test_vllm_request_fails_over_to_another_endpoint(monkeypatch): + class _Completions: + def __init__(self, error=None): + self.error = error + self.calls = 0 + + def create(self, **kwargs): + self.calls += 1 + if self.error is not None: + raise self.error + return SimpleNamespace(choices=[]) + + failed = _Completions(RuntimeError("endpoint down")) + healthy = _Completions() + replayer = TargetFeatureReplayer.__new__(TargetFeatureReplayer) + replayer.rank = 0 + replayer.vllm_timeout = 1 + replayer.vllm_max_retries = 1 + replayer.vllm_endpoint_cooldown = 5 + replayer.vllm_requests = 0 + replayer.vllm_request_seconds = 0.0 + replayer._metrics_lock = threading.Lock() + replayer._endpoint_lock = threading.Lock() + replayer._vllm_clients_initialized = True + replayer._vllm_endpoint_states = [ + _VllmEndpointState( + index=0, + url="http://host1:8000/v1", + client=SimpleNamespace(completions=failed), + model="target", + ), + _VllmEndpointState( + index=1, + url="http://host2:8000/v1", + client=SimpleNamespace(completions=healthy), + model="target", + ), + ] + monkeypatch.setattr( + "verl_speco.trainer.target_feature_replay.time.sleep", lambda _: None + ) + + response = replayer._request_vllm_response([1, 2, 3]) + + assert response.choices == [] + assert failed.calls == 1 + assert healthy.calls == 1 + assert replayer._vllm_endpoint_states[0].failures == 1 + assert replayer._vllm_endpoint_states[1].requests == 1 + + +def test_bounded_replay_cache_roundtrip(tmp_path): + cache = BoundedReplayCache( + tmp_path, + max_size_gb=0.01, + rank=1, + world_size=2, + ) + + assert cache.put("sample", _feature_sample()) is True + loaded = cache.get("sample") + + assert loaded is not None + assert torch.equal(loaded.input_ids, torch.arange(8)) + assert cache.metrics()["replay/cache_size_gb"] > 0 + assert (tmp_path / "rank00001" / "sample.pt").exists() + + +def test_bounded_replay_cache_disables_zero_budget(tmp_path): + cache = BoundedReplayCache( + tmp_path, + max_size_gb=0, + rank=0, + world_size=1, + ) + + assert cache.put("sample", _feature_sample()) is False + assert cache.get("sample") is None + + +def test_token_replay_algorithm_mismatch_is_warning_not_error(caplog): + replayer = TargetFeatureReplayer.__new__(TargetFeatureReplayer) + replayer.rank = 0 + replayer.algorithm = "DFLASH" + replayer.target_layer_ids = [1, 3] + replayer.hidden_layout = "dflash_aux" + replayer.strict_target_model_path = False + replayer._warned_replay_algorithm_mismatch = False + replayer._warned_replay_layer_mismatch = False + replayer._warned_replay_layout_mismatch = False + sample = DraftReplaySample( + algorithm="DSPARK", + input_ids=torch.arange(8), + loss_mask=torch.ones(8), + attention_mask=torch.ones(8, dtype=torch.bool), + position_ids=torch.arange(8), + feature_positions=torch.arange(2, 6), + draft_position_ids=torch.arange(3, 7), + metadata={ + "target_layer_ids": [2, 4], + "hidden_states_layout": "dflash_aux_plus_last", + }, + ) + + replayer._validate_target_path(sample) + + assert "token replay algorithm differs" in caplog.text + assert "token replay target layers differ" in caplog.text + assert "token replay hidden layout differs" in caplog.text + + +def test_vllm_payload_maps_suffix_hidden_rows_to_absolute_positions(): + replayer = TargetFeatureReplayer.__new__(TargetFeatureReplayer) + replayer.rank = 0 + replayer.target_layer_ids = [1, 3] + replayer.hidden_layout = "dflash_aux_plus_last" + replayer.dtype = torch.float32 + replayer.model_path = "/target" + replayer.target_revision = None + replayer.target_config_fingerprint = "unit" + replayer.use_logits = False + replayer.vllm_final_norm = torch.nn.RMSNorm(4, eps=1e-6) + + sample = DraftReplaySample( + algorithm="DSPARK", + input_ids=torch.arange(10, dtype=torch.long), + loss_mask=torch.ones(10, dtype=torch.float32), + attention_mask=torch.ones(10, dtype=torch.bool), + position_ids=torch.arange(10, dtype=torch.long), + feature_positions=torch.arange(4, 10, dtype=torch.long), + draft_position_ids=torch.arange(5, 11, dtype=torch.long), + metadata={"global_step": 1}, + ) + hidden = torch.arange(5 * 3 * 4, dtype=torch.float32).reshape(5, 3, 4) + payload = { + "token_ids": torch.arange(10, dtype=torch.long), + "hidden_states": hidden, + } + + feature = replayer._feature_from_vllm_payload( + sample, + payload, + prompt_ids=list(range(10)), + source="token_replay_vllm_file", + ) + + assert torch.equal(feature.input_ids, torch.arange(5, 10)) + assert torch.equal(feature.position_ids, torch.arange(6, 11)) + assert feature.hidden_states.shape == (5, 12) + assert feature.metadata["feature_start"] == 5 + assert feature.metadata["feature_end"] == 10 + assert feature.metadata["vllm_hidden_position_offset"] == 5 + torch.testing.assert_close(feature.hidden_states[:, :8], hidden[:, :2].flatten(1)) + torch.testing.assert_close( + feature.hidden_states[:, 8:], replayer.vllm_final_norm(hidden[:, 2]) + ) + + +@pytest.mark.parametrize("model_type", ["llama", "qwen2", "qwen3", "qwen3_moe"]) +@pytest.mark.parametrize("sharded", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_vllm_final_norm_matches_target_forward( + tmp_path, monkeypatch, model_type, sharded, dtype +): + transformers = pytest.importorskip("transformers") + from transformers.models.auto.configuration_auto import CONFIG_MAPPING + from verl_speco import checkpoint_tensor + + if model_type not in CONFIG_MAPPING: + pytest.skip(f"Installed Transformers does not include {model_type}") + config = transformers.AutoConfig.for_model( + model_type, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=32, + rms_norm_eps=1e-5, + head_dim=4, + moe_intermediate_size=8, + num_experts=2, + num_experts_per_tok=1, + ) + model = transformers.AutoModelForCausalLM.from_config(config).to(dtype).eval() + with torch.no_grad(): + model.model.norm.weight.copy_(torch.linspace(0.5, 2.0, 8)) + model.save_pretrained(tmp_path, max_shard_size="1KB" if sharded else "1GB") + loaded_keys = [] + original_load = checkpoint_tensor._load_checkpoint_tensor + + def load_one(path, key): + loaded_keys.append(key) + return original_load(path, key) + + monkeypatch.setattr(checkpoint_tensor, "_load_checkpoint_tensor", load_one) + norm = load_vllm_final_norm(str(tmp_path), dtype=dtype) + assert loaded_keys == ["model.norm.weight"] + assert all( + p.device.type == "cpu" and not p.requires_grad for p in norm.parameters() + ) + + captured = {} + handle = model.model.norm.register_forward_pre_hook( + lambda module, args: captured.update(final_input=args[0].detach().clone()) + ) + ids = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + output = model(ids, output_hidden_states=True) + handle.remove() + # A connector-style payload: auxiliary layer output + final PRE-norm output. + raw = torch.stack([output.hidden_states[1][0], captured["final_input"][0]], dim=1) + original_raw = raw.clone() + request = DraftReplaySample( + input_ids=ids[0], + loss_mask=torch.ones(4), + attention_mask=torch.ones(4), + position_ids=torch.arange(4), + feature_positions=torch.arange(4), + draft_position_ids=torch.arange(1, 5), + ) + for algorithm, layout in [ + ("DSPARK", "dflash_aux_plus_last"), + ("EAGLE3", "eagle3_aux_plus_last"), + ]: + contract = FeatureContract( + algorithm=algorithm, + target_layer_ids=[0], + hidden_states_layout=layout, + dtype=dtype, + target_model_id=str(tmp_path), + target_model_revision=None, + tokenizer_fingerprint="test", + ) + payload = {"token_ids": ids[0], "hidden_states": raw} + feature = feature_from_vllm_payload(payload, request, contract, final_norm=norm) + torch.testing.assert_close( + feature.hidden_states[:, :8], raw[:, 0], rtol=0, atol=0 + ) + torch.testing.assert_close( + feature.hidden_states[:, 8:], output.hidden_states[-1][0] + ) + torch.testing.assert_close(raw, original_raw, rtol=0, atol=0) + assert not feature.hidden_states.requires_grad + # Re-reading the same raw payload must not apply norm to an already-mutated tensor. + again = feature_from_vllm_payload(payload, request, contract, final_norm=norm) + torch.testing.assert_close( + again.hidden_states, feature.hidden_states, rtol=0, atol=0 + ) + with pytest.raises(ValueError, match="require the target final norm"): + feature_from_vllm_payload(payload, request, contract) + aux_only = feature_from_vllm_payload( + payload, + request, + replace(contract, hidden_states_layout="dflash_aux", algorithm="DFLASH"), + ) + torch.testing.assert_close(aux_only.hidden_states, raw[:, 0], rtol=0, atol=0) + + +def test_final_norm_loader_requires_checkpoint_weight(tmp_path): + transformers = pytest.importorskip("transformers") + from safetensors import SafetensorError + from safetensors.torch import save_file + + config = transformers.LlamaConfig( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=32, + ) + config.save_pretrained(tmp_path) + save_file({"unrelated.weight": torch.ones(8)}, str(tmp_path / "model.safetensors")) + with pytest.raises(SafetensorError, match="model.norm.weight"): + load_vllm_final_norm(str(tmp_path), dtype=torch.float32) + + +def test_vllm_replay_initializes_norm_and_invalidates_old_cache(tmp_path, monkeypatch): + from omegaconf import OmegaConf + import transformers + + transformers.LlamaConfig( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=32, + ).save_pretrained(tmp_path) + calls = [] + norm = torch.nn.RMSNorm(8) + + def loader(*args, **kwargs): + calls.append(args) + return norm + + monkeypatch.setattr( + "verl_speco.trainer.target_feature_replay.load_vllm_final_norm", loader + ) + config = OmegaConf.create( + { + "actor_rollout_ref": { + "model": {"path": str(tmp_path)}, + "rollout": { + "drafter": { + "speculative_algorithm": "DSPARK", + "target_layer_ids": [0], + "training": { + "use_logits": False, + "dspark_l1_loss_alpha": 0.9, + "target_feature_replay": { + "backend": "vllm_file", + "dtype": "float32", + }, + }, + } + }, + } + } + ) + replayer = TargetFeatureReplayer( + config, rank=0, world_size=1, device=torch.device("cpu") + ) + assert calls == [(str(tmp_path),)] + assert replayer.model is None # No full target model was loaded for replay. + assert replayer.vllm_final_norm is norm + sample = DraftReplaySample( + input_ids=torch.arange(4), + loss_mask=torch.ones(4), + attention_mask=torch.ones(4), + position_ids=torch.arange(4), + feature_positions=torch.arange(4), + draft_position_ids=torch.arange(1, 5), + ) + new_key = replayer._cache_key(sample) + replayer.backend = "torch" + assert new_key != replayer._cache_key(sample) + config.actor_rollout_ref.rollout.drafter.training.target_feature_replay.backend = ( + "torch" + ) + calls.clear() + torch_replayer = TargetFeatureReplayer( + config, rank=0, world_size=1, device=torch.device("cpu") + ) + assert calls == [] + assert torch_replayer.vllm_final_norm is None diff --git a/tests/unit/test_tq_consumer.py b/tests/unit/test_tq_consumer.py new file mode 100644 index 00000000..92f33770 --- /dev/null +++ b/tests/unit/test_tq_consumer.py @@ -0,0 +1,255 @@ +# Copyright 2026 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. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from verl_speco.trainer.feature_store import ( + DraftFeatureSample, + build_feature_store_from_config, +) +from verl_speco.trainer.tq_feature_store import ReadyEntry, TQFeatureStore +from verl_speco.trainer.tq_sample_source import ( + TQFeatureDataLoader, + build_assignments, +) +from verl_speco.transport.drafter_sample_protocol import ( + PROTOCOL_SCHEMA_VERSION, + SampleMetadata, + encode_sample, + make_ready_tag, + make_sample_key, +) + + +def _config() -> dict: + return { + "enable": True, + "ray": {"address": "ray-head:6379", "namespace": "speco-drafter"}, + "partition_id": "speco_drafter_features", + "run_id": "run-a", + "schema_version": PROTOCOL_SCHEMA_VERSION, + } + + +def _metadata(sequence_no: int = 0) -> SampleMetadata: + return SampleMetadata( + schema_version=PROTOCOL_SCHEMA_VERSION, + run_id="run-a", + sample_id=f"sample-{sequence_no}", + sequence_no=sequence_no, + ) + + +def _sample() -> DraftFeatureSample: + return DraftFeatureSample( + algorithm="DSPARK", + input_ids=torch.tensor([10, 11, 12]), + loss_mask=torch.tensor([1.0, 1.0, 1.0]), + position_ids=torch.tensor([0, 1, 2]), + hidden_states=torch.arange(12, dtype=torch.float32).reshape(3, 4), + metadata={"target_model_revision": "producer-revision"}, + ) + + +def _entry(sequence_no: int) -> ReadyEntry: + meta = _metadata(sequence_no) + return ReadyEntry(key=make_sample_key(meta), tag=make_ready_tag(meta)) + + +def test_feature_store_factory_builds_tq_without_path() -> None: + store = build_feature_store_from_config( + {"type": "tq", "path": None}, + read_only=True, + transfer_queue_cfg=_config(), + ) + assert isinstance(store, TQFeatureStore) + assert store.run_id == "run-a" + + +def test_tq_store_connect_filter_sort_and_minimal_decode(monkeypatch) -> None: + import verl_speco.trainer.tq_feature_store as module + + calls: list[tuple] = [] + monkeypatch.setattr(module, "configure_transfer_queue", lambda cfg: True) + monkeypatch.setattr( + module, + "connect_ray_cluster", + lambda address, namespace: calls.append(("ray", address, namespace)), + ) + monkeypatch.setattr( + module, + "connect_transfer_queue_client", + lambda: calls.append(("tq",)), + ) + entries = [_entry(2), _entry(1)] + unrelated = ReadyEntry( + key="other", + tag={ + **entries[0].tag, + "run_id": "another-run", + }, + ) + monkeypatch.setattr( + module, + "list_samples", + lambda: { + entries[0].key: entries[0].tag, + unrelated.key: unrelated.tag, + entries[1].key: entries[1].tag, + "control:v2:run-a:eos": { + "record_type": "control", + "status": "eos", + "schema_version": PROTOCOL_SCHEMA_VERSION, + "run_id": "run-a", + "total_samples": 2, + }, + }, + ) + fields_by_key = { + entry.key: encode_sample(_sample(), _metadata(int(entry.tag["sequence_no"]))) + for entry in entries + } + monkeypatch.setattr( + module, + "get_samples", + lambda keys: [(key, fields_by_key[key]) for key in keys], + ) + + store = TQFeatureStore.from_config(_config()) + store.connect() + ready = store.list_ready() + samples = store.get_many(ready) + + assert calls == [("ray", "ray-head:6379", "speco-drafter"), ("tq",)] + assert [entry.tag["sequence_no"] for entry in ready] == [1, 2] + assert len(samples) == 2 + # Model/revision/tokenizer/layers are intentionally not in the first + # Consumer contract; tensor and run/protocol checks still execute. + assert samples[0].metadata["target_model_revision"] == "producer-revision" + eos = store.read_eos() + assert eos is not None and eos.total_samples == 2 + + +def test_build_assignments_is_disjoint_and_complete() -> None: + entries = [_entry(index) for index in range(4)] + assignments = build_assignments(entries, batch_size=2, world_size=2) + assert [[entry.key for entry in rank] for rank in assignments] == [ + [entries[0].key, entries[1].key], + [entries[2].key, entries[3].key], + ] + + +class _FakeStore: + def __init__(self, entries: list[ReadyEntry], *, eos: bool = True): + self.entries = list(entries) + self.eos = eos + self.connected = False + self.get_calls: list[list[str]] = [] + self.clear_calls: list[list[str]] = [] + + def connect(self) -> None: + self.connected = True + + def owner_ready(self) -> bool: + return True + + def list_ready(self): + return list(self.entries) + + def read_eos(self): + return SimpleNamespace(total_samples=len(self.entries)) if self.eos else None + + def get_many(self, entries): + self.get_calls.append([entry.key for entry in entries]) + return [_sample() for _ in entries] + + def clear_many(self, keys): + normalized = list(keys) + self.clear_calls.append(normalized) + selected = set(normalized) + self.entries = [entry for entry in self.entries if entry.key not in selected] + + +def test_world_size_one_streams_trains_then_clears_and_stops() -> None: + entries = [_entry(0), _entry(1)] + store = _FakeStore(entries) + loader = TQFeatureDataLoader( + store, batch_size=2, rank=0, world_size=1, poll_interval_seconds=0.01 + ) + iterator = iter(loader) + batch = next(iterator) + assert batch.local_keys == [entry.key for entry in entries] + assert batch.global_keys == batch.local_keys + assert batch.global_sequence_nos == [0, 1] + assert store.get_calls == [batch.local_keys] + + loader.clear_completed_batch(batch.global_keys) + with pytest.raises(StopIteration): + next(iterator) + assert store.clear_calls == [batch.local_keys] + + +def test_eos_drops_incomplete_global_tail() -> None: + entry = _entry(0) + store = _FakeStore([entry]) + loader = TQFeatureDataLoader(store, batch_size=2, rank=0, world_size=1) + assert list(loader) == [] + assert store.clear_calls == [[entry.key]] + + +def test_rank0_discovery_failure_is_raised_without_clearing() -> None: + store = _FakeStore([_entry(0)]) + + def _fail_list(): + raise RuntimeError("list failed") + + store.list_ready = _fail_list + loader = TQFeatureDataLoader(store, batch_size=1, rank=0, world_size=1) + with pytest.raises(RuntimeError, match="list failed"): + next(iter(loader)) + assert store.clear_calls == [] + + +def test_nonzero_rank_uses_broadcast_assignment_without_listing(monkeypatch) -> None: + import verl_speco.trainer.tq_sample_source as module + + entries = [_entry(0), _entry(1)] + store = _FakeStore(entries, eos=False) + + def _broadcast(payload, src): + assert src == 0 + payload[0] = { + "kind": "batch", + "global_keys": [entry.key for entry in entries], + "assignments": [ + [{"key": entries[0].key, "tag": entries[0].tag}], + [{"key": entries[1].key, "tag": entries[1].tag}], + ], + } + + monkeypatch.setattr(module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(module.dist, "get_world_size", lambda: 2) + monkeypatch.setattr(module.dist, "broadcast_object_list", _broadcast) + store.list_ready = lambda: pytest.fail("nonzero rank must not list TQ keys") + + loader = TQFeatureDataLoader(store, batch_size=1, rank=1, world_size=2) + batch = next(iter(loader)) + assert batch.local_keys == [entries[1].key] + assert batch.global_keys is None + assert store.get_calls == [[entries[1].key]] diff --git a/tests/unit/test_tq_producer.py b/tests/unit/test_tq_producer.py new file mode 100644 index 00000000..b8d3543d --- /dev/null +++ b/tests/unit/test_tq_producer.py @@ -0,0 +1,508 @@ +# Copyright 2026 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. + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import Any + +import pytest +import torch + +from verl_speco.producer.vllm_feature_client import RawVllmFeature +from verl_speco.standalone_tq_producer import run_producer, validate_producer_config +from verl_speco.trainer.standalone_resume import save_standalone_resume +from verl_speco.transport.drafter_sample_protocol import PROTOCOL_SCHEMA_VERSION +from verl_speco.transport.drafter_sample_protocol import decode_sample + + +@pytest.fixture(autouse=True) +def target_final_norm(monkeypatch): + # These pipeline tests use a fake /target checkpoint. Loader accuracy is + # covered separately with real tiny HF checkpoints. + norm = torch.nn.RMSNorm(2, eps=1e-6).requires_grad_(False) + with torch.no_grad(): + norm.weight.copy_(torch.tensor([2.0, 3.0])) + monkeypatch.setattr( + "verl_speco.standalone_tq_producer.load_vllm_final_norm", + lambda *args, **kwargs: norm, + ) + return norm + + +def _config(input_path: Path) -> dict[str, Any]: + return { + "speco": { + "standalone_tq_producer": { + "input_path": str(input_path), + "tokenizer_path": "/target", + "tokenizer_fingerprint": "sha256:tokenizer", + "target_model_id": "/target", + "target_model_revision": "rev-a", + "target_layer_ids": [2, 8], + "hidden_dtype": "float32", + "trust_remote_code": False, + "vllm_endpoints": ["http://vllm:8000/v1"], + "vllm_model": "/target", + "request_timeout": 10, + "max_inflight_requests": 2, + "per_endpoint_concurrency": 1, + "input_queue_size": 2, + "publish_queue_size": 2, + "max_pending_samples": 8, + "pending_poll_interval_seconds": 0.01, + "owner_ready_timeout_seconds": 1, + "max_sequence_length": 16, + "max_feature_length": 8, + "generation_max_tokens": 4, + } + }, + "actor_rollout_ref": { + "rollout": { + "drafter": { + "speculative_algorithm": "DSPARK", + "training": { + "use_logits": False, + "dspark_l1_loss_alpha": 0.9, + "transfer_queue": { + "enable": True, + "package_version": "0.1.10", + "ray": { + "address": "ray-head:6379", + "namespace": "speco-drafter", + }, + "partition_id": "speco_drafter_features", + "run_id": "run-a", + "schema_version": PROTOCOL_SCHEMA_VERSION, + }, + }, + } + } + }, + } + + +class _Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool) -> dict[str, list[int]]: + assert add_special_tokens is False + values = { + "Q1: ": [1, 2], + "Q1: A1": [1, 2, 3, 4], + "Q2: ": [5, 6], + "Q2: A2": [5, 6, 7, 8], + } + return {"input_ids": values[text]} + + +class _ChatTokenizer: + def apply_chat_template(self, messages, *, tokenize, add_generation_prompt): + assert messages == [{"role": "user", "content": "Q3"}] + assert tokenize is True + assert add_generation_prompt is True + return [9, 10] + + +class _Transport: + def __init__(self, *, fail_sample_put: bool = False): + self.fail_sample_put = fail_sample_put + self.records: dict[str, dict[str, Any]] = { + "control:v2:run-a:owner-ready": { + "record_type": "control", + "status": "owner_ready", + "schema_version": PROTOCOL_SCHEMA_VERSION, + "run_id": "run-a", + } + } + self.payloads: dict[str, dict[str, torch.Tensor]] = {} + self.closed = False + + def configure_transfer_queue(self, config: dict[str, Any]) -> bool: + return bool(config["enable"]) + + def connect_ray_cluster(self, address: str, namespace: str | None) -> None: + assert (address, namespace) == ("ray-head:6379", "speco-drafter") + + def connect_transfer_queue_client(self) -> None: + pass + + def list_samples(self) -> dict[str, dict[str, Any]]: + return dict(self.records) + + def put_sample( + self, + key: str, + fields: dict[str, torch.Tensor], + *, + tag: dict[str, Any], + ) -> None: + if self.fail_sample_put and tag.get("record_type") == "sample": + raise RuntimeError("put failed") + self.records[key] = dict(tag) + self.payloads[key] = fields + + def close_transfer_queue_client(self) -> None: + self.closed = True + + +class _Pool: + def __init__(self, root: Path, *, close_error: BaseException | None = None): + self.root = root + self.close_error = close_error + self.paths: list[Path] = [] + self.started = False + self.closed = False + self.generate_calls = 0 + self.prefill_calls = 0 + + async def start(self) -> None: + self.started = True + + async def prefill(self, request: Any) -> RawVllmFeature: + self.prefill_calls += 1 + path = self.root / f"{request.sample_id}.safetensors" + path.write_bytes(b"temporary") + self.paths.append(path) + token_ids = torch.tensor(request.prompt_token_ids, dtype=torch.int64) + hidden = torch.arange(token_ids.numel() * 3 * 2, dtype=torch.float32).reshape( + token_ids.numel(), 3, 2 + ) + return RawVllmFeature( + payload={"token_ids": token_ids, "hidden_states": hidden}, + temporary_path=str(path), + endpoint_url="http://vllm:8000/v1", + byte_size=path.stat().st_size, + ) + + async def generate(self, request: Any) -> RawVllmFeature: + self.generate_calls += 1 + path = self.root / f"{request.sample_id}.safetensors" + path.write_bytes(b"temporary") + self.paths.append(path) + # ExampleHiddenStatesConnector excludes the final generated token because + # it was never consumed by a model forward pass. + token_ids = torch.tensor([*request.prompt_token_ids, 11], dtype=torch.int64) + hidden = torch.arange(token_ids.numel() * 3 * 2, dtype=torch.float32).reshape( + token_ids.numel(), 3, 2 + ) + return RawVllmFeature( + payload={"token_ids": token_ids, "hidden_states": hidden}, + temporary_path=str(path), + endpoint_url="http://vllm:8000/v1", + byte_size=path.stat().st_size, + generated_token_ids=(11, 12), + ) + + async def close(self) -> None: + self.closed = True + if self.close_error is not None: + raise self.close_error + + +class _MisalignedPool(_Pool): + async def prefill(self, request: Any) -> RawVllmFeature: + raw = await super().prefill(request) + raw.payload["hidden_states"] = raw.payload["hidden_states"][-1:] + return raw + + +def _write_input(path: Path) -> None: + records = [ + {"sample_id": "sample-1", "prompt": "Q1: ", "response": "A1"}, + {"sample_id": "sample-2", "prompt": "Q2: ", "response": "A2"}, + ] + path.write_text( + "".join(json.dumps(record) + "\n" for record in records), + encoding="utf-8", + ) + + +def test_run_producer_publishes_samples_then_eos( + tmp_path: Path, target_final_norm +) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + transport = _Transport() + pool = _Pool(tmp_path) + + stats = asyncio.run( + run_producer( + _config(input_path), + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + sample_keys = [ + key + for key, tag in transport.records.items() + if tag.get("record_type") == "sample" + ] + eos_tags = [tag for tag in transport.records.values() if tag.get("status") == "eos"] + assert stats.input_count == stats.published_count == 2 + assert stats.failed_count == stats.dropped_count == stats.pending_bytes == 0 + assert len(sample_keys) == 2 + assert eos_tags == [ + { + "record_type": "control", + "status": "eos", + "schema_version": PROTOCOL_SCHEMA_VERSION, + "run_id": "run-a", + "total_samples": 2, + } + ] + assert all(not path.exists() for path in pool.paths) + assert pool.started and pool.closed and transport.closed + first_fields = transport.payloads[sorted(sample_keys)[0]] + assert tuple(first_fields["sample__hidden_states"].shape) == (3, 6) + # The existing response feature window starts at prompt_length - 1 = 1. + raw = torch.arange(24, dtype=torch.float32).reshape(4, 3, 2)[1:] + torch.testing.assert_close( + first_fields["sample__hidden_states"][:, :4], raw[:, :2].flatten(1) + ) + torch.testing.assert_close( + first_fields["sample__hidden_states"][:, 4:], target_final_norm(raw[:, 2]) + ) + first_key = sorted(sample_keys)[0] + sample = decode_sample( + first_key, transport.records[first_key], first_fields, {"run_id": "run-a"} + ) + torch.testing.assert_close( + sample.hidden_states[:, 4:], target_final_norm(raw[:, 2]) + ) + assert sample.metadata["last_hidden_state_norm"] == "target_final_norm" + + +@pytest.mark.parametrize("algorithm", ["DFLASH", "DSPARK"]) +def test_aux_only_producer_does_not_load_or_apply_final_norm( + tmp_path, monkeypatch, algorithm +): + def unexpected_load(*args, **kwargs): + raise AssertionError("aux-only features must not load a target final norm") + + monkeypatch.setattr( + "verl_speco.standalone_tq_producer.load_vllm_final_norm", unexpected_load + ) + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + config = _config(input_path) + drafter = config["actor_rollout_ref"]["rollout"]["drafter"] + drafter["speculative_algorithm"] = algorithm + drafter["training"]["dspark_l1_loss_alpha"] = 0.0 + transport = _Transport() + asyncio.run( + run_producer( + config, + transport=transport, + tokenizer=_Tokenizer(), + client_pool=_Pool(tmp_path), + ) + ) + sample_key = next( + key + for key, tag in transport.records.items() + if tag.get("record_type") == "sample" + ) + sample = decode_sample( + sample_key, + transport.records[sample_key], + transport.payloads[sample_key], + {"run_id": "run-a"}, + ) + raw_aux = torch.arange(24, dtype=torch.float32).reshape(4, 3, 2)[1:, :2].flatten(1) + torch.testing.assert_close(sample.hidden_states, raw_aux) + + +def test_run_producer_restarts_input_until_max_samples(tmp_path: Path) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + config = _config(input_path) + config["speco"]["standalone_tq_producer"]["max_samples"] = 5 + transport = _Transport() + pool = _Pool(tmp_path) + + stats = asyncio.run( + run_producer( + config, + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + sample_tags = [ + tag for tag in transport.records.values() if tag.get("record_type") == "sample" + ] + eos_tags = [tag for tag in transport.records.values() if tag.get("status") == "eos"] + assert stats.input_count == stats.published_count == 5 + assert sorted(tag["sequence_no"] for tag in sample_tags) == [0, 1, 2, 3, 4] + assert pool.prefill_calls == 5 + assert eos_tags[0]["total_samples"] == 5 + + +def test_run_producer_skips_consumed_sequences_before_vllm(tmp_path: Path) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + checkpoint_path = tmp_path / "draft_step_1" + save_standalone_resume( + checkpoint_path, + [0], + optimizer_step=1, + input_path=input_path, + ) + config = _config(input_path) + producer_cfg = config["speco"]["standalone_tq_producer"] + producer_cfg["resume_checkpoint_path"] = str(checkpoint_path) + producer_cfg["max_samples"] = 2 + transport = _Transport() + pool = _Pool(tmp_path) + + stats = asyncio.run( + run_producer( + config, + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + sequence_nos = sorted( + int(tag["sequence_no"]) + for tag in transport.records.values() + if tag.get("record_type") == "sample" + ) + assert stats.input_count == 2 + assert pool.prefill_calls == 2 + assert sequence_nos == [1, 2] + + +def test_run_producer_generates_response_for_verl_chat_prompt(tmp_path: Path) -> None: + input_path = tmp_path / "dapo.jsonl" + input_path.write_text( + json.dumps( + { + "prompt": [{"role": "user", "content": "Q3"}], + "reward_model": {"ground_truth": "42"}, + "extra_info": {"index": "dapo-row"}, + } + ) + + "\n", + encoding="utf-8", + ) + transport = _Transport() + pool = _Pool(tmp_path) + + stats = asyncio.run( + run_producer( + _config(input_path), + transport=transport, + tokenizer=_ChatTokenizer(), + client_pool=pool, + ) + ) + + sample_keys = [ + key + for key, tag in transport.records.items() + if tag.get("record_type") == "sample" + ] + assert stats.input_count == stats.published_count == 1 + assert pool.generate_calls == 1 + assert pool.prefill_calls == 1 + assert len(sample_keys) == 1 + fields = transport.payloads[sample_keys[0]] + assert fields["sample__input_ids"].tolist() == [10, 11] + assert fields["sample__loss_mask"].tolist() == [0.0, 1.0] + + +def test_run_producer_drops_misaligned_hidden_states_and_writes_eos( + tmp_path: Path, +) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + transport = _Transport() + pool = _MisalignedPool(tmp_path) + + stats = asyncio.run( + run_producer( + _config(input_path), + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + assert stats.input_count == 2 + assert stats.published_count == 0 + assert stats.dropped_count == 2 + assert not any( + tag.get("record_type") == "sample" for tag in transport.records.values() + ) + eos = next(tag for tag in transport.records.values() if tag.get("status") == "eos") + assert eos["total_samples"] == 0 + assert all(not path.exists() for path in pool.paths) + + +def test_run_producer_put_failure_keeps_temporary_file_and_omits_eos( + tmp_path: Path, +) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + transport = _Transport(fail_sample_put=True) + pool = _Pool(tmp_path) + + with pytest.raises(RuntimeError, match="put failed"): + asyncio.run( + run_producer( + _config(input_path), + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + assert any(path.exists() for path in pool.paths) + assert not any(tag.get("status") == "eos" for tag in transport.records.values()) + assert pool.closed and transport.closed + + +def test_validate_producer_rejects_consumer_partition_mismatch(tmp_path: Path) -> None: + config = _config(tmp_path / "input.jsonl") + config["actor_rollout_ref"]["rollout"]["drafter"]["training"]["transfer_queue"][ + "partition_id" + ] = "other" + + with pytest.raises(ValueError, match="partition_id"): + validate_producer_config(config) + + +def test_pool_close_failure_does_not_skip_transport_close(tmp_path: Path) -> None: + input_path = tmp_path / "input.jsonl" + _write_input(input_path) + transport = _Transport() + pool = _Pool(tmp_path, close_error=RuntimeError("pool close failed")) + + with pytest.raises(RuntimeError, match="pool close failed"): + asyncio.run( + run_producer( + _config(input_path), + transport=transport, + tokenizer=_Tokenizer(), + client_pool=pool, + ) + ) + + assert pool.closed and transport.closed diff --git a/tests/unit/test_transferqueue_bridge.py b/tests/unit/test_transferqueue_bridge.py new file mode 100644 index 00000000..ff6fa694 --- /dev/null +++ b/tests/unit/test_transferqueue_bridge.py @@ -0,0 +1,188 @@ +# Copyright 2026 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. +from __future__ import annotations + +import sys + +import pytest +import torch + +from verl_speco.integration import transferqueue_bridge as bridge + + +class _FakeClient: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class _FakeTQ: + def __init__(self) -> None: + self.init_calls = [] + self.close_calls = 0 + self.clear_calls = [] + self.records: dict[str, tuple[dict, dict]] = {} + self.client = _FakeClient() + + def init(self, config=None): + self.init_calls.append(config) + + def kv_put(self, *, key, partition_id, fields, tag): + self.records[key] = (dict(fields), dict(tag)) + + def kv_list(self, *, partition_id): + return {key: tag for key, (_, tag) in self.records.items()} + + def kv_batch_get(self, *, keys, partition_id): + return {key: self.records[key][0] for key in keys} + + def kv_clear(self, *, keys, partition_id): + self.clear_calls.append((list(keys), partition_id)) + for key in keys: + self.records.pop(key, None) + + def get_client(self): + return self.client + + def close(self): + self.close_calls += 1 + + +class _FakeRay: + def __init__(self) -> None: + self.initialized = False + self.init_calls = [] + self.shutdown_calls = 0 + + def is_initialized(self): + return self.initialized + + def init(self, **kwargs): + self.initialized = True + self.init_calls.append(kwargs) + + def shutdown(self): + self.initialized = False + self.shutdown_calls += 1 + + +@pytest.fixture +def fake_runtime(monkeypatch): + fake_tq = _FakeTQ() + fake_ray = _FakeRay() + monkeypatch.setattr(bridge, "tq", fake_tq) + monkeypatch.setattr(bridge, "_TQ_IMPORTABLE", True) + monkeypatch.setitem(sys.modules, "ray", fake_ray) + monkeypatch.setattr( + bridge, + "_state", + { + "enabled": False, + "configured": False, + "initialized": False, + "config": None, + "owner": False, + "ray_initialized_here": False, + "ray_address": None, + "ray_namespace": None, + }, + ) + return fake_tq, fake_ray + + +def _config(): + return { + "enable": True, + "package_version": "0.1.10", + "partition_id": "speco_drafter_features", + "run_id": "run-a", + "schema_version": 1, + "ray": {"address": "ray-head:6379", "namespace": "speco-drafter"}, + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": 16, "num_data_storage_units": 1}, + }, + } + + +def test_owner_connects_ray_and_receives_only_native_tq_config(fake_runtime) -> None: + fake_tq, fake_ray = fake_runtime + config = _config() + assert bridge.configure_transfer_queue(config) + bridge.connect_ray_cluster("ray-head:6379", "speco-drafter") + bridge.start_transfer_queue_owner(config) + + assert fake_ray.init_calls == [ + {"address": "ray-head:6379", "namespace": "speco-drafter"} + ] + native = bridge._to_plain_dict(fake_tq.init_calls[0]) + assert set(native) == {"controller", "backend"} + assert bridge._state["owner"] is True + + bridge.close_transfer_queue_owner() + assert fake_tq.close_calls == 1 + assert fake_ray.shutdown_calls == 1 + + +def test_client_put_list_get_many_clear_and_local_close(fake_runtime) -> None: + fake_tq, fake_ray = fake_runtime + assert bridge.configure_transfer_queue(_config()) + bridge.connect_ray_cluster("ray-head:6379", "speco-drafter") + bridge.connect_transfer_queue_client() + native = bridge._to_plain_dict(fake_tq.init_calls[0]) + assert set(native) == {"controller", "backend"} + assert native["backend"]["storage_backend"] == "SimpleStorage" + assert native["backend"]["SimpleStorage"] == { + "total_storage_size": 16, + "num_data_storage_units": 1, + } + + bridge.put_sample( + "k0", + {"hidden_states": torch.ones(2, 4), "ignored": None}, + tag={"status": "ready"}, + ) + bridge.put_sample( + "k1", + {"hidden_states": torch.zeros(3, 4)}, + tag={"status": "ready"}, + ) + + assert bridge.list_samples() == { + "k0": {"status": "ready"}, + "k1": {"status": "ready"}, + } + records = bridge.get_samples(["k1", "k0"]) + assert [key for key, _ in records] == ["k1", "k0"] + assert records[0][1]["hidden_states"].shape == (3, 4) + + bridge.clear_samples(["k0", "k1"]) + assert bridge.list_samples() == {} + bridge.close_transfer_queue_client() + assert fake_tq.client.closed is True + assert fake_tq.close_calls == 0 + assert fake_ray.shutdown_calls == 1 + + +def test_client_close_cannot_be_used_by_owner(fake_runtime) -> None: + _, _ = fake_runtime + config = _config() + bridge.configure_transfer_queue(config) + bridge.connect_ray_cluster("ray-head:6379", "speco-drafter") + bridge.start_transfer_queue_owner(config) + with pytest.raises(RuntimeError, match="owner"): + bridge.close_transfer_queue_client() diff --git a/tests/unit/test_vllm_feature_client.py b/tests/unit/test_vllm_feature_client.py new file mode 100644 index 00000000..d748f1de --- /dev/null +++ b/tests/unit/test_vllm_feature_client.py @@ -0,0 +1,57 @@ +# Copyright 2026 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. + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from verl_speco.producer.vllm_feature_client import ( + VllmEndpoint, + request_generate, +) + + +def test_request_generate_only_requests_generated_token_ids() -> None: + calls = [] + + class Completions: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + choices=[ + SimpleNamespace( + prompt_token_ids=[1, 2], + token_ids=[3, 4], + ) + ], + kv_transfer_params={"hidden_states_path": "/tmp/result.safetensors"}, + ) + + client = SimpleNamespace(completions=Completions()) + + response = asyncio.run( + request_generate( + VllmEndpoint("http://vllm:8000/v1", 1), + client, + [1, 2], + model="target", + max_tokens=128, + timeout=30, + ) + ) + + assert response.generated_token_ids == (3, 4) + assert calls[0]["max_tokens"] == 128 + assert calls[0]["extra_body"] == {"return_token_ids": True} diff --git a/tools/run_qwen3-8b_drafter_hidden_state_vllm.sh b/tools/run_qwen3-8b_drafter_hidden_state_vllm.sh new file mode 100644 index 00000000..239e4b18 --- /dev/null +++ b/tools/run_qwen3-8b_drafter_hidden_state_vllm.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# Copyright 2026 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. +set -euo pipefail +set -x + +# Start only the target-model vLLM used by the standalone TQ Producer. +# Run this script in its own terminal before starting the training script. +# +# Ascend example: +# Set DEVICE_ENV=ASCEND_RT_VISIBLE_DEVICES, VLLM_DEVICES and VLLM_TP below, +# then run: bash tools/run_qwen3-8b_drafter_hidden_state_vllm.sh +# CUDA example: +# Set DEVICE_ENV=CUDA_VISIBLE_DEVICES, VLLM_DEVICES and VLLM_TP below, +# then run: bash tools/run_qwen3-8b_drafter_hidden_state_vllm.sh + +MODEL_PATH=${MODEL_PATH:-/path/to/Qwen3-8B} +DEVICE_ENV=${DEVICE_ENV:-ASCEND_RT_VISIBLE_DEVICES} +# Devices assigned to target-model vLLM. The script splits this list into +# consecutive groups of VLLM_TP devices and starts one service per group. +VLLM_DEVICES=${VLLM_DEVICES:-0,1,2,3,4,5} +VLLM_TP=${VLLM_TP:-1} +VLLM_HOST=${VLLM_HOST:-127.0.0.1} +VLLM_BASE_PORT=${VLLM_BASE_PORT:-8000} +VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION:-0.8} +VLLM_MAX_NUM_SEQS=${VLLM_MAX_NUM_SEQS:-256} +# Auxiliary training layers followed by the target model's final hidden-state +# layer. Keep the auxiliary prefix aligned with DSPARK_TARGET_LAYER_IDS in the +# standalone training script. DSpark L1 loss consumes the final entry. +VLLM_HIDDEN_STATE_LAYER_IDS=${VLLM_HIDDEN_STATE_LAYER_IDS:-'[1,9,17,25,33,36]'} +HIDDEN_STATES_DIR=${HIDDEN_STATES_DIR:-/tmp/speco-vllm-hidden-states} + +if ! [[ "${VLLM_TP}" =~ ^[1-9][0-9]*$ ]]; then + echo "VLLM_TP must be a positive integer, got: ${VLLM_TP}" >&2 + exit 2 +fi +if ! [[ "${VLLM_BASE_PORT}" =~ ^[0-9]+$ ]]; then + echo "VLLM_BASE_PORT must be an integer, got: ${VLLM_BASE_PORT}" >&2 + exit 2 +fi + +visible_devices=${VLLM_DEVICES} + +IFS=',' read -r -a DEVICE_IDS <<< "${visible_devices}" +for index in "${!DEVICE_IDS[@]}"; do + DEVICE_IDS[index]=${DEVICE_IDS[index]//[[:space:]]/} + if [[ -z "${DEVICE_IDS[index]}" ]]; then + echo "Visible device list contains an empty item: ${visible_devices}" >&2 + exit 2 + fi +done + +device_count=${#DEVICE_IDS[@]} +if (( device_count % VLLM_TP != 0 )); then + echo "Visible device count (${device_count}) must be divisible by VLLM_TP (${VLLM_TP}): ${visible_devices}" >&2 + exit 2 +fi +service_count=$((device_count / VLLM_TP)) + +SPECULATIVE_CONFIG=$(printf '{"method":"extract_hidden_states","num_speculative_tokens":1,"draft_model_config":{"hf_config":{"eagle_aux_hidden_state_layer_ids":%s}}}' "${VLLM_HIDDEN_STATE_LAYER_IDS}") + +start_vllm() { + local devices=$1 + local port=$2 + local hidden_states_dir=$3 + shift 3 + local kv_transfer_config + kv_transfer_config=$(printf '{"kv_connector":"ExampleHiddenStatesConnector","kv_role":"kv_producer","kv_connector_extra_config":{"shared_storage_path":"%s","use_synchronization_lock":true}}' "${hidden_states_dir}") + env "${DEVICE_ENV}=${devices}" vllm serve "${MODEL_PATH}" \ + --host "${VLLM_HOST}" \ + --port "${port}" \ + --tensor-parallel-size "${VLLM_TP}" \ + --gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}" \ + --max-num-seqs "${VLLM_MAX_NUM_SEQS}" \ + --speculative-config "${SPECULATIVE_CONFIG}" \ + --kv-transfer-config "${kv_transfer_config}" \ + --no-enable-chunked-prefill \ + "$@" & + STARTED_PID=$! +} + +PIDS=() +ENDPOINTS=() +cleanup() { + if (( ${#PIDS[@]} > 0 )); then + kill "${PIDS[@]}" 2>/dev/null || true + wait "${PIDS[@]}" 2>/dev/null || true + fi +} +trap cleanup EXIT INT TERM + +for ((service_index = 0; service_index < service_count; service_index++)); do + first_device=$((service_index * VLLM_TP)) + service_devices=${DEVICE_IDS[first_device]} + for ((tp_index = 1; tp_index < VLLM_TP; tp_index++)); do + service_devices+=",${DEVICE_IDS[first_device + tp_index]}" + done + + service_port=$((VLLM_BASE_PORT + service_index)) + service_hidden_states_dir="${HIDDEN_STATES_DIR}/service-${service_index}" + mkdir -p "${service_hidden_states_dir}" + start_vllm \ + "${service_devices}" \ + "${service_port}" \ + "${service_hidden_states_dir}" \ + "$@" + PIDS+=("${STARTED_PID}") + ENDPOINTS+=("http://${VLLM_HOST}:${service_port}/v1") +done + +endpoint_list=$(IFS=,; echo "[${ENDPOINTS[*]}]") +echo "VLLM_SERVICES_STARTED count=${service_count} tp=${VLLM_TP} devices=${visible_devices} endpoints=${endpoint_list} pids=${PIDS[*]}" +echo "Use this for standalone training: SPECO_VLLM_ENDPOINTS='${endpoint_list}'" +set +e +wait -n "${PIDS[@]}" +status=$? +set -e +exit "${status}" diff --git a/tools/wait_for_vllm_endpoints.py b/tools/wait_for_vllm_endpoints.py new file mode 100644 index 00000000..a7a3aac4 --- /dev/null +++ b/tools/wait_for_vllm_endpoints.py @@ -0,0 +1,92 @@ +# Copyright 2026 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. +"""Wait until every OpenAI-compatible vLLM endpoint is ready.""" + +from __future__ import annotations + +import argparse +import time +from urllib.error import HTTPError, URLError +from urllib.request import urlopen + + +def parse_endpoint_list(value: str) -> list[str]: + """Parse the Hydra-style ``[url0,url1]`` used by the training launcher.""" + + raw = value.strip() + if not (raw.startswith("[") and raw.endswith("]")): + raise ValueError("endpoints must use [url0,url1] syntax") + endpoints = [ + item.strip().strip("'\"").rstrip("/") + for item in raw[1:-1].split(",") + if item.strip() + ] + if not endpoints: + raise ValueError("endpoints must contain at least one URL") + return endpoints + + +def wait_for_endpoints( + endpoints: list[str], + *, + timeout_seconds: float, + poll_interval_seconds: float, + request_timeout_seconds: float, +) -> None: + deadline = time.monotonic() + timeout_seconds + pending = set(endpoints) + while pending: + for endpoint in list(pending): + try: + with urlopen( + f"{endpoint}/models", timeout=request_timeout_seconds + ) as response: + if 200 <= int(response.status) < 300: + print(f"EXTERNAL_VLLM_READY endpoint={endpoint}", flush=True) + pending.remove(endpoint) + except (HTTPError, OSError, URLError): + pass + if not pending: + return + if time.monotonic() >= deadline: + raise TimeoutError( + "external hidden-state vLLM is not ready at: " + + ", ".join(sorted(pending)) + ) + time.sleep(poll_interval_seconds) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoints", required=True) + parser.add_argument("--timeout-seconds", type=float, default=120.0) + parser.add_argument("--poll-interval-seconds", type=float, default=1.0) + parser.add_argument("--request-timeout-seconds", type=float, default=2.0) + args = parser.parse_args() + + try: + endpoints = parse_endpoint_list(args.endpoints) + wait_for_endpoints( + endpoints, + timeout_seconds=args.timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + request_timeout_seconds=args.request_timeout_seconds, + ) + except (TimeoutError, ValueError) as exc: + parser.error(str(exc)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/verl_speco/backends/dflash_trainer_backend.py b/verl_speco/backends/dflash_trainer_backend.py index 1a4dd902..ddce232c 100644 --- a/verl_speco/backends/dflash_trainer_backend.py +++ b/verl_speco/backends/dflash_trainer_backend.py @@ -652,9 +652,12 @@ def forward( loss_sum_per_position = ( loss_per_token.view(bsz, n_blocks, self.block_size) * binary_weights ).sum(dim=(0, 1)) - correct_per_position = ( - correct.view(bsz, n_blocks, self.block_size).float().sum(dim=(0, 1)) - ) + correct_3d = correct.view(bsz, n_blocks, self.block_size) + pred_valid_3d = binary_weights[:, :, 1:].bool() + pred_correct_3d = correct_3d[:, :, 1:] & pred_valid_3d + simulated_accept_length_sum = pred_correct_3d.float().cumprod(dim=-1).sum() + simulated_accept_block_count = pred_valid_3d.any(dim=-1).float().sum() + correct_per_position = correct_3d.float().sum(dim=(0, 1)) loss_per_position = loss_sum_per_position / count_per_pos acc_per_position = correct_per_position / count_per_pos # Prefix acceptance per block; the per-position accuracies above are @@ -681,6 +684,8 @@ def forward( "quality_token_count": quality_token_count, "valid_token_count": binary_eval_mask.sum().float(), "weighted_token_count": flat_weights.sum().float(), + "simulated_accept_length_sum": simulated_accept_length_sum, + "simulated_accept_block_count": simulated_accept_block_count, "sanitized_rows": sanitized_rows, "masked_rows": masked_rows, "loss_sum_per_position": loss_sum_per_position, @@ -784,6 +789,7 @@ def _build_fallback_config(self, target_hf_config): if mask_token_id_cfg is not None else target_text_config.vocab_size - 1 ) + target_head_dim = getattr(target_text_config, "head_dim", None) target_layer_ids = training_cfg.get("dflash_target_layer_ids", None) if target_layer_ids is None: target_layer_ids = build_target_layer_ids( @@ -803,6 +809,7 @@ def _build_fallback_config(self, target_hf_config): getattr(target_text_config, "num_attention_heads"), ) ), + head_dim=int(target_head_dim) if target_head_dim is not None else None, vocab_size=int(target_text_config.vocab_size), rms_norm_eps=float(getattr(target_text_config, "rms_norm_eps", 1e-6)), max_position_embeddings=int( @@ -818,6 +825,19 @@ def _build_fallback_config(self, target_hf_config): architectures=["DFlashDraftModel"], ) + @staticmethod + def _target_rope_theta(target_text_config) -> float: + rope_theta = getattr(target_text_config, "rope_theta", None) + if rope_theta is not None: + return float(rope_theta) + rope_parameters = getattr(target_text_config, "rope_parameters", None) + if ( + isinstance(rope_parameters, dict) + and rope_parameters.get("rope_theta") is not None + ): + return float(rope_parameters["rope_theta"]) + return 10000.0 + def _load_state_file(self, path: str) -> dict: if path.endswith(".safetensors"): with safe_open(path, framework="pt", device="cpu") as f: diff --git a/verl_speco/backends/dspark_trainer_backend.py b/verl_speco/backends/dspark_trainer_backend.py index f8e42e02..e25eef2e 100644 --- a/verl_speco/backends/dspark_trainer_backend.py +++ b/verl_speco/backends/dspark_trainer_backend.py @@ -744,6 +744,7 @@ def _build_fallback_config(self, target_hf_config): target_layer_ids = self._training_value( training_cfg, "dspark_target_layer_ids", "dflash_target_layer_ids", None ) + target_head_dim = getattr(target_text_config, "head_dim", None) if target_layer_ids is None: from verl_speco.models.dflash import build_target_layer_ids @@ -771,6 +772,7 @@ def _build_fallback_config(self, target_hf_config): getattr(target_text_config, "num_attention_heads"), ) ), + head_dim=int(target_head_dim) if target_head_dim is not None else None, vocab_size=int(target_text_config.vocab_size), rms_norm_eps=float(getattr(target_text_config, "rms_norm_eps", 1e-6)), max_position_embeddings=int( diff --git a/verl_speco/backends/eagle3_trainer_backend.py b/verl_speco/backends/eagle3_trainer_backend.py index 31aa9c16..f4468263 100644 --- a/verl_speco/backends/eagle3_trainer_backend.py +++ b/verl_speco/backends/eagle3_trainer_backend.py @@ -13,7 +13,6 @@ # limitations under the License. import logging import os -from copy import deepcopy from typing import Any, Optional, cast import torch @@ -23,7 +22,11 @@ from verl.utils.device import get_device_id, get_device_name from verl_speco.backends.lr_scheduler import build_drafter_lr_scheduler -from verl_speco.models.auto import AutoDraftModelConfig, AutoEagle3DraftModel +from verl_speco.models.auto import ( + AutoDraftModelConfig, + AutoEagle3DraftModel, + eagle3_draft_config_from_target, +) from verl_speco.models.eagle.llama_eagle import resolve_eagle3_num_aux_hidden_states from verl_speco.models.target.target_head import TargetHead from verl_speco.trainer.checkpoint import log_drafter_checkpoint_step @@ -678,17 +681,20 @@ def build_model(self): spec_model_path = self.config.rollout.drafter.model_path config_path = os.path.join(spec_model_path, "config.json") target_hf_config = self._get_target_hf_config() + training_cfg = self.config.rollout.drafter.training # 1. Load config if os.path.exists(config_path): drafter_config = AutoDraftModelConfig.from_file(config_path) else: - drafter_config = deepcopy(target_hf_config) - drafter_config.num_hidden_layers = 1 - drafter_config.torch_dtype = torch.bfloat16 - drafter_config.tie_word_embeddings = False - drafter_config.architectures = ["LlamaForCausalLMEagle3"] + drafter_config = eagle3_draft_config_from_target( + target_hf_config, + training_cfg.get("eagle3_target_layer_ids"), + ) + drafter_config.dtype = torch.bfloat16 + if not hasattr(drafter_config, "pretraining_tp"): + drafter_config.pretraining_tp = 1 if not hasattr(drafter_config, "draft_vocab_size"): drafter_config.draft_vocab_size = drafter_config.vocab_size if not hasattr(drafter_config, "target_hidden_size"): @@ -738,7 +744,6 @@ def build_model(self): drafter_module.load_embedding(target_model_path) drafter_module.freeze_embedding() - training_cfg = self.config.rollout.drafter.training if drafter_module.draft_vocab_size != drafter_module.vocab_size: if checkpoint_has_vocab_mapping and self._has_valid_vocab_mapping( drafter_module @@ -1168,6 +1173,22 @@ def compute_loss(self, model, batch, _current_pad_size): quality_tokens = torch.tensor(0.0, device=input_ids.device, dtype=torch.float32) quality_topk = min(5, int(all_step_logits[0].size(-1))) quality_step_stats = [] + collect_diagnostics = bool( + getattr(self, "enable_standalone_training_metrics", False) + ) + loss_sum_per_position = None + correct_per_position = None + count_per_position = None + if collect_diagnostics: + loss_sum_per_position = torch.zeros( + length, device=input_ids.device, dtype=torch.float32 + ) + correct_per_position = torch.zeros( + length, device=input_ids.device, dtype=torch.float32 + ) + count_per_position = torch.zeros( + length, device=input_ids.device, dtype=torch.float32 + ) sparse_base_tokens = torch.tensor( 0.0, device=input_ids.device, dtype=torch.float32 ) @@ -1280,6 +1301,9 @@ def compute_loss(self, model, batch, _current_pad_size): quality_topk_correct += step_topk_correct step_tokens = valid_position.float().sum() quality_tokens += step_tokens + if collect_diagnostics: + correct_per_position[idx] = step_top1_correct + count_per_position[idx] = step_tokens quality_step_stats.append( { "step": idx, @@ -1305,6 +1329,8 @@ def compute_loss(self, model, batch, _current_pad_size): } ) step_loss_sum = per_token_ploss.sum() + if collect_diagnostics: + loss_sum_per_position[idx] = step_loss_sum # Apply EAGLE3 step-wise temporal decay total_local_ploss += (gamma**idx) * step_loss_sum @@ -1350,13 +1376,27 @@ def compute_loss(self, model, batch, _current_pad_size): quality_step_stats, ) - return { + result = { "total_local_vloss": torch.tensor(0.0, device=input_ids.device), "total_local_ploss": total_local_ploss, "local_num_tokens": total_local_tokens, "v_weight": 0.0, "p_weight": 1.0, } + if collect_diagnostics: + result["diagnostics"] = { + "correct_count": quality_top1_correct.detach(), + "eval_token_count": quality_tokens.detach(), + "top1_correct_count": quality_top1_correct.detach(), + "top5_correct_count": quality_topk_correct.detach(), + "quality_token_count": quality_tokens.detach(), + "valid_token_count": quality_tokens.detach(), + "weighted_token_count": total_local_tokens.detach(), + "loss_sum_per_position": loss_sum_per_position.detach(), + "correct_per_position": correct_per_position.detach(), + "count_per_position": count_per_position.detach(), + } + return result def _compute_target_p_padded(self, target_scores, t2d, loss_mask, length): with torch.no_grad(): diff --git a/verl_speco/backends/lr_scheduler.py b/verl_speco/backends/lr_scheduler.py index ed811035..2980885f 100644 --- a/verl_speco/backends/lr_scheduler.py +++ b/verl_speco/backends/lr_scheduler.py @@ -83,6 +83,51 @@ def get_lr(self) -> list[float]: return [base_lr * ratio for base_lr in self.base_lrs] +class LinearWarmupDecayLR(LRScheduler): + """Linear warmup followed by linear decay over successful optimizer steps.""" + + def __init__( + self, + optimizer: Optimizer, + *, + decay_steps: int, + min_lr_ratio: float = 0.0, + warmup_steps: int = 0, + last_epoch: int = -1, + ) -> None: + self.decay_steps = int(decay_steps) + self.min_lr_ratio = float(min_lr_ratio) + self.warmup_steps = int(warmup_steps) + if self.decay_steps <= 0: + raise ValueError(f"lr_decay_steps must be positive, got {self.decay_steps}") + if self.warmup_steps < 0: + raise ValueError( + f"lr_warmup_steps must be non-negative, got {self.warmup_steps}" + ) + if self.warmup_steps >= self.decay_steps: + raise ValueError( + "lr_warmup_steps must be smaller than lr_decay_steps, " + f"got warmup={self.warmup_steps}, decay={self.decay_steps}" + ) + if not 0.0 <= self.min_lr_ratio <= 1.0: + raise ValueError(f"min_lr_ratio must be in [0, 1], got {self.min_lr_ratio}") + super().__init__(optimizer, last_epoch=last_epoch) + + def _lr_ratio(self, step: int) -> float: + step = max(int(step), 0) + if self.warmup_steps > 0 and step < self.warmup_steps: + return float(step) / float(self.warmup_steps) + + decay_span = self.decay_steps - self.warmup_steps + progress = min(max(step - self.warmup_steps, 0) / decay_span, 1.0) + linear_ratio = 1.0 - progress + return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * linear_ratio + + def get_lr(self) -> list[float]: + ratio = self._lr_ratio(self.last_epoch) + return [base_lr * ratio for base_lr in self.base_lrs] + + def build_drafter_lr_scheduler(optimizer: Optimizer, train_cfg: Any) -> LRScheduler: """Build a drafter scheduler while retaining legacy warmup_style overrides.""" @@ -116,6 +161,16 @@ def build_drafter_lr_scheduler(optimizer: Optimizer, train_cfg: Any) -> LRSchedu num_cycles=float(0.5 if num_cycles is None else num_cycles), last_epoch=last_epoch, ) + if scheduler_type in {"linear", "linear_decay"}: + decay_steps = train_cfg.get("lr_decay_steps", train_cfg.get("step", 0)) + min_lr_ratio = train_cfg.get("min_lr_ratio", 0.0) + return LinearWarmupDecayLR( + optimizer, + decay_steps=int(decay_steps or 0), + min_lr_ratio=float(0.0 if min_lr_ratio is None else min_lr_ratio), + warmup_steps=warmup_steps, + last_epoch=last_epoch, + ) if scheduler_type in {"global_cosine", "clamped_global_cosine"}: decay_steps = train_cfg.get("lr_decay_steps", 100) min_lr_ratio = train_cfg.get("min_lr_ratio", 0.1) diff --git a/verl_speco/checkpoint_tensor.py b/verl_speco/checkpoint_tensor.py new file mode 100644 index 00000000..83cd373e --- /dev/null +++ b/verl_speco/checkpoint_tensor.py @@ -0,0 +1,58 @@ +# Copyright 2026 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. +"""Read a target checkpoint tensor without importing drafter model implementations.""" + +import glob +import json +import os + +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open + + +def _load_checkpoint_tensor(model_path: str, key: str) -> torch.Tensor: + if not os.path.exists(model_path): + model_path = snapshot_download(repo_id=model_path) + + index_paths = glob.glob(os.path.join(model_path, "*.index.json")) + if len(index_paths) > 1: + raise FileNotFoundError(f"Multiple index.json files found in {model_path}") + + if index_paths: + with open(index_paths[0], encoding="utf-8") as f: + index_json = json.load(f) + weight_map = index_json.get("weight_map", {}) + if key not in weight_map: + raise KeyError( + f"Tensor {key!r} is not present in checkpoint index for {model_path}" + ) + ckpt_file = os.path.join(model_path, weight_map[key]) + if ckpt_file.endswith(".safetensors"): + with safe_open(ckpt_file, framework="pt", device="cpu") as f: + return f.get_tensor(key) + return torch.load(ckpt_file, map_location="cpu", weights_only=True)[key] + + safetensors_path = os.path.join(model_path, "model.safetensors") + if os.path.exists(safetensors_path): + with safe_open(safetensors_path, framework="pt", device="cpu") as f: + return f.get_tensor(key) + + pytorch_path = os.path.join(model_path, "pytorch_model.bin") + if os.path.exists(pytorch_path): + return torch.load(pytorch_path, map_location="cpu", weights_only=True)[key] + + raise FileNotFoundError( + f"No index.json, model.safetensors or pytorch_model.bin found in {model_path}" + ) diff --git a/verl_speco/config/draft_trainer.yaml b/verl_speco/config/draft_trainer.yaml index c536d352..1a600607 100644 --- a/verl_speco/config/draft_trainer.yaml +++ b/verl_speco/config/draft_trainer.yaml @@ -44,3 +44,43 @@ actor_rollout_ref: master_addr: ${speco.draft_training.master_addr} master_port: ${speco.draft_training.master_port} standalone: ${speco.draft_training.standalone} + # Standalone/offline-only reconstruction for compact token_replay stores. + # The frozen target model is loaded lazily after the draft trainer starts. + target_feature_replay: + backend: torch + model_path: null + target_revision: null + dtype: bfloat16 + trust_remote_code: false + strict_target_model_path: false + logits_chunk_rows: 32 + vllm_endpoint: http://localhost:8000/v1 + # Optional endpoint pool. When non-null, this takes precedence over + # vllm_endpoint and requests use least-inflight routing with failover. + vllm_endpoints: null + vllm_model: null + request_timeout: 120 + max_retries: 3 + endpoint_cooldown: 5 + on_generate: delete + require_arange_positions: true + offline_generation: + input_type: token_replay + input_path: null + output_path: null + max_samples: 0 + batch_size: 1 + shuffle: false + cache: + enabled: false + path: null + max_size_gb: 0 + # Overlap vLLM file replay with FSDP training. + # This pipeline is standalone-only. + target_feature_pipeline: + enabled: false + # Global budgets. Standalone divides them across torchrun ranks. + concurrency: 16 + producer_prefetch_depth: 4 + prefetch_depth: 2 + queue_timeout: 300 diff --git a/verl_speco/config/speco_base.yaml b/verl_speco/config/speco_base.yaml index 133193e0..8c222f2b 100644 --- a/verl_speco/config/speco_base.yaml +++ b/verl_speco/config/speco_base.yaml @@ -21,6 +21,38 @@ speco: task_runner: verl_speco.integration.task_runner.SpecoTaskRunner ray_trainer: verl_speco.trainer.speco_ray_trainer.SpecoRayPPOTrainer + # One-process standalone Producer. Prompt-only verl rows are generated by the + # target vLLM; rows with a response are replayed directly for hidden states. + standalone_tq_producer: + input_path: null + # Set internally by the unified launcher when model_path is a resumable + # standalone checkpoint. Direct Producer users may leave it null. + resume_checkpoint_path: null + tokenizer_path: null + tokenizer_fingerprint: null + target_model_id: null + target_model_revision: null + target_layer_ids: null + hidden_dtype: bfloat16 + trust_remote_code: false + vllm_endpoints: + - http://localhost:8000/v1 + vllm_model: null + request_timeout: 120 + max_inflight_requests: 16 + per_endpoint_concurrency: 4 + input_queue_size: 32 + publish_queue_size: 16 + max_pending_samples: 1024 + # Zero means one input-file pass. The unified launcher sets this to the + # exact number of samples required by max_steps and the global batch size. + max_samples: 0 + pending_poll_interval_seconds: 0.5 + owner_ready_timeout_seconds: 120 + max_sequence_length: 8192 + max_feature_length: 512 + generation_max_tokens: 511 + actor_rollout_ref: rollout: drafter: @@ -86,6 +118,8 @@ actor_rollout_ref: train_batches_per_cycle: 4 lr: 1e-5 lr_warmup_steps: 0 + # Supported standalone drafter scheduler types: + # constant, cosine, linear, global_cosine. lr_scheduler_type: constant lr_decay_steps: 100 min_lr_ratio: 0.1 @@ -137,6 +171,8 @@ actor_rollout_ref: dspark_debug_log: false dspark_debug_log_first_n: 2 dspark_debug_log_interval: 100 + # EAGLE3 target layers whose hidden states are concatenated as drafter input. + eagle3_target_layer_ids: null # P-EAGLE (parallel drafting: COD-subsampled multi-depth forward + KL loss). peagle_num_draft_layers: 4 peagle_num_aux_hidden_states: 3 @@ -190,6 +226,7 @@ actor_rollout_ref: data_buffer_max_size: 1024 hidden_state_clip_value: 1.0e4 feature_store: + # `tq` selects the streaming standalone Consumer and does not use path. type: torch_shard path: null max_samples_per_shard: 1024 @@ -198,5 +235,35 @@ actor_rollout_ref: repeat: true prefetch_depth: 2 strict_schema: true + max_seq_len: 512 + window_mode: loss + tokenizer_path: null + trust_remote_code: false + train_on: last_assistant min_sample_step: null max_sample_step: null + # TransferQueue transport for standalone drafter training. The owner, + # vLLM Producer, and every torchrun Consumer rank share this connection + # configuration. Default off; requires TransferQueue==0.1.10. + transfer_queue: + enable: false + package_version: "0.1.10" + # TQ 0.1.10 discovers its named TransferQueueController through Ray. + # Standalone owner, Producer, and every torchrun rank must use the same + # address and namespace. + ray: + address: null + namespace: speco-drafter + partition_id: speco_drafter_features + run_id: null + schema_version: 2 + connect_timeout_seconds: 120 + poll_interval_seconds: 0.5 + drop_last: true + controller: + polling_mode: true + backend: + storage_backend: SimpleStorage + SimpleStorage: + total_storage_size: 100000 + num_data_storage_units: 8 diff --git a/verl_speco/draft_train_launcher.py b/verl_speco/draft_train_launcher.py index 1f6b1a6d..1cde487a 100644 --- a/verl_speco/draft_train_launcher.py +++ b/verl_speco/draft_train_launcher.py @@ -57,6 +57,14 @@ "speco.draft_training.standalone", "actor_rollout_ref.rollout.drafter.training.standalone", ) +_FEATURE_STORE_TYPE_KEY = ( + "actor_rollout_ref.rollout.drafter.training.feature_store.type" +) +_TQ_ENABLE_KEY = "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable" +_TQ_RAY_ADDRESS_KEY = ( + "actor_rollout_ref.rollout.drafter.training.transfer_queue.ray.address" +) +_TQ_RUN_ID_KEY = "actor_rollout_ref.rollout.drafter.training.transfer_queue.run_id" _LAUNCH_OVERRIDE_KEYS = frozenset( _NPROC_KEYS @@ -166,6 +174,29 @@ def normalize_training_args( return normalized +def validate_tq_launch_config(overrides: list[str]) -> None: + """Fail early when the standalone TQ Consumer lacks connection identity.""" + + store_type = _find_override(overrides, (_FEATURE_STORE_TYPE_KEY,)) + if str(store_type or "").strip().lower() != "tq": + return + enabled = _find_override(overrides, (_TQ_ENABLE_KEY,)) + if not _parse_bool(enabled, default=False): + raise ValueError( + "feature_store.type=tq requires training.transfer_queue.enable=true" + ) + ray_address = str(_find_override(overrides, (_TQ_RAY_ADDRESS_KEY,)) or "").strip() + if not ray_address or ray_address.lower() in {"null", "none"}: + raise ValueError( + "feature_store.type=tq requires training.transfer_queue.ray.address" + ) + run_id = str(_find_override(overrides, (_TQ_RUN_ID_KEY,)) or "").strip() + if not run_id or run_id.lower() in {"null", "none"}: + raise ValueError( + "feature_store.type=tq requires training.transfer_queue.run_id" + ) + + def build_torch_distributed_command( config: DraftTrainLaunchConfig, training_args: list[str], @@ -220,6 +251,7 @@ def main(argv: list[str] | None = None) -> int: ) args, training_args = parser.parse_known_args(argv) + validate_tq_launch_config(training_args) launch_config = resolve_launch_config(training_args, module=args.module) normalized_training_args = normalize_training_args(training_args, launch_config) command = build_torch_distributed_command( diff --git a/verl_speco/inspect_feature_store.py b/verl_speco/inspect_feature_store.py index 4c09cba9..3e841ddf 100644 --- a/verl_speco/inspect_feature_store.py +++ b/verl_speco/inspect_feature_store.py @@ -28,7 +28,7 @@ def main() -> int: parser = argparse.ArgumentParser( - description="Inspect a torch_shard draft feature store." + description="Inspect a torch_shard or token_replay draft feature store." ) parser.add_argument( "path", @@ -133,15 +133,21 @@ def _sample_summary(sample: dict[str, Any]) -> dict[str, str]: keys = [ "input_ids", "loss_mask", + "attention_mask", "hidden_states", "last_hidden_states", "target_logprobs", "position_ids", + "feature_positions", + "draft_position_ids", ] return {key: _shape(sample[key]) for key in keys if key in sample} def _sample_issues(sample: dict[str, Any]) -> list[str]: + if sample.get("sample_type") == "token_replay" or "feature_positions" in sample: + return _replay_sample_issues(sample) + issues: list[str] = [] input_ids = _tensor(sample.get("input_ids"), "input_ids", issues) loss_mask = _tensor(sample.get("loss_mask"), "loss_mask", issues) @@ -232,5 +238,54 @@ def _sample_issues(sample: dict[str, Any]) -> list[str]: return issues +def _replay_sample_issues(sample: dict[str, Any]) -> list[str]: + issues: list[str] = [] + input_ids = _tensor(sample.get("input_ids"), "input_ids", issues) + loss_mask = _tensor(sample.get("loss_mask"), "loss_mask", issues) + attention_mask = _tensor(sample.get("attention_mask"), "attention_mask", issues) + position_ids = _tensor(sample.get("position_ids"), "position_ids", issues) + feature_positions = _tensor( + sample.get("feature_positions"), "feature_positions", issues + ) + draft_position_ids = _tensor( + sample.get("draft_position_ids"), "draft_position_ids", issues + ) + if input_ids is None: + return issues + + seq_len = int(input_ids.numel()) + for name, value in ( + ("loss_mask", loss_mask), + ("attention_mask", attention_mask), + ("position_ids", position_ids), + ): + if value is not None and int(value.numel()) != seq_len: + issues.append( + f"{name} length {int(value.numel())} does not match " + f"input_ids length {seq_len}" + ) + if feature_positions is None or draft_position_ids is None: + return issues + feature_positions = feature_positions.reshape(-1).long() + draft_position_ids = draft_position_ids.reshape(-1).long() + if int(feature_positions.numel()) == 0: + issues.append("feature_positions is empty") + return issues + if int(draft_position_ids.numel()) != int(feature_positions.numel()): + issues.append( + "draft_position_ids length does not match feature_positions length" + ) + if ( + int(feature_positions.min().item()) < 0 + or int(feature_positions.max().item()) >= seq_len + ): + issues.append(f"feature_positions fall outside input_ids length {seq_len}") + if int(feature_positions.numel()) > 1 and not bool( + torch.all(feature_positions[1:] == feature_positions[:-1] + 1).item() + ): + issues.append("feature_positions are not contiguous and increasing") + return issues + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/verl_speco/inspect_jsonl_samples.py b/verl_speco/inspect_jsonl_samples.py new file mode 100644 index 00000000..4af85183 --- /dev/null +++ b/verl_speco/inspect_jsonl_samples.py @@ -0,0 +1,265 @@ +# Copyright 2026 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. +"""Inspect JSONL draft-training samples without loading model dependencies.""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + + +TOKEN_REPLAY_REQUIRED_KEYS = { + "input_ids", + "loss_mask", + "attention_mask", + "position_ids", + "feature_positions", + "draft_position_ids", +} +FEATURE_REQUIRED_KEYS = {"input_ids", "loss_mask", "hidden_states"} +INPUT_LOSS_REQUIRED_KEYS = {"input_ids", "loss_mask"} +VLLM_SAFETENSORS_MANIFEST_KEYS = {"path", "num_samples", "sample"} + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Inspect a JSONL file and report whether it matches SPECO sample schemas." + ) + parser.add_argument("path", help="JSONL file to inspect.") + parser.add_argument( + "--max-lines", + type=int, + default=20, + help="Maximum number of JSONL rows to inspect.", + ) + parser.add_argument( + "--show-first", + action="store_true", + help="Print the first JSON object with long arrays summarized.", + ) + parser.add_argument( + "--strict-exit", + action="store_true", + help="Exit with code 1 when inspected rows do not match a known schema.", + ) + args = parser.parse_args() + + path = Path(args.path) + summaries: list[dict[str, Any]] = [] + key_counts: Counter[str] = Counter() + schema_counts: Counter[str] = Counter() + issues_by_schema: dict[str, list[str]] = defaultdict(list) + first_obj: dict[str, Any] | None = None + + with path.open(encoding="utf-8") as jsonl_file: + for line_number, line in enumerate(jsonl_file, start=1): + if len(summaries) >= int(args.max_lines): + break + line = line.strip() + if not line: + continue + try: + obj = json.loads(line) + except json.JSONDecodeError as exc: + summaries.append({"line": line_number, "schema": "invalid_json"}) + schema_counts["invalid_json"] += 1 + issues_by_schema["invalid_json"].append(f"line {line_number}: {exc}") + continue + if first_obj is None and isinstance(obj, dict): + first_obj = obj + if not isinstance(obj, dict): + summaries.append({"line": line_number, "schema": type(obj).__name__}) + schema_counts[type(obj).__name__] += 1 + continue + keys = set(obj) + key_counts.update(keys) + schema, issues = _classify(obj) + schema_counts[schema] += 1 + issues_by_schema[schema].extend( + f"line {line_number}: {issue}" for issue in issues + ) + summaries.append( + { + "line": line_number, + "schema": schema, + "keys": sorted(keys), + "shapes": _shape_summary(obj), + } + ) + + print(f"jsonl={path}") + print(f"inspected_lines={len(summaries)}") + print("schema_counts:") + for schema, count in schema_counts.most_common(): + print(f" {schema}: {count}") + print("key_counts:") + for key, count in key_counts.most_common(): + print(f" {key}: {count}") + print("sample_summaries:") + for summary in summaries[: min(len(summaries), 5)]: + print(json.dumps(summary, ensure_ascii=False, sort_keys=True)) + if issues_by_schema: + print("issues:") + for schema, issues in issues_by_schema.items(): + for issue in issues[:10]: + print(f" [{schema}] {issue}") + if args.show_first and first_obj is not None: + print("first_object:") + print( + json.dumps( + _compact_json(first_obj), ensure_ascii=False, indent=2, sort_keys=True + ) + ) + + invalid = any( + schema + not in { + "token_replay_jsonl", + "feature_jsonl", + "input_loss_jsonl", + "vllm_safetensors_manifest", + } + for schema in schema_counts + ) + return 1 if invalid and args.strict_exit else 0 + + +def _classify(obj: dict[str, Any]) -> tuple[str, list[str]]: + keys = set(obj) + if VLLM_SAFETENSORS_MANIFEST_KEYS.issubset(keys): + return "vllm_safetensors_manifest", [] + if TOKEN_REPLAY_REQUIRED_KEYS.issubset(keys): + return "token_replay_jsonl", _token_replay_issues(obj) + if FEATURE_REQUIRED_KEYS.issubset(keys): + return "feature_jsonl", _feature_issues(obj) + if INPUT_LOSS_REQUIRED_KEYS.issubset(keys): + return "input_loss_jsonl", _input_loss_issues(obj) + missing_token = sorted(TOKEN_REPLAY_REQUIRED_KEYS - keys) + missing_feature = sorted(FEATURE_REQUIRED_KEYS - keys) + return ( + "unknown", + [ + f"missing token_replay keys={missing_token}", + f"missing feature keys={missing_feature}", + ], + ) + + +def _token_replay_issues(obj: dict[str, Any]) -> list[str]: + issues = [] + input_len = _flat_len(obj.get("input_ids")) + for key in ("loss_mask", "attention_mask", "position_ids"): + value_len = _flat_len(obj.get(key)) + if input_len is not None and value_len != input_len: + issues.append(f"{key} length {value_len} != input_ids length {input_len}") + feature_len = _flat_len(obj.get("feature_positions")) + draft_len = _flat_len(obj.get("draft_position_ids")) + if feature_len is not None and draft_len != feature_len: + issues.append( + f"draft_position_ids length {draft_len} != feature_positions length {feature_len}" + ) + return issues + + +def _feature_issues(obj: dict[str, Any]) -> list[str]: + issues = [] + input_len = _flat_len(obj.get("input_ids")) + loss_len = _flat_len(obj.get("loss_mask")) + if input_len is not None and loss_len != input_len: + issues.append(f"loss_mask length {loss_len} != input_ids length {input_len}") + hidden_shape = _shape(obj.get("hidden_states")) + if not hidden_shape or hidden_shape[0] in {"scalar", "dict", "str", "none"}: + issues.append("hidden_states is not an array-like value") + return issues + + +def _input_loss_issues(obj: dict[str, Any]) -> list[str]: + issues = [] + input_len = _flat_len(obj.get("input_ids")) + loss_len = _flat_len(obj.get("loss_mask")) + if input_len is None: + issues.append("input_ids is not a JSON list") + if loss_len is None: + issues.append("loss_mask is not a JSON list") + if input_len is not None and loss_len is not None and loss_len != input_len: + issues.append(f"loss_mask length {loss_len} != input_ids length {input_len}") + return issues + + +def _shape_summary(obj: dict[str, Any]) -> dict[str, Any]: + return { + key: _shape(value) + for key, value in obj.items() + if key + in { + "input_ids", + "loss_mask", + "attention_mask", + "position_ids", + "feature_positions", + "draft_position_ids", + "hidden_states", + "target_logprobs", + "sample", + "path", + } + } + + +def _shape(value: Any) -> list[Any]: + if value is None: + return ["none"] + if isinstance(value, dict): + return ["dict", sorted(value)[:20]] + if isinstance(value, str): + return ["str", len(value)] + if not isinstance(value, list): + return ["scalar", type(value).__name__] + shape = [] + current: Any = value + while isinstance(current, list): + shape.append(len(current)) + current = current[0] if current else None + return shape + + +def _flat_len(value: Any) -> int | None: + if not isinstance(value, list): + return None + current = value + while isinstance(current, list) and len(current) == 1: + current = current[0] + return len(current) if isinstance(current, list) else None + + +def _compact_json(value: Any) -> Any: + if isinstance(value, dict): + return {key: _compact_json(item) for key, item in value.items()} + if isinstance(value, list): + if len(value) > 16: + return { + "__list__": True, + "shape": _shape(value), + "head": [_compact_json(item) for item in value[:4]], + "tail": [_compact_json(item) for item in value[-4:]], + } + return [_compact_json(item) for item in value] + return value + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/verl_speco/integration/transferqueue_bridge.py b/verl_speco/integration/transferqueue_bridge.py new file mode 100644 index 00000000..15bca804 --- /dev/null +++ b/verl_speco/integration/transferqueue_bridge.py @@ -0,0 +1,587 @@ +# Copyright 2026 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. + +"""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.10. 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 collections.abc import Mapping, Sequence +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.10` or disable " + "actor_rollout_ref.rollout.drafter.training.transfer_queue.enable." + ) + + return _raise + + tq = _MockTQ() # type: ignore[assignment] + + +# --------------------------------------------------------------------------- +# State +# --------------------------------------------------------------------------- + +_state_lock = threading.Lock() +_state: dict[str, Any] = { + "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 + "ray_initialized_here": False, + "ray_address": None, + "ray_namespace": None, +} + + +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: + plain = _to_plain_dict(training_cfg) + if isinstance(plain, dict) and any( + key in plain for key in ("enable", "backend", "controller", "ray") + ): + transfer_queue_cfg = plain + else: + 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 connect_ray_cluster( + ray_address: str | None, + namespace: str | None = None, +) -> None: + """Connect this ordinary process to the Ray cluster hosting TQ. + + PR #48 workers are already Ray actors and therefore do not call this + function. Standalone owner, Producer and torchrun ranks must call it + before ``tq.init`` so TQ 0.1.10 can discover its named Controller actor. + """ + + try: + import ray + except ImportError as exc: # pragma: no cover - depends on optional package + raise RuntimeError( + "Ray is required by TransferQueue 0.1.10. Install TransferQueue==0.1.10 " + "and connect all standalone processes to the same Ray cluster." + ) from exc + + if ray.is_initialized(): + return + address = str(ray_address or "auto").strip() or "auto" + kwargs: dict[str, Any] = {"address": address} + if namespace: + kwargs["namespace"] = str(namespace) + ray.init(**kwargs) + with _state_lock: + _state["ray_initialized_here"] = True + _state["ray_address"] = address + _state["ray_namespace"] = namespace + + +def start_transfer_queue_owner(tq_config: Any) -> None: + """Create the task-level named TQ Controller in the current Ray cluster.""" + + plain = _extract_tq_config(tq_config) + if plain is None: + raise ValueError("TransferQueue owner configuration is missing") + if not bool(plain.get("enable", True)): + raise ValueError("TransferQueue owner requires transfer_queue.enable=true") + if not _TQ_IMPORTABLE: + raise RuntimeError("TransferQueue==0.1.10 is required to start the TQ owner") + with _state_lock: + if _state["initialized"]: + raise RuntimeError("TransferQueue is already initialized in this process") + tq.init(_as_tq_config(_native_tq_config(plain))) + with _state_lock: + _state["config"] = plain + _state["enabled"] = True + _state["configured"] = True + _state["initialized"] = True + _state["owner"] = True + logger.info("[SpeCo TQ] standalone owner started (partition=%s)", _partition_id()) + + +def connect_transfer_queue_client() -> None: + """Attach this process to the named TQ Controller on its Ray cluster.""" + + if not _TQ_IMPORTABLE: + raise RuntimeError("TransferQueue==0.1.10 is required to connect a TQ client") + if not bool(_state["enabled"]): + raise RuntimeError( + "configure_transfer_queue() must enable TQ before client connect" + ) + _ensure_initialized() + + +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(config)`` with the same native configuration 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(_as_tq_config(_native_tq_config(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(config)`` once per worker process. + + TransferQueue 0.1.10 first tries to discover the named controller and ignores + the supplied configuration when one already exists. Supplying the same + native configuration in every process is therefore safe for ordinary + clients and also prevents an unexpectedly early client from creating a + default-configured controller. + """ + + if _state["initialized"]: + return + with _state_lock: + if _state["initialized"]: + return + configured = _state.get("config") + if not isinstance(configured, Mapping): + raise RuntimeError( + "configure_transfer_queue() must provide TQ configuration before init" + ) + tq.init(_as_tq_config(_native_tq_config(configured))) + _state["initialized"] = True + + +def _native_tq_config(tq_cfg: Mapping[str, Any]) -> dict[str, Any]: + """Remove SPECO-only connection/protocol fields before ``tq.init``.""" + + project_keys = { + "enable", + "package_version", + "ray", + "partition_id", + "run_id", + "schema_version", + "connect_timeout_seconds", + "poll_interval_seconds", + "drop_last", + } + return {key: value for key, value in tq_cfg.items() if key not in project_keys} + + +def _as_tq_config(value: Mapping[str, Any]) -> Any: + """TQ 0.1.10 annotates its config as DictConfig; keep tests dependency-light.""" + + try: + from omegaconf import OmegaConf + + return OmegaConf.create(dict(value)) + except ImportError: # pragma: no cover - project normally depends on OmegaConf + return dict(value) + + +def _partition_id() -> str: + config = _state.get("config") or {} + value = config.get("partition_id") if isinstance(config, Mapping) else None + return str(value or _SPECO_TQ_PARTITION) + + +# --------------------------------------------------------------------------- +# 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.10 on first run.) + tq.kv_put( + key=key, + partition_id=_partition_id(), + 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=_partition_id()) + value = _extract_value(result, key) + if value is None: + return {} + return _tensordict_to_dict(value) + + +def list_samples() -> dict[str, dict[str, Any]]: + """Return key -> tag for the configured partition without fetching fields.""" + + if not is_transfer_queue_enabled(): + raise RuntimeError("list_samples called while TransferQueue is not enabled.") + _ensure_initialized() + result = tq.kv_list(partition_id=_partition_id()) + if result is None: + return {} + if not isinstance(result, Mapping): + raise TypeError(f"tq.kv_list returned unsupported type {type(result)!r}") + # 0.1.10 returns key -> tag when partition_id is supplied. Accept the + # partition -> (key -> tag) wrapper as well to keep the bridge version-safe. + nested = result.get(_partition_id()) + if isinstance(nested, Mapping) and all( + isinstance(v, Mapping) for v in nested.values() + ): + result = nested + records: dict[str, dict[str, Any]] = {} + for key, tag in result.items(): + if tag is None: + records[str(key)] = {} + elif isinstance(tag, Mapping): + records[str(key)] = dict(tag) + else: + raise TypeError( + f"TQ tag for key {key!r} must be a mapping, got {type(tag)!r}" + ) + return records + + +def get_samples(keys: Sequence[str]) -> list[tuple[str, dict[str, Any]]]: + """Batch-fetch records and return one plain field dict per input key.""" + + normalized_keys = [str(key) for key in keys] + if not normalized_keys: + return [] + if len(set(normalized_keys)) != len(normalized_keys): + raise ValueError("get_samples keys must be unique") + if not is_transfer_queue_enabled(): + raise RuntimeError("get_samples called while TransferQueue is not enabled.") + _ensure_initialized() + result = tq.kv_batch_get(keys=normalized_keys, partition_id=_partition_id()) + values = _split_batch_result(result, normalized_keys) + return [ + (key, _tensordict_to_dict(value)) + for key, value in zip(normalized_keys, values, strict=True) + ] + + +def clear_samples(keys: Sequence[str]) -> None: + """Delete consumed records from the configured partition.""" + + normalized_keys = [str(key) for key in keys] + if not normalized_keys: + return + if not is_transfer_queue_enabled(): + raise RuntimeError("clear_samples called while TransferQueue is not enabled.") + _ensure_initialized() + tq.kv_clear(keys=normalized_keys, partition_id=_partition_id()) + + +def close_transfer_queue_client() -> None: + """Close only this process's TQ client; never kill the shared Controller.""" + + with _state_lock: + if _state["owner"]: + raise RuntimeError("TQ owner must use close_transfer_queue_owner()") + initialized = bool(_state["initialized"]) + _state["initialized"] = False + if initialized and _TQ_IMPORTABLE: + try: + client = tq.get_client() + if client is not None: + client.close() + except Exception: # noqa: BLE001 + logger.debug("[SpeCo TQ] local client close raised", exc_info=True) + _shutdown_local_ray_connection() + + +def close_transfer_queue_owner() -> None: + """Close global TQ resources. Only the process that started them may call.""" + + close_transfer_queue() + _shutdown_local_ray_connection() + + +def _shutdown_local_ray_connection() -> None: + with _state_lock: + initialized_here = bool(_state["ray_initialized_here"]) + _state["ray_initialized_here"] = False + if not initialized_here: + return + try: + import ray + + if ray.is_initialized(): + ray.shutdown() + except ImportError: # pragma: no cover + return + + +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 _split_batch_result(result: Any, keys: Sequence[str]) -> list[Any]: + if result is None: + raise KeyError(f"TQ returned no payload for keys={list(keys)!r}") + # TensorDict exposes a batch_size and indexes rows with ``result[index]``. + # Check this before the generic Mapping branch because some TensorDict + # versions also satisfy mapping-like protocols for their field columns. + if getattr(result, "batch_size", None) is not None: + return _index_batch_rows(result, len(keys)) + if isinstance(result, Mapping): + if all(key in result for key in keys): + return [result[key] for key in keys] + if isinstance(result, (list, tuple)): + if len(result) != len(keys): + raise RuntimeError( + f"TQ returned {len(result)} rows for {len(keys)} requested keys" + ) + return list(result) + if len(keys) == 1: + return [result] + return _index_batch_rows(result, len(keys)) + + +def _index_batch_rows(result: Any, expected_rows: int) -> list[Any]: + try: + rows = [result[index] for index in range(expected_rows)] + except Exception as exc: # noqa: BLE001 + raise TypeError( + f"Unable to split TQ batch result of type {type(result)!r} into {expected_rows} rows" + ) from exc + if len(rows) != expected_rows: + raise RuntimeError( + f"TQ returned {len(rows)} rows for {expected_rows} requested keys" + ) + return rows + + +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", + "connect_ray_cluster", + "connect_transfer_queue_client", + "start_transfer_queue_owner", + "close_transfer_queue", + "close_transfer_queue_client", + "close_transfer_queue_owner", + "clear_samples", + "init_transfer_queue", + "is_transfer_queue_enabled", + "list_samples", + "make_sample_key", + "put_sample", + "get_sample", + "get_samples", +] diff --git a/verl_speco/models/auto.py b/verl_speco/models/auto.py index 6a7acaf1..12b735a4 100644 --- a/verl_speco/models/auto.py +++ b/verl_speco/models/auto.py @@ -165,6 +165,37 @@ class AutoEagle3DraftModel(AutoDraftModel): } +def eagle3_draft_config_from_target( + target_config: PretrainedConfig, target_layer_ids=None +) -> LlamaConfig: + """Convert a target-model config to the Llama-compatible EAGLE3 config.""" + if not isinstance(target_config, PretrainedConfig): + raise TypeError( + "EAGLE3 target config must be a transformers.PretrainedConfig, got " + f"{type(target_config)!r}" + ) + + config = target_config.to_dict() + config.update( + { + "architectures": ["LlamaForCausalLMEagle3"], + "model_type": "llama", + "num_hidden_layers": 1, + "pretraining_tp": int(config.get("pretraining_tp") or 1), + "target_hidden_size": int(target_config.hidden_size), + "tie_word_embeddings": False, + } + ) + if target_layer_ids is not None: + layer_ids = _normalize_int_list(target_layer_ids) + if not layer_ids: + raise ValueError("EAGLE3 target_layer_ids must not be empty") + config["target_hidden_layer_ids"] = layer_ids + config["eagle_aux_hidden_state_layer_ids"] = layer_ids + + return LlamaConfig.from_dict(_normalize_eagle3_config_dict(config)) + + class AutoDraftModelConfig: _config_mapping = { "LlamaForCausalLMEagle3": LlamaConfig, diff --git a/verl_speco/models/dflash/configuration_dflash.py b/verl_speco/models/dflash/configuration_dflash.py index 4c2f0a46..5d56e600 100644 --- a/verl_speco/models/dflash/configuration_dflash.py +++ b/verl_speco/models/dflash/configuration_dflash.py @@ -114,6 +114,7 @@ def __init__( num_hidden_layers: int = 1, num_attention_heads: int = 32, num_key_value_heads: int = 8, + head_dim: Optional[int] = None, vocab_size: int = 152064, rms_norm_eps: float = 1e-6, max_position_embeddings: int = 32768, @@ -133,6 +134,8 @@ def __init__( self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads + if head_dim is not None: + self.head_dim = int(head_dim) self.vocab_size = vocab_size self.rms_norm_eps = rms_norm_eps self.max_position_embeddings = max_position_embeddings diff --git a/verl_speco/models/eagle/llama_eagle.py b/verl_speco/models/eagle/llama_eagle.py index b363cce1..b59d12be 100644 --- a/verl_speco/models/eagle/llama_eagle.py +++ b/verl_speco/models/eagle/llama_eagle.py @@ -1325,31 +1325,26 @@ def __init__(self, config): self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): - if self.config.pretraining_tp > 1: - slice = self.intermediate_size // self.config.pretraining_tp + pretraining_tp = int(getattr(self.config, "pretraining_tp", 1)) + if pretraining_tp > 1: + slice = self.intermediate_size // pretraining_tp gate_proj_slices = self.gate_proj.weight.split(slice, dim=0) up_proj_slices = self.up_proj.weight.split(slice, dim=0) down_proj_slices = self.down_proj.weight.split(slice, dim=1) gate_proj = torch.cat( - [ - F.linear(x, gate_proj_slices[i]) - for i in range(self.config.pretraining_tp) - ], + [F.linear(x, gate_proj_slices[i]) for i in range(pretraining_tp)], dim=-1, ) up_proj = torch.cat( - [ - F.linear(x, up_proj_slices[i]) - for i in range(self.config.pretraining_tp) - ], + [F.linear(x, up_proj_slices[i]) for i in range(pretraining_tp)], dim=-1, ) intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2) down_proj = [ F.linear(intermediate_states[i], down_proj_slices[i]) - for i in range(self.config.pretraining_tp) + for i in range(pretraining_tp) ] down_proj = sum(down_proj) else: diff --git a/verl_speco/models/target/target_head.py b/verl_speco/models/target/target_head.py index abc0830d..990d679e 100644 --- a/verl_speco/models/target/target_head.py +++ b/verl_speco/models/target/target_head.py @@ -13,50 +13,10 @@ # limitations under the License. """Minimal target lm-head loader for SPECO drafter training.""" -import glob -import json -import os - import torch -from huggingface_hub import snapshot_download -from safetensors import safe_open from torch import nn - -def _load_checkpoint_tensor(model_path: str, key: str) -> torch.Tensor: - if not os.path.exists(model_path): - model_path = snapshot_download(repo_id=model_path) - - index_paths = glob.glob(os.path.join(model_path, "*.index.json")) - if len(index_paths) > 1: - raise FileNotFoundError(f"Multiple index.json files found in {model_path}") - - if index_paths: - with open(index_paths[0], encoding="utf-8") as f: - index_json = json.load(f) - weight_map = index_json.get("weight_map", {}) - if key not in weight_map: - raise KeyError( - f"Tensor {key!r} is not present in checkpoint index for {model_path}" - ) - ckpt_file = os.path.join(model_path, weight_map[key]) - if ckpt_file.endswith(".safetensors"): - with safe_open(ckpt_file, framework="pt", device="cpu") as f: - return f.get_tensor(key) - return torch.load(ckpt_file, map_location="cpu", weights_only=True)[key] - - safetensors_path = os.path.join(model_path, "model.safetensors") - if os.path.exists(safetensors_path): - with safe_open(safetensors_path, framework="pt", device="cpu") as f: - return f.get_tensor(key) - - pytorch_path = os.path.join(model_path, "pytorch_model.bin") - if os.path.exists(pytorch_path): - return torch.load(pytorch_path, map_location="cpu", weights_only=True)[key] - - raise FileNotFoundError( - f"No index.json, model.safetensors or pytorch_model.bin found in {model_path}" - ) +from verl_speco.checkpoint_tensor import _load_checkpoint_tensor class TargetHead(nn.Module): diff --git a/verl_speco/producer/__init__.py b/verl_speco/producer/__init__.py new file mode 100644 index 00000000..00265cdf --- /dev/null +++ b/verl_speco/producer/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2026 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. +"""Standalone target-feature Producer components.""" diff --git a/verl_speco/producer/input_reader.py b/verl_speco/producer/input_reader.py new file mode 100644 index 00000000..882506dc --- /dev/null +++ b/verl_speco/producer/input_reader.py @@ -0,0 +1,605 @@ +# Copyright 2026 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. +"""Streaming verl/JSONL input and token preparation for the Producer.""" + +from __future__ import annotations + +import importlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterator, Mapping + +import torch + + +@dataclass(frozen=True) +class InputRecord: + sequence_no: int + sample_id: str + prompt: str | tuple[dict[str, str], ...] + response: str | None + source_metadata: dict[str, Any] + + +@dataclass(frozen=True) +class TokenizedRequest: + sequence_no: int + sample_id: str + input_ids: torch.Tensor + loss_mask: torch.Tensor + position_ids: torch.Tensor + feature_positions: torch.Tensor + draft_position_ids: torch.Tensor + source_metadata: dict[str, Any] + vllm_prompt_token_ids: tuple[int, ...] + + @property + def prompt_token_ids(self) -> list[int]: + return list(self.vllm_prompt_token_ids) + + +@dataclass(frozen=True) +class GenerationRequest: + """One prompt-only row that needs target-model response generation.""" + + sequence_no: int + sample_id: str + prompt_token_ids: tuple[int, ...] + max_tokens: int + source_metadata: dict[str, Any] + + +def iter_input_records(path: str | os.PathLike[str]) -> Iterator[InputRecord]: + """Yield strict prompt/response records from one JSONL or Parquet file.""" + + input_path = Path(path) + if not input_path.is_file(): + raise FileNotFoundError(f"Producer input file not found: {input_path}") + + sequence_no = 0 + for location, payload in _iter_payloads(input_path): + if not isinstance(payload, dict): + raise ValueError( + f"Producer input at {location} must be a JSON-style object" + ) + prompt_value, response = _prompt_response_from_payload(payload, location) + prompt = _normalize_prompt(prompt_value, location) + if response is not None and (not isinstance(response, str) or not response): + raise ValueError( + f"Producer input at {location} field 'response' must be a non-empty " + "string when present" + ) + sample_id = payload.get("sample_id") or _extra_info_index(payload) + if sample_id is None: + sample_id = f"train-{sequence_no:06d}" + if not isinstance(sample_id, str) or not sample_id: + raise ValueError(f"Producer input at {location} has invalid sample_id") + source_metadata = { + key: value + for key, value in payload.items() + if key + not in {"prompt", "response", "conversation", "conversations", "sample_id"} + } + yield InputRecord( + sequence_no=sequence_no, + sample_id=sample_id, + prompt=prompt, + response=response, + source_metadata=source_metadata, + ) + sequence_no += 1 + + +def _prompt_response_from_payload( + payload: Mapping[str, Any], location: str +) -> tuple[Any, Any]: + """Normalize supported row schemas to Producer ``prompt``/``response``.""" + + if "prompt" in payload: + return payload.get("prompt"), payload.get("response") + + conversation = payload.get("conversation") + if conversation is not None: + messages = _normalize_conversation_messages( + conversation, + location, + role_key="role", + content_key="content", + ) + return _split_final_assistant(messages) + + conversations = payload.get("conversations") + if conversations is not None: + messages = _normalize_conversation_messages( + conversations, + location, + role_key="from", + content_key="value", + ) + return _split_final_assistant(messages) + + return None, payload.get("response") + + +def _normalize_conversation_messages( + value: Any, + location: str, + *, + role_key: str, + content_key: str, +) -> tuple[dict[str, str], ...]: + if not isinstance(value, (list, tuple)) or not value: + raise ValueError(f"Producer input at {location} conversation must be non-empty") + role_mapping = {"human": "user", "gpt": "assistant"} + messages: list[dict[str, str]] = [] + for index, item in enumerate(value): + if not isinstance(item, Mapping): + raise ValueError( + f"Producer input at {location} conversation item {index} must be an object" + ) + role = item.get(role_key) + content = item.get(content_key) + if not isinstance(role, str) or not role: + raise ValueError( + f"Producer input at {location} conversation item {index} requires " + f"string field {role_key!r}" + ) + if not isinstance(content, str) or not content: + raise ValueError( + f"Producer input at {location} conversation item {index} requires " + f"non-empty string field {content_key!r}" + ) + messages.append( + {"role": role_mapping.get(role.strip().lower(), role), "content": content} + ) + return tuple(messages) + + +def _split_final_assistant( + messages: tuple[dict[str, str], ...], +) -> tuple[tuple[dict[str, str], ...], str | None]: + if messages[-1]["role"] != "assistant": + return messages, None + prompt = messages[:-1] + if not prompt: + raise ValueError("Conversation cannot contain only an assistant response") + return prompt, messages[-1]["content"] + + +def _normalize_prompt(value: Any, location: str) -> str | tuple[dict[str, str], ...]: + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)) and value: + messages: list[dict[str, str]] = [] + for index, message in enumerate(value): + if not isinstance(message, Mapping): + raise ValueError( + f"Producer input at {location} prompt message {index} must be " + "an object" + ) + role = message.get("role") + content = message.get("content") + if not isinstance(role, str) or not role: + raise ValueError( + f"Producer input at {location} prompt message {index} requires " + "string field 'role'" + ) + if not isinstance(content, str): + raise ValueError( + f"Producer input at {location} prompt message {index} requires " + "string field 'content'" + ) + messages.append({"role": role, "content": content}) + return tuple(messages) + raise ValueError( + f"Producer input at {location} requires 'prompt' as a string or " + "chat-message list" + ) + + +def _extra_info_index(payload: Mapping[str, Any]) -> str | None: + extra_info = payload.get("extra_info") + if not isinstance(extra_info, Mapping): + return None + value = extra_info.get("index") + return value if isinstance(value, str) and value else None + + +def _iter_payloads(input_path: Path) -> Iterator[tuple[str, Any]]: + if _is_parquet(input_path): + yield from _iter_parquet_payloads(input_path) + return + yield from _iter_jsonl_payloads(input_path) + + +def _is_parquet(input_path: Path) -> bool: + if input_path.suffix.lower() in {".parquet", ".pq"}: + return True + with input_path.open("rb") as input_file: + return input_file.read(4) == b"PAR1" + + +def _iter_jsonl_payloads(input_path: Path) -> Iterator[tuple[str, Any]]: + try: + with input_path.open("r", encoding="utf-8") as input_file: + for line_number, line in enumerate(input_file, start=1): + if not line.strip(): + continue + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"Invalid JSON object at {input_path}:{line_number}: {exc.msg}" + ) from exc + yield f"{input_path}:{line_number}", payload + except UnicodeDecodeError as exc: + raise ValueError( + f"Producer input {input_path} is not UTF-8 JSONL or a Parquet file" + ) from exc + + +def _iter_parquet_payloads(input_path: Path) -> Iterator[tuple[str, Any]]: + try: + parquet = importlib.import_module("pyarrow.parquet") + except ImportError as exc: + raise RuntimeError( + "Reading a Parquet training file requires pyarrow; install the normal " + "verl data dependencies in the training environment" + ) from exc + + parquet_file = parquet.ParquetFile(input_path) + row_number = 0 + for batch in parquet_file.iter_batches(): + for payload in batch.to_pylist(): + row_number += 1 + yield f"{input_path}:row {row_number}", payload + + +def build_loss_mask(input_ids: torch.Tensor, prompt_length: int) -> torch.Tensor: + sequence_length = int(input_ids.numel()) + if prompt_length < 0 or prompt_length > sequence_length: + raise ValueError( + f"prompt_length must be within [0, {sequence_length}], got {prompt_length}" + ) + mask = torch.ones(sequence_length, dtype=torch.float32) + mask[:prompt_length] = 0 + return mask + + +def tokenize_record( + record: InputRecord, + tokenizer: Any, + config: Mapping[str, Any] | Any, +) -> TokenizedRequest: + """Tokenize one row that already contains a response.""" + + if record.response is None: + raise ValueError( + f"Producer sample {record.sample_id!r} has no response; prepare it for " + "target-model generation instead" + ) + prompt_ids = _prompt_ids(record.prompt, tokenizer) + if isinstance(record.prompt, str): + full_ids = _token_ids( + tokenizer(record.prompt + record.response, add_special_tokens=False) + ) + if full_ids[: len(prompt_ids)] != prompt_ids: + # Some tokenizers merge text across the prompt/response boundary. + # Keep that boundary explicit so the response-only loss mask and the + # exact token IDs sent to vLLM remain aligned. + response_ids = _token_ids( + tokenizer(record.response, add_special_tokens=False) + ) + full_ids = [*prompt_ids, *response_ids] + else: + full_ids = _token_ids( + tokenizer.apply_chat_template( + [*record.prompt, {"role": "assistant", "content": record.response}], + tokenize=True, + add_generation_prompt=False, + ) + ) + if full_ids[: len(prompt_ids)] != prompt_ids: + prompt_ids, full_ids = _tokenize_chat_response_with_explicit_boundary( + record.prompt, + record.response, + tokenizer, + ) + if full_ids[: len(prompt_ids)] != prompt_ids: + raise ValueError( + f"Producer sample {record.sample_id!r} has an unstable tokenizer boundary " + "between prompt and response; prompt token IDs are not a prefix of full " + "token IDs" + ) + if len(full_ids) <= len(prompt_ids): + raise ValueError( + f"Producer sample {record.sample_id!r} produced no response tokens" + ) + return _build_tokenized_request( + sequence_no=record.sequence_no, + sample_id=record.sample_id, + prompt_length=len(prompt_ids), + full_ids=full_ids, + source_metadata=record.source_metadata, + config=config, + ) + + +def prepare_generation_request( + record: InputRecord, + tokenizer: Any, + config: Mapping[str, Any] | Any, +) -> GenerationRequest: + """Tokenize a prompt-only row and bound target-model generation.""" + + if record.response is not None: + raise ValueError( + f"Producer sample {record.sample_id!r} already contains a response" + ) + prompt_ids = _prompt_ids(record.prompt, tokenizer) + max_sequence_length = int(_config_value(config, "max_sequence_length", 0) or 0) + max_tokens = int(_config_value(config, "generation_max_tokens", 0) or 0) + if max_tokens <= 0: + raise ValueError("generation_max_tokens must be positive") + if max_sequence_length > 0: + max_tokens = min(max_tokens, max_sequence_length - len(prompt_ids)) + if max_tokens <= 0: + raise ValueError( + f"Producer sample {record.sample_id!r} prompt has {len(prompt_ids)} tokens " + "and leaves no generation capacity within " + f"max_sequence_length={max_sequence_length}" + ) + return GenerationRequest( + sequence_no=record.sequence_no, + sample_id=record.sample_id, + prompt_token_ids=tuple(prompt_ids), + max_tokens=max_tokens, + source_metadata=dict(record.source_metadata), + ) + + +def finalize_generated_request( + request: GenerationRequest, + hidden_state_token_ids: Any, + config: Mapping[str, Any] | Any, + *, + expected_response_token_ids: Any | None = None, +) -> TokenizedRequest: + """Build a training request from vLLM generation and hidden-state tokens.""" + + hidden_ids = _token_ids(hidden_state_token_ids) + prompt_ids = list(request.prompt_token_ids) + if hidden_ids[: len(prompt_ids)] != prompt_ids: + raise ValueError( + f"vLLM hidden-state token sequence for sample {request.sample_id!r} " + "does not " + "start with the rendered prompt token IDs" + ) + if expected_response_token_ids is not None: + response_ids = _token_ids(expected_response_token_ids) + full_ids = [*prompt_ids, *response_ids] + # ExampleHiddenStatesConnector deliberately excludes the final sampled + # token: that token was emitted by the model but was never fed through a + # subsequent forward pass, so no hidden state exists for it. + expected_hidden_ids = full_ids[:-1] + if hidden_ids != expected_hidden_ids: + raise ValueError( + f"vLLM hidden-state token IDs for sample {request.sample_id!r} do not " + "match the prompt plus generated completion excluding its final token " + f"(hidden={len(hidden_ids)}, expected={len(expected_hidden_ids)}, " + f"completion={len(response_ids)})" + ) + else: + # Prefilled records already contain their complete response, and their + # caller supplies the full token sequence directly. + full_ids = hidden_ids + if len(full_ids) <= len(prompt_ids): + raise ValueError( + f"vLLM generated no response tokens for sample {request.sample_id!r}; " + "the hidden-state server must enable include_output_tokens" + ) + return _build_tokenized_request( + sequence_no=request.sequence_no, + sample_id=request.sample_id, + prompt_length=len(prompt_ids), + full_ids=full_ids, + source_metadata=request.source_metadata, + config=config, + vllm_prompt_token_ids=hidden_ids, + feature_end_limit=len(hidden_ids), + ) + + +def prepare_generated_prefill_request( + request: GenerationRequest, + response_token_ids: Any, + config: Mapping[str, Any] | Any, +) -> TokenizedRequest: + """Build the full-sequence prefill request after target generation. + + The final sampled token has not itself passed through a model forward, so + target features are requested for ``prompt + completion[:-1]`` while the + full completion remains in ``input_ids`` as the next-token label sequence. + This matches the existing non-TQ vLLM replay path and does not require the + connector to capture decode-step hidden states. + """ + + response_ids = _token_ids(response_token_ids) + if not response_ids: + raise ValueError( + f"vLLM generated no response tokens for sample {request.sample_id!r}" + ) + prompt_ids = list(request.prompt_token_ids) + full_ids = [*prompt_ids, *response_ids] + hidden_input_ids = full_ids[:-1] + return _build_tokenized_request( + sequence_no=request.sequence_no, + sample_id=request.sample_id, + prompt_length=len(prompt_ids), + full_ids=full_ids, + source_metadata=request.source_metadata, + config=config, + vllm_prompt_token_ids=hidden_input_ids, + feature_end_limit=len(hidden_input_ids), + ) + + +def _prompt_ids(prompt: str | tuple[dict[str, str], ...], tokenizer: Any) -> list[int]: + if isinstance(prompt, str): + return _token_ids(tokenizer(prompt, add_special_tokens=False)) + apply_chat_template = getattr(tokenizer, "apply_chat_template", None) + if not callable(apply_chat_template): + raise RuntimeError( + "Chat-message prompts require a tokenizer with apply_chat_template()" + ) + return _token_ids( + apply_chat_template( + list(prompt), + tokenize=True, + add_generation_prompt=True, + ) + ) + + +def _tokenize_chat_response_with_explicit_boundary( + prompt: tuple[dict[str, str], ...], + response: str, + tokenizer: Any, +) -> tuple[list[int], list[int]]: + """Render a chat response while keeping its loss boundary deterministic. + + Qwen-family templates may render a generation prompt differently from an + existing assistant message (for example by inserting a thinking preamble). + A marker lets us retain the template's assistant header and suffix while + tokenizing the response as a separate loss-bearing region. + """ + + marker = "__VERL_SPECO_ASSISTANT_RESPONSE_BOUNDARY_8F7C2D91__" + while marker in response or any(marker in message["content"] for message in prompt): + marker += "_" + rendered = tokenizer.apply_chat_template( + [*prompt, {"role": "assistant", "content": marker}], + tokenize=False, + add_generation_prompt=False, + ) + if not isinstance(rendered, str) or rendered.count(marker) != 1: + raise ValueError( + "Tokenizer chat template did not preserve the assistant response marker" + ) + prompt_text, suffix_text = rendered.split(marker, 1) + explicit_prompt_ids = _token_ids(tokenizer(prompt_text, add_special_tokens=False)) + response_ids = _token_ids(tokenizer(response, add_special_tokens=False)) + suffix_ids = _token_ids(tokenizer(suffix_text, add_special_tokens=False)) + return explicit_prompt_ids, [*explicit_prompt_ids, *response_ids, *suffix_ids] + + +def _build_tokenized_request( + *, + sequence_no: int, + sample_id: str, + prompt_length: int, + full_ids: list[int], + source_metadata: Mapping[str, Any], + config: Mapping[str, Any] | Any, + vllm_prompt_token_ids: list[int] | None = None, + feature_end_limit: int | None = None, +) -> TokenizedRequest: + input_ids = torch.tensor(full_ids, dtype=torch.int64) + if int(input_ids.numel()) <= 0: + raise ValueError(f"Producer sample {sample_id!r} produced no input tokens") + + loss_mask = build_loss_mask(input_ids, prompt_length) + position_ids = torch.arange(int(input_ids.numel()), dtype=torch.int64) + + feature_start = max(prompt_length - 1, 0) + feature_end = int(input_ids.numel()) + if feature_end_limit is not None: + feature_end = min(feature_end, int(feature_end_limit)) + max_feature_length = int(_config_value(config, "max_feature_length", 0) or 0) + if max_feature_length == 1: + raise ValueError("max_feature_length must be 0 or at least 2") + if max_feature_length > 1: + feature_end = min(feature_start + max_feature_length, feature_end) + request_prompt_token_ids = ( + list(vllm_prompt_token_ids) + if vllm_prompt_token_ids is not None + else full_ids[:feature_end] + ) + max_sequence_length = int(_config_value(config, "max_sequence_length", 0) or 0) + if max_sequence_length > 0 and len(request_prompt_token_ids) > max_sequence_length: + raise ValueError( + f"Producer sample {sample_id!r} requires a vLLM prefill of " + f"{len(request_prompt_token_ids)} tokens after selecting its training " + f"window, exceeding max_sequence_length={max_sequence_length} " + f"(full_sequence_length={int(input_ids.numel())}, " + f"prompt_length={prompt_length})" + ) + feature_positions = torch.arange(feature_start, feature_end, dtype=torch.int64) + if int(feature_positions.numel()) <= 0: + raise ValueError(f"Producer sample {sample_id!r} has an empty feature window") + draft_position_ids = position_ids[feature_start:feature_end] + 1 + return TokenizedRequest( + sequence_no=sequence_no, + sample_id=sample_id, + input_ids=input_ids, + loss_mask=loss_mask, + position_ids=position_ids, + feature_positions=feature_positions, + draft_position_ids=draft_position_ids, + source_metadata=dict(source_metadata), + vllm_prompt_token_ids=tuple(request_prompt_token_ids), + ) + + +def _token_ids(encoding: Any) -> list[int]: + if isinstance(encoding, (list, tuple)) or torch.is_tensor(encoding): + value = encoding + else: + value = ( + encoding.get("input_ids") + if isinstance(encoding, Mapping) + else encoding.input_ids + ) + if value is None: + raise ValueError("Tokenizer result is missing input_ids") + if torch.is_tensor(value): + value = value.detach().cpu().reshape(-1).tolist() + if not isinstance(value, (list, tuple)): + raise TypeError("Tokenizer input_ids must be a tensor, list, or tuple") + if value and isinstance(value[0], (list, tuple)): + if len(value) != 1: + raise ValueError("Tokenizer returned more than one sequence for one input") + value = value[0] + return [int(token_id) for token_id in value] + + +def _config_value(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, Mapping): + return config.get(key, default) + return getattr(config, key, default) + + +__all__ = [ + "GenerationRequest", + "InputRecord", + "TokenizedRequest", + "build_loss_mask", + "finalize_generated_request", + "iter_input_records", + "prepare_generation_request", + "prepare_generated_prefill_request", + "tokenize_record", +] diff --git a/verl_speco/producer/vllm_feature_client.py b/verl_speco/producer/vllm_feature_client.py new file mode 100644 index 00000000..96519100 --- /dev/null +++ b/verl_speco/producer/vllm_feature_client.py @@ -0,0 +1,351 @@ +# Copyright 2026 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. +"""Bounded asynchronous vLLM hidden-state requests for the Producer.""" + +from __future__ import annotations + +import asyncio +import errno +import importlib +import inspect +import logging +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + + +logger = logging.getLogger(__name__) + +_DEFAULT_MAX_RETRIES = 3 +_RETRY_BACKOFF_BASE_SECONDS = 2 + + +@dataclass(frozen=True) +class VllmEndpoint: + base_url: str + max_concurrency: int + + def __post_init__(self) -> None: + if not self.base_url: + raise ValueError("VllmEndpoint.base_url must not be empty") + if self.max_concurrency <= 0: + raise ValueError("VllmEndpoint.max_concurrency must be positive") + + +@dataclass(frozen=True) +class VllmResponse: + hidden_states_path: str + endpoint_url: str + generated_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class RawVllmFeature: + payload: dict[str, Any] + temporary_path: str + endpoint_url: str + byte_size: int + generated_token_ids: tuple[int, ...] = () + + +@dataclass +class _EndpointState: + endpoint: VllmEndpoint + client: Any + semaphore: asyncio.Semaphore + inflight: int = 0 + requests: int = 0 + + +async def request_prefill( + endpoint: VllmEndpoint, + client: Any, + prompt_token_ids: list[int], + *, + model: str, + timeout: float, +) -> VllmResponse: + response = await client.completions.create( + model=model, + prompt=prompt_token_ids, + max_tokens=1, + extra_body={"return_token_ids": True}, + timeout=timeout, + ) + choices = getattr(response, "choices", None) or [] + if choices: + actual = getattr(choices[0], "prompt_token_ids", None) + if actual is not None and list(actual) != prompt_token_ids: + raise ValueError("vLLM prompt_token_ids mismatch") + params = getattr(response, "kv_transfer_params", None) + if not isinstance(params, Mapping): + raise ValueError("vLLM response missing kv_transfer_params") + path = params.get("hidden_states_path") + if not path: + raise ValueError("vLLM response missing hidden_states_path") + return VllmResponse(os.fspath(path), endpoint.base_url) + + +async def request_generate( + endpoint: VllmEndpoint, + client: Any, + prompt_token_ids: list[int], + *, + model: str, + max_tokens: int, + timeout: float, +) -> VllmResponse: + """Generate response tokens; hidden states are collected by a later prefill.""" + + response = await client.completions.create( + model=model, + prompt=prompt_token_ids, + max_tokens=max_tokens, + extra_body={"return_token_ids": True}, + timeout=timeout, + ) + choices = getattr(response, "choices", None) or [] + if not choices: + raise ValueError("vLLM generation response has no choices") + actual_prompt = getattr(choices[0], "prompt_token_ids", None) + if actual_prompt is not None and list(actual_prompt) != prompt_token_ids: + raise ValueError("vLLM generation prompt_token_ids mismatch") + generated = getattr(choices[0], "token_ids", None) + if not isinstance(generated, (list, tuple)) or not generated: + raise ValueError( + "vLLM generation response missing token_ids; enable return_token_ids support" + ) + params = getattr(response, "kv_transfer_params", None) + if not isinstance(params, Mapping): + raise ValueError("vLLM generation response missing kv_transfer_params") + path = params.get("hidden_states_path") + if not path: + raise ValueError("vLLM generation response missing hidden_states_path") + return VllmResponse( + os.fspath(path), + endpoint.base_url, + tuple(int(token_id) for token_id in generated), + ) + + +def load_hidden_state_result(response: VllmResponse) -> RawVllmFeature: + try: + from safetensors.torch import load_file + except ImportError as exc: + raise RuntimeError("vLLM Producer requires safetensors") from exc + path = Path(response.hidden_states_path) + _wait_for_lock(Path(f"{path}.lock")) + if not path.is_file(): + raise FileNotFoundError(f"vLLM hidden-states file not found: {path}") + return RawVllmFeature( + payload=dict(load_file(str(path), device="cpu")), + temporary_path=str(path), + endpoint_url=response.endpoint_url, + byte_size=int(path.stat().st_size), + generated_token_ids=response.generated_token_ids, + ) + + +def delete_temporary_result(raw: RawVllmFeature) -> None: + path = Path(raw.temporary_path) + path.unlink(missing_ok=True) + Path(f"{path}.lock").unlink(missing_ok=True) + + +def choose_endpoint(states: Sequence[_EndpointState]) -> _EndpointState: + if not states: + raise RuntimeError("No vLLM endpoints are configured") + return min(states, key=lambda state: (state.inflight, state.requests)) + + +class VllmFeatureClientPool: + def __init__( + self, + endpoints: Sequence[VllmEndpoint], + *, + model: str, + max_inflight_requests: int, + request_timeout: float, + ) -> None: + if not endpoints: + raise ValueError("At least one vLLM endpoint is required") + if max_inflight_requests <= 0: + raise ValueError("max_inflight_requests must be positive") + if not model: + raise ValueError("vllm_model must not be empty") + self.endpoints = list(endpoints) + self.model = model + self.request_timeout = float(request_timeout) + self._global_semaphore = asyncio.Semaphore(max_inflight_requests) + self._states: list[_EndpointState] = [] + + async def start(self) -> None: + if self._states: + return + try: + from openai import AsyncOpenAI + except ImportError as exc: + raise RuntimeError("vLLM Producer requires the openai package") from exc + self._states = [ + _EndpointState( + endpoint=endpoint, + client=AsyncOpenAI( + base_url=endpoint.base_url, + api_key="EMPTY", + max_retries=0, + ), + semaphore=asyncio.Semaphore(endpoint.max_concurrency), + ) + for endpoint in self.endpoints + ] + + async def prefill(self, request: Any) -> RawVllmFeature: + return await self._request(request, generate=False) + + async def generate(self, request: Any) -> RawVllmFeature: + return await self._request(request, generate=True) + + async def _request(self, request: Any, *, generate: bool) -> RawVllmFeature: + if not self._states: + raise RuntimeError("VllmFeatureClientPool.start() must be called first") + state = choose_endpoint(self._states) + state.inflight += 1 + try: + async with self._global_semaphore, state.semaphore: + response = await self._request_with_retries( + state, + request, + generate=generate, + ) + raw = await asyncio.to_thread(load_hidden_state_result, response) + state.requests += 1 + return raw + finally: + state.inflight = max(state.inflight - 1, 0) + + async def _request_with_retries( + self, + state: _EndpointState, + request: Any, + *, + generate: bool, + ) -> VllmResponse: + total_attempts = _DEFAULT_MAX_RETRIES + 1 + for attempt in range(1, total_attempts + 1): + try: + if generate: + return await request_generate( + state.endpoint, + state.client, + list(request.prompt_token_ids), + model=self.model, + max_tokens=int(request.max_tokens), + timeout=self.request_timeout, + ) + return await request_prefill( + state.endpoint, + state.client, + list(request.prompt_token_ids), + model=self.model, + timeout=self.request_timeout, + ) + except ValueError: + # Response validation failures are deterministic protocol/data + # errors, equivalent to speculators' InvalidResponseError. + raise + except Exception as exc: + if attempt >= total_attempts: + logger.error( + "vLLM request failed after %s attempts endpoint=%s " + "sample_id=%s error=%s", + total_attempts, + state.endpoint.base_url, + getattr(request, "sample_id", None), + exc, + ) + raise + backoff = _RETRY_BACKOFF_BASE_SECONDS**attempt + logger.warning( + "vLLM request aborted attempt=%s/%s endpoint=%s " + "sample_id=%s error=%s; retrying in %ss", + attempt, + total_attempts, + state.endpoint.base_url, + getattr(request, "sample_id", None), + exc, + backoff, + ) + await asyncio.sleep(backoff) + raise RuntimeError( + "unreachable: vLLM request retry loop exhausted without returning" + ) + + async def close(self) -> None: + states, self._states = self._states, [] + for state in states: + close = getattr(state.client, "close", None) + if not callable(close): + continue + result = close() + if inspect.isawaitable(result): + await result + + +def _wait_for_lock(lock_path: Path, timeout: float = 30.0) -> None: + if not lock_path.exists(): + return + try: + fcntl: Any = importlib.import_module("fcntl") + except ImportError: + # vLLM's file connector is Linux-only. Keep the old existence-based + # fallback for dependency-light tests on other platforms. + deadline = time.monotonic() + timeout + while lock_path.exists(): + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for vLLM hidden-state lock: {lock_path}" + ) + time.sleep(0.01) + return + + deadline = time.monotonic() + timeout + with lock_path.open("rb") as lock_file: + while True: + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH | fcntl.LOCK_NB) + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + return + except OSError as exc: + if exc.errno not in {errno.EACCES, errno.EAGAIN}: + raise + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for vLLM hidden-state lock: {lock_path}" + ) from exc + time.sleep(0.01) + + +__all__ = [ + "RawVllmFeature", + "VllmEndpoint", + "VllmFeatureClientPool", + "VllmResponse", + "choose_endpoint", + "delete_temporary_result", + "load_hidden_state_result", + "request_prefill", + "request_generate", +] diff --git a/verl_speco/standalone_tq_producer.py b/verl_speco/standalone_tq_producer.py new file mode 100644 index 00000000..0a22542f --- /dev/null +++ b/verl_speco/standalone_tq_producer.py @@ -0,0 +1,593 @@ +# Copyright 2026 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. +"""Standalone vLLM target-feature Producer writing directly to TransferQueue.""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, replace +from typing import Any, Mapping + +import torch + +from verl_speco.integration import transferqueue_bridge as default_transport +from verl_speco.integration.oldlogprob_layer_ids import ( + resolve_drafter_hidden_states_layout, +) +from verl_speco.producer.input_reader import ( + GenerationRequest, + TokenizedRequest, + iter_input_records, + prepare_generation_request, + prepare_generated_prefill_request, + tokenize_record, +) +from verl_speco.producer.vllm_feature_client import ( + RawVllmFeature, + VllmEndpoint, + VllmFeatureClientPool, + delete_temporary_result, +) +from verl_speco.trainer.feature_store import DraftFeatureSample +from verl_speco.trainer.standalone_resume import load_standalone_resume +from verl_speco.trainer.target_feature_replay import ( + FeatureContract, + HiddenStateAlignmentError, + feature_from_vllm_payload, + load_vllm_final_norm, +) +from verl_speco.transport.drafter_sample_protocol import ( + DRAFTER_TQ_PARTITION, + PROTOCOL_SCHEMA_VERSION, + SampleMetadata, + encode_sample, + is_ready_sample_tag, + make_eos_record, + make_ready_tag, + make_sample_key, +) + + +logger = logging.getLogger(__name__) +_INPUT_DONE = object() +_PUBLISH_DONE = object() + + +@dataclass +class ProducerStats: + input_count: int = 0 + published_count: int = 0 + failed_count: int = 0 + dropped_count: int = 0 + pending_bytes: int = 0 + + +@dataclass(frozen=True) +class PreparedFeature: + request: TokenizedRequest + raw: RawVllmFeature + sample: DraftFeatureSample + metadata: SampleMetadata + + +async def publish_one(result: PreparedFeature, transport: Any) -> str: + """Publish one sample and delete its temporary file only after TQ succeeds.""" + + key = make_sample_key(result.metadata) + fields = encode_sample(result.sample, result.metadata) + tag = make_ready_tag(result.metadata) + await asyncio.to_thread(transport.put_sample, key, fields, tag=tag) + delete_temporary_result(result.raw) + return key + + +def validate_producer_config(config: Any) -> None: + producer_cfg, training_cfg, tq_cfg = _config_sections(config) + required = ( + "input_path", + "tokenizer_path", + "tokenizer_fingerprint", + "target_model_id", + "target_model_revision", + "vllm_model", + ) + missing = [name for name in required if not producer_cfg.get(name)] + if missing: + raise ValueError(f"standalone_tq_producer missing required fields: {missing}") + endpoints = producer_cfg.get("vllm_endpoints") + if not isinstance(endpoints, list) or not endpoints or not all(endpoints): + raise ValueError( + "standalone_tq_producer.vllm_endpoints must be a non-empty list" + ) + target_layer_ids = producer_cfg.get("target_layer_ids") + if not isinstance(target_layer_ids, list) or not target_layer_ids: + raise ValueError( + "standalone_tq_producer.target_layer_ids must be a non-empty list" + ) + algorithm = str(training_cfg.get("speculative_algorithm", "") or "").strip() + if not algorithm: + raise ValueError("drafter.speculative_algorithm must not be empty") + if bool(training_cfg.get("use_logits", False)): + raise ValueError("Standalone TQ Producer does not support use_logits=true") + if int(tq_cfg.get("schema_version", 0)) != PROTOCOL_SCHEMA_VERSION: + raise ValueError( + f"transfer_queue.schema_version must be {PROTOCOL_SCHEMA_VERSION}" + ) + if tq_cfg.get("package_version") != "0.1.10": + raise ValueError("transfer_queue.package_version must be '0.1.10'") + if tq_cfg.get("partition_id") != DRAFTER_TQ_PARTITION: + raise ValueError( + f"transfer_queue.partition_id must be {DRAFTER_TQ_PARTITION!r}" + ) + if not tq_cfg.get("run_id"): + raise ValueError("transfer_queue.run_id must not be empty") + ray_cfg = tq_cfg.get("ray") or {} + if not isinstance(ray_cfg, Mapping) or not ray_cfg.get("address"): + raise ValueError( + "transfer_queue.ray.address must point to a running Ray cluster" + ) + positive_fields = ( + "max_inflight_requests", + "per_endpoint_concurrency", + "input_queue_size", + "publish_queue_size", + "max_pending_samples", + "generation_max_tokens", + ) + invalid = [name for name in positive_fields if int(producer_cfg.get(name, 0)) <= 0] + if invalid: + raise ValueError(f"standalone_tq_producer fields must be positive: {invalid}") + + +def _should_log_sample_progress(count: int) -> bool: + return count <= 3 or count % 50 == 0 + + +async def run_producer( + config: Any, + *, + transport: Any = default_transport, + tokenizer: Any | None = None, + client_pool: Any | None = None, +) -> ProducerStats: + """Run the bounded input -> vLLM -> TQ pipeline and publish EOS on success.""" + + validate_producer_config(config) + producer_cfg, drafter_cfg, tq_cfg = _config_sections(config) + run_id = str(tq_cfg["run_id"]) + stats = ProducerStats() + connected = False + pool = client_pool + try: + logger.info( + "Standalone TQ Producer starting run_id=%s input=%s endpoints=%s", + run_id, + producer_cfg["input_path"], + producer_cfg["vllm_endpoints"], + ) + if not transport.configure_transfer_queue(tq_cfg): + raise RuntimeError("Standalone TQ Producer requires TransferQueue==0.1.10") + ray_cfg = tq_cfg["ray"] + logger.info( + "Standalone TQ Producer connecting Ray address=%s namespace=%s", + ray_cfg["address"], + ray_cfg.get("namespace"), + ) + transport.connect_ray_cluster( + str(ray_cfg["address"]), + str(ray_cfg["namespace"]) if ray_cfg.get("namespace") else None, + ) + logger.info("Standalone TQ Producer connected Ray; initializing TQ client") + transport.connect_transfer_queue_client() + connected = True + logger.info("Standalone TQ Producer connected TQ; waiting for owner_ready") + await _wait_for_owner_ready( + transport, + run_id, + timeout=float(producer_cfg["owner_ready_timeout_seconds"]), + poll_interval=float(producer_cfg["pending_poll_interval_seconds"]), + ) + logger.info("Standalone TQ Producer observed owner_ready run_id=%s", run_id) + + consumed_sequence_nos, resume_metadata = load_standalone_resume( + producer_cfg.get("resume_checkpoint_path"), + input_path=str(producer_cfg["input_path"]), + ) + logger.info( + "Standalone TQ Producer resume progress checkpoint=%s consumed=%s step=%s", + producer_cfg.get("resume_checkpoint_path"), + len(consumed_sequence_nos), + None if resume_metadata is None else resume_metadata.get("optimizer_step"), + ) + + if tokenizer is None: + logger.info( + "Standalone TQ Producer loading tokenizer path=%s", + producer_cfg["tokenizer_path"], + ) + tokenizer = await asyncio.to_thread(_load_tokenizer, producer_cfg) + logger.info("Standalone TQ Producer tokenizer loaded") + if pool is None: + endpoint_concurrency = int(producer_cfg["per_endpoint_concurrency"]) + pool = VllmFeatureClientPool( + [ + VllmEndpoint(str(url).rstrip("/"), endpoint_concurrency) + for url in producer_cfg["vllm_endpoints"] + ], + model=str(producer_cfg["vllm_model"]), + max_inflight_requests=int(producer_cfg["max_inflight_requests"]), + request_timeout=float(producer_cfg["request_timeout"]), + ) + await pool.start() + logger.info("Standalone TQ Producer vLLM client pool started") + + algorithm = str(drafter_cfg["speculative_algorithm"]).strip().upper() + feature_contract = FeatureContract( + algorithm=algorithm, + target_layer_ids=[int(value) for value in producer_cfg["target_layer_ids"]], + hidden_states_layout=resolve_drafter_hidden_states_layout( + algorithm, drafter_cfg + ), + dtype=_parse_dtype(producer_cfg["hidden_dtype"]), + target_model_id=str(producer_cfg["target_model_id"]), + target_model_revision=str(producer_cfg["target_model_revision"]), + tokenizer_fingerprint=str(producer_cfg["tokenizer_fingerprint"]), + use_logits=False, + require_full_alignment=True, + ) + final_norm = None + if feature_contract.hidden_states_layout.endswith("_plus_last"): + final_norm = await asyncio.to_thread( + load_vllm_final_norm, + feature_contract.target_model_id, + dtype=feature_contract.dtype, + trust_remote_code=bool(producer_cfg.get("trust_remote_code", False)), + ) + worker_count = int(producer_cfg["max_inflight_requests"]) + input_queue: asyncio.Queue[Any] = asyncio.Queue( + maxsize=int(producer_cfg["input_queue_size"]) + ) + publish_queue: asyncio.Queue[Any] = asyncio.Queue( + maxsize=int(producer_cfg["publish_queue_size"]) + ) + + async def read_inputs() -> None: + max_samples = int(producer_cfg.get("max_samples", 0) or 0) + epoch = 0 + source_sequence_no = 0 + while True: + epoch_count = 0 + scanned_count = 0 + for source_record in iter_input_records( + str(producer_cfg["input_path"]) + ): + if max_samples > 0 and stats.input_count >= max_samples: + break + sequence_no = source_sequence_no + source_sequence_no += 1 + scanned_count += 1 + if sequence_no in consumed_sequence_nos: + continue + # iter_input_records restarts sequence_no at zero on every + # pass. TQ keys require a run-global sequence number so a + # repeated sample never overwrites an earlier pending copy. + record = replace( + source_record, + sequence_no=sequence_no, + ) + request = ( + prepare_generation_request(record, tokenizer, producer_cfg) + if record.response is None + else tokenize_record(record, tokenizer, producer_cfg) + ) + await input_queue.put(request) + stats.input_count += 1 + epoch_count += 1 + if _should_log_sample_progress(stats.input_count): + logger.info( + "Standalone TQ Producer queued input count=%s epoch=%s " + "sample_id=%s has_response=%s", + stats.input_count, + epoch, + record.sample_id, + record.response is not None, + ) + if scanned_count == 0 and stats.input_count == 0: + raise ValueError("Standalone TQ Producer input contains no samples") + if max_samples <= 0 or stats.input_count >= max_samples: + break + epoch += 1 + logger.info( + "Standalone TQ Producer restarting input epoch=%s " + "queued=%s target=%s", + epoch, + stats.input_count, + max_samples, + ) + for _ in range(worker_count): + await input_queue.put(_INPUT_DONE) + logger.info( + "Standalone TQ Producer input exhausted total=%s", stats.input_count + ) + + async def request_worker() -> None: + while True: + request = await input_queue.get() + if request is _INPUT_DONE: + await publish_queue.put(_PUBLISH_DONE) + return + await _wait_for_pending_capacity( + transport, + run_id, + max_pending_samples=int(producer_cfg["max_pending_samples"]), + poll_interval=float(producer_cfg["pending_poll_interval_seconds"]), + ) + if _should_log_sample_progress(int(request.sequence_no) + 1): + logger.info( + "Standalone TQ Producer requesting vLLM sequence_no=%s " + "sample_id=%s mode=%s", + request.sequence_no, + request.sample_id, + "generate_then_prefill" + if isinstance(request, GenerationRequest) + else "prefill", + ) + if isinstance(request, GenerationRequest): + generated = await pool.generate(request) + try: + request = prepare_generated_prefill_request( + request, + generated.generated_token_ids, + producer_cfg, + ) + finally: + # The generation request may still produce a prompt-only + # connector file. It is not the training payload; the + # following full-sequence prefill produces that payload. + await asyncio.to_thread(delete_temporary_result, generated) + raw = await pool.prefill(request) + else: + raw = await pool.prefill(request) + stats.pending_bytes += int(raw.byte_size) + try: + sample = feature_from_vllm_payload( + raw, request, feature_contract, final_norm=final_norm + ) + except HiddenStateAlignmentError as exc: + stats.dropped_count += 1 + stats.pending_bytes = max( + stats.pending_bytes - int(raw.byte_size), 0 + ) + await asyncio.to_thread(delete_temporary_result, raw) + logger.warning( + "Standalone TQ Producer dropped misaligned sample " + "sequence_no=%s sample_id=%s dropped=%s reason=%s", + request.sequence_no, + request.sample_id, + stats.dropped_count, + exc, + ) + continue + await publish_queue.put( + PreparedFeature( + request=request, + raw=raw, + sample=sample, + metadata=_sample_metadata( + request, sample, feature_contract, run_id, tq_cfg + ), + ) + ) + + async def publish_results() -> None: + finished_workers = 0 + while finished_workers < worker_count: + result = await publish_queue.get() + if result is _PUBLISH_DONE: + finished_workers += 1 + continue + await publish_one(result, transport) + stats.published_count += 1 + if _should_log_sample_progress(stats.published_count): + logger.info( + "Standalone TQ Producer published count=%s sequence_no=%s " + "sample_id=%s", + stats.published_count, + result.request.sequence_no, + result.request.sample_id, + ) + stats.pending_bytes = max( + stats.pending_bytes - int(result.raw.byte_size), 0 + ) + + tasks = [asyncio.create_task(read_inputs())] + tasks.extend(asyncio.create_task(request_worker()) for _ in range(worker_count)) + tasks.append(asyncio.create_task(publish_results())) + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) + failure = next( + (task.exception() for task in done if task.exception() is not None), None + ) + if failure is not None: + stats.failed_count += 1 + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + raise failure + await asyncio.gather(*pending) + + eos_key, eos_fields, eos_tag = make_eos_record(run_id, stats.published_count) + await asyncio.to_thread(transport.put_sample, eos_key, eos_fields, tag=eos_tag) + logger.info( + "Standalone TQ Producer completed inputs=%s published=%s dropped=%s", + stats.input_count, + stats.published_count, + stats.dropped_count, + ) + return stats + finally: + try: + if pool is not None: + await pool.close() + finally: + if connected: + transport.close_transfer_queue_client() + + +async def _wait_for_owner_ready( + transport: Any, + run_id: str, + *, + timeout: float, + poll_interval: float, +) -> None: + deadline = asyncio.get_running_loop().time() + timeout + while True: + records = await asyncio.to_thread(transport.list_samples) + if any( + tag.get("record_type") == "control" + and tag.get("status") == "owner_ready" + and tag.get("run_id") == run_id + for tag in records.values() + ): + return + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError( + f"Timed out waiting for TQ owner_ready for run_id={run_id!r}" + ) + await asyncio.sleep(poll_interval) + + +async def _wait_for_pending_capacity( + transport: Any, + run_id: str, + *, + max_pending_samples: int, + poll_interval: float, +) -> None: + while True: + records = await asyncio.to_thread(transport.list_samples) + ready_count = sum( + 1 + for tag in records.values() + if is_ready_sample_tag( + tag, + run_id=run_id, + schema_version=PROTOCOL_SCHEMA_VERSION, + ) + ) + if ready_count < max_pending_samples: + return + await asyncio.sleep(poll_interval) + + +def _sample_metadata( + request: TokenizedRequest, + sample: DraftFeatureSample, + contract: FeatureContract, + run_id: str, + tq_cfg: Mapping[str, Any], +) -> SampleMetadata: + del sample, contract + return SampleMetadata( + schema_version=int(tq_cfg["schema_version"]), + run_id=run_id, + sample_id=request.sample_id, + sequence_no=request.sequence_no, + ) + + +def _config_sections( + config: Any, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + plain = _plain_config(config) + try: + producer_cfg = plain["speco"]["standalone_tq_producer"] + drafter = plain["actor_rollout_ref"]["rollout"]["drafter"] + training_cfg = drafter["training"] + tq_cfg = training_cfg["transfer_queue"] + except (KeyError, TypeError) as exc: + raise ValueError(f"Producer configuration missing section {exc}") from exc + if not all( + isinstance(value, dict) for value in (producer_cfg, training_cfg, tq_cfg) + ): + raise TypeError("Producer configuration sections must resolve to mappings") + return ( + producer_cfg, + {**training_cfg, "speculative_algorithm": drafter.get("speculative_algorithm")}, + tq_cfg, + ) + + +def _plain_config(config: Any) -> dict[str, Any]: + value = config + try: + from omegaconf import OmegaConf + + if OmegaConf.is_config(config): + value = OmegaConf.to_container(config, resolve=True) + except ImportError: + pass + if not isinstance(value, Mapping): + raise TypeError("Producer configuration must be a mapping") + return dict(value) + + +def _load_tokenizer(config: Mapping[str, Any]) -> Any: + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise RuntimeError("Standalone TQ Producer requires transformers") from exc + return AutoTokenizer.from_pretrained( + str(config["tokenizer_path"]), + trust_remote_code=bool(config.get("trust_remote_code", False)), + ) + + +def _parse_dtype(value: Any) -> torch.dtype: + name = str(value).strip().lower().removeprefix("torch.") + aliases = {"fp32": "float32", "fp16": "float16", "bf16": "bfloat16"} + dtype = getattr(torch, aliases.get(name, name), None) + if not isinstance(dtype, torch.dtype): + raise ValueError(f"Unsupported standalone_tq_producer.hidden_dtype={value!r}") + return dtype + + +def _hydra_main(config: Any) -> None: + logging.basicConfig(level=logging.INFO) + asyncio.run(run_producer(config)) + + +def main() -> None: + try: + import hydra + except ImportError as exc: + raise RuntimeError("Standalone TQ Producer requires hydra-core") from exc + hydra.main(config_path="config", config_name="speco_base", version_base=None)( + _hydra_main + )() + + +if __name__ == "__main__": + main() + + +__all__ = [ + "PreparedFeature", + "ProducerStats", + "main", + "publish_one", + "run_producer", + "validate_producer_config", +] diff --git a/verl_speco/standalone_tq_training_launcher.py b/verl_speco/standalone_tq_training_launcher.py new file mode 100644 index 00000000..3e23da83 --- /dev/null +++ b/verl_speco/standalone_tq_training_launcher.py @@ -0,0 +1,797 @@ +# Copyright 2026 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. +"""Single-entry launcher for Producer -> TransferQueue -> draft training. + +The example script keeps the ordinary standalone-training interface. This +module owns the internal Ray/TQ identity and the owner, Producer and Consumer +process lifecycle so transport-specific overrides do not leak into examples. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import json +import logging +import os +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +import subprocess +import sys +import tempfile +import time +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.parse import urlparse +from urllib.request import urlopen +import uuid + +from verl_speco.trainer.standalone_resume import load_standalone_resume + + +logger = logging.getLogger(__name__) + +_MODEL_PATH_KEY = "actor_rollout_ref.model.path" +_DRAFTER_PATH_KEY = "actor_rollout_ref.rollout.drafter.model_path" +_ALGORITHM_KEY = "actor_rollout_ref.rollout.drafter.speculative_algorithm" +_TRAIN_FILES_KEY = "data.train_files" +_TOKENIZER_PATH_KEY = ( + "actor_rollout_ref.rollout.drafter.training.feature_store.tokenizer_path" +) +_PRODUCER_TARGET_LAYER_IDS_KEY = "speco.standalone_tq_producer.target_layer_ids" +_ALGORITHM_TARGET_LAYER_IDS_KEYS = { + "DFLASH": "actor_rollout_ref.rollout.drafter.training.dflash_target_layer_ids", + "DSPARK": "actor_rollout_ref.rollout.drafter.training.dspark_target_layer_ids", + "DOMINO": "actor_rollout_ref.rollout.drafter.training.domino_target_layer_ids", +} +_MAX_STEPS_KEY = "actor_rollout_ref.rollout.drafter.training.max_steps" +_BATCH_SIZE_PER_GPU_KEY = ( + "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu" +) +_NPROC_KEYS = ( + "speco.draft_training.nproc_per_node", + "speco.draft_training.num_gpus_per_node", + "actor_rollout_ref.rollout.drafter.training.nproc_per_node", + "actor_rollout_ref.rollout.drafter.training.num_gpus_per_node", +) +_NNODES_KEYS = ( + "speco.draft_training.nnodes", + "speco.draft_training.num_nodes", + "actor_rollout_ref.rollout.drafter.training.nnodes", + "actor_rollout_ref.rollout.drafter.training.num_nodes", +) + +_TQ_PREFIX = "actor_rollout_ref.rollout.drafter.training.transfer_queue" +_FEATURE_STORE_PREFIX = "actor_rollout_ref.rollout.drafter.training.feature_store" +_PRODUCER_PREFIX = "speco.standalone_tq_producer" +_PRODUCER_TUNING_KEYS = frozenset( + { + f"{_PRODUCER_PREFIX}.request_timeout", + f"{_PRODUCER_PREFIX}.max_inflight_requests", + f"{_PRODUCER_PREFIX}.per_endpoint_concurrency", + f"{_PRODUCER_PREFIX}.input_queue_size", + f"{_PRODUCER_PREFIX}.publish_queue_size", + f"{_PRODUCER_PREFIX}.max_pending_samples", + f"{_PRODUCER_PREFIX}.pending_poll_interval_seconds", + f"{_PRODUCER_PREFIX}.max_sequence_length", + f"{_PRODUCER_PREFIX}.max_feature_length", + f"{_PRODUCER_PREFIX}.generation_max_tokens", + } +) +_INTERNAL_OVERRIDE_KEYS = frozenset( + { + f"{_FEATURE_STORE_PREFIX}.type", + f"{_FEATURE_STORE_PREFIX}.path", + f"{_FEATURE_STORE_PREFIX}.shuffle", + f"{_FEATURE_STORE_PREFIX}.repeat", + f"{_TQ_PREFIX}.enable", + f"{_TQ_PREFIX}.ray.address", + f"{_TQ_PREFIX}.ray.namespace", + f"{_TQ_PREFIX}.partition_id", + f"{_TQ_PREFIX}.run_id", + f"{_TQ_PREFIX}.drop_last", + f"{_TQ_PREFIX}.backend.storage_backend", + f"{_TQ_PREFIX}.backend.SimpleStorage.total_storage_size", + f"{_TQ_PREFIX}.backend.SimpleStorage.num_data_storage_units", + } +) + +_DEFAULT_TARGET_LAYER_IDS = (1, 9, 17, 25, 33) +_DEFAULT_VLLM_ENDPOINT = "http://127.0.0.1:8000/v1" +_DEFAULT_VLLM_GPU_MEMORY_UTILIZATION = "0.4" +_VLLM_HIDDEN_STATES_DIR = "__SPECO_HIDDEN_STATES_DIR__" +_TQ_NAMESPACE = "speco-drafter" +_TQ_PARTITION = "speco_drafter_features" + + +@dataclass(frozen=True) +class PipelineConfig: + input_path: str + model_path: str + tokenizer_path: str + algorithm: str + target_layer_ids: tuple[int, ...] + vllm_endpoints: tuple[str, ...] + run_id: str + + +@dataclass(frozen=True) +class PipelineCommands: + vllm: list[str] | None + vllm_endpoints: tuple[str, ...] + owner: list[str] + producer: list[str] + consumer: list[str] + + +@dataclass(frozen=True) +class RaySession: + module: Any + address: str + + def close(self) -> None: + self.module.shutdown() + + +def _split_override(item: str) -> tuple[str, str] | None: + if "=" not in item or item.startswith("-"): + return None + key, value = item.split("=", 1) + return key, value + + +def _find_override(overrides: Sequence[str], key: str) -> str | None: + for item in reversed(overrides): + parsed = _split_override(item) + if parsed is not None and parsed[0] == key: + return parsed[1] + return None + + +def _find_first_override(overrides: Sequence[str], keys: Sequence[str]) -> str | None: + for key in keys: + value = _find_override(overrides, key) + if value is not None: + return value + return None + + +def _strip_quotes(value: str) -> str: + value = value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + return value[1:-1] + return value + + +def _single_train_file(value: str | None) -> str: + if value is None: + raise ValueError(f"Standalone TQ training requires {_TRAIN_FILES_KEY}") + text = _strip_quotes(value) + if text.startswith("[") and text.endswith("]"): + items = [_strip_quotes(item) for item in text[1:-1].split(",") if item.strip()] + if len(items) != 1: + raise ValueError("Standalone TQ Producer requires exactly one train file") + text = items[0] + if not text: + raise ValueError("Standalone TQ Producer train file must not be empty") + return text + + +def _parse_layer_ids(value: str | None, *, config_key: str) -> tuple[int, ...]: + if value is None or _strip_quotes(value).lower() in {"", "null", "none"}: + return _DEFAULT_TARGET_LAYER_IDS + text = _strip_quotes(value) + if not (text.startswith("[") and text.endswith("]")): + raise ValueError(f"{config_key} must be a Hydra integer list") + try: + result = tuple(int(item.strip()) for item in text[1:-1].split(",")) + except ValueError as exc: + raise ValueError(f"{config_key} must contain only integers") from exc + if not result or any(value < 0 for value in result): + raise ValueError(f"{config_key} must contain non-negative IDs") + return result + + +def _resolve_target_layer_ids( + training_args: Sequence[str], algorithm: str +) -> tuple[int, ...]: + """Resolve Producer layers without making the launcher DSpark-specific.""" + + algorithm_key = _ALGORITHM_TARGET_LAYER_IDS_KEYS.get(algorithm) + candidate_keys = ( + (_PRODUCER_TARGET_LAYER_IDS_KEY, algorithm_key) + if algorithm_key is not None + else (_PRODUCER_TARGET_LAYER_IDS_KEY,) + ) + raw = _find_first_override(training_args, candidate_keys) + return _parse_layer_ids(raw, config_key=candidate_keys[0]) + + +def _parse_vllm_endpoints(env: Mapping[str, str]) -> tuple[str, ...]: + """Read a Hydra-style endpoint list while preserving the singular fallback.""" + + configured = str(env.get("SPECO_VLLM_ENDPOINTS", "")).strip() + if configured: + text = _strip_quotes(configured) + if not (text.startswith("[") and text.endswith("]")): + raise ValueError( + "SPECO_VLLM_ENDPOINTS must be a Hydra-style list, for example " + "[http://127.0.0.1:8000/v1,http://127.0.0.1:8001/v1]" + ) + endpoints = tuple( + _strip_quotes(item).rstrip("/") + for item in text[1:-1].split(",") + if item.strip() + ) + else: + endpoint = str(env.get("SPECO_VLLM_ENDPOINT", _DEFAULT_VLLM_ENDPOINT)).strip() + endpoints = (endpoint.rstrip("/"),) if endpoint else () + if not endpoints: + raise ValueError("At least one hidden-state vLLM endpoint is required") + for endpoint in endpoints: + parsed = urlparse(endpoint) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError(f"Invalid vLLM endpoint: {endpoint!r}") + return endpoints + + +def _required_override(overrides: Sequence[str], key: str) -> str: + value = _find_override(overrides, key) + normalized = _strip_quotes(value or "") + if not normalized or normalized.startswith("/path/to/"): + raise ValueError(f"Standalone TQ training requires a real {key}") + return normalized + + +def _positive_int_override( + overrides: Sequence[str], keys: Sequence[str], *, default: int +) -> int: + raw = _find_first_override(overrides, keys) + value = int(_strip_quotes(raw)) if raw is not None else int(default) + if value <= 0: + raise ValueError(f"{keys[0]} must be positive, got {value}") + return value + + +def _producer_max_samples( + training_args: Sequence[str], *, resumed_optimizer_step: int = 0 +) -> int: + """Return samples needed for exactly max_steps complete global batches.""" + + raw_max_steps = _find_override(training_args, _MAX_STEPS_KEY) + max_steps = int(_strip_quotes(raw_max_steps)) if raw_max_steps is not None else 1000 + if max_steps <= 0: + # An unbounded training run cannot have a finite Producer target. Keep + # the direct Producer's one-pass behavior instead of looping forever. + return 0 + batch_size = _positive_int_override( + training_args, (_BATCH_SIZE_PER_GPU_KEY,), default=4 + ) + nproc = _positive_int_override(training_args, _NPROC_KEYS, default=1) + nnodes = _positive_int_override(training_args, _NNODES_KEYS, default=1) + remaining_steps = max(max_steps - int(resumed_optimizer_step), 0) + return remaining_steps * batch_size * nproc * nnodes + + +def _stable_path_identity(kind: str, path: str) -> str: + digest = hashlib.sha256(path.encode("utf-8")).hexdigest() + return f"{kind}-path-sha256-{digest}" + + +def _target_final_layer_id(model_path: str, target_layer_ids: Sequence[int]) -> int: + """Resolve the final transformer-layer output ID from a local HF config.""" + + config_path = Path(model_path) / "config.json" + if config_path.is_file(): + try: + model_config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError(f"Cannot read target model config: {config_path}") from exc + for candidate in ( + model_config.get("num_hidden_layers"), + (model_config.get("text_config") or {}).get("num_hidden_layers"), + ): + if candidate is not None and int(candidate) > 0: + return int(candidate) + raise ValueError(f"Target model config has no num_hidden_layers: {config_path}") + # Keep dry-run and model-registry IDs usable. The formal Qwen3-4B/8B + # defaults select layer 33 and use transformer output 36 as the final state. + return max(int(layer_id) for layer_id in target_layer_ids) + 3 + + +def resolve_pipeline_config( + training_args: Sequence[str], + *, + environ: Mapping[str, str] | None = None, +) -> PipelineConfig: + """Derive all Producer/TQ settings from ordinary training arguments.""" + + env = os.environ if environ is None else environ + model_path = _required_override(training_args, _MODEL_PATH_KEY) + input_path = _single_train_file(_find_override(training_args, _TRAIN_FILES_KEY)) + tokenizer_path = _strip_quotes( + _find_override(training_args, _TOKENIZER_PATH_KEY) or model_path + ) + algorithm = _strip_quotes( + _find_override(training_args, _ALGORITHM_KEY) or "DSPARK" + ).upper() + if not algorithm: + raise ValueError(f"{_ALGORITHM_KEY} must not be empty") + target_layer_ids = _resolve_target_layer_ids(training_args, algorithm) + endpoints = _parse_vllm_endpoints(env) + return PipelineConfig( + input_path=input_path, + model_path=model_path, + tokenizer_path=tokenizer_path, + algorithm=algorithm, + target_layer_ids=target_layer_ids, + vllm_endpoints=endpoints, + run_id=f"{algorithm.lower()}-{uuid.uuid4().hex}", + ) + + +def start_ray_session( + *, + environ: Mapping[str, str] | None = None, + ray_module: Any | None = None, +) -> RaySession: + """Create a task-local Ray control plane for the hidden TQ pipeline.""" + + env = os.environ if environ is None else environ + ray_runtime = ray_module + if ray_runtime is None: + try: + ray_runtime = importlib.import_module("ray") + except ImportError as exc: + raise RuntimeError( + "Standalone TQ training requires Ray and TransferQueue==0.1.10" + ) from exc + # ``ray.init()`` consults RAY_ADDRESS when no explicit address is supplied. + # This launcher owns the complete Producer/TQ/Consumer lifetime, so an + # inherited address (often left by another job) must never select its + # control plane. ``local`` explicitly starts this task's Ray runtime. + init_kwargs: dict[str, Any] = { + "address": "local", + "namespace": _TQ_NAMESPACE, + "include_dashboard": False, + } + num_cpus = str(env.get("SPECO_RAY_NUM_CPUS", "")).strip() + if num_cpus: + init_kwargs["num_cpus"] = int(num_cpus) + ray_runtime.init(**init_kwargs) + address = str(ray_runtime.get_runtime_context().gcs_address).strip() + if not address: + ray_runtime.shutdown() + raise RuntimeError("Ray did not report a GCS address for TQ clients") + return RaySession(module=ray_runtime, address=address) + + +def _hydra_list(values: Sequence[Any]) -> str: + return "[" + ",".join(str(value) for value in values) + "]" + + +def _replace_internal_overrides( + training_args: Sequence[str], internal: Sequence[str] +) -> list[str]: + cleaned: list[str] = [] + for item in training_args: + parsed = _split_override(item) + if parsed is not None and parsed[0] in _INTERNAL_OVERRIDE_KEYS: + continue + cleaned.append(item) + return [*cleaned, *internal] + + +def build_pipeline_commands( + config: PipelineConfig, + training_args: Sequence[str], + *, + ray_address: str, + python_executable: str = sys.executable, +) -> PipelineCommands: + """Build the internal commands without exposing transport options.""" + + drafter_path = _strip_quotes(_find_override(training_args, _DRAFTER_PATH_KEY) or "") + _, resume_metadata = load_standalone_resume( + drafter_path or None, + input_path=config.input_path, + ) + resumed_optimizer_step = ( + int(resume_metadata.get("optimizer_step", 0)) + if resume_metadata is not None + else 0 + ) + tq_overrides = [ + f"{_TQ_PREFIX}.enable=true", + f"{_TQ_PREFIX}.ray.address={ray_address}", + f"{_TQ_PREFIX}.ray.namespace={_TQ_NAMESPACE}", + f"{_TQ_PREFIX}.partition_id={_TQ_PARTITION}", + f"{_TQ_PREFIX}.run_id={config.run_id}", + f"{_TQ_PREFIX}.drop_last=true", + f"{_TQ_PREFIX}.backend.storage_backend=SimpleStorage", + f"{_TQ_PREFIX}.backend.SimpleStorage.total_storage_size=17179869184", + f"{_TQ_PREFIX}.backend.SimpleStorage.num_data_storage_units=8", + ] + parsed_endpoint = urlparse(config.vllm_endpoints[0]) + vllm_port = parsed_endpoint.port or ( + 443 if parsed_endpoint.scheme == "https" else 80 + ) + # extract_hidden_states uses the model's layer-output convention. Qwen3-4B/8B + # have 36 transformer layers; the default DSpark auxiliary selection ends at + # 33 and requests the final layer output as 36. + final_layer_id = _target_final_layer_id(config.model_path, config.target_layer_ids) + speculative_config = { + "method": "extract_hidden_states", + "num_speculative_tokens": 1, + "draft_model_config": { + "hf_config": { + "eagle_aux_hidden_state_layer_ids": [ + *config.target_layer_ids, + final_layer_id, + ] + } + }, + } + kv_transfer_config = { + "kv_connector": "ExampleHiddenStatesConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "shared_storage_path": _VLLM_HIDDEN_STATES_DIR, + "use_synchronization_lock": True, + }, + } + vllm = None + if len(config.vllm_endpoints) == 1 and parsed_endpoint.hostname in { + "127.0.0.1", + "localhost", + "0.0.0.0", + }: + vllm = [ + "vllm", + "serve", + config.model_path, + "--host", + "127.0.0.1", + "--port", + str(vllm_port), + "--gpu-memory-utilization", + _DEFAULT_VLLM_GPU_MEMORY_UTILIZATION, + "--speculative-config", + json.dumps(speculative_config, separators=(",", ":")), + "--kv-transfer-config", + json.dumps(kv_transfer_config, separators=(",", ":")), + "--no-enable-chunked-prefill", + ] + owner = [ + python_executable, + "-m", + "verl_speco.tq_owner", + *tq_overrides, + ] + producer_tuning_overrides = [ + item + for item in training_args + if (parsed := _split_override(item)) is not None + and parsed[0] in _PRODUCER_TUNING_KEYS + ] + producer = [ + python_executable, + "-m", + "verl_speco.standalone_tq_producer", + f"{_ALGORITHM_KEY}={config.algorithm}", + *tq_overrides, + *producer_tuning_overrides, + f"speco.standalone_tq_producer.input_path={config.input_path}", + "speco.standalone_tq_producer.resume_checkpoint_path=" + + (drafter_path if resume_metadata is not None else "null"), + f"speco.standalone_tq_producer.tokenizer_path={config.tokenizer_path}", + "speco.standalone_tq_producer.tokenizer_fingerprint=" + + _stable_path_identity("tokenizer", config.tokenizer_path), + f"speco.standalone_tq_producer.target_model_id={config.model_path}", + "speco.standalone_tq_producer.target_model_revision=" + + _stable_path_identity("target", config.model_path), + "speco.standalone_tq_producer.target_layer_ids=" + + _hydra_list(config.target_layer_ids), + "speco.standalone_tq_producer.vllm_endpoints=" + + _hydra_list(config.vllm_endpoints), + f"speco.standalone_tq_producer.vllm_model={config.model_path}", + "speco.standalone_tq_producer.max_samples=" + + str( + _producer_max_samples( + training_args, + resumed_optimizer_step=resumed_optimizer_step, + ) + ), + ] + consumer_internal = [ + f"{_FEATURE_STORE_PREFIX}.type=tq", + f"{_FEATURE_STORE_PREFIX}.path=null", + f"{_FEATURE_STORE_PREFIX}.shuffle=false", + f"{_FEATURE_STORE_PREFIX}.repeat=false", + *tq_overrides, + ] + algorithm_layer_ids_key = _ALGORITHM_TARGET_LAYER_IDS_KEYS.get(config.algorithm) + if algorithm_layer_ids_key is not None: + consumer_internal.append( + f"{algorithm_layer_ids_key}={_hydra_list(config.target_layer_ids)}" + ) + consumer = [ + python_executable, + "-m", + "verl_speco.draft_train_launcher", + *_replace_internal_overrides(training_args, consumer_internal), + ] + return PipelineCommands( + vllm=vllm, + vllm_endpoints=config.vllm_endpoints, + owner=owner, + producer=producer, + consumer=consumer, + ) + + +def _wait_for_owner_ready( + owner: subprocess.Popen[Any], + ready_file: Path, + *, + timeout_seconds: float, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> None: + deadline = monotonic() + timeout_seconds + while not ready_file.is_file(): + returncode = owner.poll() + if returncode is not None: + raise RuntimeError(f"TransferQueue owner exited early ({returncode})") + if monotonic() >= deadline: + raise TimeoutError("Timed out waiting for TransferQueue owner readiness") + sleep(0.1) + + +def _vllm_is_ready(endpoint: str, *, timeout_seconds: float = 1.0) -> bool: + try: + with urlopen( + f"{endpoint.rstrip('/')}/models", timeout=timeout_seconds + ) as response: + return 200 <= int(response.status) < 300 + except (HTTPError, URLError, OSError, TimeoutError): + return False + + +def _wait_for_vllm_ready( + process: subprocess.Popen[Any], + endpoint: str, + *, + timeout_seconds: float, + endpoint_ready: Callable[[str], bool] = _vllm_is_ready, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> None: + deadline = monotonic() + timeout_seconds + while not endpoint_ready(endpoint): + returncode = process.poll() + if returncode is not None: + raise RuntimeError( + f"hidden-state vLLM exited before becoming ready ({returncode})" + ) + if monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for hidden-state vLLM at {endpoint}") + sleep(1.0) + + +def _stop_process(process: subprocess.Popen[Any] | None) -> None: + if process is None or process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +def run_pipeline( + commands: PipelineCommands, + *, + ray_address: str, + environ: Mapping[str, str] | None = None, + popen: Callable[..., subprocess.Popen[Any]] = subprocess.Popen, + poll_interval_seconds: float = 0.2, + owner_ready_timeout_seconds: float = 120, + vllm_ready_timeout_seconds: float = 900, + endpoint_ready: Callable[[str], bool] = _vllm_is_ready, +) -> int: + """Run the internal processes and return the training Consumer status.""" + + base_env = dict(os.environ if environ is None else environ) + # Ray itself also reads RAY_ADDRESS. Pin every child (including the + # torchrun ranks created by the Consumer launcher) to the control plane + # created above instead of allowing a stale inherited value to win. + base_env["RAY_ADDRESS"] = ray_address + owner: subprocess.Popen[Any] | None = None + producer: subprocess.Popen[Any] | None = None + consumer: subprocess.Popen[Any] | None = None + vllm: subprocess.Popen[Any] | None = None + hidden_states_temp: tempfile.TemporaryDirectory[str] | None = None + try: + with tempfile.TemporaryDirectory(prefix="speco-tq-launch-") as temp_dir: + ready_file = Path(temp_dir) / "owner.ready" + hidden_states_temp = tempfile.TemporaryDirectory( + prefix="speco-vllm-hidden-states-" + ) + hidden_states_dir = Path(hidden_states_temp.name) + unavailable_endpoints = [ + endpoint + for endpoint in commands.vllm_endpoints + if not endpoint_ready(endpoint) + ] + if unavailable_endpoints: + if commands.vllm is None: + raise RuntimeError( + "The configured hidden-state vLLM endpoints are unavailable: " + + ", ".join(unavailable_endpoints) + ) + config_endpoint = commands.vllm_endpoints[0] + vllm_command = [ + part.replace(_VLLM_HIDDEN_STATES_DIR, str(hidden_states_dir)) + for part in commands.vllm + ] + logger.info("Starting hidden-state vLLM at %s", config_endpoint) + vllm = popen(vllm_command, env=base_env) + owner_env = {**base_env, "SPECO_TQ_OWNER_READY_FILE": str(ready_file)} + logger.info("Starting TransferQueue owner") + owner = popen(commands.owner, env=owner_env) + _wait_for_owner_ready( + owner, + ready_file, + timeout_seconds=owner_ready_timeout_seconds, + ) + if vllm is not None: + _wait_for_vllm_ready( + vllm, + config_endpoint, + timeout_seconds=vllm_ready_timeout_seconds, + endpoint_ready=endpoint_ready, + ) + else: + unavailable_endpoints = [ + endpoint + for endpoint in commands.vllm_endpoints + if not endpoint_ready(endpoint) + ] + if unavailable_endpoints: + raise RuntimeError( + "hidden-state vLLM became unavailable at: " + + ", ".join(unavailable_endpoints) + ) + logger.info("Starting standalone DSpark Consumer") + consumer = popen(commands.consumer, env=base_env) + logger.info("Starting standalone vLLM Producer") + producer = popen(commands.producer, env=base_env) + + while True: + owner_status = owner.poll() + vllm_status = None if vllm is None else vllm.poll() + producer_status = producer.poll() + consumer_status = consumer.poll() + if owner_status is not None: + raise RuntimeError( + f"TransferQueue owner exited during training ({owner_status})" + ) + if producer_status is not None and producer_status != 0: + return int(producer_status) + if vllm_status is not None: + raise RuntimeError( + f"hidden-state vLLM exited during training ({vllm_status})" + ) + if consumer_status is not None: + return int(consumer_status) + time.sleep(poll_interval_seconds) + except KeyboardInterrupt: + logger.warning("Standalone DSpark training interrupted") + return 130 + finally: + _stop_process(producer) + _stop_process(consumer) + _stop_process(owner) + _stop_process(vllm) + if hidden_states_temp is not None: + hidden_states_temp.cleanup() + + +def _format_command(command: Sequence[str]) -> str: + import shlex + + return " ".join(shlex.quote(part) for part in command) + + +def _preflight_input_file(input_path: str) -> None: + path = Path(input_path) + if not path.is_file(): + raise ValueError(f"Training file does not exist: {input_path}") + from verl_speco.producer.input_reader import iter_input_records + + try: + next(iter_input_records(path)) + except StopIteration as exc: + raise ValueError(f"Training file contains no samples: {input_path}") from exc + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Launch standalone DSpark training through Producer/TQ/Consumer." + ) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--python-executable", default=sys.executable) + args, training_args = parser.parse_known_args(argv) + logging.basicConfig(level=logging.INFO) + + try: + config = resolve_pipeline_config(training_args) + if args.dry_run: + commands = build_pipeline_commands( + config, + training_args, + ray_address="127.0.0.1:6379", + python_executable=args.python_executable, + ) + printable_commands = [ + ("vllm", commands.vllm), + ("owner", commands.owner), + ("producer", commands.producer), + ("consumer", commands.consumer), + ] + for role, command in printable_commands: + if command is None: + print( + f"{role}: external services at " + + ", ".join(commands.vllm_endpoints) + ) + continue + print(f"{role}: {_format_command(command)}") + return 0 + _preflight_input_file(config.input_path) + ray_session = start_ray_session() + try: + commands = build_pipeline_commands( + config, + training_args, + ray_address=ray_session.address, + python_executable=args.python_executable, + ) + logger.info("Using task-local Ray control plane at %s", ray_session.address) + return run_pipeline(commands, ray_address=ray_session.address) + finally: + ray_session.close() + except (OSError, RuntimeError, TimeoutError, ValueError) as exc: + logger.error("Standalone TQ training failed: %s", exc) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = [ + "PipelineCommands", + "PipelineConfig", + "RaySession", + "build_pipeline_commands", + "main", + "resolve_pipeline_config", + "run_pipeline", + "start_ray_session", +] diff --git a/verl_speco/tq_owner.py b/verl_speco/tq_owner.py new file mode 100644 index 00000000..16b1b10b --- /dev/null +++ b/verl_speco/tq_owner.py @@ -0,0 +1,121 @@ +# Copyright 2026 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. +"""Long-lived TransferQueue owner for standalone Producer/Consumer jobs.""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +import signal +import threading +from typing import Any + +import hydra +import torch +from omegaconf import OmegaConf + +from verl_speco.integration.transferqueue_bridge import ( + close_transfer_queue_owner, + configure_transfer_queue, + connect_ray_cluster, + put_sample, + start_transfer_queue_owner, +) +from verl_speco.transport.drafter_sample_protocol import PROTOCOL_SCHEMA_VERSION + + +logger = logging.getLogger(__name__) + + +def install_signal_handlers(stop_event: threading.Event) -> None: + def _request_stop(signum: int, _frame: Any) -> None: + logger.info("TQ owner received signal %s", signum) + stop_event.set() + + signal.signal(signal.SIGINT, _request_stop) + signal.signal(signal.SIGTERM, _request_stop) + + +def publish_owner_ready(run_id: str, schema_version: int) -> str: + if not run_id: + raise ValueError("transfer_queue.run_id must be set for standalone owner") + key = f"control:v{int(schema_version)}:{run_id}:owner-ready" + put_sample( + key, + {"marker": torch.tensor([1], dtype=torch.uint8)}, + tag={ + "record_type": "control", + "status": "owner_ready", + "schema_version": int(schema_version), + "run_id": run_id, + }, + ) + return key + + +def wait_until_stopped(stop_event: threading.Event) -> None: + stop_event.wait() + + +def run_owner(config: Any, *, stop_event: threading.Event | None = None) -> int: + training_cfg = config.actor_rollout_ref.rollout.drafter.training + tq_cfg = OmegaConf.to_container(training_cfg.transfer_queue, resolve=True) + if not isinstance(tq_cfg, dict): + raise TypeError("transfer_queue configuration must resolve to a mapping") + # Invoking the dedicated owner entrypoint is itself the request to enable + # TQ. Keep speco_base.yaml disabled by default for ordinary training jobs, + # and enable only this process's copied configuration. + tq_cfg["enable"] = True + if not configure_transfer_queue(tq_cfg): + raise RuntimeError("Standalone TQ owner requires TransferQueue==0.1.10") + ray_cfg = tq_cfg.get("ray", {}) + ray_address = ray_cfg.get("address") + if not ray_address: + raise ValueError( + "transfer_queue.ray.address must point to a running Ray cluster" + ) + namespace = ray_cfg.get("namespace") + event = stop_event or threading.Event() + if stop_event is None: + install_signal_handlers(event) + + started = False + try: + connect_ray_cluster(str(ray_address), str(namespace) if namespace else None) + start_transfer_queue_owner(tq_cfg) + started = True + ready_key = publish_owner_ready( + str(tq_cfg.get("run_id") or ""), + int(tq_cfg.get("schema_version", PROTOCOL_SCHEMA_VERSION)), + ) + logger.info("TQ owner ready key=%s", ready_key) + ready_file = os.environ.get("SPECO_TQ_OWNER_READY_FILE") + if ready_file: + Path(ready_file).touch() + wait_until_stopped(event) + return 0 + finally: + if started: + close_transfer_queue_owner() + + +@hydra.main(config_path="config", config_name="speco_base", version_base=None) +def main(config: Any) -> None: + logging.basicConfig(level=logging.INFO) + raise SystemExit(run_owner(config)) + + +if __name__ == "__main__": + main() diff --git a/verl_speco/trainer/base_trainer.py b/verl_speco/trainer/base_trainer.py index f504102a..6a108b09 100644 --- a/verl_speco/trainer/base_trainer.py +++ b/verl_speco/trainer/base_trainer.py @@ -827,7 +827,7 @@ def _is_block_drafter_backend(self) -> bool: def _block_drafter_metric_prefix(self) -> str: model_type = str(getattr(self.backend, "model_type", "dflash") or "dflash") - if model_type in {"dspark", "domino", "dflash2"}: + if model_type in {"dspark", "domino", "eagle3", "dflash2"}: return model_type return "dflash" @@ -873,6 +873,14 @@ def get_training_metrics(self) -> dict[str, float]: metrics[f"{prefix}/top5_acc"] = ( sums.get(f"{prefix}/top5_correct_count", 0.0) / quality_tokens ) + simulated_accept_blocks = sums.get( + f"{prefix}/simulated_accept_block_count", 0.0 + ) + if simulated_accept_blocks > 0: + metrics[f"{prefix}/simulated_acc_len"] = ( + sums.get(f"{prefix}/simulated_accept_length_sum", 0.0) + / simulated_accept_blocks + ) ce_tokens = sums.get(f"{prefix}/ce_weighted_token_count", 0.0) if ce_tokens > 0: metrics[f"{prefix}/ce_loss"] = ( @@ -903,7 +911,17 @@ def get_training_metrics(self) -> dict[str, float]: if self.optimizer is not None and self.optimizer.param_groups: metrics["drafter/current_lr"] = float(self.optimizer.param_groups[0]["lr"]) - for pos in range(int(self._block_drafter_config_value("block_size", 16))): + count_prefix = f"{prefix}/count_per_position/" + positions = sorted( + int(key.removeprefix(count_prefix)) + for key in sums + if key.startswith(count_prefix) and key.removeprefix(count_prefix).isdigit() + ) + if not positions: + positions = list( + range(int(self._block_drafter_config_value("block_size", 16))) + ) + for pos in positions: count_key = f"{prefix}/count_per_position/{pos}" count = sums.get(count_key, 0.0) if count <= 0: @@ -960,6 +978,8 @@ def _record_dflash_training_metrics(self, loss_dict: dict[str, Any]) -> None: "quality_token_count": f"{prefix}/quality_token_count", "valid_token_count": f"{prefix}/valid_token_count", "weighted_token_count": f"{prefix}/weighted_token_count", + "simulated_accept_length_sum": f"{prefix}/simulated_accept_length_sum", + "simulated_accept_block_count": f"{prefix}/simulated_accept_block_count", "ce_loss_sum": f"{prefix}/ce_loss_sum", "ce_weighted_token_count": f"{prefix}/ce_weighted_token_count", "l1_loss_sum": f"{prefix}/l1_loss_sum", @@ -4621,10 +4641,12 @@ async def training_step_from_batch( self, batch: dict[str, torch.Tensor], step: int ) -> bool: """Execute one optimizer step from a pre-built standalone batch.""" + self.last_standalone_training_error = None try: with torch.enable_grad(): return await self._training_step_on_batch(batch, step) except Exception as e: # noqa: BLE001 + self.last_standalone_training_error = e logger.exception(f"Standalone training step {step} failed with error: {e}") return False diff --git a/verl_speco/trainer/draft_dataset.py b/verl_speco/trainer/draft_dataset.py index 09edf7b0..315b4508 100644 --- a/verl_speco/trainer/draft_dataset.py +++ b/verl_speco/trainer/draft_dataset.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from typing import Iterator -from verl_speco.trainer.feature_store import DraftFeatureSample, DraftFeatureStore +from verl_speco.trainer.feature_store import DraftFeatureStore, DraftStoredSample @dataclass(frozen=True) @@ -54,7 +54,7 @@ def __init__(self, store: DraftFeatureStore, config: DraftFeatureDataLoaderConfi f"Invalid rank/world_size configuration: rank={rank}, world_size={world_size}" ) - def _sample_in_step_window(self, sample: DraftFeatureSample) -> bool: + def _sample_in_step_window(self, sample: DraftStoredSample) -> bool: if self.config.min_sample_step is None and self.config.max_sample_step is None: return True step = self._sample_step(sample) @@ -71,7 +71,7 @@ def _sample_in_step_window(self, sample: DraftFeatureSample) -> bool: return True @staticmethod - def _sample_step(sample: DraftFeatureSample) -> int | None: + def _sample_step(sample: DraftStoredSample) -> int | None: raw_step = sample.metadata.get("global_step", sample.metadata.get("step")) if raw_step is None: return None @@ -87,7 +87,7 @@ def _uses_step_window(self) -> bool: or self.config.max_sample_step is not None ) - def __iter__(self) -> Iterator[list[DraftFeatureSample]]: + def __iter__(self) -> Iterator[list[DraftStoredSample]]: epoch = 0 while True: keys = list( @@ -116,7 +116,7 @@ def __iter__(self) -> Iterator[list[DraftFeatureSample]]: if world_size > 1: rank_keys = rank_keys[: len(keys) // world_size] rank_samples = [self.store.read(key) for key in rank_keys] - batch: list[DraftFeatureSample] = [] + batch: list[DraftStoredSample] = [] for sample in rank_samples: batch.append(sample) if len(batch) >= int(self.config.batch_size): diff --git a/verl_speco/trainer/draft_training_loop.py b/verl_speco/trainer/draft_training_loop.py index d761d61b..9684c0cc 100644 --- a/verl_speco/trainer/draft_training_loop.py +++ b/verl_speco/trainer/draft_training_loop.py @@ -19,7 +19,9 @@ from copy import deepcopy import json import logging +import math import os +import time from typing import Any, cast import torch @@ -34,11 +36,34 @@ DraftFeatureDataLoader, DraftFeatureDataLoaderConfig, ) -from verl_speco.trainer.feature_store import build_feature_store_from_config +from verl_speco.trainer.feature_store import ( + DraftReplaySample, + build_feature_store_from_config, +) +from verl_speco.trainer.standalone_resume import ( + load_standalone_resume, + save_standalone_resume, +) +from verl_speco.trainer.tq_sample_source import TQFeatureDataLoader, TQLocalBatch logger = logging.getLogger(__name__) +def _should_log_batch_progress(attempted_batches: int) -> bool: + return attempted_batches <= 3 or attempted_batches % 100 == 0 + + +def _is_out_of_memory_error(error: BaseException) -> bool: + message = str(error).lower() + if "out of memory" in message or "oom" in message: + return True + return error.__class__.__name__ in {"OutOfMemoryError", "CudaOutOfMemoryError"} + + +def _contains_replay_samples(samples: list[Any]) -> bool: + return any(isinstance(sample, DraftReplaySample) for sample in samples) + + def run_standalone_draft_training(config) -> dict[str, Any]: """Run independent draft training from a feature store.""" return asyncio.run(_run_standalone_draft_training_async(config)) @@ -46,18 +71,43 @@ def run_standalone_draft_training(config) -> dict[str, Any]: async def _run_standalone_draft_training_async(config) -> dict[str, Any]: rank, local_rank, world_size = _init_distributed() + logger.info( + "[standalone rank=%s] distributed runtime initialized local_rank=%s world_size=%s", + rank, + local_rank, + world_size, + ) draft_config = config.actor_rollout_ref drafter_cfg = draft_config.rollout.drafter training_cfg = drafter_cfg.training feature_store_cfg = training_cfg.feature_store - if not feature_store_cfg.get("path"): + feature_store_type = ( + str(feature_store_cfg.get("type", "torch_shard") or "torch_shard") + .strip() + .lower() + ) + training_mode = ( + str(training_cfg.get("mode", "offline") or "offline").strip().lower() + ) + replay_feature_store_types = {"token_replay", "jsonl_token_replay", "jsonl"} + if feature_store_type != "tq" and not feature_store_cfg.get("path"): raise ValueError( "actor_rollout_ref.rollout.drafter.training.feature_store.path is required" ) + if feature_store_type == "tq" and training_mode != "offline": + raise ValueError( + "feature_store.type=tq requires standalone training.mode=offline" + ) + if feature_store_type in replay_feature_store_types and training_mode != "offline": + raise ValueError( + f"feature_store.type={feature_store_type} is supported only by " + "standalone training.mode=offline" + ) _disable_standalone_sequence_parallel(draft_config) _configure_device(local_rank) backend = _build_backend(draft_config) + setattr(backend, "enable_standalone_training_metrics", True) training_device_mesh = _build_training_device_mesh(draft_config, world_size) trainer = DrafterBaseTrainer( config=draft_config, @@ -77,7 +127,6 @@ async def _run_standalone_draft_training_async(config) -> dict[str, Any]: data_parallel_process_group=None, backend=backend, ) - max_steps = int(training_cfg.get("max_steps", training_cfg.get("step", 1000)) or 0) save_interval = int(training_cfg.get("save_interval_steps", 0) or 0) successful_steps = 0 @@ -86,41 +135,270 @@ async def _run_standalone_draft_training_async(config) -> dict[str, Any]: attempted_batches = 0 last_save_result: dict[str, Any] | None = None last_saved_step = 0 + consumed_sequence_nos: set[int] = set() + standalone_input_path = _standalone_input_path(config) store = None + feature_replayer = None + feature_producer = None + current_stage = "activate_training_model" try: + stage_started = time.perf_counter() + logger.info( + "[standalone rank=%s] activating drafter model algorithm=%s", + rank, + drafter_cfg.speculative_algorithm, + ) activated = await trainer.activate_training_model() if not activated: raise RuntimeError( f"Failed to activate standalone drafter trainer on rank={rank}" ) + logger.info( + "[standalone rank=%s] drafter model activated elapsed=%.3fs", + rank, + time.perf_counter() - stage_started, + ) initial_optimizer_step = int(trainer.optimizer_steps_total) optimizer_step = initial_optimizer_step last_saved_step = optimizer_step + if feature_store_type == "tq": + consumed_sequence_nos, resume_metadata = load_standalone_resume( + drafter_cfg.get("model_path"), + input_path=standalone_input_path, + ) + if initial_optimizer_step > 0 and resume_metadata is None: + raise ValueError( + "Standalone TQ checkpoint restored optimizer state but has no " + "standalone_resume.json; exact data resume is unavailable" + ) + logger.info( + "[standalone rank=%s] resume progress optimizer_step=%s consumed=%s", + rank, + initial_optimizer_step, + len(consumed_sequence_nos), + ) - store = build_feature_store_from_config(feature_store_cfg, read_only=True) - loader = DraftFeatureDataLoader( - store, - DraftFeatureDataLoaderConfig( + current_stage = "open_feature_store" + stage_started = time.perf_counter() + logger.info( + "[standalone rank=%s] opening feature store type=%s path=%s", + rank, + feature_store_type, + feature_store_cfg.get("path"), + ) + if feature_store_type in {"jsonl_token_replay", "jsonl"} and not ( + feature_store_cfg.get("tokenizer_path") + ): + tokenizer_path = draft_config.model.path + try: + feature_store_cfg.tokenizer_path = tokenizer_path + except AttributeError: + feature_store_cfg["tokenizer_path"] = tokenizer_path + store = build_feature_store_from_config( + feature_store_cfg, + read_only=True, + transfer_queue_cfg=training_cfg.get("transfer_queue"), + ) + if feature_store_type == "tq": + current_stage = "connect_tq_feature_store" + _connect_tq_store_across_ranks( + store, + rank=rank, + device=trainer.runtime_device, + ) + logger.info( + "[standalone rank=%s] feature store opened elapsed=%.3fs", + rank, + time.perf_counter() - stage_started, + ) + if feature_store_type in replay_feature_store_types: + # Keep the large target model entirely outside online training imports + # and lifetime. The standalone loop materializes ordinary feature + # samples before handing them to the shared trainer. + from verl_speco.trainer.target_feature_replay import ( + TargetFeatureReplayer, + ) + + current_stage = "initialize_target_feature_replayer" + stage_started = time.perf_counter() + logger.info( + "[standalone rank=%s] initializing target feature replayer", + rank, + ) + feature_replayer = TargetFeatureReplayer( + config, + rank=rank, + world_size=world_size, + device=trainer.runtime_device, + ) + logger.info( + "[standalone rank=%s] target feature replayer initialized " + "backend=%s elapsed=%.3fs", + rank, + feature_replayer.backend, + time.perf_counter() - stage_started, + ) + current_stage = "create_dataloader" + loader: Any + if feature_store_type == "tq": + tq_cfg = training_cfg.get("transfer_queue") or {} + loader = TQFeatureDataLoader( + store, batch_size=int(training_cfg.get("batch_size_per_gpu", 4)), rank=rank, world_size=world_size, - shuffle=bool(feature_store_cfg.get("shuffle", True)), - repeat=bool(feature_store_cfg.get("repeat", True)), - seed=int(training_cfg.get("seed", 0) or 0), - min_sample_step=_optional_int(feature_store_cfg.get("min_sample_step")), - max_sample_step=_optional_int(feature_store_cfg.get("max_sample_step")), - ), + poll_interval_seconds=float( + tq_cfg.get("poll_interval_seconds", 0.5) or 0.5 + ), + drop_last=bool(tq_cfg.get("drop_last", True)), + ) + else: + loader = DraftFeatureDataLoader( + store, + DraftFeatureDataLoaderConfig( + batch_size=int(training_cfg.get("batch_size_per_gpu", 4)), + rank=rank, + world_size=world_size, + shuffle=bool(feature_store_cfg.get("shuffle", True)), + repeat=bool(feature_store_cfg.get("repeat", True)), + seed=int(training_cfg.get("seed", 0) or 0), + min_sample_step=_optional_int( + feature_store_cfg.get("min_sample_step") + ), + max_sample_step=_optional_int( + feature_store_cfg.get("max_sample_step") + ), + ), + ) + logger.info( + "[standalone rank=%s] dataloader ready batch_size_per_gpu=%s " + "shuffle=%s repeat=%s", + rank, + int(training_cfg.get("batch_size_per_gpu", 4)), + bool(feature_store_cfg.get("shuffle", True)), + bool(feature_store_cfg.get("repeat", True)), ) - for samples in loader: - if max_steps > 0 and successful_steps >= max_steps: + sample_source = loader + pipeline_cfg = training_cfg.get("target_feature_pipeline", {}) or {} + pipeline_enabled = bool(pipeline_cfg.get("enabled", False)) + if feature_store_type == "tq" and pipeline_enabled: + raise ValueError( + "feature_store.type=tq already contains target hidden states and cannot be " + "combined with target_feature_pipeline.enabled=true" + ) + if pipeline_enabled: + if feature_replayer is None or not feature_replayer.backend.startswith( + "vllm_" + ): + raise ValueError( + "target_feature_pipeline.enabled=true requires a vLLM replay " + "backend (vllm_file)" + ) + from verl_speco.trainer.target_feature_pipeline import ( + TargetFeatureProducer, + ) + + feature_producer = TargetFeatureProducer( + loader, + feature_replayer, + rank=rank, + concurrency=max( + math.ceil( + int(pipeline_cfg.get("concurrency", 16) or 16) / world_size + ), + 1, + ), + producer_prefetch_depth=int( + pipeline_cfg.get("producer_prefetch_depth", 4) or 4 + ), + prefetch_depth=int(pipeline_cfg.get("prefetch_depth", 2) or 2), + queue_timeout=float(pipeline_cfg.get("queue_timeout", 300.0) or 300.0), + ) + sample_source = feature_producer + sample_iterator = iter(sample_source) + while max_steps <= 0 or optimizer_step < max_steps: + current_stage = "load_next_batch" + loaded_batch = _next_batch_across_ranks( + sample_iterator, + rank=rank, + device=trainer.runtime_device, + ) + if loaded_batch is None: break + tq_local_batch = ( + loaded_batch if isinstance(loaded_batch, TQLocalBatch) else None + ) + samples = ( + tq_local_batch.local_samples + if tq_local_batch is not None + else loaded_batch + ) + step_started = time.perf_counter() attempted_batches += 1 + log_batch_progress = _should_log_batch_progress(attempted_batches) + if log_batch_progress: + logger.info( + "[standalone rank=%s] batch=%s loaded samples=%s " + "successful_steps=%s", + rank, + attempted_batches, + len(samples), + successful_steps, + ) + current_stage = "materialize_target_features" + if feature_replayer is None and _contains_replay_samples(samples): + logger.warning( + "[standalone rank=%s] feature store type=%s yielded replay " + "samples without an initialized target feature replayer; " + "initializing replayer lazily", + rank, + feature_store_type, + ) + from verl_speco.trainer.target_feature_replay import ( + TargetFeatureReplayer, + ) + + feature_replayer = TargetFeatureReplayer( + config, + rank=rank, + world_size=world_size, + device=trainer.runtime_device, + ) + if feature_replayer is not None and feature_producer is None: + materialize_started = time.perf_counter() + if log_batch_progress: + logger.info( + "[standalone rank=%s] batch=%s materializing target features " + "backend=%s", + rank, + attempted_batches, + feature_replayer.backend, + ) + materialized_samples = feature_replayer.materialize(samples) + if log_batch_progress: + logger.info( + "[standalone rank=%s] batch=%s target features materialized " + "samples=%s elapsed=%.3fs", + rank, + attempted_batches, + len(materialized_samples), + time.perf_counter() - materialize_started, + ) + else: + materialized_samples = samples + current_stage = "prepare_training_batch" batch = trainer.prepare_training_batch_from_samples( - cast(list[Any], samples), + cast(list[Any], materialized_samples), step=optimizer_step, ) has_batch = batch is not None + current_stage = "synchronize_batch_readiness" if not _all_ranks_true(has_batch, trainer.runtime_device): + if tq_local_batch is not None: + raise RuntimeError( + "TQ Consumer could not prepare a valid batch on every rank; " + "the TQ keys were intentionally not cleared" + ) if rank == 0: logger.warning( "Skipping standalone drafter batch: at least one rank has no valid batch" @@ -128,31 +406,125 @@ async def _run_standalone_draft_training_async(config) -> dict[str, Any]: continue if batch is None: continue + trainer.reset_training_metrics() + current_stage = "training_step" + if log_batch_progress: + logger.info( + "[standalone rank=%s] batch=%s starting drafter training step " + "optimizer_step=%s", + rank, + attempted_batches, + optimizer_step, + ) ok = await trainer.training_step_from_batch(batch, optimizer_step) + step_error = getattr(trainer, "last_standalone_training_error", None) + if step_error is not None and _is_out_of_memory_error(step_error): + raise RuntimeError( + "Standalone drafter training hit an unrecoverable OOM during " + f"batch={attempted_batches} optimizer_step={optimizer_step}. " + "Reduce batch_size_per_gpu, feature_store.max_seq_len, " + "dspark_num_anchors/block_size or disable DSpark L1 loss." + ) from step_error + current_stage = "synchronize_training_step" if not _all_ranks_true(ok, trainer.runtime_device): + if tq_local_batch is not None: + raise RuntimeError( + "TQ Consumer training_step_from_batch failed on at least one rank; " + "the TQ keys were intentionally not cleared" + ) continue + if tq_local_batch is not None: + current_stage = "clear_tq_batch" + _clear_tq_batch_across_ranks( + cast(TQFeatureDataLoader, loader), + tq_local_batch.global_keys, + rank=rank, + device=trainer.runtime_device, + ) + if rank == 0: + consumed_sequence_nos.update( + tq_local_batch.global_sequence_nos or [] + ) successful_steps += 1 optimizer_step = int(trainer.optimizer_steps_total) if optimizer_step <= initial_optimizer_step: optimizer_step = initial_optimizer_step + successful_steps + step_metrics = _standalone_step_metrics( + trainer, + successful_steps=successful_steps, + attempted_batches=attempted_batches, + step_elapsed_sec=time.perf_counter() - step_started, + ) + if feature_replayer is not None: + step_metrics.update(feature_replayer.metrics()) + if feature_producer is not None: + step_metrics.update(feature_producer.metrics()) + _log_standalone_step_metrics(step_metrics, rank=rank) if save_interval > 0 and optimizer_step % save_interval == 0: - last_save_result = _save_standalone_checkpoint(trainer, optimizer_step) + current_stage = "save_checkpoint" + last_save_result = _save_standalone_checkpoint( + trainer, + optimizer_step, + consumed_sequence_nos=consumed_sequence_nos, + input_path=( + standalone_input_path if feature_store_type == "tq" else None + ), + ) if _sync_any_rank_saved_checkpoint(last_save_result.get("saved")): last_saved_step = optimizer_step _barrier() + current_stage = "load_next_batch" final_save = bool(training_cfg.get("save_final_checkpoint", True)) if final_save and successful_steps > 0 and optimizer_step != last_saved_step: + current_stage = "save_final_checkpoint" last_save_result = _save_standalone_checkpoint( - trainer, optimizer_step, wait=True + trainer, + optimizer_step, + wait=True, + consumed_sequence_nos=consumed_sequence_nos, + input_path=( + standalone_input_path if feature_store_type == "tq" else None + ), ) _barrier() + except Exception: + logger.exception( + "[standalone rank=%s] training failed stage=%s attempted_batches=%s " + "successful_steps=%s optimizer_step=%s", + rank, + current_stage, + attempted_batches, + successful_steps, + optimizer_step, + ) + raise finally: + logger.info( + "[standalone rank=%s] cleanup starting stage=%s attempted_batches=%s " + "successful_steps=%s", + rank, + current_stage, + attempted_batches, + successful_steps, + ) + if feature_producer is not None: + feature_producer.close() + if feature_replayer is not None: + feature_replayer.close() if store is not None: store.close() + logger.info("[standalone rank=%s] cleaning trainer resources", rank) await trainer.cleanup_training(clear_data=True) if dist.is_initialized(): + logger.info( + "[standalone rank=%s] entering final process-group barrier", rank + ) dist.barrier() + logger.info( + "[standalone rank=%s] final process-group barrier complete", rank + ) dist.destroy_process_group() + logger.info("[standalone rank=%s] cleanup complete", rank) return { "rank": rank, @@ -170,8 +542,16 @@ def _build_backend(draft_config): def _save_standalone_checkpoint( - trainer: DrafterBaseTrainer, step: int, *, wait: bool = False + trainer: DrafterBaseTrainer, + step: int, + *, + wait: bool = False, + consumed_sequence_nos: set[int] | None = None, + input_path: str | None = None, ) -> dict[str, Any]: + consumed_snapshot = torch.tensor( + sorted(consumed_sequence_nos or ()), dtype=torch.int64 + ) save_checkpoint = getattr(trainer, "save_checkpoint", None) if callable(save_checkpoint): result = save_checkpoint(int(step), wait=wait) @@ -180,6 +560,12 @@ def _save_standalone_checkpoint( if result.get("saved") and checkpoint_path and is_export_leader: if wait: _rewrite_standalone_block_runtime_config(trainer, checkpoint_path) + _save_resume_sidecar( + checkpoint_path, + consumed_snapshot, + step=step, + input_path=input_path, + ) else: future = getattr(trainer, "_pending_full_checkpoint_future", None) if future is not None: @@ -188,6 +574,9 @@ def _save_standalone_checkpoint( trainer, checkpoint_path, completed, + consumed_snapshot=consumed_snapshot, + step=step, + input_path=input_path, ) ) return result @@ -218,29 +607,60 @@ def _save_standalone_checkpoint( future.result() trainer._pending_full_checkpoint_future = None _rewrite_standalone_block_runtime_config(trainer, checkpoint_path) + _save_resume_sidecar( + checkpoint_path, + consumed_snapshot, + step=step, + input_path=input_path, + ) elif future is not None: future.add_done_callback( - lambda completed: _rewrite_standalone_block_runtime_config( + lambda completed: _finalize_standalone_checkpoint( trainer, checkpoint_path, completed, + consumed_snapshot=consumed_snapshot, + step=step, + input_path=input_path, ) ) return { "saved": future is not None, "path": checkpoint_path, - "reason": "saved" - if future is not None and wait - else "scheduled" - if future is not None - else "not_checkpoint_leader", + "reason": ( + "saved" + if future is not None and wait + else "scheduled" + if future is not None + else "not_checkpoint_leader" + ), } +def _load_tensor_from_safetensors( + path: str, keys: tuple[str, ...] +) -> tuple[str, torch.Tensor] | None: + try: + from safetensors import safe_open + + with safe_open(path, framework="pt", device="cpu") as f: + available_keys = set(f.keys()) + for key in keys: + if key in available_keys: + return key, f.get_tensor(key) + except Exception as exc: # noqa: BLE001 + logger.warning("Failed to load any of %s from %s: %s", keys, path, exc) + return None + + def _finalize_standalone_checkpoint( trainer: DrafterBaseTrainer, checkpoint_path: str, completed_future, + *, + consumed_snapshot: torch.Tensor | None = None, + step: int | None = None, + input_path: str | None = None, ) -> None: try: completed_future.result() @@ -251,6 +671,41 @@ def _finalize_standalone_checkpoint( return _rewrite_standalone_block_runtime_config(trainer, checkpoint_path) + _save_resume_sidecar( + checkpoint_path, + consumed_snapshot, + step=step, + input_path=input_path, + ) + + +def _save_resume_sidecar( + checkpoint_path: str, + consumed_snapshot: torch.Tensor | None, + *, + step: int | None, + input_path: str | None, +) -> None: + if consumed_snapshot is None or step is None or input_path is None: + return + save_standalone_resume( + checkpoint_path, + consumed_snapshot, + optimizer_step=step, + input_path=input_path, + ) + + +def _standalone_input_path(config: Any) -> str | None: + data_cfg = getattr(config, "data", None) + train_files = getattr(data_cfg, "train_files", None) + if train_files is None and isinstance(data_cfg, dict): + train_files = data_cfg.get("train_files") + if isinstance(train_files, str): + return train_files + if train_files is not None and len(train_files) == 1: + return str(train_files[0]) + return None def _ensure_dict_child(config: dict[str, Any], key: str) -> dict[str, Any]: @@ -491,6 +946,145 @@ def _build_training_device_mesh(draft_config, world_size: int) -> DeviceMesh | N ) +def _block_metric_prefix(trainer: DrafterBaseTrainer) -> str | None: + model_type = str(getattr(getattr(trainer, "backend", None), "model_type", "") or "") + if model_type in {"dflash", "dspark", "eagle3"}: + return model_type + return None + + +def _current_learning_rate(trainer: DrafterBaseTrainer) -> float: + optimizer = getattr(trainer, "optimizer", None) + param_groups = getattr(optimizer, "param_groups", None) + if not param_groups: + return 0.0 + return float(param_groups[0].get("lr", 0.0)) + + +def _position_metric_series( + metrics: dict[str, float], prefix: str, name: str +) -> list[float]: + values: list[float] = [] + pos = 0 + while True: + key = f"{prefix}/{name}/{pos}" + if key not in metrics: + break + values.append(float(metrics[key])) + pos += 1 + return values + + +def _weighted_average(values: list[float], counts: list[float]) -> float | None: + if not values or not counts: + return None + total_count = sum(counts[: len(values)]) + if total_count <= 0: + return None + return ( + sum(value * count for value, count in zip(values, counts, strict=False)) + / total_count + ) + + +def _simulated_accept_length(accuracies: list[float]) -> float: + cumulative = 1.0 + simulated = 0.0 + for accuracy in accuracies: + cumulative *= max(0.0, min(1.0, float(accuracy))) + simulated += cumulative + return simulated + + +def _standalone_step_metrics( + trainer: DrafterBaseTrainer, + *, + successful_steps: int, + attempted_batches: int, + step_elapsed_sec: float, +) -> dict[str, float]: + raw_metrics = trainer.get_training_metrics() + metrics: dict[str, float] = { + key: float(value) for key, value in raw_metrics.items() + } + prefix = _block_metric_prefix(trainer) + if prefix is not None: + anchor_offset = 1 if prefix == "dflash" else 0 + losses = _position_metric_series(raw_metrics, prefix, "loss_per_position") + accuracies = _position_metric_series( + raw_metrics, prefix, "accuracy_per_position" + ) + counts = _position_metric_series(raw_metrics, prefix, "count_per_position") + pred_losses = losses[anchor_offset:] + pred_accuracies = accuracies[anchor_offset:] + pred_counts = counts[anchor_offset:] + + avg_loss = _weighted_average(pred_losses, pred_counts) + avg_acc = _weighted_average(pred_accuracies, pred_counts) + if avg_loss is not None: + metrics["train/avg_loss"] = avg_loss + if avg_acc is not None: + metrics["train/avg_acc"] = avg_acc + if f"{prefix}/simulated_acc_len" in raw_metrics: + metrics["train/simulated_acc_len"] = float( + raw_metrics[f"{prefix}/simulated_acc_len"] + ) + elif pred_accuracies: + metrics["train/simulated_acc_len"] = _simulated_accept_length( + pred_accuracies + ) + if f"{prefix}/top1_acc" in raw_metrics: + metrics["train/top1_acc"] = float(raw_metrics[f"{prefix}/top1_acc"]) + if f"{prefix}/top5_acc" in raw_metrics: + metrics["train/top5_acc"] = float(raw_metrics[f"{prefix}/top5_acc"]) + for idx, value in enumerate(pred_losses): + metrics[f"train/ploss_{idx}"] = float(value) + for idx, value in enumerate(pred_accuracies): + metrics[f"train/acc_{idx}"] = float(value) + metrics["train/step"] = float(successful_steps) + metrics["train/global_step"] = float( + getattr(trainer, "training_steps", successful_steps) + ) + metrics["train/lr"] = _current_learning_rate(trainer) + metrics["drafter/train_successful_steps"] = float(successful_steps) + metrics["drafter/train_attempted_batches"] = float(attempted_batches) + metrics["perf/step_time"] = float(step_elapsed_sec) + return metrics + + +def _log_standalone_step_metrics(metrics: dict[str, float], *, rank: int) -> None: + if rank != 0: + return + fields = [f"step={int(metrics.get('train/step', 0.0))}"] + for key, label in ( + ("train/avg_loss", "avg_loss"), + ("train/avg_acc", "avg_acc"), + ("train/top1_acc", "top1"), + ("train/top5_acc", "top5"), + ("train/simulated_acc_len", "sim_acc_len"), + ("train/lr", "lr"), + ("perf/step_time", "step_time"), + ("replay/cache_hit_ratio", "cache_hit"), + ("replay/target_forward_time_total", "target_forward_total"), + ("replay/vllm_request_time_total", "vllm_request_total"), + ("producer/consumer_wait_time_total", "producer_wait_total"), + ("producer/ready_queue_size", "ready_batches"), + ): + if key not in metrics: + continue + value = float(metrics[key]) + if key == "train/lr": + fields.append(f"{label}={value:.3e}") + elif key.endswith("_time_total") or key in { + "perf/step_time", + "replay/target_forward_time_total", + }: + fields.append(f"{label}={value:.3f}s") + else: + fields.append(f"{label}={value:.4f}") + logger.warning("[standalone drafter metrics] %s", " ".join(fields)) + + def _init_distributed() -> tuple[int, int, int]: rank = int(os.environ.get("RANK", "0")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) @@ -539,6 +1133,114 @@ def _all_ranks_true(value: bool, device: torch.device) -> bool: return bool(ready.item()) +def _clear_tq_batch_across_ranks( + loader: TQFeatureDataLoader, + global_keys: list[str] | None, + *, + rank: int, + device: torch.device, +) -> None: + """Clear once on rank 0 and report a clear failure to every training rank.""" + + local_error: BaseException | None = None + if rank == 0: + try: + loader.clear_completed_batch(global_keys) + except BaseException as exc: # noqa: BLE001 + local_error = exc + failed = torch.tensor( + 1 if local_error is not None else 0, + dtype=torch.int32, + device=device, + ) + if dist.is_initialized() and dist.get_world_size() > 1: + dist.all_reduce(failed, op=dist.ReduceOp.MAX) + if bool(failed.item()): + if local_error is not None: + raise RuntimeError( + "rank 0 failed to clear a completed TQ batch" + ) from local_error + raise RuntimeError("rank 0 failed to clear a completed TQ batch") + + +def _connect_tq_store_across_ranks(store, *, rank: int, device: torch.device) -> None: + """Connect every rank before any rank enters TQ key-discovery broadcasts.""" + + local_error: BaseException | None = None + try: + store.connect() + except BaseException as exc: # noqa: BLE001 + local_error = exc + connected = _all_ranks_true(local_error is None, device) + if connected: + return + if local_error is not None: + raise RuntimeError( + f"TQ Consumer failed to connect on rank={rank}" + ) from local_error + raise RuntimeError( + f"TQ Consumer failed to connect on another rank; rank={rank} is stopping" + ) + + +def _next_batch_across_ranks( + source, + *, + rank: int, + device: torch.device, +) -> Any | None: + """Fetch one batch and make producer failures visible to every rank. + + Producer and replay errors happen before the FSDP training step. Every + rank therefore reports its fetch result through the same collective before + any rank is allowed to enter model collectives. This prevents healthy + ranks from waiting in FSDP after another rank has already started cleanup. + """ + samples: Any | None = None + local_error: BaseException | None = None + exhausted = False + try: + samples = next(source) + except StopIteration: + exhausted = True + except BaseException as exc: # noqa: BLE001 + local_error = exc + + state = torch.tensor( + [1 if local_error is not None else 0, 1 if exhausted else 0], + dtype=torch.int32, + device=device, + ) + if dist.is_initialized() and dist.get_world_size() > 1: + dist.all_reduce(state, op=dist.ReduceOp.MAX) + + any_failed = bool(state[0].item()) + any_exhausted = bool(state[1].item()) + if any_failed: + if local_error is not None: + raise RuntimeError( + f"Standalone target-feature producer failed on rank={rank}; " + "all ranks are stopping before the next training collective" + ) from local_error + raise RuntimeError( + "Standalone target-feature producer failed on another rank; " + f"rank={rank} is stopping before the next training collective" + ) + if any_exhausted: + if not exhausted: + logger.warning( + "[standalone rank=%s] discarding a prefetched batch because " + "another rank exhausted its data source", + rank, + ) + return None + if samples is None: + raise RuntimeError( + f"Rank={rank} reported a successful batch fetch without samples" + ) + return samples + + def _sync_any_rank_saved_checkpoint(saved: Any) -> bool: if not dist.is_initialized(): return bool(saved) diff --git a/verl_speco/trainer/feature_store.py b/verl_speco/trainer/feature_store.py index 2f500c9f..cf1ef0df 100644 --- a/verl_speco/trainer/feature_store.py +++ b/verl_speco/trainer/feature_store.py @@ -202,12 +202,132 @@ def to_training_item(self) -> dict[str, Any]: return item +@dataclass +class DraftReplaySample: + """Compact token sample used to reconstruct target features offline.""" + + input_ids: torch.Tensor + loss_mask: torch.Tensor + attention_mask: torch.Tensor + position_ids: torch.Tensor + feature_positions: torch.Tensor + draft_position_ids: torch.Tensor + algorithm: str = "EAGLE3" + schema_version: int = SCHEMA_VERSION + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict( + cls, payload: dict[str, Any], *, strict: bool = True + ) -> "DraftReplaySample": + sample = cls( + schema_version=int(payload.get("schema_version", SCHEMA_VERSION)), + algorithm=str( + payload.get( + "algorithm", payload.get("metadata", {}).get("algorithm", "EAGLE3") + ) + ), + input_ids=payload["input_ids"], + loss_mask=payload["loss_mask"], + attention_mask=payload["attention_mask"], + position_ids=payload["position_ids"], + feature_positions=payload["feature_positions"], + draft_position_ids=payload["draft_position_ids"], + metadata=dict(payload.get("metadata") or {}), + ) + sample.validate(strict=strict) + return sample + + def validate(self, *, strict: bool = True) -> None: + if self.schema_version != SCHEMA_VERSION and strict: + raise ValueError( + f"Unsupported DraftReplaySample schema_version={self.schema_version}" + ) + tensor_fields = ( + "input_ids", + "loss_mask", + "attention_mask", + "position_ids", + "feature_positions", + "draft_position_ids", + ) + for name in tensor_fields: + value = getattr(self, name) + if not torch.is_tensor(value): + raise TypeError(f"DraftReplaySample.{name} must be a torch.Tensor") + if value.dim() > 1: + setattr(self, name, value.reshape(-1)) + + sequence_length = int(self.input_ids.numel()) + for name in ("loss_mask", "attention_mask", "position_ids"): + value = cast(torch.Tensor, getattr(self, name)) + if strict and int(value.numel()) != sequence_length: + raise ValueError( + f"DraftReplaySample input_ids/{name} length mismatch: " + f"{sequence_length} vs {int(value.numel())}" + ) + if strict and int(self.feature_positions.numel()) <= 0: + raise ValueError("DraftReplaySample.feature_positions must not be empty") + if strict and int(self.draft_position_ids.numel()) != int( + self.feature_positions.numel() + ): + raise ValueError( + "DraftReplaySample feature_positions/draft_position_ids length mismatch: " + f"{int(self.feature_positions.numel())} vs " + f"{int(self.draft_position_ids.numel())}" + ) + if int(self.feature_positions.numel()) > 0: + positions = self.feature_positions.detach().cpu().long() + if strict and ( + int(positions.min().item()) < 0 + or int(positions.max().item()) >= sequence_length + ): + raise ValueError( + "DraftReplaySample.feature_positions are outside input_ids: " + f"min={int(positions.min().item())} " + f"max={int(positions.max().item())} sequence_length={sequence_length}" + ) + if strict and int(positions.numel()) > 1: + deltas = positions[1:] - positions[:-1] + if not bool((deltas == 1).all().item()): + raise ValueError( + "DraftReplaySample.feature_positions must be contiguous and increasing" + ) + + def to_dict(self) -> dict[str, Any]: + self.validate(strict=False) + return { + "schema_version": self.schema_version, + "sample_type": "token_replay", + "algorithm": self.algorithm, + "input_ids": self.input_ids.detach().cpu().to(torch.int32).contiguous(), + "loss_mask": self.loss_mask.detach().cpu().to(torch.float16).contiguous(), + "attention_mask": self.attention_mask.detach().cpu().bool().contiguous(), + "position_ids": self.position_ids.detach() + .cpu() + .to(torch.int32) + .contiguous(), + "feature_positions": self.feature_positions.detach() + .cpu() + .to(torch.int32) + .contiguous(), + "draft_position_ids": self.draft_position_ids.detach() + .cpu() + .to(torch.int32) + .contiguous(), + "metadata": dict(self.metadata), + } + + +DraftStoredSample = DraftFeatureSample | DraftReplaySample + + class DraftFeatureStore(Protocol): def write_many( - self, samples: list[DraftFeatureSample | dict[str, Any]] + self, samples: list[DraftStoredSample | dict[str, Any]] ) -> list[str]: ... - def read(self, key: str) -> DraftFeatureSample: ... + def read(self, key: str) -> DraftStoredSample: ... def iter_keys(self, *, shuffle: bool = False, seed: int = 0) -> Iterator[str]: ... @@ -291,7 +411,7 @@ def __init__( self._write_metadata() def write_many( - self, samples: list[DraftFeatureSample | dict[str, Any]] + self, samples: list[DraftStoredSample | dict[str, Any]] ) -> list[str]: if self.read_only: raise RuntimeError("Cannot write to a read-only TorchShardFeatureStore") @@ -342,7 +462,7 @@ def flush_on_step(self, global_step: int | None, interval_steps: int) -> list[st return [] return self.flush() - def read(self, key: str) -> DraftFeatureSample: + def read(self, key: str) -> DraftStoredSample: shard_name, sample_index = _parse_key(key) shard = self._load_shard(shard_name) samples = shard.get("samples") or [] @@ -432,22 +552,560 @@ def _load_shard(self, shard_name: str) -> dict[str, Any]: return torch.load(path, map_location="cpu") +class TokenReplayFeatureStore(TorchShardFeatureStore): + """Compact shard store containing tokens and replay alignment metadata.""" + + def __init__( + self, + path: str | os.PathLike[str], + *, + max_samples_per_shard: int = 1024, + metadata: dict[str, Any] | None = None, + strict_schema: bool = True, + read_only: bool = False, + shard_prefix: str = "shard", + ): + replay_metadata = dict(metadata or {}) + replay_metadata["format"] = "token_replay" + super().__init__( + path, + max_samples_per_shard=max_samples_per_shard, + metadata=replay_metadata, + strict_schema=strict_schema, + read_only=read_only, + shard_prefix=shard_prefix, + ) + + def write_many( + self, samples: list[DraftStoredSample | dict[str, Any]] + ) -> list[str]: + if self.read_only: + raise RuntimeError("Cannot write to a read-only TokenReplayFeatureStore") + keys: list[str] = [] + for sample_like in samples: + sample = _coerce_replay_sample(sample_like, strict=self.strict_schema) + self._pending.append(sample.to_dict()) + keys.append(f"pending:{len(self._pending) - 1}") + if len(self._pending) >= self.max_samples_per_shard: + self.flush() + return keys + + def read(self, key: str) -> DraftReplaySample: + shard_name, sample_index = _parse_key(key) + shard = self._load_shard(shard_name) + samples = shard.get("samples") or [] + sample = samples[int(sample_index)] + return DraftReplaySample.from_dict(sample, strict=self.strict_schema) + + +class JsonlTokenReplayFeatureStore: + """Read token replay JSONL rows as compact replay samples. + + Supported row formats: + - ``{"input_ids": [...], "loss_mask": [...]}`` + - ``{"conversations": [{"from": "human", "value": "..."}, ...]}`` + """ + + def __init__( + self, + path: str | os.PathLike[str], + *, + max_samples_per_shard: int = 1024, + metadata: dict[str, Any] | None = None, + strict_schema: bool = True, + read_only: bool = False, + shard_prefix: str = "shard", + max_seq_len: int = 512, + window_mode: str = "loss", + tokenizer_path: str | os.PathLike[str] | None = None, + trust_remote_code: bool = False, + train_on: str = "last_assistant", + ): + if path is None: + raise ValueError("JsonlTokenReplayFeatureStore requires a non-empty path") + if not read_only: + raise RuntimeError("JsonlTokenReplayFeatureStore is read-only") + self.path = Path(path) + self.max_samples_per_shard = max(int(max_samples_per_shard), 1) + self.metadata = { + "schema_version": SCHEMA_VERSION, + "format": "jsonl_token_replay", + "created_by": "verl_speco", + "created_at": time.time(), + } + if metadata: + self.metadata.update(metadata) + self.strict_schema = bool(strict_schema) + self.read_only = bool(read_only) + self.shard_prefix = str(shard_prefix or "shard") + self.max_seq_len = int(max_seq_len or 0) + self.window_mode = str(window_mode or "loss").strip().lower() + self.tokenizer_path = os.fspath(tokenizer_path) if tokenizer_path else None + self.trust_remote_code = bool(trust_remote_code) + self.train_on = str(train_on or "last_assistant").strip().lower() + self._tokenizer: Any | None = None + if self.window_mode not in {"loss", "front", "full"}: + raise ValueError( + "feature_store.window_mode for jsonl_token_replay must be " + "'loss', 'front' or 'full'" + ) + if self.train_on != "last_assistant": + raise ValueError( + "feature_store.train_on for jsonl_token_replay currently supports " + "only 'last_assistant'" + ) + self._files = self._resolve_jsonl_files() + self._line_offsets: dict[str, list[int]] = {} + self._keys = self._build_keys() + + def write_many( + self, samples: list[DraftStoredSample | dict[str, Any]] + ) -> list[str]: + raise RuntimeError("JsonlTokenReplayFeatureStore is read-only") + + def read(self, key: str) -> DraftReplaySample: + file_name, row_index = _parse_key(key) + file_path = self.path / file_name if self.path.is_dir() else self.path + offsets = self._line_offsets.get(file_name) + if offsets is None: + self._keys = self._build_keys() + offsets = self._line_offsets.get(file_name) + if offsets is None or int(row_index) >= len(offsets): + raise IndexError(f"JSONL row {row_index} not found in {file_path}") + payload = _load_jsonl_offset(file_path, offsets[int(row_index)]) + return self._row_to_replay_sample(payload, file_name, int(row_index)) + + def iter_keys(self, *, shuffle: bool = False, seed: int = 0) -> Iterator[str]: + keys = list(self._keys) + if shuffle: + random.Random(int(seed)).shuffle(keys) + yield from keys + + def get_metadata(self) -> dict[str, Any]: + metadata = dict(self.metadata) + metadata.update( + { + "num_files": len(self._files), + "num_samples": len(self._keys), + "max_seq_len": self.max_seq_len, + "window_mode": self.window_mode, + "train_on": self.train_on, + "tokenizer_path": self.tokenizer_path, + } + ) + return metadata + + def close(self) -> None: + return + + def _resolve_jsonl_files(self) -> list[Path]: + if self.path.is_file(): + return [self.path] + if self.path.is_dir(): + files = sorted(self.path.glob("*.jsonl")) + if files: + return files + raise FileNotFoundError(f"No JSONL file found at {self.path}") + + def _build_keys(self) -> list[str]: + keys: list[str] = [] + self._line_offsets = {} + for file_path in self._files: + file_key = ( + file_path.relative_to(self.path).as_posix() + if self.path.is_dir() + else file_path.name + ) + offsets: list[int] = [] + with file_path.open("rb") as jsonl_file: + while True: + offset = int(jsonl_file.tell()) + line = jsonl_file.readline() + if not line: + break + if line.strip(): + row_index = len(offsets) + offsets.append(offset) + keys.append(f"{file_key}:{row_index}") + self._line_offsets[file_key] = offsets + return keys + + def _row_to_replay_sample( + self, payload: dict[str, Any], file_name: str, line_index: int + ) -> DraftReplaySample: + row_source = "jsonl_token_replay" + if "input_ids" in payload or "loss_mask" in payload: + input_ids = _json_list_tensor(payload, "input_ids", dtype=torch.long) + loss_mask = _json_list_tensor(payload, "loss_mask", dtype=torch.float32) + elif "conversations" in payload: + input_ids, loss_mask = self._conversation_to_input_ids_and_loss_mask( + payload + ) + row_source = "jsonl_conversations" + else: + raise KeyError( + "jsonl_token_replay sample must contain input_ids/loss_mask or " + "conversations" + ) + if int(input_ids.numel()) != int(loss_mask.numel()): + raise ValueError( + "jsonl_token_replay input_ids/loss_mask length mismatch: " + f"{int(input_ids.numel())} vs {int(loss_mask.numel())}" + ) + attention_mask = _optional_json_list_tensor( + payload, "attention_mask", dtype=torch.bool + ) + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + position_ids = _optional_json_list_tensor( + payload, "position_ids", dtype=torch.long + ) + if position_ids is None: + position_ids = attention_mask.long().cumsum(dim=0).sub(1).clamp_min(0) + + sequence_length = int(input_ids.numel()) + start, end = self._feature_window(loss_mask) + feature_positions = torch.arange(start, end, dtype=torch.long) + draft_position_ids = position_ids[start:end].long() + 1 + metadata = { + "source": row_source, + "jsonl_path": file_name, + "jsonl_line": line_index, + "sequence_length": end - start, + "full_sequence_length": sequence_length, + "feature_start": start, + "feature_end": end, + "loss_tokens": int(loss_mask[start:end].sum().item()), + } + for key in ("id", "hash", "primary_id", "finish_reason"): + if key in payload: + metadata[key] = payload[key] + return DraftReplaySample( + algorithm=str(payload.get("algorithm", "EAGLE3")), + input_ids=input_ids, + loss_mask=loss_mask, + attention_mask=attention_mask, + position_ids=position_ids, + feature_positions=feature_positions, + draft_position_ids=draft_position_ids, + metadata=metadata, + ) + + def _ensure_tokenizer(self) -> Any: + if self._tokenizer is not None: + return self._tokenizer + if not self.tokenizer_path: + raise ValueError( + "feature_store.tokenizer_path is required when " + "jsonl_token_replay reads conversations rows" + ) + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise RuntimeError( + "jsonl_token_replay conversations rows require transformers" + ) from exc + self._tokenizer = AutoTokenizer.from_pretrained( + self.tokenizer_path, + trust_remote_code=self.trust_remote_code, + ) + return self._tokenizer + + def _conversation_to_input_ids_and_loss_mask( + self, payload: dict[str, Any] + ) -> tuple[torch.Tensor, torch.Tensor]: + conversations = payload.get("conversations") + if not isinstance(conversations, list) or not conversations: + raise ValueError( + "jsonl_token_replay conversations must be a non-empty list" + ) + messages = [_conversation_item_to_message(item) for item in conversations] + if not messages or messages[-1]["role"] != "assistant": + raise ValueError( + "jsonl_token_replay train_on=last_assistant requires the final " + "conversation item to be assistant" + ) + + tokenizer = self._ensure_tokenizer() + prompt_ids = tokenizer.apply_chat_template( + messages[:-1], + tokenize=True, + add_generation_prompt=True, + ) + full_ids = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=False, + ) + prompt_ids = _token_ids_to_list(prompt_ids) + full_ids = _token_ids_to_list(full_ids) + if len(full_ids) <= len(prompt_ids): + raise ValueError( + "jsonl_token_replay conversations produced no assistant tokens" + ) + loss_mask = [0.0] * len(prompt_ids) + [1.0] * (len(full_ids) - len(prompt_ids)) + return ( + torch.tensor(full_ids, dtype=torch.long).reshape(-1), + torch.tensor(loss_mask, dtype=torch.float32).reshape(-1), + ) + + def _feature_window(self, loss_mask: torch.Tensor) -> tuple[int, int]: + sequence_length = int(loss_mask.numel()) + if sequence_length <= 0: + raise ValueError("jsonl_token_replay input_ids must not be empty") + max_seq_len = self.max_seq_len if self.max_seq_len > 0 else sequence_length + max_seq_len = min(max_seq_len, sequence_length) + if self.window_mode == "full": + return 0, sequence_length + if self.window_mode == "front": + return 0, max_seq_len + + active = torch.nonzero(loss_mask.float() > 0, as_tuple=False).reshape(-1) + if int(active.numel()) <= 0: + return 0, max_seq_len + first_loss = int(active[0].item()) + start = max(first_loss - 1, 0) + end = min(start + max_seq_len, sequence_length) + if end <= start: + end = min(start + 1, sequence_length) + return start, end + + +class VllmSafetensorsFeatureStore(TorchShardFeatureStore): + """Feature store for vLLM-extracted hidden states saved as safetensors. + + Each sample is stored in an individual ``.safetensors`` file with the + non-tensor schema/metadata recorded in the shared manifest. This mirrors + the vLLM/speculators hidden-state extraction flow while exposing the same + ``DraftFeatureStore`` interface used by standalone training. + """ + + def __init__( + self, + path: str | os.PathLike[str], + *, + max_samples_per_shard: int = 1024, + metadata: dict[str, Any] | None = None, + strict_schema: bool = True, + read_only: bool = False, + shard_prefix: str = "hs", + ): + safetensors_metadata = dict(metadata or {}) + safetensors_metadata["format"] = "vllm_safetensors" + super().__init__( + path, + max_samples_per_shard=max_samples_per_shard, + metadata=safetensors_metadata, + strict_schema=strict_schema, + read_only=read_only, + shard_prefix=shard_prefix, + ) + + def write_many( + self, samples: list[DraftStoredSample | dict[str, Any]] + ) -> list[str]: + if self.read_only: + raise RuntimeError( + "Cannot write to a read-only VllmSafetensorsFeatureStore" + ) + try: + from safetensors.torch import save_file + except ImportError as exc: + raise RuntimeError( + "feature_store.type=vllm_safetensors requires safetensors" + ) from exc + + keys: list[str] = [] + for sample_like in samples: + sample = _coerce_sample(sample_like, strict=self.strict_schema) + sample_index = self._infer_next_shard_index() + file_name = f"{self.shard_prefix}_{sample_index:06d}.safetensors" + file_path = self.path / file_name + tensor_payload = self._sample_to_safetensors(sample) + with tempfile.NamedTemporaryFile( + prefix=file_path.name, + suffix=".tmp", + dir=file_path.parent, + delete=False, + ) as tmp_file: + tmp_name = tmp_file.name + try: + save_file(tensor_payload, tmp_name) + os.replace(tmp_name, file_path) + finally: + if os.path.exists(tmp_name): + os.remove(tmp_name) + + entry = { + "path": file_name, + "num_samples": 1, + "num_tokens": _sample_token_count(sample.to_dict()), + "sample": self._sample_manifest(sample), + } + with self.manifest_path.open("a", encoding="utf-8") as manifest_file: + manifest_file.write( + json.dumps(entry, ensure_ascii=True, sort_keys=True) + "\n" + ) + self._manifest.append(entry) + self._next_shard_index = sample_index + 1 + keys.append(f"{file_name}:0") + return keys + + def flush(self) -> list[str]: + return [] + + def read(self, key: str) -> DraftFeatureSample: + try: + from safetensors.torch import load_file + except ImportError as exc: + raise RuntimeError( + "feature_store.type=vllm_safetensors requires safetensors" + ) from exc + + file_name, sample_index = _parse_key(key) + if int(sample_index) != 0: + raise ValueError( + f"Invalid vllm_safetensors key {key!r}; expected sample index 0" + ) + entries = {str(entry.get("path")): entry for entry in self._load_manifest()} + entry = entries.get(file_name) + if entry is None: + raise KeyError(f"Missing vllm_safetensors manifest entry for {file_name}") + tensors = load_file(str(self.path / file_name), device="cpu") + manifest_sample = dict(entry.get("sample") or {}) + payload: dict[str, Any] = { + "schema_version": int( + manifest_sample.get("schema_version", SCHEMA_VERSION) + ), + "algorithm": manifest_sample.get("algorithm", "EAGLE3"), + "input_ids": tensors["input_ids"], + "loss_mask": tensors["loss_mask"], + "hidden_states": tensors["hidden_states"], + "metadata": dict(manifest_sample.get("metadata") or {}), + } + if "metadata.hidden_positions" in tensors: + payload["metadata"]["hidden_positions"] = tensors[ + "metadata.hidden_positions" + ].long() + for optional_key in ( + "last_hidden_states", + "target", + "target_logprobs", + "position_ids", + ): + if optional_key in tensors: + payload[optional_key] = tensors[optional_key] + return DraftFeatureSample.from_dict(payload, strict=self.strict_schema) + + def iter_keys(self, *, shuffle: bool = False, seed: int = 0) -> Iterator[str]: + keys = [f"{entry['path']}:0" for entry in self._load_manifest()] + if shuffle: + random.Random(int(seed)).shuffle(keys) + yield from keys + + def _sample_to_safetensors( + self, sample: DraftFeatureSample + ) -> dict[str, torch.Tensor]: + payload = sample.to_dict() + hidden_states = payload["hidden_states"] + if not torch.is_tensor(hidden_states): + raise TypeError( + "feature_store.type=vllm_safetensors requires tensor hidden_states" + ) + tensors = { + "input_ids": payload["input_ids"].long().contiguous(), + "loss_mask": payload["loss_mask"].float().contiguous(), + "hidden_states": hidden_states.contiguous(), + } + for optional_key in ( + "last_hidden_states", + "target", + "target_logprobs", + "position_ids", + ): + value = payload.get(optional_key) + if value is not None and torch.is_tensor(value): + tensors[optional_key] = value.contiguous() + metadata = payload.get("metadata") or {} + hidden_positions = metadata.get("hidden_positions") + if hidden_positions is not None and torch.is_tensor(hidden_positions): + tensors["metadata.hidden_positions"] = hidden_positions.long().contiguous() + return tensors + + def _sample_manifest(self, sample: DraftFeatureSample) -> dict[str, Any]: + return { + "schema_version": sample.schema_version, + "algorithm": sample.algorithm, + "metadata": _json_safe_metadata(sample.metadata), + } + + def build_feature_store_from_config( - feature_store_cfg, *, read_only: bool = False -) -> TorchShardFeatureStore: - store_type = str(feature_store_cfg.get("type", "torch_shard") or "torch_shard") - if store_type != "torch_shard": + feature_store_cfg, + *, + read_only: bool = False, + metadata: dict[str, Any] | None = None, + shard_prefix: str = "shard", + transfer_queue_cfg: Any | None = None, +) -> Any: + store_type = ( + str(feature_store_cfg.get("type", "torch_shard") or "torch_shard") + .strip() + .lower() + ) + if store_type == "tq": + if not read_only: + raise ValueError( + "feature_store.type=tq is a read-only Consumer data source" + ) + from verl_speco.trainer.tq_feature_store import TQFeatureStore + + tq_cfg = transfer_queue_cfg + if tq_cfg is None: + tq_cfg = feature_store_cfg.get("tq") + if tq_cfg is None: + raise ValueError( + "feature_store.type=tq requires the sibling training.transfer_queue configuration" + ) + return TQFeatureStore.from_config(tq_cfg) + if store_type == "torch_shard": + store_cls: type[TorchShardFeatureStore] = TorchShardFeatureStore + elif store_type == "token_replay": + store_cls = TokenReplayFeatureStore + elif store_type in {"vllm_safetensors", "safetensors"}: + store_cls = VllmSafetensorsFeatureStore + elif store_type in {"jsonl_token_replay", "jsonl"}: + return JsonlTokenReplayFeatureStore( + feature_store_cfg.get("path"), + max_samples_per_shard=int( + feature_store_cfg.get("max_samples_per_shard", 1024) + ), + metadata=metadata, + strict_schema=bool(feature_store_cfg.get("strict_schema", True)), + read_only=read_only, + shard_prefix=shard_prefix, + max_seq_len=int(feature_store_cfg.get("max_seq_len", 512) or 0), + window_mode=str(feature_store_cfg.get("window_mode", "loss") or "loss"), + tokenizer_path=feature_store_cfg.get("tokenizer_path"), + trust_remote_code=bool(feature_store_cfg.get("trust_remote_code", False)), + train_on=str( + feature_store_cfg.get("train_on", "last_assistant") or "last_assistant" + ), + ) + else: raise NotImplementedError(f"Unsupported draft feature store type: {store_type}") - return TorchShardFeatureStore( + return store_cls( feature_store_cfg.get("path"), max_samples_per_shard=int(feature_store_cfg.get("max_samples_per_shard", 1024)), + metadata=metadata, strict_schema=bool(feature_store_cfg.get("strict_schema", True)), read_only=read_only, + shard_prefix=shard_prefix, ) def _coerce_sample( - sample_like: DraftFeatureSample | dict[str, Any], *, strict: bool + sample_like: DraftStoredSample | dict[str, Any], *, strict: bool ) -> DraftFeatureSample: if isinstance(sample_like, DraftFeatureSample): sample_like.validate(strict=strict) @@ -457,6 +1115,17 @@ def _coerce_sample( raise TypeError(f"Unsupported draft feature sample type: {type(sample_like)!r}") +def _coerce_replay_sample( + sample_like: DraftStoredSample | dict[str, Any], *, strict: bool +) -> DraftReplaySample: + if isinstance(sample_like, DraftReplaySample): + sample_like.validate(strict=strict) + return sample_like + if isinstance(sample_like, dict): + return DraftReplaySample.from_dict(sample_like, strict=strict) + raise TypeError(f"Unsupported draft replay sample type: {type(sample_like)!r}") + + def _cpu_tensor_tree(value: Any) -> Any: if torch.is_tensor(value): return value.detach().cpu().contiguous() @@ -518,3 +1187,102 @@ def _atomic_torch_save(payload: dict[str, Any], path: Path) -> None: finally: if os.path.exists(tmp_name): os.remove(tmp_name) + + +def _load_jsonl_offset(path: Path, offset: int) -> dict[str, Any]: + with path.open("rb") as jsonl_file: + jsonl_file.seek(int(offset)) + line = jsonl_file.readline().decode("utf-8").strip() + if not line: + raise ValueError(f"JSONL offset {offset} in {path} points to an empty line") + payload = json.loads(line) + if not isinstance(payload, dict): + raise TypeError(f"JSONL offset {offset} in {path} must contain a JSON object") + return payload + + +def _json_list_tensor( + payload: dict[str, Any], key: str, *, dtype: torch.dtype +) -> torch.Tensor: + if key not in payload: + raise KeyError(f"jsonl_token_replay sample missing required key {key!r}") + value = payload[key] + if not isinstance(value, list): + raise TypeError( + f"jsonl_token_replay {key} must be a JSON list, got {type(value).__name__}" + ) + return torch.tensor(value, dtype=dtype).reshape(-1) + + +def _optional_json_list_tensor( + payload: dict[str, Any], key: str, *, dtype: torch.dtype +) -> torch.Tensor | None: + if key not in payload or payload[key] is None: + return None + return _json_list_tensor(payload, key, dtype=dtype) + + +def _normalize_conversation_role(value: Any) -> str: + role = str(value or "").strip().lower() + if role in {"human", "user"}: + return "user" + if role in {"assistant", "gpt"}: + return "assistant" + if role == "system": + return "system" + return role + + +def _conversation_item_to_message(item: Any) -> dict[str, str]: + if not isinstance(item, dict): + raise TypeError("jsonl_token_replay conversations entries must be JSON objects") + role = _normalize_conversation_role(item.get("role", item.get("from"))) + content = item.get("content", item.get("value", "")) + if role not in {"system", "user", "assistant"}: + raise ValueError( + f"Unsupported conversation role for jsonl_token_replay: {role!r}" + ) + return {"role": role, "content": str(content)} + + +def _token_ids_to_list(value: Any) -> list[int]: + if hasattr(value, "data") and isinstance(getattr(value, "data", None), dict): + data = getattr(value, "data") + if "input_ids" in data: + value = data["input_ids"] + elif isinstance(value, dict) and "input_ids" in value: + value = value["input_ids"] + if torch.is_tensor(value): + value = value.detach().cpu().tolist() + elif hasattr(value, "tolist"): + value = value.tolist() + if ( + isinstance(value, list) + and len(value) == 1 + and isinstance(value[0], (list, tuple)) + ): + value = value[0] + if not isinstance(value, (list, tuple)): + raise TypeError( + "tokenizer.apply_chat_template must return token ids as a list, " + f"tuple, tensor, or tolist()-compatible value; got {type(value).__name__}" + ) + return [int(token_id) for token_id in value] + + +def _json_safe_metadata(value: Any) -> Any: + if torch.is_tensor(value): + if value.numel() <= 128: + return value.detach().cpu().tolist() + return { + "__tensor__": True, + "shape": list(value.shape), + "dtype": str(value.dtype), + } + if isinstance(value, dict): + return {str(key): _json_safe_metadata(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe_metadata(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) diff --git a/verl_speco/trainer/standalone_resume.py b/verl_speco/trainer/standalone_resume.py new file mode 100644 index 00000000..206388e9 --- /dev/null +++ b/verl_speco/trainer/standalone_resume.py @@ -0,0 +1,129 @@ +# Copyright 2026 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. +"""Standalone-only data progress stored beside a drafter checkpoint.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Iterable + +import torch + + +RESUME_METADATA_NAME = "standalone_resume.json" +CONSUMED_SEQUENCE_NAME = "consumed_sequence_nos.pt" +RESUME_SCHEMA_VERSION = 1 + + +def build_input_fingerprint(path: str | os.PathLike[str]) -> dict[str, Any]: + input_path = Path(path).resolve() + stat = input_path.stat() + return { + "path": str(input_path), + "size_bytes": int(stat.st_size), + "mtime_ns": int(stat.st_mtime_ns), + } + + +def save_standalone_resume( + checkpoint_path: str | os.PathLike[str], + consumed_sequence_nos: Iterable[int] | torch.Tensor, + *, + optimizer_step: int, + input_path: str | os.PathLike[str], +) -> None: + checkpoint_dir = Path(checkpoint_path) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + if isinstance(consumed_sequence_nos, torch.Tensor): + values = consumed_sequence_nos.detach().to(device="cpu", dtype=torch.int64) + else: + values = torch.tensor( + sorted({int(value) for value in consumed_sequence_nos}), + dtype=torch.int64, + ) + values = torch.unique(values.flatten(), sorted=True) + if values.numel() and int(values[0]) < 0: + raise ValueError("consumed sequence numbers must be non-negative") + + tensor_path = checkpoint_dir / CONSUMED_SEQUENCE_NAME + tensor_temporary = tensor_path.with_suffix(tensor_path.suffix + ".incomplete") + torch.save(values, tensor_temporary) + os.replace(tensor_temporary, tensor_path) + + metadata = { + "schema_version": RESUME_SCHEMA_VERSION, + "optimizer_step": int(optimizer_step), + "consumed_count": int(values.numel()), + "consumed_sequence_file": CONSUMED_SEQUENCE_NAME, + "input_fingerprint": build_input_fingerprint(input_path), + } + metadata_path = checkpoint_dir / RESUME_METADATA_NAME + metadata_temporary = metadata_path.with_suffix(metadata_path.suffix + ".incomplete") + metadata_temporary.write_text( + json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + os.replace(metadata_temporary, metadata_path) + + +def load_standalone_resume( + checkpoint_path: str | os.PathLike[str] | None, + *, + input_path: str | os.PathLike[str] | None = None, +) -> tuple[set[int], dict[str, Any] | None]: + if not checkpoint_path: + return set(), None + checkpoint_dir = Path(checkpoint_path) + metadata_path = checkpoint_dir / RESUME_METADATA_NAME + if not metadata_path.is_file(): + return set(), None + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if int(metadata.get("schema_version", 0)) != RESUME_SCHEMA_VERSION: + raise ValueError(f"Unsupported standalone resume metadata: {metadata_path}") + tensor_name = metadata.get("consumed_sequence_file") + if tensor_name != CONSUMED_SEQUENCE_NAME: + raise ValueError(f"Invalid consumed sequence file in {metadata_path}") + tensor_path = checkpoint_dir / tensor_name + try: + values = torch.load(tensor_path, map_location="cpu", weights_only=True) + except TypeError: + values = torch.load(tensor_path, map_location="cpu") + if not isinstance(values, torch.Tensor) or values.dtype != torch.int64: + raise ValueError(f"Invalid consumed sequence tensor: {tensor_path}") + values = values.flatten() + if values.numel() and int(values[0]) < 0: + raise ValueError(f"Negative consumed sequence number in {tensor_path}") + consumed = {int(value) for value in values.tolist()} + if len(consumed) != int(metadata.get("consumed_count", -1)): + raise ValueError(f"Consumed sequence count mismatch in {checkpoint_dir}") + if input_path is not None: + saved_fingerprint = metadata.get("input_fingerprint") + current_fingerprint = build_input_fingerprint(input_path) + if saved_fingerprint != current_fingerprint: + raise ValueError( + "Standalone resume input file changed since the checkpoint was saved: " + f"saved={saved_fingerprint!r}, current={current_fingerprint!r}" + ) + return consumed, metadata + + +__all__ = [ + "CONSUMED_SEQUENCE_NAME", + "RESUME_METADATA_NAME", + "build_input_fingerprint", + "load_standalone_resume", + "save_standalone_resume", +] diff --git a/verl_speco/trainer/target_feature_pipeline.py b/verl_speco/trainer/target_feature_pipeline.py new file mode 100644 index 00000000..284d87bf --- /dev/null +++ b/verl_speco/trainer/target_feature_pipeline.py @@ -0,0 +1,168 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +"""Bounded producer/prefetch pipeline for standalone target features.""" + +from __future__ import annotations + +import logging +import queue +import threading +import time +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any, Iterable, Iterator + +from verl_speco.trainer.feature_store import DraftFeatureSample + +logger = logging.getLogger(__name__) + + +@dataclass +class _PipelineFailure: + error: BaseException + + +_END = object() + + +class TargetFeatureProducer: + """Materialize future batches while the current FSDP step is running. + + The coordinator owns the source iterator and puts complete batches into a + bounded ready queue. Sample requests inside each batch are concurrent. + Consequently a training rank never observes a partially materialized batch. + """ + + def __init__( + self, + source: Iterable[list[Any]], + replayer: Any, + *, + rank: int, + concurrency: int, + producer_prefetch_depth: int, + prefetch_depth: int, + queue_timeout: float, + ): + self.source = iter(source) + self.replayer = replayer + self.rank = int(rank) + self.concurrency = max(int(concurrency), 1) + self.producer_prefetch_depth = max(int(producer_prefetch_depth), 1) + self.prefetch_depth = max(int(prefetch_depth), 1) + self.queue_timeout = max(float(queue_timeout), 1.0) + self._ready: queue.Queue[Any] = queue.Queue(maxsize=self.prefetch_depth) + self._stop = threading.Event() + self._request_executor = ThreadPoolExecutor( + max_workers=self.concurrency, + thread_name_prefix=f"speco-request-r{self.rank}", + ) + self._thread = threading.Thread( + target=self._run, + name=f"speco-target-producer-r{self.rank}", + daemon=True, + ) + self.produced_batches = 0 + self.produced_samples = 0 + self.producer_seconds = 0.0 + self.queue_wait_seconds = 0.0 + self.consumer_wait_seconds = 0.0 + self.transfer_seconds = 0.0 + self.failed_batches = 0 + self._thread.start() + logger.info( + "[target producer rank=%s] started request_concurrency=%s " + "producer_prefetch_depth=%s prefetch_depth=%s", + self.rank, + self.concurrency, + self.producer_prefetch_depth, + self.prefetch_depth, + ) + + def _run(self) -> None: + try: + pending: deque[tuple[float, list[Future[Any]]]] = deque() + + def submit_next() -> bool: + if self._stop.is_set(): + return False + try: + samples = next(self.source) + except StopIteration: + return False + started = time.perf_counter() + futures = [ + self._request_executor.submit(self.replayer.materialize, [sample]) + for sample in samples + ] + pending.append((started, futures)) + return True + + for _ in range(self.producer_prefetch_depth): + if not submit_next(): + break + + while pending and not self._stop.is_set(): + started, request_futures = pending.popleft() + produced = [future.result() for future in request_futures] + self.producer_seconds += time.perf_counter() - started + submit_next() + transfer_started = time.perf_counter() + batch = [item for group in produced for item in group] + self.transfer_seconds += time.perf_counter() - transfer_started + self.produced_batches += 1 + self.produced_samples += len(batch) + self._put(batch) + self._put(_END) + except BaseException as exc: # noqa: BLE001 + self.failed_batches += 1 + self._put(_PipelineFailure(exc)) + + def _put(self, value: Any) -> None: + started = time.perf_counter() + while not self._stop.is_set(): + try: + self._ready.put(value, timeout=0.2) + self.queue_wait_seconds += time.perf_counter() - started + return + except queue.Full: + continue + + def __iter__(self) -> Iterator[list[DraftFeatureSample]]: + return self + + def __next__(self) -> list[DraftFeatureSample]: + started = time.perf_counter() + try: + value = self._ready.get(timeout=self.queue_timeout) + except queue.Empty as exc: + raise TimeoutError( + "Timed out waiting for target-feature producer; inspect the vLLM " + "logs for a stalled request or missing hidden-state file" + ) from exc + self.consumer_wait_seconds += time.perf_counter() - started + if value is _END: + raise StopIteration + if isinstance(value, _PipelineFailure): + raise RuntimeError("Target-feature producer failed") from value.error + return value + + def metrics(self) -> dict[str, float]: + return { + "producer/batches_total": float(self.produced_batches), + "producer/samples_total": float(self.produced_samples), + "producer/materialize_time_total": float(self.producer_seconds), + "producer/transfer_time_total": float(self.transfer_seconds), + "producer/queue_block_time_total": float(self.queue_wait_seconds), + "producer/consumer_wait_time_total": float(self.consumer_wait_seconds), + "producer/ready_queue_size": float(self._ready.qsize()), + "producer/ready_queue_capacity": float(self.prefetch_depth), + "producer/failed_batches_total": float(self.failed_batches), + } + + def close(self) -> None: + self._stop.set() + self._request_executor.shutdown(wait=True, cancel_futures=True) + self._thread.join(timeout=5.0) diff --git a/verl_speco/trainer/target_feature_replay.py b/verl_speco/trainer/target_feature_replay.py new file mode 100644 index 00000000..b5d4b9a8 --- /dev/null +++ b/verl_speco/trainer/target_feature_replay.py @@ -0,0 +1,1537 @@ +# Copyright 2026 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. +"""Reconstruct standalone draft-training features from compact token samples.""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import logging +import os +import tempfile +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, cast + +import torch +from torch import nn + +from verl_speco.integration.oldlogprob_layer_ids import ( + resolve_oldlogprob_aux_layer_ids, +) +from verl_speco.trainer.feature_store import DraftFeatureSample, DraftReplaySample + +logger = logging.getLogger(__name__) + + +class HiddenStateAlignmentError(ValueError): + """The returned hidden states cannot cover the requested training sample.""" + + +@dataclass(frozen=True) +class FeatureContract: + """Explicit inputs for converting one vLLM payload into a training sample.""" + + algorithm: str + target_layer_ids: list[int] + hidden_states_layout: str + dtype: torch.dtype + target_model_id: str + target_model_revision: str | None + tokenizer_fingerprint: str + use_logits: bool = False + target_config_fingerprint: str | None = None + source: str = "standalone_tq_producer" + require_full_alignment: bool = False + + +@dataclass +class _VllmEndpointState: + index: int + url: str + client: Any | None = None + model: str | None = None + inflight: int = 0 + requests: int = 0 + failures: int = 0 + consecutive_failures: int = 0 + request_seconds: float = 0.0 + cooldown_until: float = 0.0 + + +def _config_value(config: Any, key: str, default: Any = None) -> Any: + if config is None: + return default + if hasattr(config, "get"): + return config.get(key, default) + return getattr(config, key, default) + + +def _normalize_vllm_endpoints(config: Any) -> list[str]: + configured = _config_value(config, "vllm_endpoints", None) + if configured is None: + configured = [ + _config_value(config, "vllm_endpoint", "http://localhost:8000/v1") + ] + elif isinstance(configured, str): + configured = [configured] + endpoints: list[str] = [] + for value in configured: + endpoint = str(value or "").strip().rstrip("/") + if endpoint and endpoint not in endpoints: + endpoints.append(endpoint) + if not endpoints: + raise ValueError( + "target_feature_replay.vllm_endpoints must contain at least one URL" + ) + return endpoints + + +def _parse_dtype(value: Any) -> torch.dtype: + normalized = str(value or "bfloat16").strip().lower() + dtypes = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, + "fp32": torch.float32, + "float32": torch.float32, + } + if normalized not in dtypes: + raise ValueError( + f"Unsupported target feature replay dtype {value!r}; " + "expected bfloat16, float16, or float32" + ) + return dtypes[normalized] + + +def _tensor_from_module_output(value: Any) -> torch.Tensor: + if torch.is_tensor(value): + return cast(torch.Tensor, value) + if isinstance(value, (tuple, list)) and value and torch.is_tensor(value[0]): + return cast(torch.Tensor, value[0]) + last_hidden_state = getattr(value, "last_hidden_state", None) + if torch.is_tensor(last_hidden_state): + return cast(torch.Tensor, last_hidden_state) + raise TypeError( + f"Target feature replay expected tensor-like module output, got {type(value)!r}" + ) + + +def _get_module_by_path(root: Any, path: str) -> Any: + current = root + for part in path.split("."): + if not part: + continue + current = getattr(current, part, None) + if current is None: + return None + return current + + +def _find_layers_and_final_norm(model: nn.Module) -> tuple[list[nn.Module], nn.Module]: + roots: list[nn.Module] = [model] + base_model = getattr(model, "base_model", None) + if isinstance(base_model, nn.Module) and base_model is not model: + roots.append(base_model) + + candidates = ( + ("model.layers", "model.norm"), + ("base_model.model.layers", "base_model.model.norm"), + ("model.decoder.layers", "model.decoder.final_layer_norm"), + ("transformer.h", "transformer.ln_f"), + ("gpt_neox.layers", "gpt_neox.final_layer_norm"), + ) + for root in roots: + for layers_path, norm_path in candidates: + layers = _get_module_by_path(root, layers_path) + norm = _get_module_by_path(root, norm_path) + if ( + isinstance(layers, (nn.ModuleList, list, tuple)) + and len(layers) > 0 + and isinstance(norm, nn.Module) + ): + return list(layers), norm + + for name, child in root.named_modules(): + if not isinstance(child, nn.ModuleList) or len(child) <= 0: + continue + if not (name.endswith("layers") or name.endswith("h")): + continue + parent_path = name.rsplit(".", 1)[0] if "." in name else "" + for norm_name in ("norm", "final_layer_norm", "ln_f"): + norm_path = f"{parent_path}.{norm_name}" if parent_path else norm_name + norm = _get_module_by_path(root, norm_path) + if isinstance(norm, nn.Module): + return list(child), norm + + raise RuntimeError( + "Target feature replay could not find transformer layers and final norm" + ) + + +def _hidden_capture_target(layer_id: int, num_layers: int) -> tuple[str, int | None]: + hidden_state_index = ( + int(layer_id) + 1 if int(layer_id) >= 0 else num_layers + 1 + int(layer_id) + ) + if hidden_state_index <= 0 or hidden_state_index > num_layers: + raise IndexError( + f"Target replay layer id {layer_id} resolved to hidden-state index " + f"{hidden_state_index}, but the model has {num_layers} layers" + ) + if hidden_state_index == num_layers: + return "final", None + return "layer", hidden_state_index - 1 + + +def load_vllm_final_norm( + model_path: str, + *, + dtype: torch.dtype, + trust_remote_code: bool = False, + target_config: Any = None, +) -> nn.Module: + """Load only the target's final norm, using its actual HF implementation. + + extract_hidden_states collects layer residuals BEFORE the target final norm. + Constructing the architecture on meta discovers the norm without allocating + target weights; only that module's checkpoint tensors are read into CPU RAM. + The checkpoint must be the same frozen target served by the vLLM endpoint. + """ + from transformers import AutoConfig, AutoModelForCausalLM + + from verl_speco.checkpoint_tensor import _load_checkpoint_tensor + + if target_config is None: + target_config = AutoConfig.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) + with torch.device("meta"): + model = AutoModelForCausalLM.from_config( + target_config, + trust_remote_code=trust_remote_code, + attn_implementation="eager", + ) + _, norm = _find_layers_and_final_norm(model) + norm_name = next(name for name, module in model.named_modules() if module is norm) + state = { + key: _load_checkpoint_tensor(model_path, f"{norm_name}.{key}") + for key in norm.state_dict() + } + norm.load_state_dict(state, strict=True, assign=True) + norm = norm.to(device="cpu", dtype=dtype).eval().requires_grad_(False) + logger.info("Loaded vLLM target final norm %s from %s", norm_name, model_path) + return norm + + +def _load_json_config(path: Any) -> dict[str, Any] | None: + if not path: + return None + config_path = os.path.join(os.fspath(path), "config.json") + try: + with open(config_path, encoding="utf-8") as config_file: + value = json.load(config_file) + except (OSError, json.JSONDecodeError): + return None + return value if isinstance(value, dict) else None + + +def _wait_for_lock(lock_path: Path, timeout: float = 30.0) -> None: + if not lock_path.exists(): + return + try: + import fcntl + except ImportError: + deadline = time.monotonic() + float(timeout) + while lock_path.exists(): + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for hidden-states lock: {lock_path}" + ) + time.sleep(0.1) + return + + fd = os.open(lock_path, os.O_RDONLY) + try: + deadline = time.monotonic() + float(timeout) + while True: + try: + fcntl.flock(fd, fcntl.LOCK_SH | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for hidden-states lock: {lock_path}" + ) from None + time.sleep(0.1) + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + try: + lock_path.unlink() + except OSError: + pass + + +class BoundedReplayCache: + """Per-rank disk cache with a hard least-recently-used size budget.""" + + def __init__( + self, + path: str | os.PathLike[str], + *, + max_size_gb: float, + rank: int, + world_size: int, + ): + global_max_bytes = max(int(float(max_size_gb) * 1024**3), 0) + self.max_bytes = global_max_bytes // max(int(world_size), 1) + self.path = Path(path) / f"rank{int(rank):05d}" + self.path.mkdir(parents=True, exist_ok=True) + self._entries: dict[Path, tuple[int, float]] = {} + self._total_bytes = 0 + self._scan() + + @property + def enabled(self) -> bool: + return self.max_bytes > 0 + + def _scan(self) -> None: + self._entries = {} + self._total_bytes = 0 + for path in self.path.glob("*.pt"): + try: + stat = path.stat() + except OSError: + continue + size = int(stat.st_size) + self._entries[path] = (size, float(stat.st_mtime)) + self._total_bytes += size + + def get(self, key: str) -> DraftFeatureSample | None: + if not self.enabled: + return None + path = self.path / f"{key}.pt" + if not path.exists(): + return None + try: + try: + payload = torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + payload = torch.load(path, map_location="cpu") + sample = DraftFeatureSample.from_dict(payload, strict=True) + now = time.time() + os.utime(path, (now, now)) + size = int(path.stat().st_size) + self._entries[path] = (size, now) + return sample + except Exception as exc: # noqa: BLE001 + logger.warning( + "Discard invalid target replay cache entry %s: %s", path, exc + ) + self._forget(path) + try: + path.unlink() + except OSError: + pass + return None + + def put(self, key: str, sample: DraftFeatureSample) -> bool: + if not self.enabled: + return False + path = self.path / f"{key}.pt" + if path.exists(): + return True + with tempfile.NamedTemporaryFile( + prefix=path.name, + suffix=".tmp", + dir=self.path, + delete=False, + ) as tmp_file: + tmp_path = Path(tmp_file.name) + try: + torch.save(sample.to_dict(), tmp_path) + size = int(tmp_path.stat().st_size) + if size > self.max_bytes: + return False + self._evict_until_fits(size) + os.replace(tmp_path, path) + now = time.time() + self._entries[path] = (size, now) + self._total_bytes += size + return True + except Exception as exc: # noqa: BLE001 + logger.warning( + "Failed to write target replay cache entry %s: %s", path, exc + ) + return False + finally: + if tmp_path.exists(): + try: + tmp_path.unlink() + except OSError: + pass + + def _evict_until_fits(self, incoming_size: int) -> None: + entries = sorted(self._entries.items(), key=lambda item: item[1][1]) + for path, _ in entries: + if self._total_bytes + incoming_size <= self.max_bytes: + break + try: + path.unlink() + except OSError: + continue + self._forget(path) + + def _forget(self, path: Path) -> None: + previous = self._entries.pop(path, None) + if previous is not None: + self._total_bytes = max(self._total_bytes - int(previous[0]), 0) + + def metrics(self) -> dict[str, float]: + return { + "replay/cache_size_gb": self._total_bytes / float(1024**3), + "replay/cache_budget_gb_per_rank": self.max_bytes / float(1024**3), + } + + +def feature_from_vllm_payload( + payload: Mapping[str, Any] | Any, + request: DraftReplaySample | Any, + feature_config: FeatureContract, + *, + final_norm: nn.Module | None = None, +) -> DraftFeatureSample: + """Pure vLLM payload conversion shared by replay and standalone Producer.""" + + values = getattr(payload, "payload", payload) + if not isinstance(values, Mapping): + raise TypeError("vLLM hidden-states payload must be a mapping") + token_ids = values.get("token_ids") + hidden = values.get("hidden_states") + if token_ids is None or hidden is None: + raise ValueError( + "vLLM hidden-states payload must contain token_ids and hidden_states" + ) + if not torch.is_tensor(token_ids) or not torch.is_tensor(hidden): + raise ValueError( + "vLLM hidden-states payload must contain token_ids and hidden_states" + ) + feature_positions = request.feature_positions.detach().cpu().long() + if int(feature_positions.numel()) <= 0: + raise ValueError("vLLM feature positions must not be empty") + feature_end_for_request = int(feature_positions[-1].item()) + 1 + expected_prompt_ids = ( + list(request.prompt_token_ids) + if hasattr(request, "prompt_token_ids") + else request.input_ids[:feature_end_for_request].detach().cpu().long().tolist() + ) + if token_ids.detach().cpu().long().tolist() != expected_prompt_ids: + raise HiddenStateAlignmentError( + "vLLM hidden-states token_ids do not match replay input" + ) + if hidden.dim() != 3: + raise ValueError( + "vLLM hidden_states must have shape [seq, layers, hidden], " + f"got {tuple(hidden.shape)}" + ) + + algorithm = str(feature_config.algorithm).strip().upper() + if algorithm not in {"EAGLE3", "DFLASH", "DSPARK"}: + raise ValueError(f"Unsupported vLLM feature algorithm {algorithm!r}") + target_layer_ids = [int(layer_id) for layer_id in feature_config.target_layer_ids] + if not target_layer_ids: + raise ValueError("FeatureContract.target_layer_ids must not be empty") + hidden_layout = str(feature_config.hidden_states_layout) + if hidden_layout not in { + "eagle3_aux_plus_last", + "dflash_aux", + "dflash_aux_plus_last", + }: + raise ValueError(f"Unsupported vLLM hidden_states_layout {hidden_layout!r}") + + hidden_position_offset = max(len(expected_prompt_ids) - int(hidden.size(0)), 0) + include_final = hidden_layout in { + "eagle3_aux_plus_last", + "dflash_aux_plus_last", + } + required_layers = len(target_layer_ids) + (1 if include_final else 0) + if int(hidden.size(1)) < required_layers: + raise HiddenStateAlignmentError( + "vLLM hidden_states layer count is too small: " + f"got {int(hidden.size(1))}, need at least {required_layers}. " + "Start vLLM with target layer ids plus the final layer when the " + "training layout needs last hidden states." + ) + relative_positions = feature_positions - hidden_position_offset + keep_mask = (relative_positions >= 0) & (relative_positions < int(hidden.size(0))) + filtered = not bool(keep_mask.all().item()) + if filtered: + if feature_config.require_full_alignment: + raise HiddenStateAlignmentError( + "vLLM hidden-state rows do not cover the complete feature window: " + f"hidden_rows={int(hidden.size(0))}, " + f"hidden_position_offset={hidden_position_offset}, " + f"dropped={int((~keep_mask).sum().item())}, " + f"feature_min={int(feature_positions.min().item())}, " + f"feature_max={int(feature_positions.max().item())}" + ) + logger.warning( + "Dropping vLLM feature positions outside hidden rows dropped=%s " + "hidden_rows=%s hidden_offset=%s feature_min=%s feature_max=%s", + int((~keep_mask).sum().item()), + int(hidden.size(0)), + hidden_position_offset, + int(feature_positions.min().item()), + int(feature_positions.max().item()), + ) + feature_positions = feature_positions[keep_mask] + relative_positions = relative_positions[keep_mask] + if int(feature_positions.numel()) <= 0: + raise HiddenStateAlignmentError( + "vLLM hidden_states contain no rows for replay feature positions: " + f"hidden_rows={int(hidden.size(0))}, " + f"hidden_position_offset={hidden_position_offset}" + ) + + selected = hidden.index_select(0, relative_positions).to(dtype=feature_config.dtype) + aux_hidden = selected[:, : len(target_layer_ids), :].flatten(1) + if include_final: + if final_norm is None: + raise ValueError("vLLM plus_last features require the target final norm") + # Auxiliary layers stay raw. Only the final supervision block goes + # through the frozen target norm, exactly once, before storage/transport. + with torch.no_grad(): + final_hidden = final_norm(selected[:, required_layers - 1, :]) + output_hidden = torch.cat([aux_hidden, final_hidden], dim=-1) + else: + output_hidden = aux_hidden + selected_input_ids = request.input_ids.index_select(0, feature_positions).long() + selected_loss_mask = request.loss_mask.index_select(0, feature_positions).float() + draft_position_ids = request.draft_position_ids.detach().cpu().long() + if filtered: + draft_position_ids = draft_position_ids[keep_mask] + + source_metadata = getattr(request, "source_metadata", None) + if source_metadata is None: + source_metadata = getattr(request, "metadata", {}) + metadata = dict(source_metadata or {}) + feature_start = int(feature_positions[0].item()) + feature_end = int(feature_positions[-1].item()) + 1 + metadata.update( + { + "source": feature_config.source, + "target_model_path": feature_config.target_model_id, + "target_revision": feature_config.target_model_revision, + "target_config_fingerprint": feature_config.target_config_fingerprint, + "tokenizer_fingerprint": feature_config.tokenizer_fingerprint, + "target_layer_ids": target_layer_ids, + "vllm_hidden_layers": int(hidden.size(1)), + "vllm_hidden_rows": int(hidden.size(0)), + "vllm_hidden_position_offset": hidden_position_offset, + "hidden_states_layout": hidden_layout, + "feature_start": feature_start, + "feature_end": feature_end, + "hidden_position_start": feature_start, + "hidden_position_end": feature_end, + "hidden_positions": feature_positions, + "sequence_length": int(selected_input_ids.numel()), + "full_sequence_length": int(request.input_ids.numel()), + "use_logits": feature_config.use_logits, + } + ) + if include_final: + metadata["last_hidden_state_norm"] = "target_final_norm" + return DraftFeatureSample( + algorithm=algorithm, + input_ids=selected_input_ids, + loss_mask=selected_loss_mask, + hidden_states=output_hidden.cpu().contiguous(), + position_ids=draft_position_ids, + metadata=metadata, + ) + + +class TargetFeatureReplayer: + """Materialize target hidden states only for standalone token replay.""" + + def __init__( + self, + config: Any, + *, + rank: int, + world_size: int, + device: torch.device, + ): + self.config = config + self.rank = int(rank) + self.world_size = int(world_size) + self.device = torch.device(device) + self.draft_config = config.actor_rollout_ref + self.drafter_cfg = self.draft_config.rollout.drafter + self.training_cfg = self.drafter_cfg.training + self.replay_cfg = self.training_cfg.get("target_feature_replay", {}) or {} + self.backend = ( + str(_config_value(self.replay_cfg, "backend", "torch") or "torch") + .strip() + .lower() + ) + if self.backend not in {"torch", "vllm_file"}: + raise ValueError( + f"Unsupported target_feature_replay.backend={self.backend!r}; " + "expected 'torch' or 'vllm_file'" + ) + configured_model_path = _config_value(self.replay_cfg, "model_path", None) + model_path = configured_model_path or self.draft_config.model.path + if not model_path: + raise ValueError( + "Token replay requires target_feature_replay.model_path or " + "actor_rollout_ref.model.path" + ) + self.model_path = os.fspath(model_path) + self.target_revision = str( + _config_value(self.replay_cfg, "target_revision", None) or self.model_path + ) + self.dtype = _parse_dtype(_config_value(self.replay_cfg, "dtype", "bfloat16")) + self.trust_remote_code = bool( + _config_value(self.replay_cfg, "trust_remote_code", False) + ) + self.strict_target_model_path = bool( + _config_value(self.replay_cfg, "strict_target_model_path", False) + ) + self.algorithm = str(self.drafter_cfg.speculative_algorithm).upper() + if self.algorithm not in {"EAGLE3", "DFLASH", "DSPARK"}: + raise ValueError( + f"Token replay does not support drafter algorithm {self.algorithm!r}" + ) + self.use_logits = bool(self.training_cfg.get("use_logits", False)) + self.logits_topk = int(self.training_cfg.get("logits_topk", 128) or 128) + self.logits_chunk_rows = max( + int(_config_value(self.replay_cfg, "logits_chunk_rows", 32) or 32), 1 + ) + self.vllm_endpoints = _normalize_vllm_endpoints(self.replay_cfg) + self.vllm_endpoint = self.vllm_endpoints[0] + self.vllm_model = _config_value(self.replay_cfg, "vllm_model", None) + self.vllm_timeout = float( + _config_value(self.replay_cfg, "request_timeout", 120.0) or 120.0 + ) + self.vllm_max_retries = max( + int(_config_value(self.replay_cfg, "max_retries", 3) or 0), 0 + ) + self.vllm_endpoint_cooldown = max( + float(_config_value(self.replay_cfg, "endpoint_cooldown", 5.0) or 0.0), + 0.0, + ) + self.vllm_on_generate = ( + str(_config_value(self.replay_cfg, "on_generate", "delete") or "delete") + .strip() + .lower() + ) + if self.vllm_on_generate not in {"delete", "keep"}: + raise ValueError( + "target_feature_replay.on_generate must be 'delete' or 'keep'" + ) + self.vllm_require_arange_positions = bool( + _config_value(self.replay_cfg, "require_arange_positions", True) + ) + + from transformers import AutoConfig + + self.target_config = AutoConfig.from_pretrained( + self.model_path, + trust_remote_code=self.trust_remote_code, + ) + self.target_num_hidden_layers = int( + getattr( + getattr(self.target_config, "text_config", self.target_config), + "num_hidden_layers", + ) + ) + model_configs = [ + value + for value in ( + _load_json_config(self.drafter_cfg.get("model_path", None)), + _load_json_config(self.drafter_cfg.get("checkpoint_path", None)), + ) + if value is not None + ] + layer_ids = resolve_oldlogprob_aux_layer_ids( + self.drafter_cfg, + target_num_hidden_layers=self.target_num_hidden_layers, + model_configs=model_configs, + ) + if not layer_ids: + raise RuntimeError( + "Token replay could not resolve target auxiliary layer ids" + ) + self.target_layer_ids = [int(layer_id) for layer_id in layer_ids] + dspark_l1_enabled = ( + self.algorithm == "DSPARK" + and float(self.training_cfg.get("dspark_l1_loss_alpha", 0.9) or 0.0) > 0 + ) + self.hidden_layout = ( + "dflash_aux_plus_last" + if dspark_l1_enabled + else "dflash_aux" + if self.algorithm in {"DFLASH", "DSPARK"} + else "eagle3_aux_plus_last" + ) + self.vllm_final_norm = None + if self.backend == "vllm_file" and self.hidden_layout.endswith("_plus_last"): + self.vllm_final_norm = load_vllm_final_norm( + self.model_path, + dtype=self.dtype, + trust_remote_code=self.trust_remote_code, + target_config=self.target_config, + ) + config_json = json.dumps( + self.target_config.to_dict(), sort_keys=True, default=str + ).encode() + self.target_config_fingerprint = hashlib.sha256(config_json).hexdigest() + + self.cache: BoundedReplayCache | None = None + self._cache_lock = threading.Lock() + cache_cfg = _config_value(self.replay_cfg, "cache", {}) or {} + if bool(_config_value(cache_cfg, "enabled", False)): + cache_path = _config_value(cache_cfg, "path", None) + if not cache_path: + feature_path = os.fspath(self.training_cfg.feature_store.path) + cache_path = f"{feature_path}.hidden_cache" + self.cache = BoundedReplayCache( + cache_path, + max_size_gb=float(_config_value(cache_cfg, "max_size_gb", 0.0) or 0.0), + rank=self.rank, + world_size=self.world_size, + ) + + self.model: nn.Module | None = None + self.layers: list[nn.Module] = [] + self.final_norm: nn.Module | None = None + self.backbone: nn.Module | None = None + self.output_embedding: nn.Module | None = None + self.vllm_client: Any | None = None + self.vllm_resolved_model: str | None = None + self._vllm_endpoint_states = [ + _VllmEndpointState(index=index, url=endpoint) + for index, endpoint in enumerate(self.vllm_endpoints) + ] + self._vllm_clients_initialized = False + self._client_lock = threading.Lock() + self._endpoint_lock = threading.Lock() + self._metrics_lock = threading.Lock() + self.cache_hits = 0 + self.cache_misses = 0 + self.materialized_samples = 0 + self.target_forward_seconds = 0.0 + self.vllm_request_seconds = 0.0 + self.vllm_requests = 0 + self._warned_replay_algorithm_mismatch = False + self._warned_replay_layer_mismatch = False + self._warned_replay_layout_mismatch = False + logger.info( + "[target replay rank=%s] initialized backend=%s algorithm=%s " + "target_layers=%s hidden_layout=%s use_logits=%s endpoints=%s cache=%s", + self.rank, + self.backend, + self.algorithm, + self.target_layer_ids, + self.hidden_layout, + self.use_logits, + self.vllm_endpoints if self.backend.startswith("vllm_") else None, + self.cache is not None, + ) + + def materialize( + self, samples: Iterable[DraftReplaySample | DraftFeatureSample] + ) -> list[DraftFeatureSample]: + materialized: list[DraftFeatureSample] = [] + for sample_index, sample in enumerate(samples): + try: + if isinstance(sample, DraftFeatureSample): + materialized.append(sample) + continue + if not isinstance(sample, DraftReplaySample): + raise TypeError( + "Target feature replay expected DraftReplaySample, " + f"got {type(sample)!r}" + ) + self._validate_target_path(sample) + key = self._cache_key(sample) + with self._cache_lock: + cached = self.cache.get(key) if self.cache is not None else None + if cached is not None: + with self._metrics_lock: + self.cache_hits += 1 + materialized.append(cached) + continue + with self._metrics_lock: + self.cache_misses += 1 + replayed = self._materialize_one(sample) + if self.cache is not None: + with self._cache_lock: + self.cache.put(key, replayed) + materialized.append(replayed) + except Exception: + metadata = getattr(sample, "metadata", {}) or {} + logger.exception( + "[target replay rank=%s] sample materialization failed " + "sample_index=%s algorithm=%s source=%s global_step=%s", + self.rank, + sample_index, + getattr(sample, "algorithm", None), + metadata.get("source"), + metadata.get("global_step"), + ) + raise + with self._metrics_lock: + self.materialized_samples += len(materialized) + return materialized + + def _validate_target_path(self, sample: DraftReplaySample) -> None: + if sample.algorithm.upper() != self.algorithm: + if not self._warned_replay_algorithm_mismatch: + logger.warning( + "[target replay rank=%s] token replay algorithm differs from " + "training algorithm; using the training algorithm for " + "materialized features sample=%s training=%s", + self.rank, + sample.algorithm, + self.algorithm, + ) + self._warned_replay_algorithm_mismatch = True + collected_layer_ids = sample.metadata.get("target_layer_ids") + if collected_layer_ids is not None: + normalized_layer_ids = ( + [int(collected_layer_ids)] + if isinstance(collected_layer_ids, int) + else [int(value) for value in collected_layer_ids] + ) + if normalized_layer_ids != self.target_layer_ids: + if not self._warned_replay_layer_mismatch: + logger.warning( + "[target replay rank=%s] token replay target layers differ " + "from replay target layers; recomputing hidden states with " + "the training configuration collected=%s replay=%s", + self.rank, + normalized_layer_ids, + self.target_layer_ids, + ) + self._warned_replay_layer_mismatch = True + collected_layout = sample.metadata.get("hidden_states_layout") + if collected_layout and str(collected_layout) != self.hidden_layout: + if not self._warned_replay_layout_mismatch: + logger.warning( + "[target replay rank=%s] token replay hidden layout differs " + "from replay layout; recomputing hidden states with the " + "training configuration collected=%s replay=%s", + self.rank, + collected_layout, + self.hidden_layout, + ) + self._warned_replay_layout_mismatch = True + if not self.strict_target_model_path: + return + collected_path = sample.metadata.get("target_model_path") + if collected_path and os.path.normpath( + os.fspath(collected_path) + ) != os.path.normpath(self.model_path): + raise ValueError( + "Token replay target model path mismatch: " + f"collected={collected_path!r} replay={self.model_path!r}" + ) + + def _cache_key(self, sample: DraftReplaySample) -> str: + digest = hashlib.sha256() + contract = { + "target_revision": self.target_revision, + "target_config": self.target_config_fingerprint, + "algorithm": self.algorithm, + "target_layer_ids": self.target_layer_ids, + "hidden_layout": self.hidden_layout, + "dtype": str(self.dtype), + "use_logits": self.use_logits, + "logits_topk": self.logits_topk, + } + if self.backend == "vllm_file" and self.hidden_layout.endswith("_plus_last"): + # Old cache entries contain pre-norm final hidden; never reuse them. + contract["vllm_last_hidden_state_norm"] = "target_final_norm_v1" + digest.update(json.dumps(contract, sort_keys=True).encode()) + for tensor in ( + sample.input_ids, + sample.attention_mask, + sample.position_ids, + sample.feature_positions, + sample.draft_position_ids, + sample.loss_mask, + ): + contiguous = tensor.detach().cpu().contiguous() + digest.update(str(contiguous.dtype).encode()) + digest.update(str(tuple(contiguous.shape)).encode()) + digest.update(contiguous.numpy().tobytes()) + return digest.hexdigest() + + def _ensure_model(self) -> None: + if self.model is not None: + return + from transformers import AutoModelForCausalLM + + logger.warning( + "Loading frozen target model for standalone token replay: path=%s dtype=%s device=%s", + self.model_path, + self.dtype, + self.device, + ) + model = AutoModelForCausalLM.from_pretrained( + self.model_path, + torch_dtype=self.dtype, + trust_remote_code=self.trust_remote_code, + low_cpu_mem_usage=True, + ) + model.eval() + model.requires_grad_(False) + model.to(self.device) + self.layers, self.final_norm = _find_layers_and_final_norm(model) + base_model_prefix = str(getattr(model, "base_model_prefix", "") or "") + backbone = ( + getattr(model, base_model_prefix, None) if base_model_prefix else None + ) + self.backbone = backbone if isinstance(backbone, nn.Module) else model + output_embedding = model.get_output_embeddings() + self.output_embedding = ( + output_embedding if isinstance(output_embedding, nn.Module) else None + ) + self.model = model + + def _materialize_one(self, sample: DraftReplaySample) -> DraftFeatureSample: + if self.backend == "vllm_file": + return self._materialize_one_vllm_file(sample) + return self._materialize_one_torch(sample) + + def _materialize_one_torch(self, sample: DraftReplaySample) -> DraftFeatureSample: + self._ensure_model() + assert self.backbone is not None + assert self.final_norm is not None + + feature_positions = sample.feature_positions.detach().cpu().long() + feature_end = int(feature_positions[-1].item()) + 1 + input_ids = sample.input_ids[:feature_end].to( + self.device, dtype=torch.long, non_blocking=True + ) + attention_mask = sample.attention_mask[:feature_end].to( + self.device, dtype=torch.long, non_blocking=True + ) + position_ids = sample.position_ids[:feature_end].to( + self.device, dtype=torch.long, non_blocking=True + ) + captures: dict[str, torch.Tensor] = {} + handles = [] + + def capture(key: str): + def hook(_module, _inputs, output): + captures[key] = _tensor_from_module_output(output) + + return hook + + aux_keys: list[str] = [] + modules: dict[str, nn.Module] = {} + for layer_id in self.target_layer_ids: + kind, layer_index = _hidden_capture_target( + layer_id, self.target_num_hidden_layers + ) + if kind == "final": + key = "final" + module = self.final_norm + else: + assert layer_index is not None + key = f"layer:{layer_index}" + module = self.layers[layer_index] + aux_keys.append(key) + modules[key] = module + include_final = self.hidden_layout in { + "eagle3_aux_plus_last", + "dflash_aux_plus_last", + } + need_final = include_final or (self.algorithm == "EAGLE3" and self.use_logits) + if need_final: + modules["final"] = self.final_norm + for key, module in modules.items(): + handles.append(module.register_forward_hook(capture(key))) + + started = time.perf_counter() + try: + forward_kwargs = { + "input_ids": input_ids.unsqueeze(0), + "attention_mask": attention_mask.unsqueeze(0), + "position_ids": position_ids.unsqueeze(0), + "use_cache": False, + "return_dict": True, + } + forward_kwargs = _supported_forward_kwargs( + self.backbone.forward, forward_kwargs + ) + with torch.inference_mode(): + self.backbone(**forward_kwargs) + finally: + for handle in handles: + handle.remove() + self.target_forward_seconds += time.perf_counter() - started + + required_keys = list(aux_keys) + if need_final: + required_keys.append("final") + missing = [key for key in required_keys if key not in captures] + if missing: + raise RuntimeError( + f"Target feature replay missed hidden-state captures: {missing}" + ) + + device_positions = feature_positions.to(self.device) + hidden_parts = [ + captures[key].squeeze(0).index_select(0, device_positions) + for key in aux_keys + ] + selected_final = ( + captures["final"].squeeze(0).index_select(0, device_positions) + if need_final + else None + ) + if include_final: + assert selected_final is not None + hidden_parts.append(selected_final) + hidden_states = torch.cat(hidden_parts, dim=-1).to( + device="cpu", dtype=self.dtype + ) + + target_logprobs = None + if self.algorithm == "EAGLE3" and self.use_logits: + assert selected_final is not None + target_logprobs = self._build_sparse_target_logprobs(selected_final[:-1]) + + selected_input_ids = sample.input_ids.index_select(0, feature_positions).long() + selected_loss_mask = sample.loss_mask.index_select(0, feature_positions).float() + metadata = dict(sample.metadata) + feature_start = int(feature_positions[0].item()) + feature_end = int(feature_positions[-1].item()) + 1 + metadata.update( + { + "source": "token_replay", + "target_model_path": self.model_path, + "target_revision": self.target_revision, + "target_config_fingerprint": self.target_config_fingerprint, + "target_layer_ids": list(self.target_layer_ids), + "hidden_states_layout": self.hidden_layout, + "feature_start": feature_start, + "feature_end": feature_end, + "hidden_position_start": feature_start, + "hidden_position_end": feature_end, + "hidden_positions": feature_positions, + "sequence_length": int(selected_input_ids.numel()), + "full_sequence_length": int(sample.input_ids.numel()), + "use_logits": self.use_logits, + } + ) + if target_logprobs is not None: + metadata["target_logprobs_position_start"] = feature_start + 1 + metadata["target_logprobs_position_end"] = feature_end + + return DraftFeatureSample( + algorithm=self.algorithm, + input_ids=selected_input_ids, + loss_mask=selected_loss_mask, + hidden_states=hidden_states, + target_logprobs=target_logprobs, + position_ids=sample.draft_position_ids.long(), + metadata=metadata, + ) + + def _materialize_one_vllm_file( + self, sample: DraftReplaySample + ) -> DraftFeatureSample: + if self.use_logits: + raise NotImplementedError( + "target_feature_replay.backend=vllm_file does not yet support " + "training.use_logits=true; use backend=torch for EAGLE3 logits." + ) + self._validate_vllm_positions(sample) + feature_positions = sample.feature_positions.detach().cpu().long() + feature_end = int(feature_positions[-1].item()) + 1 + prompt_ids = sample.input_ids[:feature_end].detach().cpu().long().tolist() + hidden_payload = self._request_vllm_hidden_states(prompt_ids) + try: + feature = self._feature_from_vllm_payload( + sample, + hidden_payload, + prompt_ids=prompt_ids, + source="token_replay_vllm_file", + ) + finally: + path = hidden_payload.get("_path") + if self.vllm_on_generate == "delete" and path: + try: + Path(os.fspath(path)).unlink(missing_ok=True) + except OSError: + logger.warning("Failed to delete vLLM hidden-states file %s", path) + return feature + + def _validate_vllm_positions(self, sample: DraftReplaySample) -> None: + if not self.vllm_require_arange_positions: + return + feature_positions = sample.feature_positions.detach().cpu().long() + feature_end = int(feature_positions[-1].item()) + 1 + expected = torch.arange(feature_end, dtype=torch.long) + actual = sample.position_ids[:feature_end].detach().cpu().long() + if not torch.equal(actual, expected): + raise ValueError( + "vLLM target replay currently requires " + "position_ids to be contiguous arange positions for the replay prefix" + ) + + def _ensure_vllm_clients(self) -> None: + if self._vllm_clients_initialized: + return + with self._client_lock: + if self._vllm_clients_initialized: + return + started = time.perf_counter() + logger.info( + "[target replay rank=%s] initializing vLLM endpoint pool=%s " + "configured_model=%s", + self.rank, + self.vllm_endpoints, + self.vllm_model, + ) + try: + import openai + except ImportError as exc: + raise RuntimeError( + "vLLM target replay requires the openai package" + ) from exc + resolved_models: set[str] = set() + for state in self._vllm_endpoint_states: + state.client = openai.OpenAI( + base_url=state.url, + api_key="EMPTY", + max_retries=0, + ) + state.model = ( + os.fspath(self.vllm_model) if self.vllm_model else self.model_path + ) + try: + models = state.client.models.list(timeout=self.vllm_timeout) + if not self.vllm_model and models.data: + state.model = str(models.data[0].id) + resolved_models.add(str(state.model)) + logger.info( + "[target replay rank=%s] vLLM endpoint[%s] ready url=%s model=%s", + self.rank, + state.index, + state.url, + state.model, + ) + except Exception as exc: # noqa: BLE001 + state.failures += 1 + state.consecutive_failures += 1 + state.cooldown_until = ( + time.monotonic() + self.vllm_endpoint_cooldown + ) + logger.warning( + "[target replay rank=%s] vLLM endpoint[%s] health check " + "failed; requests may retry it after cooldown url=%s error=%r", + self.rank, + state.index, + state.url, + exc, + ) + self._vllm_clients_initialized = True + self.vllm_client = self._vllm_endpoint_states[0].client + self.vllm_resolved_model = self._vllm_endpoint_states[0].model + if len(resolved_models) > 1: + logger.warning( + "[target replay rank=%s] vLLM endpoints advertise different " + "models=%s; set target_feature_replay.vllm_model explicitly " + "after confirming all servers use identical target weights", + self.rank, + sorted(resolved_models), + ) + logger.info( + "[target replay rank=%s] vLLM endpoint pool initialized " + "endpoints=%s elapsed=%.3fs", + self.rank, + len(self._vllm_endpoint_states), + time.perf_counter() - started, + ) + + def _acquire_vllm_endpoint( + self, excluded: set[int] | None = None + ) -> _VllmEndpointState: + self._ensure_vllm_clients() + excluded = excluded or set() + now = time.monotonic() + with self._endpoint_lock: + candidates = [ + state + for state in self._vllm_endpoint_states + if state.client is not None + and state.model is not None + and state.index not in excluded + and state.cooldown_until <= now + ] + if not candidates: + candidates = [ + state + for state in self._vllm_endpoint_states + if state.client is not None + and state.model is not None + and state.index not in excluded + ] + if not candidates: + candidates = [ + state + for state in self._vllm_endpoint_states + if state.client is not None and state.model is not None + ] + if not candidates: + raise RuntimeError("No configured vLLM endpoint has a usable client") + state = min( + candidates, + key=lambda item: ( + item.inflight, + item.consecutive_failures, + item.cooldown_until, + item.index, + ), + ) + state.inflight += 1 + return state + + def _release_vllm_endpoint( + self, + state: _VllmEndpointState, + *, + elapsed: float, + succeeded: bool, + ) -> None: + with self._endpoint_lock: + state.inflight = max(state.inflight - 1, 0) + state.request_seconds += float(elapsed) + if succeeded: + state.requests += 1 + state.consecutive_failures = 0 + state.cooldown_until = 0.0 + else: + state.failures += 1 + state.consecutive_failures += 1 + state.cooldown_until = time.monotonic() + self.vllm_endpoint_cooldown + + def _request_vllm_response(self, prompt_ids: list[int]) -> Any: + last_error: Exception | None = None + started = time.perf_counter() + attempted_endpoints: set[int] = set() + for attempt in range(self.vllm_max_retries + 1): + state = self._acquire_vllm_endpoint(attempted_endpoints) + if state.client is None: + raise RuntimeError(f"vLLM endpoint {state.url} has no client") + attempt_started = time.perf_counter() + try: + response = state.client.completions.create( + model=state.model, + prompt=prompt_ids, + max_tokens=1, + extra_body={"return_token_ids": True}, + timeout=self.vllm_timeout, + ) + choices = getattr(response, "choices", None) or [] + if choices: + actual = getattr(choices[0], "prompt_token_ids", None) + if actual is not None and list(actual) != prompt_ids: + raise ValueError("vLLM prompt_token_ids mismatch") + with self._metrics_lock: + self.vllm_requests += 1 + self.vllm_request_seconds += time.perf_counter() - started + self._release_vllm_endpoint( + state, + elapsed=time.perf_counter() - attempt_started, + succeeded=True, + ) + return response + except Exception as exc: # noqa: BLE001 + last_error = exc + attempted_endpoints.add(state.index) + self._release_vllm_endpoint( + state, + elapsed=time.perf_counter() - attempt_started, + succeeded=False, + ) + if attempt >= self.vllm_max_retries: + break + time.sleep(float(2**attempt)) + with self._metrics_lock: + self.vllm_request_seconds += time.perf_counter() - started + raise RuntimeError( + "Failed to request vLLM hidden states after " + f"{self.vllm_max_retries + 1} attempts: {last_error}" + ) from last_error + + def _request_vllm_hidden_states(self, prompt_ids: list[int]) -> dict[str, Any]: + last_error: Exception | None = None + started = time.perf_counter() + request_index = self.vllm_requests + 1 + log_request = request_index <= 2 or request_index % 100 == 0 + attempted_endpoints: set[int] = set() + for attempt in range(self.vllm_max_retries + 1): + state = self._acquire_vllm_endpoint(attempted_endpoints) + if state.client is None: + raise RuntimeError(f"vLLM endpoint {state.url} has no client") + try: + attempt_started = time.perf_counter() + if log_request: + logger.info( + "[target replay rank=%s] vLLM request starting request=%s " + "attempt=%s/%s endpoint=%s prompt_tokens=%s", + self.rank, + request_index, + attempt + 1, + self.vllm_max_retries + 1, + state.url, + len(prompt_ids), + ) + response = state.client.completions.create( + model=state.model, + prompt=prompt_ids, + max_tokens=1, + extra_body={"return_token_ids": True}, + timeout=self.vllm_timeout, + ) + path = self._extract_hidden_states_path(response, prompt_ids) + payload = self._load_vllm_hidden_states(path) + payload["_path"] = path + with self._metrics_lock: + self.vllm_requests += 1 + self.vllm_request_seconds += time.perf_counter() - started + self._release_vllm_endpoint( + state, + elapsed=time.perf_counter() - attempt_started, + succeeded=True, + ) + if log_request: + hidden_states = payload.get("hidden_states") + hidden_shape = ( + tuple(hidden_states.shape) + if hidden_states is not None and torch.is_tensor(hidden_states) + else None + ) + logger.info( + "[target replay rank=%s] vLLM request completed request=%s " + "attempt=%s endpoint=%s path=%s hidden_shape=%s elapsed=%.3fs", + self.rank, + request_index, + attempt + 1, + state.url, + path, + hidden_shape, + time.perf_counter() - attempt_started, + ) + return payload + except Exception as exc: # noqa: BLE001 + last_error = exc + attempted_endpoints.add(state.index) + self._release_vllm_endpoint( + state, + elapsed=time.perf_counter() - attempt_started, + succeeded=False, + ) + logger.warning( + "[target replay rank=%s] vLLM request failed request=%s " + "attempt=%s/%s endpoint=%s prompt_tokens=%s elapsed=%.3fs error=%r", + self.rank, + request_index, + attempt + 1, + self.vllm_max_retries + 1, + state.url, + len(prompt_ids), + time.perf_counter() - started, + exc, + ) + if attempt >= self.vllm_max_retries: + break + time.sleep(float(2**attempt)) + with self._metrics_lock: + self.vllm_request_seconds += time.perf_counter() - started + raise RuntimeError( + f"Failed to request vLLM hidden states after " + f"{self.vllm_max_retries + 1} attempts: {last_error}" + ) from last_error + + def _extract_hidden_states_path(self, response: Any, prompt_ids: list[int]) -> str: + choices = getattr(response, "choices", None) or [] + if choices: + prompt_token_ids = getattr(choices[0], "prompt_token_ids", None) + if prompt_token_ids is not None and list(prompt_token_ids) != prompt_ids: + raise ValueError( + "vLLM prompt_token_ids mismatch while extracting hidden states" + ) + kv_transfer_params = getattr(response, "kv_transfer_params", None) + if kv_transfer_params is None: + raise ValueError("vLLM response missing kv_transfer_params") + path = kv_transfer_params.get("hidden_states_path") + if not path: + raise ValueError("vLLM response missing hidden_states_path") + return os.fspath(path) + + def _load_vllm_hidden_states(self, path: str) -> dict[str, Any]: + try: + from safetensors.torch import load_file + except ImportError as exc: + raise RuntimeError("vLLM hidden-state replay requires safetensors") from exc + file_path = Path(path) + lock_path = Path(f"{path}.lock") + if lock_path.exists(): + _wait_for_lock(lock_path) + if not file_path.exists(): + raise FileNotFoundError(f"vLLM hidden-states file not found: {path}") + return dict(load_file(str(file_path), device="cpu")) + + def _feature_from_vllm_payload( + self, + sample: DraftReplaySample, + payload: dict[str, Any], + *, + prompt_ids: list[int], + source: str, + ) -> DraftFeatureSample: + expected_prompt_ids = ( + sample.input_ids[: int(sample.feature_positions[-1].item()) + 1] + .detach() + .cpu() + .long() + .tolist() + ) + if expected_prompt_ids != prompt_ids: + raise ValueError("prompt_ids do not match replay feature window") + return feature_from_vllm_payload( + payload, + sample, + FeatureContract( + algorithm=getattr(self, "algorithm", sample.algorithm), + target_layer_ids=list(self.target_layer_ids), + hidden_states_layout=self.hidden_layout, + dtype=self.dtype, + target_model_id=self.model_path, + target_model_revision=self.target_revision, + tokenizer_fingerprint=str( + getattr(self, "tokenizer_fingerprint", "replay-unspecified") + ), + use_logits=self.use_logits, + target_config_fingerprint=self.target_config_fingerprint, + source=source, + ), + final_norm=self.vllm_final_norm, + ) + + def _build_sparse_target_logprobs( + self, final_hidden_states: torch.Tensor + ) -> torch.Tensor: + if self.output_embedding is None: + raise RuntimeError( + "EAGLE3 token replay with use_logits=true requires target output embeddings" + ) + rows: list[torch.Tensor] = [] + topk = max(self.logits_topk, 1) + with torch.inference_mode(): + for start in range( + 0, int(final_hidden_states.size(0)), self.logits_chunk_rows + ): + hidden = final_hidden_states[start : start + self.logits_chunk_rows] + logits = self.output_embedding(hidden).float() + local_topk = min(topk, int(logits.size(-1))) + values, ids = logits.topk(local_topk, dim=-1) + values = values - torch.logsumexp(logits, dim=-1, keepdim=True) + rows.append( + torch.stack((values, ids.to(dtype=values.dtype)), dim=-1).cpu() + ) + if not rows: + return torch.empty(0, topk, 2, dtype=torch.float32) + return torch.cat(rows, dim=0).contiguous() + + def metrics(self) -> dict[str, float]: + metrics = { + "replay/cache_hits_total": float(self.cache_hits), + "replay/cache_misses_total": float(self.cache_misses), + "replay/materialized_samples_total": float(self.materialized_samples), + "replay/target_forward_time_total": float(self.target_forward_seconds), + } + if self.backend.startswith("vllm_"): + metrics["replay/vllm_requests_total"] = float(self.vllm_requests) + metrics["replay/vllm_request_time_total"] = float(self.vllm_request_seconds) + with self._endpoint_lock: + metrics["replay/vllm_endpoints_total"] = float( + len(self._vllm_endpoint_states) + ) + for state in self._vllm_endpoint_states: + prefix = f"replay/vllm_endpoint_{state.index}" + metrics[f"{prefix}_inflight"] = float(state.inflight) + metrics[f"{prefix}_requests_total"] = float(state.requests) + metrics[f"{prefix}_failures_total"] = float(state.failures) + metrics[f"{prefix}_request_time_total"] = float( + state.request_seconds + ) + total = self.cache_hits + self.cache_misses + if total > 0: + metrics["replay/cache_hit_ratio"] = self.cache_hits / float(total) + if self.cache is not None: + metrics.update(self.cache.metrics()) + return metrics + + def close(self) -> None: + self.vllm_final_norm = None + for state in self._vllm_endpoint_states: + client = state.client + close = getattr(client, "close", None) + if callable(close): + try: + close() + except Exception: # noqa: BLE001 + logger.warning( + "Failed to close vLLM endpoint client %s", + state.url, + exc_info=True, + ) + state.client = None + if self.model is None: + return + try: + self.model.to("cpu") + except Exception: # noqa: BLE001 + pass + self.model = None + self.layers = [] + self.final_norm = None + self.backbone = None + self.output_embedding = None + + +def _supported_forward_kwargs(forward: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(forward) + except (TypeError, ValueError): + return kwargs + if any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in signature.parameters.values() + ): + return kwargs + return {key: value for key, value in kwargs.items() if key in signature.parameters} diff --git a/verl_speco/trainer/tq_feature_store.py b/verl_speco/trainer/tq_feature_store.py new file mode 100644 index 00000000..7e25af9a --- /dev/null +++ b/verl_speco/trainer/tq_feature_store.py @@ -0,0 +1,223 @@ +# Copyright 2026 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. +"""Streaming TransferQueue feature source for standalone drafter training.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Sequence + +from verl_speco.integration.transferqueue_bridge import ( + clear_samples, + close_transfer_queue_client, + configure_transfer_queue, + connect_ray_cluster, + connect_transfer_queue_client, + get_samples, + list_samples, +) +from verl_speco.trainer.feature_store import DraftFeatureSample +from verl_speco.transport.drafter_sample_protocol import ( + ExpectedFeatureConfig, + decode_sample, + parse_ready_tag, +) + + +@dataclass(frozen=True) +class ReadyEntry: + """One discoverable sample record; payload tensors are not loaded yet.""" + + key: str + tag: dict[str, Any] + + +@dataclass(frozen=True) +class EosMetadata: + """End-of-stream control record published by the Producer.""" + + key: str + run_id: str + schema_version: int + total_samples: int + + +class TQFeatureStore: + """Thin Consumer adapter over the shared TransferQueue bridge. + + This intentionally does not implement the static ``iter_keys/read`` feature + store protocol. TQ keys are added by the Producer and removed after a + successful optimizer step, so discovery must happen for every global batch. + """ + + def __init__(self, config: Mapping[str, Any]): + self.config = _plain_dict(config) + self.run_id = str(self.config.get("run_id") or "").strip() + if not self.run_id: + raise ValueError("transfer_queue.run_id is required for a TQ Consumer") + self.schema_version = int(self.config.get("schema_version", 1)) + ray_cfg = _plain_dict(self.config.get("ray") or {}) + self.ray_address = str(ray_cfg.get("address") or "").strip() + self.ray_namespace = str(ray_cfg.get("namespace") or "").strip() or None + if not self.ray_address: + raise ValueError("transfer_queue.ray.address is required for a TQ Consumer") + self._connected = False + # First version deliberately checks only run/protocol. Tensor + # presence, lengths, shape and dtype self-consistency remain enforced by + # decode_sample; model/tokenizer/layer identity checks stay disabled. + self.expected_config = ExpectedFeatureConfig( + run_id=self.run_id, + schema_version=self.schema_version, + ) + + @classmethod + def from_config(cls, config: Any) -> "TQFeatureStore": + return cls(_plain_dict(config)) + + def connect(self) -> None: + if self._connected: + return + if not configure_transfer_queue(self.config): + raise RuntimeError( + "TQ Consumer requires transfer_queue.enable=true and TransferQueue==0.1.10" + ) + connect_ray_cluster(self.ray_address, self.ray_namespace) + connect_transfer_queue_client() + self._connected = True + + def list_ready(self, run_id: str | None = None) -> list[ReadyEntry]: + self._require_connected() + expected_run_id = str(run_id or self.run_id) + ready: list[ReadyEntry] = [] + for key, raw_tag in list_samples().items(): + tag = dict(raw_tag) + meta = parse_ready_tag( + tag, + run_id=expected_run_id, + schema_version=self.schema_version, + ) + if meta is None: + continue + ready.append(ReadyEntry(key=str(key), tag=tag)) + ready.sort(key=lambda entry: (int(entry.tag["sequence_no"]), entry.key)) + return ready + + def owner_ready(self) -> bool: + """Whether the standalone Owner published this run's readiness marker.""" + + self._require_connected() + key = f"control:v{self.schema_version}:{self.run_id}:owner-ready" + tag = list_samples().get(key) + if not isinstance(tag, Mapping): + return False + return ( + tag.get("record_type") == "control" + and tag.get("status") == "owner_ready" + and str(tag.get("run_id") or "") == self.run_id + and int(tag.get("schema_version", -1)) == self.schema_version + ) + + def get_many(self, entries: Sequence[ReadyEntry]) -> list[DraftFeatureSample]: + self._require_connected() + if not entries: + return [] + records = get_samples([entry.key for entry in entries]) + if len(records) != len(entries): + raise RuntimeError( + f"TQ returned {len(records)} records for {len(entries)} requested entries" + ) + samples: list[DraftFeatureSample] = [] + for entry, (key, fields) in zip(entries, records, strict=True): + if key != entry.key: + raise RuntimeError( + f"TQ batch result order mismatch: got key={key!r}, expected={entry.key!r}" + ) + samples.append( + decode_sample( + key=key, + tag=entry.tag, + fields=fields, + expected_config=self.expected_config, + ) + ) + return samples + + def clear_many(self, keys: Sequence[str]) -> None: + self._require_connected() + clear_samples([str(key) for key in keys]) + + def read_eos(self, run_id: str | None = None) -> EosMetadata | None: + self._require_connected() + expected_run_id = str(run_id or self.run_id) + for key, raw_tag in list_samples().items(): + tag = dict(raw_tag) + if tag.get("record_type") != "control" or tag.get("status") != "eos": + continue + if str(tag.get("run_id") or "") != expected_run_id: + continue + if int(tag.get("schema_version", -1)) != self.schema_version: + continue + return EosMetadata( + key=str(key), + run_id=expected_run_id, + schema_version=self.schema_version, + total_samples=int(tag.get("total_samples", 0)), + ) + return None + + def close_local(self) -> None: + if not self._connected: + return + close_transfer_queue_client() + self._connected = False + + def close(self) -> None: + """Compatibility with the standalone loop's existing cleanup path.""" + + self.close_local() + + def _require_connected(self) -> None: + if not self._connected: + raise RuntimeError( + "TQFeatureStore.connect() must be called before data access" + ) + + +def _plain_dict(value: Any) -> dict[str, Any]: + if value is None: + return {} + try: + from omegaconf import DictConfig, OmegaConf + + if isinstance(value, DictConfig): + converted = OmegaConf.to_container(value, resolve=True) + if not isinstance(converted, dict): + raise TypeError("Expected a mapping configuration") + return dict(converted) + except ImportError: # pragma: no cover - the project depends on OmegaConf + pass + if isinstance(value, Mapping): + return {str(key): _plain_value(item) for key, item in value.items()} + raise TypeError(f"Expected a mapping configuration, got {type(value)!r}") + + +def _plain_value(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain_value(item) for item in value] + return value + + +__all__ = ["EosMetadata", "ReadyEntry", "TQFeatureStore"] diff --git a/verl_speco/trainer/tq_sample_source.py b/verl_speco/trainer/tq_sample_source.py new file mode 100644 index 00000000..07ab9d06 --- /dev/null +++ b/verl_speco/trainer/tq_sample_source.py @@ -0,0 +1,219 @@ +# Copyright 2026 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. +"""Distributed streaming sample source for a TQ-backed standalone Consumer.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import Any, Iterator, Sequence + +import torch.distributed as dist + +from verl_speco.trainer.feature_store import DraftFeatureSample +from verl_speco.trainer.tq_feature_store import ReadyEntry, TQFeatureStore + + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class TQLocalBatch: + """The payload owned by one rank plus rank-0's global cleanup keys.""" + + local_keys: list[str] + local_samples: list[DraftFeatureSample] + global_keys: list[str] | None + global_sequence_nos: list[int] | None = None + + +def build_assignments( + entries: Sequence[ReadyEntry], *, batch_size: int, world_size: int +) -> list[list[ReadyEntry]]: + """Split one complete global batch into disjoint contiguous rank batches.""" + + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if world_size <= 0: + raise ValueError(f"world_size must be positive, got {world_size}") + expected = batch_size * world_size + if len(entries) != expected: + raise ValueError( + f"Expected exactly {expected} ready entries for one global batch, got {len(entries)}" + ) + return [ + list(entries[rank * batch_size : (rank + 1) * batch_size]) + for rank in range(world_size) + ] + + +class TQFeatureDataLoader: + """Poll TQ on rank 0, distribute keys, and fetch payloads on owner ranks.""" + + def __init__( + self, + store: TQFeatureStore, + *, + batch_size: int, + rank: int, + world_size: int, + poll_interval_seconds: float = 0.5, + drop_last: bool = True, + ): + self.store = store + self.batch_size = int(batch_size) + self.rank = int(rank) + self.world_size = int(world_size) + self.poll_interval_seconds = max(float(poll_interval_seconds), 0.01) + self.drop_last = bool(drop_last) + if self.batch_size <= 0: + raise ValueError("TQ Consumer batch_size_per_gpu must be positive") + if self.world_size <= 0 or not (0 <= self.rank < self.world_size): + raise ValueError( + "Invalid TQ Consumer rank/world_size: " + f"rank={self.rank}, world_size={self.world_size}" + ) + if not self.drop_last: + raise ValueError( + "TQ Consumer first version requires transfer_queue.drop_last=true" + ) + if dist.is_initialized() and dist.get_world_size() != self.world_size: + raise ValueError( + "TQFeatureDataLoader world_size does not match torch.distributed world size" + ) + + def __iter__(self) -> Iterator[TQLocalBatch]: + self.store.connect() + global_batch_size = self.batch_size * self.world_size + owner_ready = False + while True: + command: dict[str, Any] | None = None + if self.rank == 0: + try: + if not owner_ready: + owner_ready = self.store.owner_ready() + if not owner_ready: + time.sleep(self.poll_interval_seconds) + continue + ready = self.store.list_ready() + if len(ready) >= global_batch_size: + selected = ready[:global_batch_size] + assignments = build_assignments( + selected, + batch_size=self.batch_size, + world_size=self.world_size, + ) + command = { + "kind": "batch", + "global_keys": [entry.key for entry in selected], + "global_sequence_nos": [ + int(entry.tag["sequence_no"]) for entry in selected + ], + "assignments": [ + [_entry_to_wire(entry) for entry in rank_entries] + for rank_entries in assignments + ], + } + else: + eos = self.store.read_eos() + if eos is not None: + tail_keys = [entry.key for entry in ready] + if tail_keys: + logger.info( + "Dropping %s TQ tail samples after EOS because one " + "global batch requires %s", + len(tail_keys), + global_batch_size, + ) + self.store.clear_many(tail_keys) + command = {"kind": "stop"} + else: + time.sleep(self.poll_interval_seconds) + continue + except BaseException as exc: # noqa: BLE001 + command = { + "kind": "error", + "message": f"rank 0 failed while discovering TQ samples: {exc}", + } + + command = self._broadcast_command(command) + if command.get("kind") == "error": + raise RuntimeError(str(command.get("message") or "TQ discovery failed")) + if command.get("kind") == "stop": + return + if command.get("kind") != "batch": + raise RuntimeError(f"Unsupported TQ loader command: {command!r}") + wire_assignments = command.get("assignments") + if ( + not isinstance(wire_assignments, list) + or len(wire_assignments) != self.world_size + ): + raise RuntimeError("TQ loader received malformed rank assignments") + local_entries = [ + _entry_from_wire(item) for item in wire_assignments[self.rank] + ] + samples = self.store.get_many(local_entries) + global_keys = ( + [str(key) for key in command.get("global_keys", [])] + if self.rank == 0 + else None + ) + global_sequence_nos = ( + [int(value) for value in command.get("global_sequence_nos", [])] + if self.rank == 0 + else None + ) + yield TQLocalBatch( + local_keys=[entry.key for entry in local_entries], + local_samples=samples, + global_keys=global_keys, + global_sequence_nos=global_sequence_nos, + ) + + def clear_completed_batch(self, global_keys: Sequence[str] | None) -> None: + if self.rank != 0: + return + if not global_keys: + raise ValueError( + "rank 0 requires global_keys to clear a completed TQ batch" + ) + self.store.clear_many(global_keys) + + def _broadcast_command(self, command: dict[str, Any] | None) -> dict[str, Any]: + if not dist.is_initialized() or self.world_size == 1: + if command is None: + raise RuntimeError("rank 0 did not create a TQ loader command") + return command + payload: list[Any] = [command if self.rank == 0 else None] + dist.broadcast_object_list(payload, src=0) + received = payload[0] + if not isinstance(received, dict): + raise RuntimeError("TQ loader broadcast did not contain a command mapping") + return received + + +def _entry_to_wire(entry: ReadyEntry) -> dict[str, Any]: + return {"key": entry.key, "tag": dict(entry.tag)} + + +def _entry_from_wire(value: Any) -> ReadyEntry: + if not isinstance(value, dict) or "key" not in value or "tag" not in value: + raise TypeError(f"Invalid serialized ReadyEntry: {value!r}") + if not isinstance(value["tag"], dict): + raise TypeError("Serialized ReadyEntry.tag must be a mapping") + return ReadyEntry(key=str(value["key"]), tag=dict(value["tag"])) + + +__all__ = ["TQFeatureDataLoader", "TQLocalBatch", "build_assignments"] diff --git a/verl_speco/transport/__init__.py b/verl_speco/transport/__init__.py new file mode 100644 index 00000000..6f8b6dd2 --- /dev/null +++ b/verl_speco/transport/__init__.py @@ -0,0 +1,42 @@ +# Copyright 2026 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. +"""Transport protocols shared by standalone SPECO producers and consumers.""" + +from verl_speco.transport.drafter_sample_protocol import ( + DRAFTER_TQ_PARTITION, + PROTOCOL_SCHEMA_VERSION, + ExpectedFeatureConfig, + SampleMetadata, + decode_sample, + encode_sample, + is_ready_sample_tag, + make_eos_record, + make_ready_tag, + make_sample_key, + parse_ready_tag, +) + +__all__ = [ + "DRAFTER_TQ_PARTITION", + "PROTOCOL_SCHEMA_VERSION", + "ExpectedFeatureConfig", + "SampleMetadata", + "decode_sample", + "encode_sample", + "is_ready_sample_tag", + "make_eos_record", + "make_ready_tag", + "make_sample_key", + "parse_ready_tag", +] diff --git a/verl_speco/transport/drafter_sample_protocol.py b/verl_speco/transport/drafter_sample_protocol.py new file mode 100644 index 00000000..2bb5d355 --- /dev/null +++ b/verl_speco/transport/drafter_sample_protocol.py @@ -0,0 +1,386 @@ +# Copyright 2026 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. +"""Algorithm-neutral TransferQueue codec for ``DraftFeatureSample``.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from typing import Any, Mapping + +import torch + +from verl_speco.trainer.feature_store import DraftFeatureSample + + +PROTOCOL_SCHEMA_VERSION = 2 +DRAFTER_TQ_PARTITION = "speco_drafter_features" +_MANIFEST_FIELD = "sample__manifest_json" +_REQUIRED_SAMPLE_FIELDS = ("input_ids", "loss_mask", "hidden_states") +_OPTIONAL_TENSOR_FIELDS = ( + "last_hidden_states", + "target", + "target_logprobs", + "position_ids", +) + + +@dataclass(frozen=True) +class SampleMetadata: + """Small control-plane envelope; training metadata belongs to the sample.""" + + schema_version: int + run_id: str + sample_id: str + sequence_no: int + + def validate(self) -> None: + if self.schema_version != PROTOCOL_SCHEMA_VERSION: + raise ValueError( + f"Unsupported drafter protocol schema_version={self.schema_version}; " + f"expected {PROTOCOL_SCHEMA_VERSION}" + ) + if not self.run_id: + raise ValueError("SampleMetadata.run_id must not be empty") + if not self.sample_id: + raise ValueError("SampleMetadata.sample_id must not be empty") + if self.sequence_no < 0: + raise ValueError("SampleMetadata.sequence_no must be non-negative") + + def to_dict(self) -> dict[str, Any]: + self.validate() + return asdict(self) + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "SampleMetadata": + try: + meta = cls( + schema_version=int(payload["schema_version"]), + run_id=str(payload["run_id"]), + sample_id=str(payload["sample_id"]), + sequence_no=int(payload["sequence_no"]), + ) + except KeyError as exc: + raise ValueError( + f"ready tag missing required field {exc.args[0]!r}" + ) from exc + meta.validate() + return meta + + +@dataclass(frozen=True) +class ExpectedFeatureConfig: + """Consumer-side transport contract.""" + + run_id: str + schema_version: int = PROTOCOL_SCHEMA_VERSION + + +def make_sample_key(meta: SampleMetadata) -> str: + meta.validate() + return ( + f"drafter:v{meta.schema_version}:{meta.run_id}:" + f"{meta.sequence_no:012d}:{meta.sample_id}" + ) + + +def make_ready_tag(meta: SampleMetadata) -> dict[str, Any]: + meta.validate() + return {"record_type": "sample", "status": "ready", **meta.to_dict()} + + +def parse_ready_tag( + tag: Mapping[str, Any], + *, + run_id: str | None = None, + schema_version: int = PROTOCOL_SCHEMA_VERSION, +) -> SampleMetadata | None: + """Return one valid ready envelope, or ``None`` when it is not consumable.""" + + if tag.get("record_type") != "sample" or tag.get("status") != "ready": + return None + try: + meta = SampleMetadata.from_dict(tag) + except (TypeError, ValueError): + return None + if meta.schema_version != int(schema_version): + return None + if run_id is not None and meta.run_id != str(run_id): + return None + return meta + + +def is_ready_sample_tag( + tag: Mapping[str, Any], + *, + run_id: str, + schema_version: int = PROTOCOL_SCHEMA_VERSION, +) -> bool: + return ( + parse_ready_tag(tag, run_id=run_id, schema_version=schema_version) is not None + ) + + +def encode_sample( + sample: DraftFeatureSample | Mapping[str, Any], meta: SampleMetadata +) -> dict[str, torch.Tensor]: + """Losslessly encode one normalized feature sample into TQ tensor fields.""" + + meta.validate() + normalized = ( + sample + if isinstance(sample, DraftFeatureSample) + else DraftFeatureSample.from_dict(dict(sample), strict=True) + ) + normalized.validate(strict=True) + payload = normalized.to_dict() + fields: dict[str, torch.Tensor] = {} + manifest: dict[str, Any] = { + "draft_feature_schema_version": int(normalized.schema_version), + "algorithm": str(normalized.algorithm), + "present_fields": [], + } + + for name in ("input_ids", "loss_mask", *_OPTIONAL_TENSOR_FIELDS): + value = payload.get(name) + if value is None: + continue + if not torch.is_tensor(value): + raise TypeError(f"DraftFeatureSample.{name} must be a torch.Tensor") + fields[f"sample__{name}"] = _cpu_contiguous(value) + manifest["present_fields"].append(name) + + hidden = payload["hidden_states"] + if torch.is_tensor(hidden): + fields["sample__hidden_states"] = _cpu_contiguous(hidden) + manifest["hidden_states_kind"] = "tensor" + elif isinstance(hidden, (list, tuple)): + hidden_fields: list[str] = [] + for index, value in enumerate(hidden): + if not torch.is_tensor(value): + raise TypeError( + f"DraftFeatureSample.hidden_states[{index}] must be a torch.Tensor" + ) + field_name = f"sample__hidden_states__{index:06d}" + fields[field_name] = _cpu_contiguous(value) + hidden_fields.append(field_name) + if not hidden_fields: + raise ValueError("DraftFeatureSample.hidden_states list must not be empty") + manifest["hidden_states_kind"] = "list" + manifest["hidden_states_fields"] = hidden_fields + else: + raise TypeError( + "DraftFeatureSample.hidden_states must be a tensor or tensor list" + ) + manifest["present_fields"].append("hidden_states") + manifest["metadata"] = _encode_metadata_tree( + payload.get("metadata", {}), fields, path="metadata" + ) + fields[_MANIFEST_FIELD] = _json_to_tensor(manifest) + return fields + + +def decode_sample( + key: str, + tag: Mapping[str, Any], + fields: Mapping[str, Any], + expected_config: ExpectedFeatureConfig | Mapping[str, Any], +) -> DraftFeatureSample: + """Validate one queue record and restore the complete training sample.""" + + expected = ( + expected_config + if isinstance(expected_config, ExpectedFeatureConfig) + else ExpectedFeatureConfig(**dict(expected_config)) + ) + meta = parse_ready_tag( + tag, run_id=expected.run_id, schema_version=expected.schema_version + ) + if meta is None: + raise ValueError(f"TQ sample {key!r} has an invalid or unexpected ready tag") + expected_key = make_sample_key(meta) + if key != expected_key: + raise ValueError( + f"TQ sample key mismatch: got {key!r}, expected {expected_key!r}" + ) + if _MANIFEST_FIELD not in fields: + raise ValueError(f"TQ sample {key!r} missing field {_MANIFEST_FIELD!r}") + manifest = _tensor_to_json(fields[_MANIFEST_FIELD], name=_MANIFEST_FIELD) + present = manifest.get("present_fields") + if not isinstance(present, list): + raise ValueError("sample manifest present_fields must be a list") + missing = [name for name in _REQUIRED_SAMPLE_FIELDS if name not in present] + if missing: + raise ValueError(f"TQ sample {key!r} missing required sample fields: {missing}") + + try: + sample_schema_version = int(manifest["draft_feature_schema_version"]) + algorithm = str(manifest["algorithm"]) + except KeyError as exc: + raise ValueError(f"sample manifest missing field {exc.args[0]!r}") from exc + payload: dict[str, Any] = { + "schema_version": sample_schema_version, + "algorithm": algorithm, + "metadata": _decode_metadata_tree( + manifest.get("metadata", {}), fields, path="metadata" + ), + } + for name in ("input_ids", "loss_mask", *_OPTIONAL_TENSOR_FIELDS): + if name in present: + payload[name] = _require_tensor(fields, f"sample__{name}") + payload["hidden_states"] = _decode_hidden_states(fields, manifest) + return DraftFeatureSample.from_dict(payload, strict=True) + + +def make_eos_record( + run_id: str, total_samples: int +) -> tuple[str, dict[str, torch.Tensor], dict[str, Any]]: + if not run_id: + raise ValueError("run_id must not be empty") + if total_samples < 0: + raise ValueError("total_samples must be non-negative") + key = f"control:v{PROTOCOL_SCHEMA_VERSION}:{run_id}:eos" + fields = {"marker": torch.tensor([1], dtype=torch.uint8)} + tag = { + "record_type": "control", + "status": "eos", + "schema_version": PROTOCOL_SCHEMA_VERSION, + "run_id": run_id, + "total_samples": int(total_samples), + } + return key, fields, tag + + +def _encode_metadata_tree( + value: Any, fields: dict[str, torch.Tensor], *, path: str +) -> Any: + if torch.is_tensor(value): + field_name = f"sample__metadata_tensor__{len(fields):06d}" + fields[field_name] = _cpu_contiguous(value) + return {"__tq_tensor_ref__": field_name} + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Mapping): + encoded: dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError( + f"DraftFeatureSample metadata key at {path} must be str" + ) + encoded[key] = _encode_metadata_tree(item, fields, path=f"{path}.{key}") + return {"__tq_mapping__": encoded} + if isinstance(value, (list, tuple)): + items = [ + _encode_metadata_tree(item, fields, path=f"{path}[{index}]") + for index, item in enumerate(value) + ] + return { + "__tq_sequence__": "tuple" if isinstance(value, tuple) else "list", + "items": items, + } + raise TypeError( + f"Unsupported DraftFeatureSample metadata value at {path}: {type(value).__name__}" + ) + + +def _decode_metadata_tree(value: Any, fields: Mapping[str, Any], *, path: str) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if not isinstance(value, Mapping): + raise ValueError(f"Invalid metadata manifest node at {path}") + if "__tq_tensor_ref__" in value: + return _require_tensor(fields, str(value["__tq_tensor_ref__"])) + if "__tq_mapping__" in value: + mapping = value["__tq_mapping__"] + if not isinstance(mapping, Mapping): + raise ValueError(f"Invalid metadata mapping node at {path}") + return { + str(key): _decode_metadata_tree(item, fields, path=f"{path}.{key}") + for key, item in mapping.items() + } + if "__tq_sequence__" in value: + items = value.get("items") + if not isinstance(items, list): + raise ValueError(f"Invalid metadata sequence node at {path}") + decoded = [ + _decode_metadata_tree(item, fields, path=f"{path}[{index}]") + for index, item in enumerate(items) + ] + if value["__tq_sequence__"] == "tuple": + return tuple(decoded) + if value["__tq_sequence__"] == "list": + return decoded + raise ValueError(f"Unknown metadata manifest node at {path}") + + +def _decode_hidden_states( + fields: Mapping[str, Any], manifest: Mapping[str, Any] +) -> torch.Tensor | list[torch.Tensor]: + kind = manifest.get("hidden_states_kind") + if kind == "tensor": + return _require_tensor(fields, "sample__hidden_states") + if kind == "list": + names = manifest.get("hidden_states_fields") + if not isinstance(names, list) or not names: + raise ValueError("sample manifest hidden_states_fields must be non-empty") + return [_require_tensor(fields, str(name)) for name in names] + raise ValueError(f"Unsupported hidden_states_kind={kind!r}") + + +def _json_to_tensor(payload: Mapping[str, Any]) -> torch.Tensor: + raw = json.dumps(dict(payload), sort_keys=True, separators=(",", ":")).encode( + "utf-8" + ) + return torch.tensor(list(raw), dtype=torch.uint8) + + +def _tensor_to_json(value: Any, *, name: str) -> dict[str, Any]: + if not torch.is_tensor(value): + raise TypeError(f"{name} must be a torch.Tensor") + tensor = value.detach().cpu().to(torch.uint8).reshape(-1) + try: + decoded = json.loads(bytes(tensor.tolist()).decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError(f"{name} is not valid UTF-8 JSON") from exc + if not isinstance(decoded, dict): + raise ValueError(f"{name} must decode to a JSON object") + return decoded + + +def _require_tensor(fields: Mapping[str, Any], name: str) -> torch.Tensor: + value = fields.get(name) + if value is None or not torch.is_tensor(value): + raise TypeError(f"TQ field {name!r} must be a torch.Tensor") + return value.detach().cpu().contiguous() + + +def _cpu_contiguous(value: torch.Tensor) -> torch.Tensor: + if not torch.is_tensor(value): + raise TypeError(f"Expected torch.Tensor, got {type(value)!r}") + return value.detach().cpu().contiguous() + + +__all__ = [ + "DRAFTER_TQ_PARTITION", + "PROTOCOL_SCHEMA_VERSION", + "ExpectedFeatureConfig", + "SampleMetadata", + "decode_sample", + "encode_sample", + "is_ready_sample_tag", + "make_eos_record", + "make_ready_tag", + "make_sample_key", + "parse_ready_tag", +] diff --git a/verl_speco/vllm_hidden_states_generate.py b/verl_speco/vllm_hidden_states_generate.py new file mode 100644 index 00000000..16ab4375 --- /dev/null +++ b/verl_speco/vllm_hidden_states_generate.py @@ -0,0 +1,143 @@ +# Copyright 2026 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. +"""Generate vLLM safetensors draft features from token replay samples.""" + +from __future__ import annotations + +import logging +import os +from copy import deepcopy +from typing import Any + +import hydra +import torch +from omegaconf import OmegaConf, open_dict + +from verl_speco.trainer.draft_dataset import ( + DraftFeatureDataLoader, + DraftFeatureDataLoaderConfig, +) +from verl_speco.trainer.feature_store import build_feature_store_from_config +from verl_speco.trainer.target_feature_replay import TargetFeatureReplayer + +logger = logging.getLogger(__name__) + + +def _plain_config(value: Any) -> dict[str, Any]: + return dict(OmegaConf.to_container(value, resolve=True) or {}) + + +def generate_vllm_safetensors_features(config) -> dict[str, Any]: + """Materialize token replay samples through vLLM and save safetensors.""" + + draft_config = config.actor_rollout_ref + training_cfg = draft_config.rollout.drafter.training + replay_cfg = training_cfg.get("target_feature_replay", {}) or {} + generation_cfg = replay_cfg.get("offline_generation", {}) or {} + feature_store_cfg = training_cfg.feature_store + + input_path = generation_cfg.get("input_path", None) + output_path = generation_cfg.get("output_path", None) or feature_store_cfg.get( + "path", None + ) + if not input_path: + raise ValueError( + "target_feature_replay.offline_generation.input_path is required" + ) + if not output_path: + raise ValueError( + "target_feature_replay.offline_generation.output_path or " + "training.feature_store.path is required" + ) + + input_type = str(generation_cfg.get("input_type", "token_replay") or "token_replay") + input_cfg = _plain_config(feature_store_cfg) + input_cfg.update({"type": input_type, "path": os.fspath(input_path)}) + + output_cfg = _plain_config(feature_store_cfg) + output_cfg.update({"type": "vllm_safetensors", "path": os.fspath(output_path)}) + + max_samples = int(generation_cfg.get("max_samples", 0) or 0) + batch_size = max(int(generation_cfg.get("batch_size", 1) or 1), 1) + shuffle = bool(generation_cfg.get("shuffle", False)) + seed = int(training_cfg.get("seed", 0) or 0) + + replay_config = deepcopy(config) + with open_dict(replay_config): + target_feature_replay = replay_config.actor_rollout_ref.rollout.drafter.training.target_feature_replay + target_feature_replay.backend = "vllm_file" + + input_store = build_feature_store_from_config(input_cfg, read_only=True) + output_store = build_feature_store_from_config( + output_cfg, + read_only=False, + metadata={ + "source_format": input_type, + "source_path": os.fspath(input_path), + "target_feature_backend": "vllm_file", + }, + ) + replayer = TargetFeatureReplayer( + replay_config, + rank=0, + world_size=1, + device=torch.device("cpu"), + ) + written = 0 + try: + loader = DraftFeatureDataLoader( + input_store, + DraftFeatureDataLoaderConfig( + batch_size=batch_size, + rank=0, + world_size=1, + shuffle=shuffle, + repeat=False, + seed=seed, + ), + ) + for samples in loader: + if max_samples > 0 and written >= max_samples: + break + if max_samples > 0: + samples = samples[: max_samples - written] + features = replayer.materialize(samples) + output_store.write_many(features) + written += len(features) + if written % 100 == 0: + logger.warning("Generated %s vLLM safetensors samples", written) + finally: + input_store.close() + output_store.close() + replayer.close() + + metrics = replayer.metrics() + result = { + "input_path": os.fspath(input_path), + "output_path": os.fspath(output_path), + "written_samples": written, + "metrics": metrics, + } + logger.warning("vLLM safetensors generation finished: %s", result) + return result + + +@hydra.main(config_path="config", config_name="draft_trainer", version_base=None) +def main(config): + logging.basicConfig(level=logging.INFO) + generate_vllm_safetensors_features(config) + + +if __name__ == "__main__": + main()