diff --git a/pipelinerl/actor.py b/pipelinerl/actor.py index 1fec56f5..2bfe762a 100644 --- a/pipelinerl/actor.py +++ b/pipelinerl/actor.py @@ -762,12 +762,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_model_version": last_trainer_version, + "trainer_completed_step": self.trainer_state.completed_step, } _t_before_stats = time.monotonic() diff --git a/pipelinerl/async_llm.py b/pipelinerl/async_llm.py index fd9e1137..eb7dc1d2 100644 --- a/pipelinerl/async_llm.py +++ b/pipelinerl/async_llm.py @@ -6,7 +6,7 @@ import aiohttp 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 @@ -174,10 +174,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, ) ) @@ -306,6 +308,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: @@ -320,6 +328,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/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 318592ff..089b99d9 100644 --- a/pipelinerl/preprocess.py +++ b/pipelinerl/preprocess.py @@ -355,6 +355,9 @@ def convert_to_fast_llm_format(entry: dict) -> dict: - loss_masking_spans: list of (start, end) spans where loss IS computed (completion only) - 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) @@ -390,6 +393,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: @@ -397,6 +405,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 diff --git a/pipelinerl/rollouts.py b/pipelinerl/rollouts.py index 1200ba23..a325317f 100644 --- a/pipelinerl/rollouts.py +++ b/pipelinerl/rollouts.py @@ -18,6 +18,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. @@ -35,6 +38,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 a9539812..69e69d6e 100644 --- a/pipelinerl/state.py +++ b/pipelinerl/state.py @@ -26,12 +26,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() @@ -105,10 +109,18 @@ def listen_events(): event_type = event.get("type") step = event.get("step") + # Fast-LLM sends the cumulative document count as the model version (to align + # staleness with DeepSpeed's document clock); fall back to `step` for older + # trainers that only send the step. + document_count = event.get("document_count") + version = document_count if document_count is not None else step if event_type == "weights_ready": - logger.info(f"Received weights_ready event: step={step}") - self.propagated_weight_version = step + logger.info( + f"Received weights_ready event: step={step}, document_count={document_count}" + ) + 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 855a992d..16dc9df8 100644 --- a/pipelinerl/vllm1.py +++ b/pipelinerl/vllm1.py @@ -61,6 +61,110 @@ 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}") + 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, token_ids, top_logprobs, *args, **kwargs): + result = original_create_chat_logprobs(self, token_ids, top_logprobs, *args, **kwargs) + if _version_tagging_disabled["value"]: + return result + try: + content = getattr(result, "content", None) + if content: + for index, item in enumerate(content): + position = top_logprobs[index] if index < len(top_logprobs) else None + sampled = position.get(token_ids[index]) if position 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 @@ -408,13 +512,17 @@ async def init_fast_llm_receiver(self): f"Fast-LLM receiver initialized (Redis {self._redis_host}:{self._redis_port})" ) - async def receive_weight_update_fast_llm(self): + async def receive_weight_update_fast_llm(self, version: int | None = None): """Run a fast-llm broadcast weight update paused-for-the-duration. Pause/resume wraps the collective RPC symmetrically with the HTTP path 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 @@ -434,8 +542,10 @@ async def receive_weight_update_fast_llm(self): 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 " + f"Fast-llm weight update processed version={version} " f"in {time.perf_counter() - update_started_at:.3f}s" ) finally: @@ -505,24 +615,38 @@ def monitor_redis_stream(): event_type = event.get("type") step = event.get("step") + # Fast-LLM sends the cumulative document count as the model version + # (aligns staleness with the trainer's document clock); fall back to + # `step` for older trainers that only send the step. + document_count = event.get("document_count") + version = document_count if document_count is not None else step if event_type == "weights_ready": if not first_weights_ready_seen: logger.info( - f"[FastLLM] weights_ready step={step} (initial broadcast — no pause wrap)" + f"[FastLLM] weights_ready step={step} document_count={document_count} " + f"(initial broadcast — no pause wrap)" ) coro = self.engine.engine_core.collective_rpc_async( "receive_weight_update_fast_llm", args=() ) first_weights_ready_seen = True + initial_broadcast = True else: logger.info( - f"[FastLLM] weights_ready step={step}, dispatching to workers" + f"[FastLLM] weights_ready step={step} document_count={document_count}, " + f"dispatching to workers" ) - coro = self.receive_weight_update_fast_llm() + 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}" ) @@ -640,6 +764,7 @@ async def create_engine( # Initialize Fast-LLM mode if enabled 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") diff --git a/tests/test_model_version.py b/tests/test_model_version.py new file mode 100644 index 00000000..0b98a3eb --- /dev/null +++ b/tests/test_model_version.py @@ -0,0 +1,51 @@ +"""Unit tests for the per-token model_version plumbing. + +Covers the two pure pieces of the (otherwise cluster-only) feature: parsing the optional +`:v` suffix off a `token_id:` string, and the model_version padding / fallback in +`convert_to_fast_llm_format`. +""" + +import pytest + +from pipelinerl.llm import parse_token_id_and_version +from pipelinerl.preprocess import convert_to_fast_llm_format + + +@pytest.mark.parametrize( + "token, expected", + [ + ("token_id:1271", (1271, None)), # no version suffix (backward compatible) + ("token_id:1271:v5", (1271, 5)), + ("token_id:1271:v0", (1271, 0)), # version 0 is a real version, not "absent" + ("token_id:50257:v123456", (50257, 123456)), + ], +) +def test_parse_token_id_and_version(token, expected): + assert parse_token_id_and_version(token) == expected + + +def test_convert_model_version_per_token_left_padded(): + # Completion versions are left-padded to the full sequence with the per-rollout scalar. + entry = {"input_ids": [10, 11, 12, 13, 14, 15], "model_version": 1, "token_versions": [2, 3]} + assert convert_to_fast_llm_format(entry)["model_version"] == [1, 1, 1, 1, 2, 3] + + +def test_convert_model_version_per_token_pads_with_first_when_no_scalar(): + entry = {"input_ids": [10, 11, 12, 13], "token_versions": [7, 8]} + assert convert_to_fast_llm_format(entry)["model_version"] == [7, 7, 7, 8] + + +def test_convert_model_version_scalar_broadcast_fallback(): + # No per-token versions: broadcast the per-rollout scalar across the sequence. + entry = {"input_ids": [10, 11, 12], "model_version": 4, "token_versions": []} + assert convert_to_fast_llm_format(entry)["model_version"] == [4, 4, 4] + + +def test_convert_model_version_absent(): + entry = {"input_ids": [10, 11, 12]} + assert "model_version" not in convert_to_fast_llm_format(entry) + + +def test_convert_model_version_full_completion_no_prompt(): + entry = {"input_ids": [10, 11, 12], "model_version": 9, "token_versions": [2, 3, 4]} + assert convert_to_fast_llm_format(entry)["model_version"] == [2, 3, 4]