adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/common/env_grpo.py

@@ -0,0 +1,361 @@
+import os
+from dataclasses import dataclass
+from typing import Callable, Sequence
+
+import numpy as np
+
+from adaptive_harmony import (
+    CosineScheduler,
+    DataSet,
+    JobNotifier,
+    Logger,
+    StageNotifier,
+    StringThread,
+    TokenizedThread,
+    TrainingModel,
+)
+from adaptive_harmony.common import RecipeCallback
+from adaptive_harmony.common.checkpointing import CheckpointManager
+from adaptive_harmony.core.utils import async_map, async_map_batch, get_minibatches, hash_hyperparams, log_args
+from adaptive_harmony.environment import EnvironmentFactory, TrajectoryScore
+from adaptive_harmony.metric_logger import StdoutLogger
+
+
+def compute_advantages(
+    scores: list[TrajectoryScore],
+    logprobs: list[list[float]],
+    samples: list[TokenizedThread],
+    num_generated_turns: list[int],
+) -> list[list[float]]:
+    def get_assistant_lengths(samples: list[TokenizedThread], num_generated_turns: list[int]) -> list[list[int]]:
+        # here the +1 is because we have a loss weight on the EOD token of a turn, which is not represented when you look at the tokenized
+        # Keep only the last num_generated_turns assistant turns (those with weight>0) since logprobs_per_token only returns logprobs for them
+        return [
+            [len(turn.content) + 1 for turn in sample.get_turns() if turn.role == "assistant"][-num_gen:]
+            for sample, num_gen in zip(samples, num_generated_turns)
+        ]
+
+    # FROM https://arxiv.org/pdf/2402.03300 -> Process Supervision RL with GRPO
+    # HERE PADDING DOES NOT PLAYS A ROLE IN ADVANTAGE COMPUTATION. SINCE nan are ignored.
+
+    mapped_scores = [[turn_score.score for turn_score in score.scores] for score in scores]
+
+    max_len = max(map(len, mapped_scores))
+
+    # pad with np.nan instead of 0
+    all_scores = np.full((len(mapped_scores), max_len), np.nan)
+    for i, s in enumerate(mapped_scores):
+        all_scores[i, : len(s)] = s
+
+    # nan-aware mean and std
+    mean = np.nanmean(all_scores)
+    std = np.nanstd(all_scores) + 1e-8
+
+    normalized_rewards = (all_scores - mean) / std
+
+    # cumulative sum per row, ignoring nans
+    score_level_advantage = np.where(
+        np.isnan(normalized_rewards),
+        np.nan,
+        np.cumsum(np.nan_to_num(normalized_rewards)[:, ::-1], axis=1)[:, ::-1],
+    )
+
+    turn_level_advantage = [
+        np.repeat(
+            adv[: len(score.scores)],
+            [turn_score.num_assistant_turns for turn_score in score.scores],
+        )
+        for adv, score in zip(score_level_advantage, scores)
+    ]
+
+    assistant_lengths = get_assistant_lengths(samples, num_generated_turns)
+    assert all([len(lp) == sum(al) for lp, al in zip(logprobs, assistant_lengths)])
+
+    token_level_advantage = [np.repeat(adv, al).tolist() for adv, al in zip(turn_level_advantage, assistant_lengths)]
+
+    return token_level_advantage
+
+
+@dataclass
+class Sample:
+    sample: TokenizedThread
+    logprobs: list[float]
+    ref_logprobs: list[float]
+    advantage: list[float]
+    kl_div: list[float]
+    # for logging
+    score: float
+    gen_len: float
+
+
+ENVGRPO_HYPERPARAMS = {
+    "max_num_grpo_steps",
+    "completions_per_sample",
+    "lr",
+    "lr_scheduler",
+    "samples_per_batch",
+    "samples_per_mini_batch",
+    "mini_epochs_per_batch",
+    "max_grad_norm",
+    "clip_range",
+    "kl_beta",
+    "weight_decays",
+    "skip_nan_gradients",
+}
+
+
+class ENVGRPO:
+    @log_args
+    @hash_hyperparams(include=ENVGRPO_HYPERPARAMS)
+    def __init__(
+        self,
+        dataset: list[StringThread],
+        model: TrainingModel,
+        environment_factory: EnvironmentFactory,
+        logger: Logger = StdoutLogger(),
+        stage_notifier: StageNotifier = JobNotifier().stage_notifier("ENVGRPO Training"),
+        callbacks: Sequence[RecipeCallback] = [],
+        validation_dataset: list[StringThread] | None = None,
+        validation_frequency: float = 0.2,
+        max_num_grpo_steps: int | None = None,
+        completions_per_sample=8,
+        lr: float = 7.5e-7,
+        lr_scheduler: Callable[[float], float] | None = None,
+        samples_per_batch=128,
+        samples_per_mini_batch=128,
+        mini_epochs_per_batch=1,
+        max_grad_norm=1.0,
+        clip_range=0.1,
+        kl_beta=0.1,
+        weight_decays: float = 0.0,
+        skip_nan_gradients: bool = False,
+        restart_from_checkpoint: str | None = None,
+        checkpoint_frequency: float = 0.2,
+    ):
+        # Core components
+        self.dataset = DataSet(dataset, allow_looping=True)
+        self.model = model
+        self.logger = logger
+        self.stage_notifier = stage_notifier
+        self.sample_index_counter = 0
+        self.skip_nan_gradients = skip_nan_gradients
+        # Validation data/params
+        self.validation_dataset = validation_dataset
+        self.validation_frequency = validation_frequency
+        self.last_validation_percentage = -1.0  # Validation will run before training starts
+        # GRPO HP's
+        self.max_num_batches = max_num_grpo_steps
+        self.completions_per_sample = completions_per_sample
+        self.lr_schedule = lr_scheduler or CosineScheduler(lr)
+        self.samples_per_batch = samples_per_batch // completions_per_sample
+        self.samples_per_mini_batch = samples_per_mini_batch
+        self.total_num_samples = (
+            self.max_num_batches * self.samples_per_batch if self.max_num_batches else len(self.dataset)
+        )
+        self.max_grad_norm = max_grad_norm
+        self.environment_factory = environment_factory
+        self.clip_range = clip_range
+        self.kl_beta = kl_beta
+        self.weight_decays = weight_decays
+        self.mini_epochs_per_batch = mini_epochs_per_batch
+
+        self.num_batches_processed = 0
+        self.callbacks = callbacks
+
+        self.checkpoint_manager = CheckpointManager(
+            recipe_name="ENVGRPO",
+            dataset=self.dataset,
+            threads_dataset=dataset,
+            callbacks=self.callbacks,
+            hyperparams_hash=self._hyperparams_hash,  # type: ignore
+            job_id=os.environ.get("HARMONY_JOB_ID"),
+            checkpoint_frequency=checkpoint_frequency,
+            restart_from_checkpoint=restart_from_checkpoint,
+        )
+
+    @property
+    def training_completion_percentage(self):
+        return (
+            self.dataset.completion_percentage()
+            if self.max_num_batches is None
+            else min(self.num_batches_processed / self.max_num_batches, 1.0)
+        )
+
+    async def gen_data(self, sample: StringThread) -> list[Sample]:
+        async def generate_trajectory(
+            prompt: StringThread,
+        ) -> tuple[TokenizedThread, TrajectoryScore, int]:
+            # this create the environment for the first turn.
+            environment = self.environment_factory.create_environment(prompt.metadata)
+            prompt = await environment.bootstrap_prompt(prompt)
+
+            # Count assistant turns in the context (before generation)
+            nb_context_assistant_turns = sum(1 for turn in prompt.get_turns() if turn.role == "assistant")
+
+            string_trajectory = await self.model.generate(prompt)  # generate the first response from the agent.
+            num_generated_turns = 1
+            # we loop until the environment returns a score.
+            # notice how the environment can return a score or a tool or user response.
+            while not isinstance(
+                environment_response := await environment.react_to(string_trajectory),
+                TrajectoryScore,
+            ):
+                for env_role, env_content in environment_response:
+                    if not isinstance(env_content, str):
+                        raise ValueError(f"env_content should be a str, got {env_content}")
+                    if env_role == "user":
+                        string_trajectory = string_trajectory.user(env_content)
+                    elif env_role == "tool":
+                        string_trajectory = string_trajectory.tool(env_content)
+                    else:
+                        raise ValueError
+                string_trajectory = await self.model.generate(string_trajectory)
+                num_generated_turns += 1
+
+            tokenized_trajectory = (
+                await self.model.tokenize_thread(string_trajectory)
+            ).with_weight_assistant_turns_from_index(nb_context_assistant_turns)
+
+            return tokenized_trajectory, environment_response, num_generated_turns
+
+        assert self.model_ref is not None, "Calling `gen_data` before reference model has been set"
+
+        trajs_and_scores = await async_map(generate_trajectory, [sample] * self.completions_per_sample)
+        all_samples = [traj for traj, _, _ in trajs_and_scores]
+        num_generated_turns_list = [num_turns for _, _, num_turns in trajs_and_scores]
+        logprobs = await async_map(self.model.logprobs_per_token, all_samples)
+        ref_logprobs = await async_map(self.model_ref.logprobs_per_token, all_samples)
+
+        all_trajectory_scores = [score for _, score, _ in trajs_and_scores]
+        advantages = compute_advantages(all_trajectory_scores, logprobs, all_samples, num_generated_turns_list)
+
+        kl = [
+            (np.array(lp, dtype=np.float32) - np.array(ref_lp, dtype=np.float32)).tolist()
+            for lp, ref_lp in zip(logprobs, ref_logprobs)
+        ]
+
+        samples = []
+        for i in range(len(logprobs)):
+            samples.append(
+                Sample(
+                    sample=all_samples[i],
+                    logprobs=logprobs[i],
+                    ref_logprobs=ref_logprobs[i],
+                    advantage=advantages[i],
+                    kl_div=kl[i],
+                    score=all_trajectory_scores[i].cumulative_score,
+                    gen_len=all_samples[i].len_last_turn(),
+                )
+            )
+        return samples
+
+    async def train_sample(self, sample: Sample):
+        await self.model.train_grpo(
+            sample.sample,
+            sample.logprobs,
+            sample.ref_logprobs,
+            sample.advantage,
+            self.clip_range,
+            self.kl_beta,
+        )
+
+    async def _recipe_specific_checkpoint_loading(self, checkpoint_data: dict) -> None:
+        self.num_batches_processed = checkpoint_data["num_batches_processed"]
+
+        model_checkpoint_name = checkpoint_data["model_checkpoint_name"]
+        model_checkpoint_name = await self.model.load(f"model_registry://{model_checkpoint_name}")
+
+        self.last_validation_percentage = checkpoint_data.get("last_validation_percentage", -1.0)
+        self.sample_index_counter = checkpoint_data.get("sample_index_counter", 0)
+
+        self.model.set_optim_step(checkpoint_data["optim_step"])
+
+    async def _recipe_specific_checkpoint_saving(self) -> dict:
+        progress_pct = int(self.training_completion_percentage * 100)
+        model_checkpoint_name = f"checkpoint-{self.checkpoint_manager.job_id}-{progress_pct}-policy"
+        await self.model.save(model_checkpoint_name, inference_only=False)
+
+        return {
+            "num_batches_processed": self.num_batches_processed,
+            "model_checkpoint_name": model_checkpoint_name,
+            "last_validation_percentage": self.last_validation_percentage,
+            "sample_index_counter": self.sample_index_counter,
+            "optim_step": self.model.get_optim_step(),
+        }
+
+    async def run(self):
+        self.model_ref = await self.model.clone_inf()
+        await self.checkpoint_manager.maybe_restore_checkpoint(self._recipe_specific_checkpoint_loading)
+
+        self.stage_notifier.report_progress(
+            tot_num_samples=self.total_num_samples,
+            processed_num_samples=self.num_batches_processed * self.samples_per_batch,
+            monitoring_link=self.logger.training_monitoring_link,
+        )
+
+        while self.training_completion_percentage < 1.0:
+            self.num_batches_processed += 1
+
+            for callback in self.callbacks:
+                if logs := await callback.maybe_call(self.training_completion_percentage):
+                    self.logger(logs)
+
+            # Generate training samples
+            data = await async_map_batch(self.gen_data, self.dataset, self.samples_per_batch)
+
+            scorer_logs = {}
+            for key, value in self.environment_factory.get_logs(clear=True).items():
+                if "/" not in key:
+                    key = f"environment/{key}"
+                scorer_logs[key] = value
+            batch_logs = {
+                **scorer_logs,
+                **self.get_train_batch_logs(data),
+            }
+
+            current_lr = self.lr_schedule(self.training_completion_percentage)
+
+            # Train on generated samples
+            flattened_data = sum([inner_list for inner_list in data], start=[])
+            minibatches = get_minibatches(flattened_data, self.samples_per_mini_batch, self.mini_epochs_per_batch)
+            for idx, mini_batch in enumerate(minibatches):
+                await async_map(self.train_sample, mini_batch)
+                optim_logs = await self.model.optim_step(
+                    current_lr, wd=0.0, max_grad_norm=self.max_grad_norm, skip_nan_gradients=self.skip_nan_gradients
+                )
+                if idx == len(minibatches) - 1:
+                    # only log tables and full batch-related logs on the final minibatch
+                    self.logger(optim_logs | batch_logs)
+                else:
+                    self.logger(optim_logs | dict(completion_percentage=self.training_completion_percentage))
+
+            self.stage_notifier.report_progress(
+                tot_num_samples=self.total_num_samples,
+                processed_num_samples=self.num_batches_processed * self.samples_per_batch,
+                monitoring_link=self.logger.training_monitoring_link,
+            )
+
+            if await self.checkpoint_manager.maybe_checkpoint(
+                self.training_completion_percentage, self._recipe_specific_checkpoint_saving
+            ):
+                break
+
+    def get_train_batch_logs(self, data: list[list[Sample]]) -> dict:
+        return {
+            **dict(
+                completion_percentage=self.training_completion_percentage,
+                score_mean=np.mean([[sample.score for sample in batch] for batch in data]).item(),
+                percentage_no_advantages=np.mean(
+                    [all(sample.advantage == batch[0].advantage for sample in batch) for batch in data]
+                ).item(),
+                score_std=np.std([[sample.score for sample in batch] for batch in data]).item(),
+                kl_div=np.mean([[np.mean(sample.kl_div) for sample in batch] for batch in data]).item(),
+                generation_length=np.mean([np.mean([sample.gen_len for sample in batch]) for batch in data]).item(),
+                logprobs=np.mean(
+                    np.concatenate([np.concatenate([sample.logprobs for sample in batch]) for batch in data])
+                ).item(),
+                ref_logprobs=np.mean(
+                    np.concatenate([np.concatenate([sample.ref_logprobs for sample in batch]) for batch in data])
+                ).item(),
+            ),
+        }

