Fast-LLM RL: per-token model-version tagging, metrics, and GSPO configs

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.

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

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

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)

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

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

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

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(

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,

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

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: