diff --git a/conf/base.yaml b/conf/base.yaml index 0d5f176b..7d01a189 100644 --- a/conf/base.yaml +++ b/conf/base.yaml @@ -116,6 +116,7 @@ debug: streams_from: null place_inference_workers: true use_existing_llms: false + log_data_pipeline: false # Fast-LLM integration: when true, fast-llm is used as the trainer. # Data flows actors -> Redis (fast_llm_streaming) -> fast-llm training loop. diff --git a/conf/counting.yaml b/conf/counting.yaml index 97f61a88..ed4122b5 100644 --- a/conf/counting.yaml +++ b/conf/counting.yaml @@ -3,6 +3,25 @@ defaults: finetune: seq_length: 4000 gradient_accumulation_passes: 1024 +vllm_config: + vllm_kwargs: + max_model_len: 4000 +fast_llm: + training: + num_workers: 1 + schedule: + depth_first_micro_batches: 256 + model: + base_model: + head: + losses: + grpo: + type: grpo + epsilon_low: 0.2 + epsilon_high: 0.2 + optimizer: + learning_rate: + base: 1e-5 llm: parameters: max_tokens: 1000 diff --git a/conf/math_qwen05_gspo_common.yaml b/conf/math_qwen05_gspo_common.yaml new file mode 100644 index 00000000..c67ff2f8 --- /dev/null +++ b/conf/math_qwen05_gspo_common.yaml @@ -0,0 +1,48 @@ +# Shared Qwen2.5-0.5B GSPO recipe — backend-agnostic parts only: sampling, topology, the +# preprocessor's RL data prep, and orchestration. Composed by the per-backend configs: +# math_qwen05_gspo_fllm (Fast-LLM trainer) / math_qwen05_gspo_ds (DeepSpeed trainer) +# Change a shared hyperparameter HERE; trainer values that must match across backends are +# interpolated from here (e.g. the FL config reads ${finetune.seq_length} / ${finetune.rl.epsilon_*}). +# Rationale for each value: Fast-LLM docs/recipes/reinforcement-learning.md. + +defaults: + - math + - _self_ + +seed: 43 +# Public HF id; override with a local snapshot path on the CLI if it can't be resolved. +model_path: Qwen/Qwen2.5-0.5B + +# Topology: 7 vLLM actors + 1 trainer on 8 GPUs. +world: + actor_fraction: 7 + preprocessor_fraction: 0 + finetune_fraction: 1 + +llm: + parameters: {max_tokens: 7000, temperature: 0.7} +test_llm: + parameters: {max_tokens: 7000, temperature: 0.7} + +actor: + llm_max_rollouts: 128 + +vllm_config: + vllm_kwargs: + max_model_len: 10000 + +# PipelineRL's preprocessor + orchestration config (runs for BOTH backends). The DeepSpeed +# *trainer* fields live in math_qwen05_gspo_ds; the Fast-LLM trainer fields in _fllm. +finetune: + seq_length: 10000 # sample packing in the preprocessor; the FL trainer's + # micro_batch_size interpolates this. + max_train_steps: 2000 + save_checkpoint_steps: 100 + attempts: 8 # rollouts per problem + rl: # RLConfig(**cfg.finetune.rl) — advantage/reward prep (preprocessor) + group_normalization: false + filter_zero_advantage_groups: true + # Canonical clip epsilon. The DeepSpeed loss uses these directly; the Fast-LLM loss + # interpolates them (so the value is defined once, shared across backends). + epsilon_low: 0.003 + epsilon_high: 0.004 diff --git a/conf/math_qwen05_gspo_ds.yaml b/conf/math_qwen05_gspo_ds.yaml new file mode 100644 index 00000000..7ac8cf1c --- /dev/null +++ b/conf/math_qwen05_gspo_ds.yaml @@ -0,0 +1,21 @@ +# Qwen2.5-0.5B GSPO — DeepSpeed trainer backend (bf16 reference arm). Launch: +# python -m pipelinerl.launch --config-name math_qwen05_gspo_ds +# Shared recipe (sampling/topology/preprocessor/orchestration) comes from _common. + +defaults: + - math_qwen05_gspo_common + - _self_ + +use_fast_llm: false +deepspeed_config: deepspeed_stage1_bf16 # single-GPU trainer -> stage is a no-op; 1 is minimal + +# DeepSpeed trainer optimizer/schedule (the Fast-LLM arm uses fast_llm.optimizer.* instead). +finetune: + learning_rate: 5.0e-7 + num_warmup_steps: 0 + lr_scheduler_type: constant + gradient_accumulation_passes: 256 # = rollouts (documents) per optimizer step + adam_beta1: 0.974004 # sqrt-rule betas (m=4), matching the Fast-LLM arm + adam_beta2: 0.999750 + rl: + policy_loss: gspo # selects the DeepSpeed GSPO loss; its clip = finetune.rl.epsilon_* (from _common) diff --git a/conf/math_qwen05_gspo_fllm.yaml b/conf/math_qwen05_gspo_fllm.yaml new file mode 100644 index 00000000..4d4dbf48 --- /dev/null +++ b/conf/math_qwen05_gspo_fllm.yaml @@ -0,0 +1,53 @@ +# Qwen2.5-0.5B GSPO — Fast-LLM trainer backend (bf16). Launch: +# python -m pipelinerl.launch --config-name math_qwen05_gspo_fllm +# Shared recipe (sampling/topology/preprocessor/orchestration) comes from _common. +# Precision variants: add on the CLI +# fast_llm.model.distributed.compute_dtype=float16 vllm_config.vllm_kwargs.dtype=float16 (fp16-matched) +# fast_llm.model.distributed.compute_dtype=float32 vllm_config.vllm_kwargs.dtype=float32 (fp32-matched) + +defaults: + - math_qwen05_gspo_common + - _self_ + +use_fast_llm: true + +# Fast-LLM trainer config (written to a YAML at launch and passed to fast-llm). Must-match +# values are interpolated from the shared _common config so they can't drift. +fast_llm: + training: + train_iters: ${finetune.max_train_steps} + checkpoint: + interval: ${finetune.save_checkpoint_steps} + export: + interval: 1000 # periodic HF-format export cadence + data: + micro_batch_size: ${finetune.seq_length} # tokens; one packed sequence per micro-batch + schedule: + depth_first_micro_batches: 48 + docs_per_step: 256 # accumulate to this many rollouts, then step + optimizer: + learning_rate: + base: 5.0e-7 + warmup_iterations: 0 + decay_style: constant + beta_1: 0.974004 # sqrt-rule, effective-batch multiplier m=4 + beta_2: 0.999750 + gradient_norm_clipping: 76.8 # 0.3 * docs_per_step + model: + distributed: + compute_dtype: bfloat16 # bf16 arm; fp16/fp32 variants override this + the vLLM dtype + sequence_data_parallel: 1 + timeout: 3600 + multi_stage: + zero_stage: 1 # single-GPU trainer (FSDP size 1) -> stage is a no-op + base_model: + head: + fp32_lm_head: true + losses: + # `losses` is a dict keyed by metric-name prefix; this entry logs as gspo_*. + gspo: + type: gspo + epsilon_low: ${finetune.rl.epsilon_low} # shared clip (defined in _common) + epsilon_high: ${finetune.rl.epsilon_high} + logits_scale_factor: 1.4285714 # = 1 / sampling temperature (0.7) + metrics: basic # staleness/clip/ratio/KL/advantage/reward diff --git a/conf/math_qwen7b_gspo_ds.yaml b/conf/math_qwen7b_gspo_ds.yaml new file mode 100644 index 00000000..31bfa93b --- /dev/null +++ b/conf/math_qwen7b_gspo_ds.yaml @@ -0,0 +1,17 @@ +# Qwen2.5-7B GSPO — DeepSpeed trainer backend (bf16 reference arm). Multi-node launch: +# python -m pipelinerl.launch --config-name math_qwen7b_gspo_ds +# Node count is set by the launcher (env WORLD_SIZE), not here. +# +# Same recipe as the 0.5B DeepSpeed config — this inherits it (including the config-driven +# sqrt-rule Adam betas) and overrides only the model and the ZeRO stage. + +defaults: + - math_qwen05_gspo_ds + - _self_ + +# Public HF id; override with a local snapshot path on the CLI if it can't be resolved. +model_path: Qwen/Qwen2.5-7B + +# ZeRO-3 shards parameters + optimizer state + gradients across the multi-GPU trainer +# (memory is tight for the 7B trainer; see the zero_stage note in the FL config). +deepspeed_config: deepspeed_stage3_bf16 diff --git a/conf/math_qwen7b_gspo_fllm.yaml b/conf/math_qwen7b_gspo_fllm.yaml new file mode 100644 index 00000000..5f2c0988 --- /dev/null +++ b/conf/math_qwen7b_gspo_fllm.yaml @@ -0,0 +1,31 @@ +# Qwen2.5-7B GSPO — Fast-LLM trainer backend (bf16). Multi-node launch: +# python -m pipelinerl.launch --config-name math_qwen7b_gspo_fllm +# Node count is set by the launcher (env WORLD_SIZE), not here. +# +# Same recipe as the 0.5B Fast-LLM config — this inherits it and overrides only what the +# larger model needs to fit in trainer memory. Change a shared hyperparameter in +# math_qwen05_gspo_fllm / _common; it flows here automatically. +# Precision variants: same CLI overrides as the 0.5B config +# fast_llm.model.distributed.compute_dtype=float16 vllm_config.vllm_kwargs.dtype=float16 (fp16-matched) + +defaults: + - math_qwen05_gspo_fllm + - _self_ + +# Public HF id; override with a local snapshot path on the CLI if it can't be resolved. +model_path: Qwen/Qwen2.5-7B + +fast_llm: + model: + multi_stage: + # ZeRO-3 also shards parameters (on top of optimizer state + gradients). MLP recompute + # below is already on, so trainer memory is tight — ZeRO-2's replicated parameters may + # not fit. Default to 3; drop to 2 only if there is headroom (2 avoids the per-step + # parameter all-gather). + zero_stage: 3 + base_model: + decoder: + block: + mlp: + # Recompute MLP activations in the backward pass to save trainer memory. + recompute_level: full diff --git a/pipelinerl/actor.py b/pipelinerl/actor.py index 319dabda..8ba8c080 100644 --- a/pipelinerl/actor.py +++ b/pipelinerl/actor.py @@ -696,12 +696,14 @@ def _run(self, dataset: list[tuple[str, dict]]): "result_queue_size": self.result_queue.qsize(), "finished_groups": finished_groups, "trainer_model_version": trainer_version_to_publish, + "trainer_completed_step": self.trainer_state.completed_step, "time_since_start": time.time() - loop_start_time, } trainer_version_to_publish = None else: loop_stats = { "trainer_model_version": last_trainer_version, + "trainer_completed_step": self.trainer_state.completed_step, } self.publish_stats( diff --git a/pipelinerl/async_llm.py b/pipelinerl/async_llm.py index 26ca3da1..327a60f7 100644 --- a/pipelinerl/async_llm.py +++ b/pipelinerl/async_llm.py @@ -6,7 +6,7 @@ import litellm import numpy as np from PIL import Image -from pipelinerl.llm import LLMCall, LLMOutput, Prompt, TokenLogprob, TrainableLLM +from pipelinerl.llm import LLMCall, LLMOutput, Prompt, TokenLogprob, TrainableLLM, parse_token_id_and_version from pipelinerl.finetune.data import MASKED_TOKEN_ID from pipelinerl.rollouts import TrainingText, apply_rollout_reward @@ -184,10 +184,12 @@ async def llm_async_generate( try: # We assume that the server was launched with --return-tokens-as-token-ids # and that the tokens are provided as: ['token_id:1271', 'token_id:1505', ' + token_id, version = parse_token_id_and_version(logprob["token"]) parsed_logprobs.append( TokenLogprob( - token_id=int(logprob["token"].split(":")[-1]), + token_id=token_id, logprob=logprob["logprob"], + version=version, generated=1, ) ) @@ -325,6 +327,12 @@ def make_training_text(llm: TrainableLLM, llm_call: LLMCall) -> TrainingText: # Apply masking to input tokens that aren't generated labels = [MASKED_TOKEN_ID] * len(prompt_token_ids) + labels logprobs = [lp.logprob for lp in llm_call.logprobs] + # Per-token model version, parallel to logprobs. Kept only when the server reported a + # version for every token; otherwise left empty so the trainer falls back to the + # per-rollout version. + token_versions = [lp.version for lp in llm_call.logprobs] + if any(version is None for version in token_versions): + token_versions = [] if finish_reason is not None: finished = finish_reason != "length" else: @@ -339,6 +347,7 @@ def make_training_text(llm: TrainableLLM, llm_call: LLMCall) -> TrainingText: input_ids=input_ids, labels=labels, logprobs=logprobs, + token_versions=token_versions, finished=finished, prompt_tokens=prompt_tokens, output_tokens=output_tokens, diff --git a/pipelinerl/launch.py b/pipelinerl/launch.py index d529ec61..e32ca6e9 100644 --- a/pipelinerl/launch.py +++ b/pipelinerl/launch.py @@ -460,6 +460,9 @@ def _run_finetune_fast_llm(cfg: DictConfig, world_map: WorldMap, gpus: list[int] fast_llm_cfg["run"]["experiment_name"] = experiment_name fast_llm_cfg["data"]["datasets"]["training"]["host"] = cfg.streams.host fast_llm_cfg["data"]["datasets"]["training"]["port"] = cfg.streams.port + if cfg.debug.log_data_pipeline: + fast_llm_cfg["data"]["datasets"]["training"]["log_data_pipeline"] = True + fast_llm_cfg.setdefault("schedule", {})["log_data_pipeline"] = True fast_llm_cfg["training"]["wandb"]["entity_name"] = cfg.wandb.wandb_entity_name fast_llm_cfg["training"]["wandb"]["project_name"] = cfg.wandb.wandb_project_name fast_llm_cfg["training"]["wandb"]["group_name"] = cfg.wandb.wandb_group diff --git a/pipelinerl/llm.py b/pipelinerl/llm.py index 42325433..4f6d9ae6 100644 --- a/pipelinerl/llm.py +++ b/pipelinerl/llm.py @@ -67,6 +67,19 @@ def __bool__(self) -> bool: class TokenLogprob(BaseModel): logprob: float token_id: int + version: int | None = None + + +def parse_token_id_and_version(token: str) -> tuple[int, int | None]: + """Parse a `--return-tokens-as-token-ids` token string into (token_id, version). + + The server emits ``token_id:`` and, when it reports a per-token weight version, + ``token_id::v``. The version suffix is optional. + """ + parts = token.split(":") + if len(parts) >= 2 and parts[-1].startswith("v") and parts[-1][1:].isdigit(): + return int(parts[-2]), int(parts[-1][1:]) + return int(parts[-1]), None class LLMCall(BaseModel): @@ -391,10 +404,12 @@ def parse_completion_logprobs(self, completion_logprobs: list[dict]) -> list[Tok try: # We assume that the server was launched with --return-tokens-as-token-ids # and that the tokens are provided as: ['token_id:1271', 'token_id:1505', ' + token_id, version = parse_token_id_and_version(logprob["token"]) logprobs.append( TokenLogprob( - token_id=int(logprob["token"].split(":")[-1]), + token_id=token_id, logprob=logprob["logprob"], + version=version, ) ) except Exception as e: diff --git a/pipelinerl/preprocess.py b/pipelinerl/preprocess.py index e7984556..03987c5c 100644 --- a/pipelinerl/preprocess.py +++ b/pipelinerl/preprocess.py @@ -3,6 +3,8 @@ os.environ["HF_DATASETS_DISABLE_PROGRESS_BARS"] = "1" +import contextlib +import json import logging import queue import threading @@ -375,6 +377,9 @@ def convert_to_fast_llm_format(entry: dict) -> dict: - loss_masking_spans: list of (start, end) spans masked out of the loss (label == -100; prompt tokens) - advantage: scalar float (per-rollout GRPO advantage) - old_log_probabilities: list of floats, full sequence length (zeros for prompt tokens) + - reward: scalar float (raw per-rollout reward, a diagnostic; distinct from advantage) + - model_version: list of ints, full sequence length (per-token weight version; prompt positions + padded and masked out on the trainer side) """ input_ids = entry["input_ids"] tokens = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) @@ -410,6 +415,11 @@ def convert_to_fast_llm_format(entry: dict) -> dict: if advantages: result["advantage"] = float(advantages[0]) + # reward: raw (un-normalized) reward, a scalar per rollout (distinct from the group-relative + # advantage). Fast-LLM logs it as a diagnostic; it does not affect the loss. + if "reward" in entry: + result["reward"] = float(entry["reward"]) + # old_log_probabilities: full sequence length, zeros for prompt tokens # (prepare_rl_fields pads with zeros on the left to match len(input_ids)) if "old_logprobs" in entry: @@ -417,6 +427,21 @@ def convert_to_fast_llm_format(entry: dict) -> dict: old_logprobs = old_logprobs.tolist() if hasattr(old_logprobs, "tolist") else list(old_logprobs) result["old_log_probabilities"] = [float(x) for x in old_logprobs] + # model_version: full sequence length per-token weight version. When the server reports a + # per-completion-token version (`token_versions`, in-flight weight swaps), left-pad it to the full + # sequence like old_log_probabilities; prompt positions are masked out on the trainer side, so the + # pad value is inert. Otherwise fall back to the per-rollout scalar broadcast across all tokens. + scalar_version = entry.get("model_version") + token_versions = entry.get("token_versions") + if token_versions is not None and hasattr(token_versions, "tolist"): + token_versions = token_versions.tolist() + if token_versions: + pad_value = int(scalar_version) if scalar_version is not None else int(token_versions[0]) + pad = [pad_value] * (len(tokens) - len(token_versions)) + result["model_version"] = pad + [int(x) for x in token_versions] + elif scalar_version is not None: + result["model_version"] = [int(scalar_version)] * len(tokens) + return result @@ -553,7 +578,14 @@ def is_trainer_finished() -> bool: # Per-trainer sample tracking (similar to finetune_loop.py) total_filtered_out = 0 # Track total filtered samples across all batches - with write_to_streams(output_stream, shared=use_shared_stream, stream_name_override=fast_llm_stream_name) as data_writer, write_to_streams(stats_streams) as stats_writer: + pipeline_log_file = None + + with write_to_streams(output_stream, shared=use_shared_stream, stream_name_override=fast_llm_stream_name) as data_writer, write_to_streams(stats_streams) as stats_writer, contextlib.ExitStack() as pipeline_log_stack: + if cfg.use_fast_llm and cfg.debug.log_data_pipeline: + # Write alongside fast-llm rank files: {exp_dir}/finetune/data_pipeline_log/ + log_dir = Path(cfg.output_dir) / "finetune" / "data_pipeline_log" + log_dir.mkdir(parents=True, exist_ok=True) + pipeline_log_file = pipeline_log_stack.enter_context(open(log_dir / "preprocessor.jsonl", "a")) with SharedMemoryManager() as smm: # Create shared memory queues without the manager parameter input_queue = SharedMemoryQueue(smm, cfg.preprocess.input_queue_size, cfg.preprocess.shared_memory_entry_size) @@ -589,6 +621,7 @@ def is_trainer_finished() -> bool: fetching_took = 0 writing_took = 0 num_filtered_out = 0 + last_backpressure_log = 0.0 while True: if is_trainer_finished(): logger.info("Trainer signalled completion; stopping preprocessor loop") @@ -656,6 +689,13 @@ def is_trainer_finished() -> bool: assert isinstance(trainer_state.samples_processed, int) if published_samples - trainer_state.samples_processed > max_unconsumed_samples: # wait for the finetune loop to finish processing data + now = time.time() + if now - last_backpressure_log >= 10.0: + last_backpressure_log = now + logger.info( + f"Back-pressure: published={published_samples} consumed={trainer_state.samples_processed}" + f" unconsumed={published_samples - trainer_state.samples_processed} > max={max_unconsumed_samples}, waiting" + ) continue batch_done = False @@ -665,10 +705,25 @@ def is_trainer_finished() -> bool: # Fast-LLM path: write individual samples directly (Fast-LLM does its own packing) if cfg.use_fast_llm: + write_start = time.time() if pipeline_log_file is not None else None + write_samples = 0 + write_tokens = 0 while len(processed_entries_queue) > 0: entry = processed_entries_queue.popleft() + if pipeline_log_file is not None: + write_samples += 1 + write_tokens += len(entry.get("input_ids", [])) data_writer.write(convert_to_fast_llm_format(entry)) published_samples += 1 + if pipeline_log_file is not None and write_samples > 0: + pipeline_log_file.write(json.dumps({ + "event": "WRITE", + "t_start": round(write_start, 3), + "t_end": round(time.time(), 3), + "samples": write_samples, + "tokens": write_tokens, + }) + "\n") + pipeline_log_file.flush() batch_done = True elif cfg.finetune.seq_packing: if samples_per_trainer[trainer_id] == target_samples_per_lead: diff --git a/pipelinerl/rollouts.py b/pipelinerl/rollouts.py index 4c71dda7..637db371 100644 --- a/pipelinerl/rollouts.py +++ b/pipelinerl/rollouts.py @@ -19,6 +19,9 @@ class TrainingText(BaseModel): n_predicted (int): The number of predicted tokens in the text. reward (float): The reward associated with the training instance. Defaults to 0.0. logprobs (List[float]): A list of log probabilities of the completion tokens from the assistant model. + token_versions (List[int]): Per-completion-token model version (parallel to logprobs). Captures + weight versions that change mid-generation when the server swaps weights in flight. Empty + when the server does not report per-token versions. ref_logprobs (List[float]): A list of reference log probabilities of the completion tokens from the reference model. input_ids (List[int]): A list of token IDs representing the input text, including the prompt and the predicted tokens. labels (List[int]): A list of token IDs that are used as labels for training. The last n_predicted tokens are set to MASKED_TOKEN_ID. @@ -36,6 +39,7 @@ class TrainingText(BaseModel): n_predicted: int reward: float = 0.0 logprobs: List[float] = Field(default_factory=list) + token_versions: List[int] = Field(default_factory=list) ref_logprobs: List[float] = Field(default_factory=list) input_ids: List[int] = Field(default_factory=list) labels: List[int] = Field(default_factory=list) diff --git a/pipelinerl/state.py b/pipelinerl/state.py index 8db2752d..233fc661 100644 --- a/pipelinerl/state.py +++ b/pipelinerl/state.py @@ -81,12 +81,16 @@ def __init__(self, exp_path: Path, use_fast_llm: bool = False, weight_broadcast: self.use_fast_llm = use_fast_llm self.weight_broadcast = weight_broadcast self.propagated_weight_version: int | None = None if weight_broadcast else 0 + # Raw trainer step behind the current weights (for logging only); the version stamped onto + # rollouts is `propagated_weight_version`, which is the document count when Fast-LLM sends it. + self.completed_step: int | None = None if weight_broadcast else 0 self.samples_processed: int | None = None if weight_broadcast else 0 self.training_done: bool = False self._training_done_event = threading.Event() def debug_mode_init(self): self.propagated_weight_version = 0 + self.completed_step = 0 self.samples_processed = 0 self.training_done = True self._training_done_event.set() @@ -134,6 +138,7 @@ def listen_events(): if event_type == "weights_ready": logger.info(f"Received weights_ready event: step={step}, documents_seen={documents_seen}") self.propagated_weight_version = version + self.completed_step = step elif event_type == "training_finished": logger.info("Received training_finished event") self.training_done = True diff --git a/pipelinerl/vllm1.py b/pipelinerl/vllm1.py index ab15a318..6540ef37 100644 --- a/pipelinerl/vllm1.py +++ b/pipelinerl/vllm1.py @@ -52,6 +52,113 @@ logger.propagate = False +# --- Per-token model version capture -------------------------------------------------- +# The active weight version is a global, serialized quantity: it changes only inside a +# weight swap, while generation is paused (`_pause_generation`). We record the version +# active when the output processor commits each token, then ride it to the client inside +# the existing per-token `token` string of the chat logprobs +# (`token_id:` -> `token_id::v`), so no response-schema change is needed. +# Both seams run in the API-server process, alongside the version-tracking monitor thread. +# +# Any missing link (unpatched vLLM build, flat logprobs, a token absent from its own +# top-logprobs) simply omits the version; the consumer then falls back to the per-rollout +# version. +_current_model_version: dict[str, int | None] = {"value": None} +# Set once if a patched seam ever raises: the annotation hooks then no-op cheaply and the +# consumer falls back to the per-rollout version. +_version_tagging_disabled: dict[str, bool] = {"value": False} + + +def _set_current_model_version(version: int | None) -> None: + _current_model_version["value"] = version + + +def _disable_version_tagging(context: str, error: Exception) -> None: + if not _version_tagging_disabled["value"]: + _version_tagging_disabled["value"] = True + logger.warning( + f"[FastLLM] Per-token model_version tagging disabled after error in {context}: {error!r}" + ) + + +def _install_model_version_patches() -> None: + """Monkeypatch the vLLM v1 output path to tag generated tokens with the model version. + + Two seams, both in the API-server process: + 1. `LogprobsProcessor.update_from_output` — annotate each newly committed position's + `Logprob` objects with `.version` = the version active at commit time. + 2. `OpenAIServingChat._create_chat_logprobs` — append `:v` to each per-token + `token` string, read back from the annotated `Logprob`. + Idempotent, and defensive: a version mismatch that moves these seams disables per-token + versions (consumer falls back to the per-rollout version) rather than crashing the server. + """ + try: + from vllm.v1.engine.logprobs import LogprobsProcessor + + try: + # Newer vLLM keeps the chat serving class in a chat_completion package; + # older builds define it in serving_chat.py. + from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat + except ImportError: + from vllm.entrypoints.openai.serving_chat import OpenAIServingChat + except ImportError as error: + logger.warning(f"[FastLLM] Per-token model_version disabled (vLLM layout changed): {error!r}") + return + + if getattr(LogprobsProcessor, "_pipelinerl_version_patched", False): + return + + original_update_from_output = LogprobsProcessor.update_from_output + + def update_from_output(self, *args, **kwargs): + previous_length = len(self.logprobs) if isinstance(self.logprobs, list) else None + original_update_from_output(self, *args, **kwargs) + if _version_tagging_disabled["value"]: + return + try: + version = _current_model_version["value"] + if version is None or previous_length is None or not isinstance(self.logprobs, list): + return + # Every logprob at a decode position shares that position's version; annotate all of + # them so the serving layer reads the right value regardless of dict ordering. + for position in self.logprobs[previous_length:]: + if isinstance(position, dict): + for logprob in position.values(): + logprob.version = version + except Exception as error: + # Best effort: never let version tagging break the output processor. + _disable_version_tagging("output processor", error) + + original_create_chat_logprobs = OpenAIServingChat._create_chat_logprobs + + def _create_chat_logprobs(self, *args, **kwargs): + result = original_create_chat_logprobs(self, *args, **kwargs) + if _version_tagging_disabled["value"]: + return result + try: + token_ids = args[0] if args else kwargs.get("token_ids") + top_logprobs = args[1] if len(args) > 1 else kwargs.get("top_logprobs") + content = getattr(result, "content", None) + if content and token_ids is not None and top_logprobs is not None: + for index, item in enumerate(content): + position = top_logprobs[index] if index < len(top_logprobs) else None + token_id = token_ids[index] if index < len(token_ids) else None + sampled = position.get(token_id) if position is not None and token_id is not None else None + version = getattr(sampled, "version", None) + # Only extend the `token_id:` form; never mangle a decoded text token. + if version is not None and item.token.startswith("token_id:"): + item.token = f"{item.token}:v{version}" + except Exception as error: + # Best effort: never let version tagging break the response. + _disable_version_tagging("chat logprobs", error) + return result + + LogprobsProcessor.update_from_output = update_from_output + OpenAIServingChat._create_chat_logprobs = _create_chat_logprobs + LogprobsProcessor._pipelinerl_version_patched = True + logger.info("[FastLLM] Per-token model_version patches installed") + + @runtime_checkable class LikeWorker(Protocol): rank: int @@ -327,6 +434,10 @@ async def receive_weight_update_fast_llm(self, version: int | None = None): so that in-flight generation cannot interleave with a mid-broadcast parameter swap (the source of logprob drift PR #137 closed). + `version` is recorded as the active model version once the new weights are + loaded but before generation resumes, so tokens sampled after the swap are + stamped with the new version and those before it keep the old one. + NOTE: this must NOT be used for the very first weights_ready event after process startup, because at that point the actor has not yet begun issuing rollouts (it's blocked in wait_for_model_version) and @@ -346,6 +457,8 @@ async def receive_weight_update_fast_llm(self, version: int | None = None): await self.engine.engine_core.collective_rpc_async( "receive_weight_update_fast_llm", args=() ) + # Weights are loaded; stamp subsequent tokens with the new version before resuming. + _set_current_model_version(version) logger.info( f"Fast-llm weight update processed version={version} " f"in {time.perf_counter() - update_started_at:.3f}s" @@ -401,15 +514,22 @@ def monitor_redis_stream(): "receive_weight_update_fast_llm", args=() ) first_weights_ready_seen = True + initial_broadcast = True else: logger.info( f"[FastLLM] weights_ready step={step} documents_seen={documents_seen}, " f"dispatching to workers" ) coro = self.receive_weight_update_fast_llm(version) + initial_broadcast = False try: future = asyncio.run_coroutine_threadsafe(coro, loop) future.result() + # The pause-wrapped path stamps the version internally (before + # resume); the initial raw path runs before the actor generates, + # so setting it here has no token to race with. + if initial_broadcast: + _set_current_model_version(version) logger.info(f"[FastLLM] Weight update complete: step={step}") except Exception as e: logger.error(f"[FastLLM] Error receiving weight update: {e}") @@ -480,6 +600,7 @@ async def create_engine(args: Any): await manager.init_actor_update_group() if weight_update_mode == "fast-llm": + _install_model_version_patches() await manager.init_fast_llm_receiver() await manager.start_fast_llm_monitoring() logger.info("Fast-LLM weight update mode enabled")