adaptive_harmony/common/grpo.py

@@ -0,0 +1,260 @@
+import os
+from dataclasses import dataclass
+from typing import Callable, Sequence, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+
+from adaptive_harmony import (
+    CosineScheduler,
+    DataSet,
+    JobNotifier,
+    Logger,
+    StageNotifier,
+    StringThread,
+    TokenizedThread,
+    TrainingModel,
+)
+from adaptive_harmony.common.callbacks import RecipeCallback
+from adaptive_harmony.common.checkpointing import CheckpointManager
+from adaptive_harmony.core.utils import async_map, async_map_batch, get_minibatches, hash_hyperparams, log_args
+from adaptive_harmony.graders import BaseGrader
+from adaptive_harmony.metric_logger import StdoutLogger
+
+FloatArray: TypeAlias = NDArray[np.float32]
+
+
+@dataclass
+class Sample:
+    sample: TokenizedThread
+    logprobs: list[float]
+    ref_logprobs: list[float]
+    advantage: float
+    score: float
+    kl_div: list[float]
+    gen_len: float
+
+
+GRPO_HYPERPARAMS = {
+    "max_num_grpo_steps",
+    "completions_per_sample",
+    "lr",
+    "lr_scheduler",
+    "samples_per_batch",
+    "samples_per_mini_batch",
+    "mini_epochs_per_batch",
+    "max_grad_norm",
+    "clip_range",
+    "kl_beta",
+    "weight_decay",
+    "skip_nan_gradients",
+}
+
+
+class GRPO:
+    @log_args
+    @hash_hyperparams(include=GRPO_HYPERPARAMS)
+    def __init__(
+        self,
+        dataset: list[StringThread],
+        model: TrainingModel,
+        grader: BaseGrader,
+        logger: Logger = StdoutLogger(),
+        stage_notifier: StageNotifier = JobNotifier().stage_notifier("GRPO Training"),
+        callbacks: Sequence[RecipeCallback] = [],
+        max_num_grpo_steps: int | None = None,
+        completions_per_sample=8,
+        lr: float = 7.5e-7,
+        lr_scheduler: Callable[[float], float] | None = None,
+        samples_per_batch=128,
+        samples_per_mini_batch=128,
+        mini_epochs_per_batch=1,
+        max_grad_norm=1.0,
+        clip_range=0.1,
+        kl_beta=0.01,
+        weight_decay=0.0,
+        skip_nan_gradients: bool = False,
+        restart_from_checkpoint: str | None = None,
+        checkpoint_frequency: float = 0.2,
+    ):
+        # Core components
+        self.dataset = DataSet(dataset, allow_looping=True)
+        self.model = model
+        self.grader = grader
+        self.scoring_fn = grader.score_float_value
+        self.logger = logger
+        self.stage_notifier = stage_notifier
+        self.callbacks = callbacks
+        self.skip_nan_gradients = skip_nan_gradients
+        # GRPO HP's
+        self.max_num_batches = max_num_grpo_steps
+        self.completions_per_sample = completions_per_sample
+        self.lr_schedule = lr_scheduler or CosineScheduler(lr)
+        self.prompts_per_batch = samples_per_batch // completions_per_sample
+        self.samples_per_mini_batch = samples_per_mini_batch
+        self.total_num_samples = (
+            self.max_num_batches * self.prompts_per_batch if self.max_num_batches else len(self.dataset)
+        )
+        self.max_grad_norm = max_grad_norm
+        self.clip_range = clip_range
+        self.kl_beta = kl_beta
+        self.weight_decay = weight_decay
+        self.mini_epochs_per_batch = mini_epochs_per_batch
+
+        self.num_batches_processed = 0
+
+        self.checkpoint_manager = CheckpointManager(
+            recipe_name="GRPO",
+            dataset=self.dataset,
+            threads_dataset=dataset,
+            callbacks=self.callbacks,
+            hyperparams_hash=self._hyperparams_hash,  # type: ignore
+            job_id=os.environ.get("HARMONY_JOB_ID"),
+            checkpoint_frequency=checkpoint_frequency,
+            restart_from_checkpoint=restart_from_checkpoint,
+        )
+
+    @property
+    def training_completion_percentage(self):
+        return (
+            self.dataset.completion_percentage()
+            if self.max_num_batches is None
+            else min(self.num_batches_processed / self.max_num_batches, 1.0)
+        )
+
+    async def gen_data(self, sample: StringThread) -> list[Sample]:
+        assert self.model_ref is not None, "Calling `gen_data` before reference model has been set"
+
+        all_samples = await async_map(self.model.generate_tokens, [sample] * self.completions_per_sample)
+        string_samples = await async_map(self.model.detokenize_thread, all_samples)
+        all_scores = np.array(await async_map(self.scoring_fn, string_samples), dtype=np.float32)
+
+        advantages: FloatArray = all_scores - all_scores.mean()
+        advantages /= advantages.std() + 1e-8
+
+        logprobs = await async_map(self.model.logprobs_per_token, all_samples)
+        ref_logprobs = await async_map(self.model_ref.logprobs_per_token, all_samples)
+        kl = [
+            (np.array(lp, dtype=np.float32) - np.array(ref_lp, dtype=np.float32)).tolist()
+            for lp, ref_lp in zip(logprobs, ref_logprobs)
+        ]
+
+        samples = []
+        for i in range(len(logprobs)):
+            samples.append(
+                Sample(
+                    sample=all_samples[i],
+                    logprobs=logprobs[i],
+                    ref_logprobs=ref_logprobs[i],
+                    advantage=advantages[i],
+                    score=all_scores[i],
+                    kl_div=kl[i],
+                    gen_len=all_samples[i].len_last_turn(),
+                )
+            )
+        return samples
+
+    async def train_sample(self, sample: Sample):
+        await self.model.train_grpo(
+            sample.sample,
+            sample.logprobs,
+            sample.ref_logprobs,
+            [sample.advantage] * len(sample.logprobs),
+            self.clip_range,
+            self.kl_beta,
+        )
+
+    async def _recipe_specific_checkpoint_loading(self, checkpoint_data: dict) -> None:
+        self.num_batches_processed = checkpoint_data["num_batches_processed"]
+
+        model_checkpoint_name = checkpoint_data["model_checkpoint_name"]
+        await self.model.load(f"model_registry://{model_checkpoint_name}")
+
+        self.model.set_optim_step(checkpoint_data["optim_step"])
+
+    async def _recipe_specific_checkpoint_saving(self) -> dict:
+        progress_pct = int(self.training_completion_percentage * 100)
+        model_checkpoint_name = f"checkpoint-{self.checkpoint_manager.job_id}-{progress_pct}-policy"
+        model_checkpoint_name = await self.model.save(model_checkpoint_name, inference_only=False)
+
+        return {
+            "num_batches_processed": self.num_batches_processed,
+            "model_checkpoint_name": model_checkpoint_name,
+            "optim_step": self.model.get_optim_step(),
+        }
+
+    async def run(self):
+        self.model_ref = await self.model.clone_inf()
+        await self.checkpoint_manager.maybe_restore_checkpoint(self._recipe_specific_checkpoint_loading)
+
+        self.stage_notifier.report_progress(
+            tot_num_samples=self.total_num_samples,
+            processed_num_samples=self.dataset.idx,
+            monitoring_link=self.logger.training_monitoring_link,
+        )
+
+        while self.training_completion_percentage < 1.0:
+            self.num_batches_processed += 1
+
+            for callback in self.callbacks:
+                if logs := await callback.maybe_call(self.training_completion_percentage):
+                    self.logger(logs)
+
+            # Generate training samples
+            data = await async_map_batch(self.gen_data, self.dataset, self.prompts_per_batch)
+            scorer_logs = self.grader.get_logs(clear=True)
+            batch_logs = {
+                **{f"rewards/{key}": value for key, value in scorer_logs.items()},
+                **self.get_train_batch_logs(data),
+            }
+
+            current_lr = self.lr_schedule(self.training_completion_percentage)
+            # Train on generated samples
+            flattened_data = sum([inner_list for inner_list in data], start=[])
+            minibatches = get_minibatches(flattened_data, self.samples_per_mini_batch, self.mini_epochs_per_batch)
+            for idx, mini_batch in enumerate(minibatches):
+                await async_map(self.train_sample, mini_batch)
+                optim_logs = await self.model.optim_step(
+                    current_lr,
+                    wd=self.weight_decay,
+                    max_grad_norm=self.max_grad_norm,
+                    skip_nan_gradients=self.skip_nan_gradients,
+                )
+                if idx == len(minibatches) - 1:
+                    # only log tables and full batch-related logs on the final minibatch
+                    self.logger(optim_logs | batch_logs)
+                else:
+                    self.logger(optim_logs | dict(completion_percentage=self.training_completion_percentage))
+
+            self.stage_notifier.report_progress(
+                tot_num_samples=self.total_num_samples,
+                processed_num_samples=self.dataset.idx,
+                monitoring_link=self.logger.training_monitoring_link,
+            )
+
+            if await self.checkpoint_manager.maybe_checkpoint(
+                self.training_completion_percentage, self._recipe_specific_checkpoint_saving
+            ):
+                break
+
+    def get_train_batch_logs(self, data: list[list[Sample]]) -> dict:
+        return {
+            **dict(
+                completion_percentage=self.training_completion_percentage,
+                percentage_no_advantages=np.mean(
+                    [all(sample.advantage == batch[0].advantage for sample in batch) for batch in data]
+                ).item(),
+                score_mean=np.mean([[sample.score for sample in batch] for batch in data]).item(),
+                score_std=np.std([[sample.score for sample in batch] for batch in data]).item(),
+                kl_div=np.mean([[np.mean(sample.kl_div) for sample in batch] for batch in data]).item(),
+                advantages=np.mean(np.concatenate([[sample.advantage for sample in batch] for batch in data])).item(),
+                generation_length=np.mean([np.mean([sample.gen_len for sample in batch]) for batch in data]).item(),
+                logprobs=np.mean(
+                    np.concatenate([np.concatenate([sample.logprobs for sample in batch]) for batch in data])
+                ).item(),
+                ref_logprobs=np.mean(
+                    np.concatenate([np.concatenate([sample.ref_logprobs for sample in batch]) for batch in data])
+                ).item(),
+            ),
+            **{"training/completion_percentage": self.training_completion_percentage},
+        }

