Forward raw reward, document-count model version, and per-token model version

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()

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,

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:<id>`` and, when it reports a per-token weight version,
+    ``token_id:<id>:v<version>``. 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:

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,13 +393,33 @@ 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:
         old_logprobs = entry["old_logprobs"]
         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
 
 

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)

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