freesolo-flash-dev 1.0.4__tar.gz → 1.0.5__tar.gz

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.4
+Version: 1.0.5
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/flash/cli/training_doc.py

@@ -109,6 +109,8 @@ epochs = 1                  # one pass over the retained train rows
 max_examples = 2            # rows to train on (the starter dataset has 2)
 # max_steps = 100           # positive values set the exact optimizer-update horizon
 # save_at_steps = [10, 50, 100]  # requires max_steps; overrides save_every
+# multi-turn GRPO defaults to one reward per rollout; choose "per_turn" for turn-level credit.
+# credit_assignment = "per_episode"
 lora_rank = 32
 lora_alpha = 64
 # All SFT/GRPO knobs live under [train]. Do not add [sft] or [grpo] tables.
@@ -339,10 +341,12 @@ def score_response(self, example, response_text) -> RewardResult:
     )
 ```
 
-`score` is what GRPO optimizes (it becomes the run's `total`). Each `RewardMetric` you
-attach is logged by name in the per-scorer breakdown — that is how the clean success
-rate becomes visible. Use the shaped `score` to confirm the model is learning *at all*,
-and judge the run on the explicit `success` metric.
+`score` is what GRPO optimizes (it becomes the run's `total`). In standard (single-turn)
+GRPO, each `RewardMetric` is averaged across scored completions and logged by name at
+the managed heartbeat cadence, which is not guaranteed to be every optimizer step. That
+is how the clean success rate becomes visible. Multi-turn scoring currently reports only
+the scalar reward. Use the shaped `score` to confirm the model is learning *at all*, and
+judge the run on the explicit `success` metric.
 
 When `thinking = true`, score the final answer unless you intentionally need the
 reasoning trace. Flash passes a string-compatible response object to `score_response`;
@@ -550,10 +554,9 @@ every run, the last two matter more the smaller the model:
   frontier one outright; a frontier `teacher_model` only earns its keep once the student is large
   enough to track it (~9B+). Early-stopping also largely neutralizes this gap, since the teacher-driven
   over-sharpening only compounds over many steps.
-- **Diagnose it in-band.** Watch the per-step **mean completion entropy** in the run's telemetry — a
-  steady decline toward zero is the collapse happening. Confirm at serving by evaluating at
-  **temperature=0** and flagging `finish_reason=length` completions that never emit your answer token,
-  and compare an early checkpoint against the final one to watch the loop emerge over steps.
+- **Diagnose it at serving.** Evaluate at **temperature=0** and flag
+  `finish_reason=length` completions that never emit your answer token. Compare an early checkpoint
+  against the final one to watch the loop emerge over steps.
 
 ### Distilling from base with no format anchor
 

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/flash/client/http.py

@@ -512,13 +512,22 @@ class ApiClient:
                 if content:
                     yield str(content)
                 return
-            while raw := resp.read(1):
-                chunk = decoder.decode(raw)
-                if chunk:
-                    yield chunk
-            tail = decoder.decode(b"", final=True)
-            if tail:
-                yield tail
+            read1 = getattr(resp, "read1", None)
+            read = read1 if read1 is not None else resp.read
+            read_size = 4096 if read1 is not None else 1
+            while raw := read(read_size):
+                state = decoder.getstate()
+                try:
+                    decoded = decoder.decode(raw)
+                except UnicodeDecodeError as exc:
+                    decoder.setstate(state)
+                    prefix_end = max(0, exc.start - len(state[0]))
+                    yield from decoder.decode(raw[:prefix_end])
+                    # bind + re-raise explicitly: the yield above clears the active exception, so a
+                    # bare `raise` here would fail with "No active exception to reraise".
+                    raise exc
+                yield from decoded
+            yield from decoder.decode(b"", final=True)
 
 
 def client_from_config(require_key: bool = True) -> ApiClient:

freesolo_flash_dev-1.0.5/flash/engine/multiturn_reward_scoring.py

@@ -0,0 +1,73 @@
+"""Typed terminal reward scoring for multi-turn GRPO rollouts."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+from flash.envs.base import BaseEnvironment, RolloutReward
+
+
+@dataclass(frozen=True)
+class RolloutScoreRequest:
+    example: dict
+    state: dict
+    turn_count: int
+
+
+def _validated_reward(reward: RolloutReward, request: RolloutScoreRequest) -> RolloutReward:
+    episode = float(reward.episode)
+    # a non-finite episode reward is unscorable and has no valid scalar fallback. canonicalize it to
+    # nan, which is trl's ONLY unscorable marker (trl uses torch.isnan to exclude the row from the
+    # group baseline and then nan_to_num zeros its advantage, matching stock grpo). forwarding a raw
+    # inf instead would NOT be recognized as unscorable and would contaminate the whole group with
+    # huge advantages. per-turn credit is disabled so the unscorable row is never revived.
+    if not math.isfinite(episode):
+        print(
+            "[grpo][warn] episode reward non-finite; rollout unscorable, per-turn credit disabled"
+        )
+        return RolloutReward(episode=float("nan"), turns=None)
+    if reward.turns is None:
+        return RolloutReward(episode=episode, turns=None)
+    if not isinstance(reward.turns, (list, tuple)):
+        # only an ordered list/tuple carries a well-defined per-turn order. a str, bytes,
+        # bytearray, mapping, or unordered set would still iterate into floats and could pass
+        # the count check while assigning rewards to the wrong turns -- reject and fall back.
+        print(
+            "[grpo][warn] per-turn rewards unavailable (turns is not an ordered list/tuple); "
+            "using episode reward"
+        )
+        return RolloutReward(episode=episode, turns=None)
+
+    reason: str | None = None
+    coerced: tuple[float, ...] | None = None
+    try:
+        coerced = tuple(float(value) for value in reward.turns)
+    except (TypeError, ValueError):
+        reason = "per-turn rewards contain a non-number"
+    else:
+        if len(coerced) != request.turn_count:
+            reason = f"received {len(coerced)} reward(s) for {request.turn_count} assistant turn(s)"
+        elif not all(math.isfinite(value) for value in coerced):
+            reason = "per-turn rewards contain a non-finite value"
+
+    if reason is not None:
+        print(f"[grpo][warn] per-turn rewards unavailable ({reason}); using episode reward")
+        return RolloutReward(episode=episode, turns=None)
+    return RolloutReward(episode=episode, turns=coerced)
+
+
+def score_rollouts(active_env, requests: list[RolloutScoreRequest]) -> list[RolloutReward]:
+    """Score terminal rollout states once and return normalized typed rewards."""
+    items = [(request.example, request.state) for request in requests]
+    rollout_rewards_many = getattr(active_env, "rollout_rewards_many", None)
+    if callable(rollout_rewards_many):
+        rewards = rollout_rewards_many(items)
+    else:
+        rewards = BaseEnvironment.rollout_rewards_many(active_env, items)
+    if len(rewards) != len(requests):
+        raise RuntimeError("env.rollout_rewards_many returned the wrong number of rewards")
+    return [
+        _validated_reward(reward, request)
+        for reward, request in zip(rewards, requests, strict=True)
+    ]

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/flash/engine/multiturn_rollout.py

@@ -16,10 +16,11 @@ import threading
 import time
 from collections import OrderedDict
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from typing import TypedDict
 
+from flash.engine.multiturn_reward_scoring import RolloutScoreRequest, score_rollouts
+
 
 class RolloutResult(TypedDict):
     """Token-aligned fields returned per rollout."""
@@ -29,6 +30,8 @@ class RolloutResult(TypedDict):
     logprobs: list[float]
     env_mask: list[int]
     reward: float
+    turn_spans: list[tuple[int, int]]
+    turn_rewards: list[float] | None
 
 
 RolloutCompletion = tuple[str, list[int], list[float], str]
@@ -71,6 +74,8 @@ _ROLLOUT_FIELDS: tuple[str, ...] = (
     "logprobs",
     "env_mask",
     "reward",
+    "turn_spans",
+    "turn_rewards",
 )
 
 
@@ -174,6 +179,7 @@ def rollout_one(
     completion_ids: list[int] = []
     logprobs: list[float] = []
     env_mask: list[int] = []
+    turn_spans: list[tuple[int, int]] = []
 
     turns = 0
     while True:
@@ -184,7 +190,9 @@ def rollout_one(
                 break
             max_new = min(max_new, remaining)
         asst_ids, asst_lp, text = generate(cur_ids, max_new)
+        turn_start = len(completion_ids)
         completion_ids.extend(asst_ids)
+        turn_spans.append((turn_start, len(completion_ids)))
         logprobs.extend(asst_lp)
         env_mask.extend([1] * len(asst_ids))
         cur_ids.extend(asst_ids)
@@ -212,13 +220,18 @@ def rollout_one(
         env_mask.extend([0] * len(glue))
         cur_ids.extend(glue)
 
-    reward = active_env.reward("", example, state)
+    score = score_rollouts(
+        active_env,
+        [RolloutScoreRequest(example=example, state=state, turn_count=len(turn_spans))],
+    )[0]
     return {
         "prompt_ids": prompt_ids,
         "completion_ids": completion_ids,
         "logprobs": logprobs,
         "env_mask": env_mask,
-        "reward": float(reward),
+        "reward": score.episode,
+        "turn_spans": turn_spans,
+        "turn_rewards": list(score.turns) if score.turns is not None else None,
     }
 
 
@@ -331,6 +344,7 @@ class _RolloutState:
         "messages",
         "prompt_ids",
         "state",
+        "turn_spans",
         "turns",
     )
 
@@ -343,18 +357,21 @@ class _RolloutState:
         self.completion_ids: list[int] = []
         self.logprobs: list[float] = []
         self.env_mask: list[int] = []
+        self.turn_spans: list[tuple[int, int]] = []
         self.state = state
         self.turns = 0
         self.budget = budget
         self.done = False
 
-    def result(self, reward: float) -> RolloutResult:
+    def result(self, reward: float, turn_rewards: list[float] | None) -> RolloutResult:
         return {
             "prompt_ids": self.prompt_ids,
             "completion_ids": self.completion_ids,
             "logprobs": self.logprobs,
             "env_mask": self.env_mask,
             "reward": float(reward),
+            "turn_spans": self.turn_spans,
+            "turn_rewards": turn_rewards,
         }
 
 
@@ -369,7 +386,9 @@ def _advance_after_turn(
     max_turns: int,
 ) -> None:
     """Fold one assistant turn into ``r`` and run its env step. Sets ``r.done`` when finished."""
+    turn_start = len(r.completion_ids)
     r.completion_ids.extend(asst_ids)
+    r.turn_spans.append((turn_start, len(r.completion_ids)))
     r.logprobs.extend(asst_lp)
     r.env_mask.extend([1] * len(asst_ids))
     r.cur_ids.extend(asst_ids)
@@ -465,31 +484,6 @@ def _turn_budget(r: _RolloutState, per_turn_max_tokens: int) -> int | None:
     return max(1, max_new)
 
 
-def _score_rollouts(active_env, rollouts: list[_RolloutState]) -> list[float]:
-    """Reward each rollout in input order, using reward_many, concurrent, or serial scoring."""
-    reward_many = getattr(active_env, "reward_many", None)
-    if callable(reward_many):
-        rewards = reward_many([(r.example, r.state) for r in rollouts])
-        if len(rewards) != len(rollouts):
-            raise RuntimeError("env.reward_many returned the wrong number of rewards")
-        return [float(x) for x in rewards]
-
-    def _score(r: _RolloutState) -> float:
-        return float(active_env.reward("", r.example, r.state))
-
-    if len(rollouts) <= 1 or not getattr(active_env, "reward_thread_safe", True):
-        return [_score(r) for r in rollouts]
-    pool = ThreadPoolExecutor(max_workers=min(16, len(rollouts)))
-    try:
-        futures = {pool.submit(_score, r): i for i, r in enumerate(rollouts)}
-        scores: list[float] = [0.0] * len(rollouts)
-        for fut in as_completed(futures):
-            scores[futures[fut]] = fut.result()  # re-raises the first failed scorer
-    finally:
-        pool.shutdown(wait=True, cancel_futures=True)
-    return scores
-
-
 _PHYSICAL_REQUEST_COUNTER = itertools.count()
 _PHYSICAL_REQUEST_COUNTER_LOCK = threading.Lock()
 _ROLLOUT_REQUEST_TIMEOUT_MIN_SECONDS = 600.0
@@ -690,8 +684,22 @@ def rollout_async(
         to_env.put(None)
         worker.join()
 
-    rewards = _score_rollouts(active_env, rollouts)
-    return [r.result(rw) for r, rw in zip(rollouts, rewards, strict=True)]
+    requests = [
+        RolloutScoreRequest(
+            example=rollout.example,
+            state=rollout.state,
+            turn_count=len(rollout.turn_spans),
+        )
+        for rollout in rollouts
+    ]
+    scores = score_rollouts(active_env, requests)
+    return [
+        rollout.result(
+            score.episode,
+            list(score.turns) if score.turns is not None else None,
+        )
+        for rollout, score in zip(rollouts, scores, strict=True)
+    ]
 
 
 def render_message_ids(tok, messages, add_generation_prompt: bool, *, thinking: bool) -> list[int]:
@@ -964,10 +972,10 @@ def build_rollout_func(
                 monotonic=monotonic,
                 request_id_factory=request_id_factory,
             )
-            out: dict[str, list] = {k: [] for k in _ROLLOUT_FIELDS}
-            for r in rollouts:
-                for k in out:
-                    out[k].append(r[k])
+            out: dict[str, list] = {key: [] for key in _ROLLOUT_FIELDS}
+            for rollout in rollouts:
+                for key in out:
+                    out[key].append(rollout[key])
             return out
         finally:
             # Abort in-flight requests on error so they don't corrupt the next GRPO step.

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/flash/engine/worker/adapter.py

@@ -2,11 +2,21 @@
 
 from __future__ import annotations
 
+import json
 import os
+import shutil
 
 from flash.engine.recipe import RECIPE
 from flash.engine.worker._pkg import W as _w
+from flash.engine.worker.hf import (
+    RetriableInfraError,
+    _has_deployable_adapter,
+    _prefetch_error_is_retriable,
+    _require_hf_deadline_allowance,
+    _sleep_with_hf_deadline,
+)
 from flash.engine.worker.lora import (
+    _read_adapter_tensor_keys,
     adapter_is_vl_warmstart,
     assert_adapter_delta_nonzero,
     assert_adapter_load_clean,
@@ -15,6 +25,9 @@ from flash.engine.worker.lora import (
 )
 from flash.engine.worker.perf import optimal_attn_impl
 
+_ADAPTER_DOWNLOAD_RETRIES = 4
+_ADAPTER_DOWNLOAD_BACKOFF_S = 5.0
+
 
 def make_lora(model_id: str | None = None):
     """Build LoRA config targeting all linear layers (VL models included: the vision tower /
@@ -196,6 +209,20 @@ def _resolve_adapter_ref(adapter_ref: str) -> tuple[str, str] | None:
     return parse_adapter_storage_ref(adapter_ref)
 
 
+def _warmstart_adapter_is_loadable(adir: str) -> bool:
+    """return true only for a structurally complete adapter config and weight file."""
+    if not _has_deployable_adapter(adir):
+        return False
+    try:
+        with open(os.path.join(adir, "adapter_config.json"), encoding="utf-8") as config_file:
+            config = json.load(config_file)
+        if not isinstance(config, dict) or str(config.get("peft_type", "")).upper() != "LORA":
+            return False
+        return bool(_read_adapter_tensor_keys(adir))
+    except Exception:
+        return False
+
+
 def _download_adapter(adapter_prefix: str | None) -> str | None:
     """Download an init_from_adapter LoRA to /tmp/evdl/<prefix>/adapter and return its dir.
 
@@ -210,19 +237,45 @@ def _download_adapter(adapter_prefix: str | None) -> str | None:
     repo, prefix = resolved
     from huggingface_hub import snapshot_download
 
-    try:
-        snapshot_download(
-            repo_id=repo,
-            repo_type="dataset",
-            allow_patterns=[f"{prefix}/adapter/*"],
-            local_dir="/tmp/evdl",
-            token=os.environ.get("HF_TOKEN"),
-            revision=(_w.JOB_SPEC.train.init_from_adapter_revision if _w.JOB_SPEC else None)
-            or None,
-        )
-    except Exception:
-        raise RuntimeError(
-            "the prepared warm-start source adapter could not be downloaded"
-        ) from None
     adir = os.path.join("/tmp/evdl", prefix, "adapter")
-    return adir if os.path.isdir(adir) else None
+    # start from a clean path so the loadable-check can only ever accept files THIS download
+    # materialized -- leftover materialization from an earlier worker subprocess, attempt, or a
+    # different run sharing the same prefix must not satisfy the post-exception loadable check and
+    # mask a terminal 404/403/429 for the current repo/revision.
+    shutil.rmtree(adir, ignore_errors=True)
+    for attempt in range(_ADAPTER_DOWNLOAD_RETRIES):
+        _require_hf_deadline_allowance()
+        try:
+            snapshot_download(
+                repo_id=repo,
+                repo_type="dataset",
+                allow_patterns=[f"{prefix}/adapter/*"],
+                local_dir="/tmp/evdl",
+                token=os.environ.get("HF_TOKEN"),
+                revision=(_w.JOB_SPEC.train.init_from_adapter_revision if _w.JOB_SPEC else None)
+                or None,
+            )
+        except Exception as error:
+            # a later nonessential sidecar may fail after the config and weights are already complete.
+            if _warmstart_adapter_is_loadable(adir):
+                return adir
+            if not _prefetch_error_is_retriable(error):
+                raise RuntimeError(
+                    "the prepared warm-start source adapter could not be downloaded"
+                ) from None
+        else:
+            # a returned snapshot can still be incomplete: an interrupted transfer, or hf falling
+            # back to a partial local_dir when a throttled metadata call cannot confirm the file set.
+            if _warmstart_adapter_is_loadable(adir):
+                return adir
+        # discard partial local_dir materialization so the next attempt cannot reuse stale files.
+        shutil.rmtree(adir, ignore_errors=True)
+        if attempt + 1 < _ADAPTER_DOWNLOAD_RETRIES:
+            try:
+                if not _sleep_with_hf_deadline(_ADAPTER_DOWNLOAD_BACKOFF_S * (attempt + 1)):
+                    break
+            except Exception:
+                break
+    raise RetriableInfraError(
+        "the prepared warm-start source adapter could not be downloaded after transient failures"
+    ) from None

freesolo_flash_dev-1.0.5/flash/engine/worker/grpo_perturn_trainer.py

@@ -0,0 +1,138 @@
+"""GRPO trainer support for group-relative per-turn credit assignment."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import cast
+
+import torch
+from trl import GRPOTrainer
+
+
+@dataclass(frozen=True)
+class TurnCreditRow:
+    spans: tuple[tuple[int, int], ...]
+    turns: tuple[float, ...] | None
+
+
+def build_per_turn_advantages(
+    turn_spans_per_completion: list[list[tuple[int, int]]],
+    turn_rewards_per_completion: list[list[float] | None],
+    num_generations: int,
+    completion_len: int,
+    *,
+    episode_advantages: torch.Tensor,
+) -> torch.Tensor:
+    """Build token-aligned advantages from per-turn rewards in consecutive GRPO groups."""
+    batch_size = len(turn_spans_per_completion)
+    if len(turn_rewards_per_completion) != batch_size:
+        raise ValueError("turn span and reward row counts must match")
+    if num_generations <= 0 or batch_size % num_generations != 0:
+        raise ValueError("batch size must be divisible by num_generations")
+    if completion_len < 0:
+        raise ValueError("completion_len must be non-negative")
+    if episode_advantages.dim() != 1 or episode_advantages.numel() != batch_size:
+        raise ValueError("episode_advantages must have shape [B]")
+
+    rows: list[TurnCreditRow] = []
+    for row_index, (spans, turns) in enumerate(
+        zip(turn_spans_per_completion, turn_rewards_per_completion, strict=True)
+    ):
+        normalized_spans = tuple((int(start), int(end)) for start, end in spans)
+        for start, end in normalized_spans:
+            if not 0 <= start <= end <= completion_len:
+                raise ValueError(
+                    f"turn span [{start}, {end}) for row {row_index} exceeds completion width "
+                    f"{completion_len}"
+                )
+        rows.append(
+            TurnCreditRow(
+                spans=normalized_spans,
+                turns=None if turns is None else tuple(turns),
+            )
+        )
+
+    advantages = episode_advantages.new_zeros((batch_size, completion_len))
+    for group_start in range(0, batch_size, num_generations):
+        group_end = group_start + num_generations
+        group = rows[group_start:group_end]
+        if any(row.turns is None for row in group):
+            for row_index in range(group_start, group_end):
+                completion_end = rows[row_index].spans[-1][1] if rows[row_index].spans else 0
+                advantages[row_index, :completion_end] = episode_advantages[row_index]
+            continue
+
+        max_turns = max(len(row.turns or ()) for row in group)
+        for turn_index in range(max_turns):
+            member_indexes = [
+                row_index
+                for row_index in range(group_start, group_end)
+                if turn_index < len(rows[row_index].turns or ())
+                and rows[row_index].spans[turn_index][1] > rows[row_index].spans[turn_index][0]
+            ]
+            if not member_indexes:
+                # every member's span for this turn is zero-width (no emitted tokens); an
+                # empty turn contributes no advantage and must not skew the group baseline.
+                continue
+            mean_reward = sum(
+                cast(tuple[float, ...], rows[row_index].turns)[turn_index]
+                for row_index in member_indexes
+            ) / len(member_indexes)
+            for row_index in member_indexes:
+                row = rows[row_index]
+                reward = cast(tuple[float, ...], row.turns)[turn_index]
+                start, end = row.spans[turn_index]
+                advantages[row_index, start:end] = reward - mean_reward
+
+    if not bool(torch.isfinite(advantages).all()):
+        raise ValueError("per-turn advantages must be finite")
+    return advantages
+
+
+class GRPOPerTurnTrainer(GRPOTrainer):
+    """Replace scalar GRPO advantages with aligned per-turn advantages when supplied."""
+
+    def _generate_and_score_completions(self, inputs: list[dict[str, object]]) -> dict[str, object]:
+        output = cast(dict[str, object], super()._generate_and_score_completions(inputs))
+        turn_rewards = cast(
+            list[list[float] | None],
+            [item.get("turn_rewards") for item in inputs],
+        )
+        if not any(rewards is not None for rewards in turn_rewards):
+            return output
+        if self.accelerator.num_processes > 1:
+            raise NotImplementedError(
+                "per-turn GRPO advantages currently support single-process training only; "
+                "distributed group centering requires gather-aligned turn metadata"
+            )
+
+        turn_spans = cast(
+            list[list[tuple[int, int]] | None],
+            [item.get("turn_spans") for item in inputs],
+        )
+        if any(spans is None for spans in turn_spans):
+            raise ValueError("per-turn rollout rows must all include turn_spans")
+        aligned_turn_spans = cast(list[list[tuple[int, int]]], turn_spans)
+        scalar_advantages = cast(torch.Tensor, output["advantages"])
+        if scalar_advantages.dim() != 1:
+            raise ValueError("expected TRL scalar advantages with shape [B]")
+        completion_ids = cast(torch.Tensor, output["completion_ids"])
+        batch_size, completion_len = completion_ids.shape
+        if len(inputs) != batch_size:
+            raise ValueError(
+                f"per-turn metadata has {len(inputs)} row(s) for output batch size {batch_size}"
+            )
+
+        output["advantages"] = build_per_turn_advantages(
+            aligned_turn_spans,
+            turn_rewards,
+            num_generations=(
+                self.num_generations if self.model.training else self.num_generations_eval
+            ),
+            completion_len=completion_len,
+            episode_advantages=scalar_advantages,
+        ).to(device=scalar_advantages.device, dtype=scalar_advantages.dtype)
+        if not getattr(self, "_per_turn_credit_logged", False):
+            print("[rl] multi-turn per-turn group-relative credit is active")
+            self._per_turn_credit_logged = True
+        return output

{freesolo_flash_dev-1.0.4 → freesolo_flash_dev-1.0.5}/flash/engine/worker/heartbeat.py

@@ -9,7 +9,9 @@ from __future__ import annotations
 import contextlib
 import faulthandler
 import json
+import math
 import os
+import re
 import sys
 import threading
 import time
@@ -235,7 +237,60 @@ def _maybe_attach_gpu_diag(payload: dict, last_gpu_diag_at: float, now: float) -
     return last_gpu_diag_at
 
 
-def make_reward_heartbeat_callback():
+_REWARD_METRIC_NAME_DISALLOWED = re.compile(r"[^A-Za-z0-9_.-]")
+_REWARD_METRIC_RESERVED_NAMES = frozenset(
+    {
+        "reward",
+        "reward_last",
+        "step",
+        "epoch",
+        "loss",
+        "grad_norm",
+        "learning_rate",
+        "stage",
+        "gpu",
+        "diag",
+    }
+)
+_REWARD_METRIC_LIMIT = 12
+# names TRAINING.md tells users to judge on: never dropped by the alphabetical cap.
+_REWARD_METRIC_PRIORITY_NAMES = ("success",)
+
+
+def _bounded_reward_metrics(metrics) -> dict[str, float]:
+    if not isinstance(metrics, dict):
+        return {}
+    surviving: dict[str, float] = {}
+    for name, value in metrics.items():
+        sanitized_name = _REWARD_METRIC_NAME_DISALLOWED.sub("", str(name))[:64]
+        if not sanitized_name or sanitized_name in _REWARD_METRIC_RESERVED_NAMES:
+            continue
+        try:
+            score = float(value)
+        except (TypeError, ValueError):
+            continue
+        if not math.isfinite(score):
+            continue
+        # distinct source names that sanitize to the same key must not silently overwrite each
+        # other; disambiguate with a numeric suffix (kept within the 64-char bound, allowed chars).
+        unique_name = sanitized_name
+        suffix = 2
+        while unique_name in surviving:
+            tail = f"_{suffix}"
+            unique_name = sanitized_name[: 64 - len(tail)] + tail
+            suffix += 1
+        surviving[unique_name] = score
+    if len(surviving) <= _REWARD_METRIC_LIMIT:
+        return dict(sorted(surviving.items()))
+    # the cap must not drop metrics users are told to judge on (e.g. success); keep those first,
+    # then fill the remaining slots alphabetically.
+    priority = [n for n in _REWARD_METRIC_PRIORITY_NAMES if n in surviving]
+    remaining = max(0, _REWARD_METRIC_LIMIT - len(priority))
+    rest = sorted(n for n in surviving if n not in priority)[:remaining]
+    return {n: surviving[n] for n in sorted(priority + rest)}
+
+
+def make_reward_heartbeat_callback(reward_metrics=None):
     """Return a TRL callback that streams per-step reward to the HF heartbeat channel."""
     from transformers import TrainerCallback
 
@@ -243,6 +298,12 @@ def make_reward_heartbeat_callback():
         def __init__(self):
             self.reward_history = []
             self.last_gpu_diag_at = 0.0
+            self.latest_reward_metrics: dict[str, float] = {}
+
+        def latest_fields(self) -> dict:
+            if not self.latest_reward_metrics:
+                return {}
+            return {"reward_metrics": dict(self.latest_reward_metrics)}
 
         def on_log(self, args, state, control, logs=None, **kwargs):
             if not logs:
@@ -261,6 +322,10 @@ def make_reward_heartbeat_callback():
                 "reward": r,
                 "reward_last": self.reward_history[-8:],
             }
+            latest_metrics = reward_metrics() if callable(reward_metrics) else reward_metrics
+            self.latest_reward_metrics = _bounded_reward_metrics(latest_metrics)
+            if self.latest_reward_metrics:
+                payload["reward_metrics"] = dict(self.latest_reward_metrics)
             now = time.monotonic()
             self.last_gpu_diag_at = _maybe_attach_gpu_diag(payload, self.last_gpu_diag_at, now)
             _w.heartbeat("rl_step", **payload)