adaptive_harmony/common/gspo.py

@@ -0,0 +1,70 @@
+from typing import Callable, Sequence
+
+from adaptive_harmony import (
+    JobNotifier,
+    Logger,
+    StageNotifier,
+    StringThread,
+    TrainingModel,
+)
+from adaptive_harmony.common.callbacks import RecipeCallback
+from adaptive_harmony.common.grpo import GRPO, Sample
+from adaptive_harmony.graders import BaseGrader
+from adaptive_harmony.metric_logger import StdoutLogger
+
+
+class GSPO(GRPO):  # grpo already logs args so we don't do it here
+    def __init__(
+        self,
+        dataset: list[StringThread],
+        model: TrainingModel,
+        grader: BaseGrader,
+        logger: Logger = StdoutLogger(),
+        stage_notifier: StageNotifier = JobNotifier().stage_notifier("GSPO Training"),
+        callbacks: Sequence[RecipeCallback] = [],
+        max_num_gspo_steps: int | None = None,
+        completions_per_sample=8,
+        lr: float = 7.5e-7,
+        lr_scheduler: Callable[[float], float] | None = None,
+        samples_per_batch=128,
+        samples_per_mini_batch=128,
+        mini_epochs_per_batch=1,
+        max_grad_norm=1.0,
+        # wildly different defaults than GRPO because we are looking at the
+        # entire sequence at once, I tried the number from the GSPO paper but
+        # it was skipping ~60% of the samples. I had good success with 0.01 but
+        # it's not properly swept yet.
+        clip_range=0.01,
+        kl_beta=0.01,
+        weight_decay=0.0,
+    ):
+        super().__init__(
+            dataset,
+            model,
+            grader,
+            logger,
+            stage_notifier,
+            callbacks,
+            max_num_grpo_steps=max_num_gspo_steps,
+            completions_per_sample=completions_per_sample,
+            lr=lr,
+            lr_scheduler=lr_scheduler,
+            samples_per_batch=samples_per_batch,
+            samples_per_mini_batch=samples_per_mini_batch,
+            mini_epochs_per_batch=mini_epochs_per_batch,
+            max_grad_norm=max_grad_norm,
+            clip_range=clip_range,
+            kl_beta=kl_beta,
+            weight_decay=weight_decay,
+        )
+
+    async def train_sample(self, sample: Sample):
+        await self.model.train_gspo(
+            sample.sample,
+            sample.logprobs,
+            sample.ref_logprobs,
+            [sample.advantage],  # only diff with grpo is we train with a single advantage
+            self.clip_range,
+            self.clip_range,
+            self.kl_beta,
+        )