Buckets:
| """Streaming GRPO metrics and identity-bound validation artifacts.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| import time | |
| from collections import Counter | |
| from pathlib import Path | |
| from shared.artifacts import file_sha256, write_json | |
| from .config import RLConfig | |
| from .distributed import DistributedContext | |
| from .pipeline import protocol_limits, score_prompt_group, score_trajectory | |
| from .rollout import GRPORolloutController | |
| from .self_score import BehaviorSelfScorer | |
| from .uniar_backend import UniARRLBackend | |
| from .scheduling import map_rank_tasks | |
| from .validation_artifacts import read_sample, save_sample, validation_directory | |
| VALIDATION_SCHEMA = "spatial_mmcot_grpo_validation_v3" | |
| def validation_rollouts_per_sample(config: RLConfig) -> int: | |
| return config.evaluation.rollouts_per_sample or config.rollout.group_size | |
| def make_backend(actor, processor, config: RLConfig) -> UniARRLBackend: | |
| return UniARRLBackend( | |
| actor, | |
| processor, | |
| text_temperature=config.rollout.text_temperature, | |
| visual_temperature=config.rollout.visual_temperature, | |
| max_text_tokens=config.rollout.max_text_tokens, | |
| max_scaffold_tokens=config.rollout.max_scaffold_tokens, | |
| min_grid_side=config.rollout.min_grid_side, | |
| max_grid_side=config.rollout.max_grid_side, | |
| context_limit=config.data.max_sequence_tokens, | |
| cache_autocast=config.rollout.cache_autocast, | |
| ) | |
| def make_self_scorer(actor, config: RLConfig): | |
| if not config.reward.information_gain.enabled: | |
| return None | |
| return BehaviorSelfScorer( | |
| actor, | |
| text_temperature=config.rollout.text_temperature, | |
| visual_temperature=config.rollout.visual_temperature, | |
| probability_floor=config.reward.information_gain.p_min, | |
| ) | |
| def _group_signature(trajectory) -> str: | |
| digest = hashlib.sha256(trajectory.input_ids.numpy().tobytes()) | |
| digest.update(trajectory.sampled_visual_codes.numpy().tobytes()) | |
| return digest.hexdigest() | |
| class GroupMetricsAccumulator: | |
| """Reduce prompt groups without retaining their trajectory tensors.""" | |
| def __init__(self) -> None: | |
| self.prompts = 0 | |
| self.rollouts = 0 | |
| self.protocol_valid = 0 | |
| self.accuracy_correct = 0 | |
| self.reward_sum = 0.0 | |
| self.reward_min = math.inf | |
| self.reward_max = -math.inf | |
| self.ig_sum = 0.0 | |
| self.zero_advantage_groups = 0 | |
| self.unique_trajectories = 0 | |
| self.generated_images = 0 | |
| self.text_actions = 0 | |
| self.visual_steps = 0 | |
| self.failure_reasons = Counter() | |
| def add(self, group) -> None: | |
| self._add( | |
| tuple(group.trajectories), | |
| zero_advantage=group.zero_advantage_reason is not None, | |
| ) | |
| def add_scored(self, values) -> None: | |
| """Add evaluation trajectories without constructing GRPO advantages.""" | |
| self._add(tuple(values), zero_advantage=False) | |
| def _add(self, values, *, zero_advantage: bool) -> None: | |
| if not values: | |
| raise ValueError("GRPO metric group has no trajectories") | |
| self.prompts += 1 | |
| self.rollouts += len(values) | |
| self.zero_advantage_groups += zero_advantage | |
| self.unique_trajectories += len( | |
| {_group_signature(value.trajectory) for value in values} | |
| ) | |
| for value in values: | |
| reward = value.reward.total | |
| self.protocol_valid += value.protocol.valid | |
| self.accuracy_correct += value.verifier_correct | |
| self.reward_sum += reward | |
| self.reward_min = min(self.reward_min, reward) | |
| self.reward_max = max(self.reward_max, reward) | |
| self.ig_sum += value.reward.ig | |
| self.generated_images += sum( | |
| span.start >= value.trajectory.prefix_length | |
| for span in value.trajectory.visual_sequence_spans | |
| ) | |
| self.text_actions += len(value.trajectory.text_action_positions) | |
| self.visual_steps += len(value.trajectory.visual_action_positions) | |
| self.failure_reasons[value.protocol.failure_reason or "valid"] += 1 | |
| def result(self) -> dict[str, float | int]: | |
| if not self.rollouts: | |
| raise ValueError("GRPO metrics require at least one rollout") | |
| result = { | |
| "prompts": self.prompts, | |
| "rollouts": self.rollouts, | |
| "protocol_valid": self.protocol_valid, | |
| "accuracy_correct": self.accuracy_correct, | |
| "reward_mean": self.reward_sum / self.rollouts, | |
| "reward_min": self.reward_min, | |
| "reward_max": self.reward_max, | |
| "information_gain_reward_mean": self.ig_sum / self.rollouts, | |
| "zero_advantage_groups": self.zero_advantage_groups, | |
| "unique_trajectories": self.unique_trajectories, | |
| "stochastic_diversity_fraction": ( | |
| self.unique_trajectories / self.rollouts | |
| ), | |
| "generated_images_mean": self.generated_images / self.rollouts, | |
| "text_actions": self.text_actions, | |
| "visual_steps": self.visual_steps, | |
| } | |
| result.update( | |
| { | |
| f"protocol_failure/{key}": count | |
| for key, count in self.failure_reasons.items() | |
| } | |
| ) | |
| return result | |
| def state_dict(self) -> dict: | |
| return { | |
| name: (dict(value) if isinstance(value, Counter) else value) | |
| for name, value in self.__dict__.items() | |
| } | |
| def merge_state_dict(self, state: dict) -> None: | |
| additive = ( | |
| "prompts", | |
| "rollouts", | |
| "protocol_valid", | |
| "accuracy_correct", | |
| "reward_sum", | |
| "ig_sum", | |
| "zero_advantage_groups", | |
| "unique_trajectories", | |
| "generated_images", | |
| "text_actions", | |
| "visual_steps", | |
| ) | |
| for name in additive: | |
| setattr(self, name, getattr(self, name) + state[name]) | |
| self.reward_min = min(self.reward_min, state["reward_min"]) | |
| self.reward_max = max(self.reward_max, state["reward_max"]) | |
| self.failure_reasons.update(state["failure_reasons"]) | |
| def group_metrics(groups) -> dict[str, float | int]: | |
| accumulator = GroupMetricsAccumulator() | |
| for group in groups: | |
| accumulator.add(group) | |
| return accumulator.result() | |
| def evaluate( | |
| *, | |
| actor, | |
| processor, | |
| cache, | |
| indices: tuple[int, ...], | |
| config: RLConfig, | |
| ) -> dict: | |
| selected = indices[: config.evaluation.samples] | |
| return _evaluate_entries( | |
| actor=actor, | |
| processor=processor, | |
| cache=cache, | |
| entries=tuple(enumerate(selected)), | |
| total=len(selected), | |
| config=config, | |
| rank=0, | |
| ).result() | |
| def _evaluate_entries( | |
| *, | |
| actor, | |
| processor, | |
| cache, | |
| entries: tuple[tuple[int, int], ...], | |
| total: int, | |
| config: RLConfig, | |
| rank: int, | |
| ) -> GroupMetricsAccumulator: | |
| actor.eval() | |
| controller = GRPORolloutController( | |
| make_backend(actor, processor, config), protocol_limits(config) | |
| ) | |
| scorer = make_self_scorer(actor, config) | |
| metrics = GroupMetricsAccumulator() | |
| started = time.monotonic() | |
| for completed, (position, index) in enumerate(entries, 1): | |
| evaluation_rollouts = config.evaluation.rollouts_per_sample | |
| if evaluation_rollouts is None: | |
| group = controller.collect_group( | |
| index, | |
| cache.load(index), | |
| group_size=config.rollout.group_size, | |
| seed=config.run.seed + 1_000_003, | |
| update=0, | |
| prompt_slot=position, | |
| ) | |
| metrics.add( | |
| score_prompt_group( | |
| group, | |
| tokenizer=processor.tokenizer, | |
| config=config, | |
| self_scorer=scorer, | |
| ) | |
| ) | |
| else: | |
| trajectories = controller.collect_group_members( | |
| index, | |
| cache.load(index), | |
| group_size=evaluation_rollouts, | |
| group_indices=range(evaluation_rollouts), | |
| seed=config.run.seed + 1_000_003, | |
| update=0, | |
| prompt_slot=position, | |
| ) | |
| metrics.add_scored( | |
| score_trajectory( | |
| trajectory, | |
| tokenizer=processor.tokenizer, | |
| config=config, | |
| self_scorer=scorer, | |
| ) | |
| for trajectory in trajectories | |
| ) | |
| print( | |
| "RL_VALIDATION_PROGRESS " | |
| f"sample={position + 1}/{total} rank={rank} " | |
| f"local={completed}/{len(entries)} " | |
| f"seconds={time.monotonic() - started:.1f}", | |
| flush=True, | |
| ) | |
| return metrics | |
| def validation_identity( | |
| *, | |
| cache, | |
| indices: tuple[int, ...], | |
| config: RLConfig, | |
| provenance: dict, | |
| sampling: dict, | |
| self_scoring: dict, | |
| phase: str, | |
| successful_update: int, | |
| actor_identity: dict, | |
| ) -> dict: | |
| selected = indices[: config.evaluation.samples] | |
| sample_ids = [cache.sample_id(index) for index in selected] | |
| return { | |
| "phase": phase, | |
| "successful_update": successful_update, | |
| "actor": actor_identity, | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance, | |
| "sampling_fingerprint": sampling["fingerprint"], | |
| "self_scoring_fingerprint": self_scoring["fingerprint"], | |
| "seed": config.run.seed + 1_000_003, | |
| "training_group_size": config.rollout.group_size, | |
| "rollouts_per_sample": validation_rollouts_per_sample(config), | |
| "expected_samples": len(selected), | |
| "sample_ids": sample_ids, | |
| "sample_ids_sha256": hashlib.sha256( | |
| "\n".join(sample_ids).encode() | |
| ).hexdigest(), | |
| } | |
| def checkpoint_actor_identity(run_dir: Path, checkpoint: Path) -> dict: | |
| resolved_run = run_dir.resolve() | |
| resolved = checkpoint.resolve(strict=True) | |
| if not resolved.is_relative_to(resolved_run): | |
| raise ValueError("validation checkpoint is outside the GRPO run") | |
| return { | |
| "kind": "rl_checkpoint_manifest", | |
| "path": str(resolved.relative_to(resolved_run)), | |
| "manifest_sha256": file_sha256(resolved / "manifest.json"), | |
| } | |
| def _completed_validation_path( | |
| run_dir: Path, stem: str, identity: dict | |
| ) -> Path | None: | |
| candidates = sorted(run_dir.glob(f"{stem}*.json")) | |
| for path in reversed(candidates): | |
| try: | |
| payload = json.loads(path.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| continue | |
| metrics = payload.get("metrics") if isinstance(payload, dict) else None | |
| if ( | |
| isinstance(payload, dict) | |
| and isinstance(metrics, dict) | |
| and payload.get("schema") == VALIDATION_SCHEMA | |
| and payload.get("identity") == identity | |
| and payload.get("completed_samples") == metrics.get("prompts") | |
| and ( | |
| "expected_samples" not in identity | |
| or payload.get("completed_samples") == identity["expected_samples"] | |
| ) | |
| ): | |
| return path | |
| return None | |
| def _unique_path(path: Path) -> Path: | |
| if not path.exists(): | |
| return path | |
| for retry in range(1, 10_000): | |
| candidate = path.with_name(f"{path.stem}-retry-{retry:04d}{path.suffix}") | |
| if not candidate.exists(): | |
| return candidate | |
| raise RuntimeError(f"too many preserved RL retry artifacts for {path.name}") | |
| def ensure_validation( | |
| *, | |
| run_dir: Path, | |
| stem: str, | |
| identity: dict, | |
| actor, | |
| processor, | |
| cache, | |
| indices: tuple[int, ...], | |
| config: RLConfig, | |
| parallel: DistributedContext | None = None, | |
| reuse_from: Path | None = None, | |
| ) -> Path: | |
| parallel = parallel or DistributedContext() | |
| existing = parallel.main_call(lambda: _completed_validation_path(run_dir, stem, identity)) | |
| if existing is not None: | |
| return existing | |
| if reuse_from is not None: | |
| from .resume import reuse_parent_validation | |
| reused = parallel.main_call(lambda: reuse_parent_validation(reuse_from, run_dir=run_dir, stem=stem, identity=identity)) | |
| if reused is not None: | |
| print(f"RL_VALIDATION_REUSED {reused}", flush=True) | |
| return reused | |
| selected = indices[: config.evaluation.samples] | |
| root = validation_directory(run_dir, stem, identity) | |
| pending = parallel.main_call(lambda: [position for position, index in enumerate(selected) | |
| if read_sample(root, position, cache.sample_id(index), identity) is None]) | |
| actor.eval() | |
| controller = GRPORolloutController(make_backend(actor, processor, config), protocol_limits(config)) | |
| scorer = make_self_scorer(actor, config) | |
| print(f"RL_VALIDATION_RESUME phase={stem} completed={len(selected) - len(pending)} " | |
| f"pending={len(pending)} rank={parallel.rank}", flush=True) | |
| def evaluate_sample(position): | |
| index = selected[position] | |
| count = validation_rollouts_per_sample(config) | |
| trajectories = controller.collect_group_members( | |
| index, cache.load(index), group_size=count, group_indices=range(count), | |
| seed=config.run.seed + 1_000_003, update=0, prompt_slot=position, | |
| ) | |
| scored = tuple(score_trajectory(trajectory, tokenizer=processor.tokenizer, | |
| config=config, self_scorer=scorer) | |
| for trajectory in trajectories) | |
| metrics = GroupMetricsAccumulator() | |
| metrics.add_scored(scored) | |
| save_sample(root, position, cache.sample_id(index), identity, scored, | |
| metrics.state_dict(), processor.tokenizer) | |
| print(f"RL_VALIDATION_PROGRESS sample={position + 1}/{len(selected)} " | |
| f"rank={parallel.rank} committed=true", flush=True) | |
| map_rank_tasks(pending, evaluate_sample, directory=run_dir, parallel=parallel, | |
| scheduling=config.evaluation.scheduling) | |
| def publish(): | |
| metrics_accumulator = GroupMetricsAccumulator() | |
| for position, index in enumerate(selected): | |
| state = read_sample(root, position, cache.sample_id(index), identity) | |
| if state is None: | |
| raise RuntimeError("validation queue omitted a sample") | |
| metrics_accumulator.merge_state_dict(state) | |
| metrics = metrics_accumulator.result() | |
| if metrics["prompts"] != identity["expected_samples"]: | |
| raise RuntimeError("GRPO validation completed an unexpected sample count") | |
| if ( | |
| metrics["rollouts"] | |
| != identity["expected_samples"] * identity["rollouts_per_sample"] | |
| ): | |
| raise RuntimeError("GRPO validation completed an unexpected rollout count") | |
| path = _unique_path(run_dir / f"{stem}.json") | |
| write_json( | |
| path, | |
| { | |
| "schema": VALIDATION_SCHEMA, | |
| "identity": identity, | |
| "completed_samples": metrics["prompts"], | |
| "metrics": metrics, | |
| }, | |
| ) | |
| return path | |
| return parallel.main_call(publish) | |
| __all__ = [ | |
| "checkpoint_actor_identity", | |
| "ensure_validation", | |
| "evaluate", | |
| "group_metrics", | |
| "GroupMetricsAccumulator", | |
| "make_backend", | |
| "make_self_scorer", | |
| "validation_rollouts_per_sample", | |
| "validation_identity", | |
| ] | |
Xet Storage Details
- Size:
- 15.6 kB
- Xet hash:
- 595a10ed2682ae0c859b76e59ad2ca7e8acb8cb4ab01258a2904ed8ed43c64d0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.